Blame for tiny_computer_agent/__main__.py

"""Command line entry point.

    python -m tiny_computer_agent                     # scripted demo, dry-run
    python -m tiny_computer_agent --doctor            # what can this machine do?
    python -m tiny_computer_agent --image example     # replay a screenshot, dry-run
    python -m tiny_computer_agent --image shot.png    # replay your own screenshot
    python -m tiny_computer_agent --live              # your real screen, dry-run
    python -m tiny_computer_agent --live --real-input "click 'OK'"

2.0 adds two subcommands, for everything after that first run:

    python -m tiny_computer_agent skills list                  # what has been learned
    python -m tiny_computer_agent skills learn "click 'Save'" --script '[{"kind":"click","label":"Save"}]' --verify-text Saved
    python -m tiny_computer_agent batch tasks/day.json          # a queue, resumable, reported

Real mouse and keyboard input additionally requires
``TINY_COMPUTER_AGENT_ALLOW_REAL_INPUT=1``. Dry run is the default, so you can read a
whole episode before anything touches your pointer.
"""

from __future__ import annotations

import argparse
import json
import os
import platform
import sys

from tiny_computer_agent.act import DryRunExecutor, PyAutoGuiExecutor
from tiny_computer_agent.diff import ChangeVerifier
from tiny_computer_agent.brain import MockBrain, PlannerBrain
from tiny_computer_agent.demo import fake_desktop
from tiny_computer_agent.desktop import (
    BRAIN_NAMES, DEFAULT_KILL_SWITCH, Check, display_available, run_checks,
)
from tiny_computer_agent.loop import AgentLoop, Step
from tiny_computer_agent.observe import (
    Element, FakeScreen, GrabScreen, ImageScreen, ScreenUnavailable,
)
from tiny_computer_agent.perception import IndexedScreen, detector_available, make_detector
from tiny_computer_agent.safety import Policy
from tiny_computer_agent.runner import Task, TaskRunner
from tiny_computer_agent.skills import DEFAULT_SKILLS_PATH, SkillBrain, SkillLibrary, record_skill
from tiny_computer_agent.trace import Trace, score_trace
from tiny_computer_agent.audit import AuditLog, verify_audit
from tiny_computer_agent.gates import Gate, gate_from_spec

ALLOW_REAL_INPUT_ENV = "TINY_COMPUTER_AGENT_ALLOW_REAL_INPUT"
EXAMPLE_IMAGE = "example"


def demo_screen() -> FakeScreen:
    """A three-frame desktop: a dialog, then an editor, then a saved file."""
    return FakeScreen(
        [
            (Element("Terminal", x=0, y=0, width=40, height=24),
             Element("Files", x=60, y=0, width=40, height=24)),
            (Element("Name", kind="input", x=100, y=120, width=200, height=24),
             Element("Save", x=320, y=120, width=80, height=24)),
            (Element("Save", x=320, y=120, width=80, height=24),),
        ],
        text="Untitled document - unsaved changes",
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="tiny-computer-agent",
        description="Drive a computer with a tiny local model (dry-run by default).",
    )
    parser.add_argument("goal", nargs="?", default="click 'Save' to save the document")
    parser.add_argument("--max-steps", type=int, default=6)
    parser.add_argument("--script", default=None,
                        help="JSON list of actions to replay instead of using a brain")
    parser.add_argument("--doctor", action="store_true",
                        help="report what this machine can and cannot do, then exit")
    parser.add_argument("--image", default=None,
                        help="replay a screenshot instead of the screen; "
                             f"'{EXAMPLE_IMAGE}' uses the built-in one")
    parser.add_argument("--brain", default="planner", choices=BRAIN_NAMES,
                        help="planner (rule-based, no weights) | text-planner | vlm")
    parser.add_argument("--detector", default="auto",
                        choices=("auto", "best", "atspi", "ocr", "grid"),
                        help="how to build the element index "
                             "(best = accessibility tree, then OCR, then grid)")
    parser.add_argument("--live", action="store_true",
                        help="capture your real screen instead of the scripted demo")
    parser.add_argument("--real-input", action="store_true",
                        help="actually move the mouse and type (requires the allow env var)")
    parser.add_argument("--monitor", type=int, default=1, help="1 is the primary monitor")
    parser.add_argument("--scale", type=float, default=None,
                        help="screen pixels per captured pixel (HiDPI, secondary monitors)")
    parser.add_argument("--no-ocr", action="store_true",
                        help="skip OCR and index a numbered grid instead")
    parser.add_argument("--max-width", type=int, default=1600,
                        help="downscale captures wider than this before indexing")
    parser.add_argument("--verify-text", default=None, metavar="TEXT",
                        help="stop with success once this text appears on screen "
                             "(otherwise a successful action that removes its own "
                             "target is indistinguishable from a failure)")
    parser.add_argument("--allow-grid", action="store_true",
                        help="permit clicks on numbered grid cells with --real-input")
    parser.add_argument("--kill-switch", default=DEFAULT_KILL_SWITCH,
                        help="create this file to abort the run before the next action")
    parser.add_argument("--trace", default=None, metavar="PATH",
                        help="write a portable trace of the run (frames, no pixels) so it "
                             "can be replayed and scored offline")
    parser.add_argument("--gate", default=None,
                        help="the operator's yes, before each action: read-only | interactive "
                             "| labels:Save,Cancel | allow-all")
    parser.add_argument("--audit", default=None, metavar="PATH",
                        help="append a hash-chained record of this run to a JSONL file")
    parser.add_argument("--verify-audit", default=None, metavar="PATH",
                        help="check an audit log's chain and exit (0 intact, 1 tampered)")
    parser.add_argument("--expected", default=None, metavar="LABEL",
                        help="the element label that is the right answer for this goal; "
                             "recorded in the trace so the run can be scored")
    return parser


def report(checks: list[Check]) -> int:
    """Print a preflight report; non-zero when something required is missing."""
    blocking = [c for c in checks if c.required and not c.ok]
    lines = [
        f"tiny-computer-agent preflight ({platform.system()} {platform.machine()}, "
        f"python {platform.python_version()})",
        "",
        *(c.render() for c in checks),
        "",
    ]
    if blocking:
        lines.append(f"NOT READY: {len(blocking)} blocking problem(s): "
                     + ", ".join(c.name for c in blocking))
    else:
        lines.append("READY: the harness can drive this desktop.")
    print("\n".join(lines))
    return 1 if blocking else 0


def _probe_live(args: argparse.Namespace) -> Check:
    """Actually grab one frame, because only a real capture proves it works."""
    try:
        observation = GrabScreen(
            monitor=args.monitor, max_width=args.max_width, scale=args.scale
        ).capture()
        frame = observation.frame
        return Check("live capture", True, f"captured {frame.width}x{frame.height} pixels")
    except ScreenUnavailable as exc:
        return Check("live capture", False, str(exc),
                     fix="run this on the desktop session itself, not a headless shell")


def _gate_from_spec_arg(args: argparse.Namespace) -> tuple[Gate | None, str | None]:
    """The operator's gate, from the one-line spelling on the command line."""
    spec = getattr(args, "gate", None)
    if not spec:
        return None, None
    try:
        return gate_from_spec(spec), None
    except ValueError as exc:
        return None, f"bad --gate: {exc}"


def _audit_from_arg(args: argparse.Namespace) -> AuditLog | None:
    """Append-only evidence for this run, if the operator asked for it."""
    path = getattr(args, "audit", None)
    return AuditLog(path) if path else None


def _build_screen(args: argparse.Namespace):
    """(screen, description). Raises ScreenUnavailable / OSError on failure."""
    detector_name = "grid" if args.no_ocr else args.detector
    if args.image is not None:
        source = ImageScreen(fake_desktop() if args.image == EXAMPLE_IMAGE else args.image)
        return IndexedScreen(source, make_detector(detector_name)), f"screenshot replay: {args.image}"
    if args.live:
        source = GrabScreen(monitor=args.monitor, max_width=args.max_width, scale=args.scale)
        return IndexedScreen(source, make_detector(detector_name)), f"real desktop (monitor {args.monitor})"
    return demo_screen(), "scripted demo"


SUBCOMMANDS = ("skills", "batch")


def _screen_arguments(parser: argparse.ArgumentParser) -> None:
    """The screen flags the subcommands share with the main parser."""
    parser.add_argument("--image", default=None,
                        help=f"replay a screenshot; '{EXAMPLE_IMAGE}' uses the built-in one")
    parser.add_argument("--live", action="store_true", help="use this machine's real screen")
    parser.add_argument("--detector", default="auto",
                        choices=("auto", "best", "atspi", "ocr", "grid"))
    parser.add_argument("--no-ocr", action="store_true", help="index a numbered grid instead")
    parser.add_argument("--monitor", type=int, default=1)
    parser.add_argument("--scale", type=float, default=None)
    parser.add_argument("--max-width", type=int, default=1600)
    parser.add_argument("--path", default=str(DEFAULT_SKILLS_PATH),
                        help="the skill library to read or write")


def build_skills_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="tiny-computer-agent skills",
        description="Skills are verified runs, written down once, replayed forever.")
    actions = parser.add_subparsers(dest="action", required=True)
    listing = actions.add_parser("list", help="what the library holds")
    show = actions.add_parser("show", help="one skill, step by step")
    show.add_argument("name")
    remove = actions.add_parser("remove", help="forget a skill")
    remove.add_argument("name")
    learn = actions.add_parser("learn", help="run a script once, and keep it if it verifies")
    learn.add_argument("goal")
    learn.add_argument("--script", required=True, help="JSON list of actions to perform once")
    learn.add_argument("--verify-text", default=None, metavar="TEXT",
                       help="what 'it worked' looks like on screen: some text appears")
    learn.add_argument("--verify-gone", default=None, metavar="LABEL",
                       help="or: this element must disappear (the dialog closed)")
    learn.add_argument("--verify-appeared", default=None, metavar="LABEL",
                       help="or: this element must appear (the menu opened)")
    learn.add_argument("--name", default=None, help="name the skill yourself")
    _screen_arguments(learn)                      # which also carries --path
    for read_only in (listing, show, remove):     # these only ever need the library
        read_only.add_argument("--path", default=str(DEFAULT_SKILLS_PATH),
                               help="the skill library to read or write")
    return parser


def skills_main(argv: list[str]) -> int:
    args = build_skills_parser().parse_args(argv)
    library = SkillLibrary.load(args.path)
    if args.action == "list":
        if not library:
            print(f"no skills in {args.path} yet. Learn one with:\n"
                  f"  python -m tiny_computer_agent skills learn \"click 'Save'\" "
                  f"--script '[{{\"kind\": \"click\", \"label\": \"Save\"}}]' --verify-text Saved")
            return 0
        rows = library.to_rows()
        widths = {key: max(len(key), *(len(str(row[key])) for row in rows)) for key in rows[0]}
        print(f"{len(rows)} skill(s) in {args.path}")
        print("  " + "  ".join(key.ljust(widths[key]) for key in rows[0]))
        for row in rows:
            print("  " + "  ".join(str(row[key]).ljust(widths[key]) for key in row))
        return 0
    if args.action == "show":
        skill = library.get(args.name)
        if skill is None:
            print(f"no skill named {args.name!r} in {args.path}", file=sys.stderr)
            return 1
        print(skill.line())
        print(f"  verify: {skill.verification()}")
        print(f"  entry:  {', '.join(skill.preconditions()) or '(none - no label to re-ground)'}")
        for index, step in enumerate(skill.steps, 1):
            print(f"  {index}. {step.describe()}")
        return 0
    if args.action == "remove":
        library = SkillLibrary.load(args.path)
        if not library.remove(args.name):
            print(f"no skill named {args.name!r} in {args.path}", file=sys.stderr)
            return 1
        print(f"removed {args.name!r}; {len(library)} left in {library.save()}")
        return 0

    # learn: the first run is scripted, and only a verified one becomes a skill
    if not (args.verify_text or args.verify_gone or args.verify_appeared):
        print("refusing to learn: pass --verify-text (or --verify-gone / --verify-appeared) so "
              "the run has evidence. A skill recorded from a brain's own 'done' would be a "
              "guess with a name.", file=sys.stderr)
        return 2
    try:
        screen, description = _build_screen(args)
    except ScreenUnavailable as exc:
        print(f"cannot see the screen: {exc}", file=sys.stderr)
        return 3
    try:
        script = json.loads(args.script)
    except json.JSONDecodeError as exc:
        print(f"--script is not valid JSON: {exc}", file=sys.stderr)
        return 2
    _, index_detail = detector_available("grid" if args.no_ocr else args.detector)
    print(f"learning {args.goal!r} on {description} ({index_detail})")
    if args.verify_text:
        needle = args.verify_text.lower()
        verify: object = lambda observation: needle in observation.text.lower()
    else:
        verify = ChangeVerifier(appeared=(args.verify_appeared,) if args.verify_appeared else (),
                                disappeared=(args.verify_gone,) if args.verify_gone else ())
    loop = AgentLoop(screen=screen, brain=MockBrain(script), executor=DryRunExecutor(),
                     policy=Policy(dry_run=True, step_budget=len(script) + 2),
                     on_step=lambda step: print("  " + step.line()),
                     verify=verify)
    episode = loop.run(args.goal)
    print()
    print(episode.render())
    skill = record_skill(episode, name=args.name, library=library, verify_text=args.verify_text or "",
                         verify_gone=args.verify_gone or "",
                         verify_appeared=args.verify_appeared or "")
    if skill is None:
        print("\nnot learned: the run did not end verified, so there is nothing worth keeping.",
              file=sys.stderr)
        return 1
    print(f"\nlearned {skill.name!r} ({len(skill.steps)} step(s)) -> {library.save()}")
    return 0


def build_batch_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="tiny-computer-agent batch",
        description="Run a queue of goals, resuming what already finished, and report on it.")
    parser.add_argument("tasks", help="JSON: a list of {goal, name?, verify_text?, budget?}")
    parser.add_argument("--state", default=None,
                        help="JSONL file appended as each task finishes (resume from it)")
    parser.add_argument("--report", default=None, help="write the report as JSON here")
    parser.add_argument("--max-steps", type=int, default=6)
    parser.add_argument("--real-input", action="store_true",
                        help="actually drive the mouse (requires the allow env var)")
    parser.add_argument("--brain", default="planner", choices=BRAIN_NAMES)
    _screen_arguments(parser)
    return parser


def batch_main(argv: list[str]) -> int:
    args = build_batch_parser().parse_args(argv)
    try:
        payload = json.loads(open(args.tasks).read())
    except (OSError, json.JSONDecodeError) as exc:
        print(f"cannot read {args.tasks}: {exc}", file=sys.stderr)
        return 2
    if not isinstance(payload, list) or not payload:
        print(f"{args.tasks} should hold a non-empty list of tasks", file=sys.stderr)
        return 2
    try:
        from_args = [Task(goal=entry["goal"], name=entry.get("name", ""),
                          budget=entry.get("budget"), verify_text=entry.get("verify_text", ""),
                          verify_gone=entry.get("verify_gone", ""),
                          verify_appeared=entry.get("verify_appeared", ""))
                     for entry in payload]
    except (TypeError, KeyError) as exc:
        print(f"a task needs at least a 'goal': {exc}", file=sys.stderr)
        return 2

    if args.real_input and os.environ.get(ALLOW_REAL_INPUT_ENV) != "1":
        print(f"refusing --real-input: set {ALLOW_REAL_INPUT_ENV}=1 first.", file=sys.stderr)
        return 2
    try:
        screen, description = _build_screen(args)
    except ScreenUnavailable as exc:
        print(f"cannot see the screen: {exc}", file=sys.stderr)
        return 3

    def make_loop(task: Task) -> AgentLoop:
        from tiny_computer_agent.models import make_model_brain
        brain = PlannerBrain() if args.brain == "planner" else make_model_brain(args.brain)
        return AgentLoop(screen=screen, brain=brain,
                         executor=PyAutoGuiExecutor() if args.real_input else DryRunExecutor(),
                         policy=Policy(dry_run=not args.real_input, step_budget=task.budget or args.max_steps))

    print(f"batch: {len(from_args)} task(s) on {description}, "
          f"{'REAL INPUT' if args.real_input else 'dry-run'}")
    report = TaskRunner(make_loop, state_path=args.state,
                        on_task=lambda outcome: print("  " + outcome.line())).run(from_args)
    print()
    print(report.render())
    if args.report:
        print(f"report: wrote {report.save(args.report)}")
    return 0 if report.failed == 0 else 1


def main(argv: list[str] | None = None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)
    if argv and argv[0] in SUBCOMMANDS:
        return skills_main(argv[1:]) if argv[0] == "skills" else batch_main(argv[1:])
    args = build_parser().parse_args(argv)
    detector_name = "grid" if args.no_ocr else args.detector

    if args.verify_audit:
        intact, detail = verify_audit(args.verify_audit)
        print(f"{'ok  ' if intact else 'FAIL'} audit chain: {detail}")
        return 0 if intact else 1

    gate, gate_error = _gate_from_spec_arg(args)
    if gate_error is not None:
        print(gate_error, file=sys.stderr)
        return 2
    audit = _audit_from_arg(args)

    if args.doctor:
        checks = run_checks(brain=args.brain, detector=detector_name,
                            real_input=args.real_input, kill_switch_path=args.kill_switch)
        if args.live:
            checks.append(_probe_live(args))
        return report(checks)

    if args.real_input and os.environ.get(ALLOW_REAL_INPUT_ENV) != "1":
        print(f"refusing --real-input: set {ALLOW_REAL_INPUT_ENV}=1 to confirm you "
              f"understand this will drive your actual mouse and keyboard.", file=sys.stderr)
        return 2

    if args.brain == "vlm" and not (args.live or args.image):
        print("the vlm brain needs --live (or --image) so it has pixels to look at; "
              "use --brain planner to exercise the harness offline.", file=sys.stderr)
        return 2

    if args.live:
        available, detail = display_available()
        if not available:
            print(f"live capture needs a display session: {detail}")
            print(f"refusing to run live: {detail}. Try --image to replay a screenshot, "
                  f"or --doctor to diagnose.", file=sys.stderr)
            return 2

    try:
        screen, description = _build_screen(args)
    except ScreenUnavailable as exc:
        print(f"cannot see the screen: {exc}", file=sys.stderr)
        return 3
    except Exception as exc:
        print(f"could not start the screen: {type(exc).__name__}: {exc}", file=sys.stderr)
        return 2

    if args.script:
        try:
            script = json.loads(args.script)
        except json.JSONDecodeError as exc:
            print(f"--script is not valid JSON: {exc}", file=sys.stderr)
            return 2
        brain = MockBrain(script)
    elif args.brain == "planner":
        brain = PlannerBrain()
    else:
        from tiny_computer_agent.models import make_model_brain

        try:
            brain = make_model_brain(args.brain)
        except Exception as exc:
            print(f"could not load the {args.brain} brain: {type(exc).__name__}: {exc}",
                  file=sys.stderr)
            return 2

    policy = Policy(
        dry_run=not args.real_input,
        step_budget=args.max_steps,
        kill_switch_path=args.kill_switch,
        allow_grid_clicks=args.allow_grid,
    )
    executor = PyAutoGuiExecutor() if args.real_input else DryRunExecutor()

    mode = "REAL INPUT (mouse and keyboard live)" if args.real_input else "dry-run"
    print(f"goal: {args.goal}")
    print(f"mode: {mode}  |  screen: {description}  |  brain: {args.brain}  |  "
          f"kill switch: create {args.kill_switch} to stop")
    if gate is not None:
        print(f"gate: {args.gate} - every action needs the operator's yes")
    if audit is not None:
        print(f"audit: recording this run to {audit.path}")
    if args.image is not None or args.live:
        _, index_detail = detector_available(detector_name)
        print(f"index: {detector_name} - {index_detail}")
    if args.image is not None:
        # One read, so the header can state how big the action space is.
        try:
            preview = screen.capture()
        except ScreenUnavailable as exc:
            print(f"cannot see the screen: {exc}", file=sys.stderr)
            return 3
        print(f"elements detected: {len(preview.elements)}")
    print()

    warned = False

    verify = None
    if args.verify_text:
        _needle = args.verify_text.lower()

        def verify(observation) -> bool:  # noqa: F811 - the loop calls this
            return _needle in observation.text.lower()

    def on_step(step: Step) -> None:
        nonlocal warned
        elements = step.observation.elements
        if not warned and elements and all(e.kind == "grid" for e in elements):
            warned = True
            _, index_detail = detector_available(detector_name)
            print(f"warning: no text was found on screen, so the index is "
                  f"{len(elements)} numbered grid cells and the agent is choosing "
                  f"blind ({index_detail})", file=sys.stderr)
        print("  " + step.line())

    try:
        episode = AgentLoop(screen=screen, brain=brain, executor=executor,
                            policy=policy, on_step=on_step, verify=verify,
                            gate=gate, audit=audit).run(args.goal)
    except ScreenUnavailable as exc:
        print(f"cannot see the screen: {exc}", file=sys.stderr)
        print("run `python -m tiny_computer_agent --doctor --live` for details.", file=sys.stderr)
        return 3

    print()
    print(episode.render())
    if audit is not None:
        print(audit.render())

    if args.trace:
        trace = Trace.record(episode, expected=args.expected, brain=args.brain,
                             scenario=args.goal)
        print(f"trace: wrote {trace.save(args.trace)} ({trace.frames} frames)")
        if args.expected:
            print("score: " + score_trace(trace).line().strip())
            print(f"       replay it later: Trace.load({str(args.trace)!r}).render()")
    return 0 if episode.success else 1


if __name__ == "__main__":
    raise SystemExit(main())