scripts/colab_a0.py

#!/usr/bin/env python3
"""Small lifecycle CLI for Agent Zero smoke tests in Google Colab."""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import signal
import subprocess
import sys
import time
from urllib.request import urlopen

OUTPUT_TAIL_CHARS = 12_000


def _paths(state_dir: Path) -> tuple[Path, Path]:
    state_dir.mkdir(parents=True, exist_ok=True)
    return state_dir / "agent-zero.pid", state_dir / "agent-zero.log"


def _pid(path: Path) -> int | None:
    try:
        value = int(path.read_text().strip())
        os.kill(value, 0)
        proc = Path(f"/proc/{value}")
        if proc.exists():
            fields = (proc / "stat").read_text().split()
            if len(fields) > 2 and fields[2] == "Z":
                return None
            command = (proc / "cmdline").read_bytes().replace(b"\0", b" ")
            if b"run_ui.py" not in command:
                return None
        return value
    except (FileNotFoundError, PermissionError, ValueError, ProcessLookupError):
        return None


def health(url: str, timeout: float) -> bool:
    try:
        with urlopen(url.rstrip("/") + "/api/health", timeout=timeout) as response:
            return 200 <= response.status < 300
    except OSError:
        return False


def start(root: Path, state_dir: Path, host: str, port: int, wait: float) -> int:
    pid_path, log_path = _paths(state_dir)
    if running := _pid(pid_path):
        print(json.dumps({"status": "running", "pid": running}))
        return 0
    log = log_path.open("ab")
    env = {**os.environ, "WEB_UI_HOST": host, "WEB_UI_PORT": str(port)}
    process = subprocess.Popen(
        [sys.executable, "run_ui.py"],
        cwd=root,
        env=env,
        stdout=log,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    log.close()
    pid_path.write_text(str(process.pid))
    url = f"http://127.0.0.1:{port}"
    deadline = time.monotonic() + wait
    while time.monotonic() < deadline:
        if process.poll() is not None:
            print(
                json.dumps(
                    {"status": "failed", "pid": process.pid, "log": str(log_path)}
                )
            )
            return 1
        if health(url, 2):
            print(json.dumps({"status": "healthy", "pid": process.pid, "url": url}))
            return 0
        time.sleep(0.5)
    print(
        json.dumps(
            {"status": "starting", "pid": process.pid, "url": url, "log": str(log_path)}
        )
    )
    return 2


def stop(state_dir: Path, wait: float) -> int:
    pid_path, _ = _paths(state_dir)
    running = _pid(pid_path)
    if running is None:
        pid_path.unlink(missing_ok=True)
        print(json.dumps({"status": "stopped"}))
        return 0
    os.killpg(running, signal.SIGTERM)
    deadline = time.monotonic() + wait
    while time.monotonic() < deadline:
        if _pid(pid_path) is None:
            pid_path.unlink(missing_ok=True)
            print(json.dumps({"status": "stopped", "pid": running}))
            return 0
        time.sleep(0.25)
    os.killpg(running, signal.SIGKILL)
    pid_path.unlink(missing_ok=True)
    print(json.dumps({"status": "killed", "pid": running}))
    return 0


def execute(root: Path, command: list[str], timeout: float) -> int:
    """Execute an argv command in the checked-out repository and emit evidence."""
    if not command:
        print(json.dumps({"result": "BLOCKED", "error": "command is required"}))
        return 2
    root = root.resolve()
    if not (root / ".git").exists():
        print(json.dumps({"result": "BLOCKED", "error": "root is not a git checkout"}))
        return 2
    started = time.monotonic()
    try:
        completed = subprocess.run(
            command,
            cwd=root,
            capture_output=True,
            text=True,
            timeout=max(1.0, timeout),
            check=False,
        )
        commit = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=root,
            capture_output=True,
            text=True,
            timeout=10,
            check=False,
        ).stdout.strip()
        evidence = {
            "execution_id": f"colab-{int(time.time())}-{os.getpid()}",
            "backend": "colab",
            "repository": str(root),
            "commit": commit,
            "python": sys.version.split()[0],
            "command": command,
            "exit_code": completed.returncode,
            "passed": None,
            "failed": None,
            "skipped": None,
            "blocked": 0,
            "duration": round(time.monotonic() - started, 3),
            "artifacts": [],
            "stdout_tail": completed.stdout[-OUTPUT_TAIL_CHARS:],
            "stderr_tail": completed.stderr[-OUTPUT_TAIL_CHARS:],
            "result": "PASS" if completed.returncode == 0 else "FAIL",
        }
        print(json.dumps(evidence))
        return completed.returncode
    except subprocess.TimeoutExpired as error:
        print(
            json.dumps(
                {
                    "execution_id": f"colab-{int(time.time())}-{os.getpid()}",
                    "backend": "colab",
                    "repository": str(root),
                    "commit": "",
                    "python": sys.version.split()[0],
                    "command": command,
                    "exit_code": None,
                    "passed": 0,
                    "failed": 0,
                    "skipped": 0,
                    "blocked": 1,
                    "duration": round(time.monotonic() - started, 3),
                    "artifacts": [],
                    "stdout_tail": str(error.stdout or "")[-OUTPUT_TAIL_CHARS:],
                    "stderr_tail": str(error.stderr or "")[-OUTPUT_TAIL_CHARS:],
                    "result": "BLOCKED",
                    "error": "timeout",
                }
            )
        )
        return 124


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=("start", "health", "stop", "exec"))
    parser.add_argument(
        "--root", type=Path, default=Path(__file__).resolve().parents[1]
    )
    parser.add_argument("--state-dir", type=Path, default=Path("/content/.agent-zero"))
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--wait", type=float, default=60)
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args()
    if args.action == "start":
        return start(args.root, args.state_dir, args.host, args.port, args.wait)
    if args.action == "stop":
        return stop(args.state_dir, args.wait)
    if args.action == "exec":
        command = args.command[1:] if args.command[:1] == ["--"] else args.command
        return execute(args.root, command, args.wait)
    ok = health(f"http://127.0.0.1:{args.port}", min(args.wait, 10))
    print(json.dumps({"status": "healthy" if ok else "unhealthy", "port": args.port}))
    return 0 if ok else 1


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