tai/taicpp/tests/capmcp_cli_diff.py

#!/usr/bin/env python3
"""Runs the same `--cap-manifest` / `--cap ...` argv through tai.py and through
the C++ test driver (build/capmcp/capmcp_cli) and diffs stdout, stderr and the
exit code. Every case runs in a throwaway HOME / TAI_HOME / cwd.

    python3 -I tests/capmcp_cli_diff.py CPP_DRIVER [--tai PATH_TO_TAI_PY] [--filter TEXT] [--verbose]

Exit 0 only when all cases are byte-identical."""
import http.server
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
import textwrap
import threading

HERE = os.path.dirname(os.path.abspath(__file__))
TAI = os.path.join(os.path.dirname(os.path.dirname(HERE)), "tai.py")

FAKE_STREAM_AGENT = r'''
import json, sys
print(json.dumps({"type": "system", "subtype": "init", "session_id": "fake-1"}), flush=True)
n = 0
for line in sys.stdin:
    try:
        frame = json.loads(line)
    except ValueError:
        continue
    n += 1
    text = frame.get("message")
    print(json.dumps({"type": "assistant", "message": {"id": "m%d" % n, "role": "assistant", "content": [
        {"type": "text", "text": "echo %d: %s" % (n, text)},
        {"type": "tool_use", "id": "t%d" % n, "name": "Read", "input": {"path": "x"}}]}}), flush=True)
    print(json.dumps({"type": "user", "message": {"role": "user", "content": [
        {"type": "tool_result", "tool_use_id": "t%d" % n, "content": "ok", "is_error": False}]}}), flush=True)
    print(json.dumps({"type": "result", "subtype": "success", "usage": {"input_tokens": 3, "output_tokens": 5}, "duration_ms": 7, "total_cost_usd": 0.5}), flush=True)
'''

FAKE_ACP_AGENT = r'''
import json, sys
def send(obj):
    print(json.dumps(obj), flush=True)
for line in sys.stdin:
    msg = json.loads(line)
    if "method" in msg and "id" in msg and msg["method"] == "initialize":
        send({"jsonrpc": "2.0", "id": msg["id"], "result": {"protocolVersion": 1, "agentCapabilities": {}}})
    elif msg.get("method") == "session/new":
        send({"jsonrpc": "2.0", "id": msg["id"], "result": {"sessionId": "s1"}})
    elif msg.get("method") == "session/prompt":
        text = msg["params"]["prompt"][0]["text"]
        send({"jsonrpc": "2.0", "method": "session/update", "params": {"sessionId": "s1", "update": {"sessionUpdate": "agent_message_chunk", "content": {"type": "text", "text": "acp says " + text}}}})
        send({"jsonrpc": "2.0", "method": "session/update", "params": {"sessionId": "s1", "update": {"sessionUpdate": "tool_call", "toolCallId": "c1", "title": "Read file"}}})
        send({"jsonrpc": "2.0", "method": "session/update", "params": {"sessionId": "s1", "update": {"sessionUpdate": "tool_call_update", "toolCallId": "c1", "status": "completed", "content": [{"type": "text", "text": "done"}]}}})
        send({"jsonrpc": "2.0", "method": "session/update", "params": {"sessionId": "s1", "update": {"sessionUpdate": "usage_update", "usage": {"inputTokens": 10, "outputTokens": 4}}}})
        send({"jsonrpc": "2.0", "id": msg["id"], "result": {"stopReason": "end_turn"}})
'''

FAKE_PTY_AGENT = r'''
import sys
sys.stdout.write("fake agent ready\nfake> ")
sys.stdout.flush()
for line in sys.stdin:
    text = line.strip()
    sys.stdout.write("you said: %s\n" % text)
    sys.stdout.write("fake> ")
    sys.stdout.flush()
'''

MANIFEST_STREAM = '''[agent]
name = "fake"
binary = "fake"
profiles = ["chat"]
[startup]
command = ["python3", "-c", "pass"]
ready_when = { pattern = "ready" }
[fast_path]
stream_json = ["python3", "{agent}"]
[pty]
cols = 80
rows = 24
[capabilities]
streaming_output = true
'''

MANIFEST_ACP = '''[agent]
name = "acpfake"
binary = "acpfake"
profiles = ["chat"]
[startup]
command = ["python3", "-c", "pass"]
ready_when = { pattern = "ready" }
[fast_path]
acp_stdio = true
[pty]
cols = 80
rows = 24
[capabilities]
streaming_output = true
'''

MANIFEST_PTY = '''[agent]
name = "ptyfake"
binary = "ptyfake"
profiles = ["chat"]
[startup]
command = ["python3", "{agent}"]
ready_when = { pattern = "fake agent ready" }
init_timeout_seconds = 20
[pty]
cols = 100
rows = 30
[parse]
idle = ["^fake> $"]
[capabilities]
streaming_output = true
'''

MANIFEST_A2A = '''[agent]
name = "a2afake"
binary = "a2afake"
profiles = ["chat"]
[startup]
command = ["python3", "-c", "pass"]
ready_when = { pattern = "ready" }
[fast_path]
a2a_serve_at = "http://127.0.0.1:{port}"
[pty]
cols = 80
rows = 24
[capabilities]
streaming_output = true
'''

FLEET_TOML = '''[fleet]
start = "cap://a"
task = "begin"
budget_usd = 100

[[sessions]]
name = "a"
agent = "fake"

[[sessions]]
name = "b"
agent = "fake"

[[routes]]
when = "cap://a.done"
route_to = "cap://b"
'''


class A2aHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, *args):
        pass

    def do_GET(self):
        if self.path == "/.well-known/agent-card.json":
            body = json.dumps({"name": "a2a-card", "extensions": [{"uri": "cap-protocol/v1"}]}).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
        else:
            self.send_response(404)
            self.send_header("Content-Length", "0")
            self.end_headers()

    def do_POST(self):
        length = int(self.headers.get("Content-Length") or 0)
        raw = self.rfile.read(length)
        if self.path != "/message/send":
            self.send_response(404)
            self.send_header("Content-Length", "0")
            self.end_headers()
            return
        data = json.loads(raw)
        text = data["message"]["parts"][0]["text"]
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.end_headers()
        frames = [
            {"task": {"id": "t1", "contextId": "c1", "status": {"state": "working"}}},
            {"message": {"messageId": "mm1", "parts": [{"kind": "text", "text": "a2a heard " + text}]}},
            {"task": {"id": "t1", "contextId": "c1", "status": {"state": "completed"}}},
        ]
        for frame in frames:
            self.wfile.write(("data: " + json.dumps(frame) + "\n\n").encode())
            self.wfile.flush()


def normalize(text, root):
    text = text.replace(root, "<ROOT>")
    text = re.sub(r"cap_[0-9a-f]{12}", "cap_<ID>", text)
    return text


def run(argv, env, cwd, timeout=60, stdin=None):
    done = subprocess.run(argv, cwd=cwd, env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout,
                          stdin=subprocess.DEVNULL if stdin is None else None, input=stdin)
    return done.returncode, done.stdout.decode("utf-8", "replace"), done.stderr.decode("utf-8", "replace")


def main():
    args = sys.argv[1:]
    if not args:
        print(__doc__)
        return 2
    cpp = os.path.abspath(args[0])
    tai = TAI
    if "--tai" in args:
        tai = args[args.index("--tai") + 1]
    only = args[args.index("--filter") + 1] if "--filter" in args else None
    root = tempfile.mkdtemp(prefix="capmcp-cli-")
    server = http.server.HTTPServer(("127.0.0.1", 0), A2aHandler)
    port = server.server_address[1]
    threading.Thread(target=server.serve_forever, daemon=True).start()
    try:
        home = os.path.join(root, "home")
        agents = os.path.join(home, ".config", "cap", "agents")
        os.makedirs(agents)
        work = os.path.join(root, "work")
        os.makedirs(work)
        files = {
            "fake_stream.py": FAKE_STREAM_AGENT, "fake_acp.py": FAKE_ACP_AGENT, "fake_pty.py": FAKE_PTY_AGENT,
        }
        for name, text in files.items():
            with open(os.path.join(root, name), "w") as handle:
                handle.write(text)
        manifests = {
            "fake.toml": MANIFEST_STREAM.replace("{agent}", os.path.join(root, "fake_stream.py")),
            "acpfake.toml": MANIFEST_ACP,
            "ptyfake.toml": MANIFEST_PTY.replace("{agent}", os.path.join(root, "fake_pty.py")),
            "a2afake.toml": MANIFEST_A2A.replace("{port}", str(port)),
            "broken.toml": "[agent]\nname = \"broken\"\n",
            "wrongname.toml": MANIFEST_STREAM.replace('name = "fake"', 'name = "other"').replace("{agent}", "x"),
            "badre.toml": MANIFEST_STREAM.replace('name = "fake"', 'name = "badre"').replace("{agent}", "x") + "\n[parse]\nidle = [\"(?<=a)b\"]\n",
            "notoml.toml": "this is not toml = = =\n",
        }
        for name, text in manifests.items():
            with open(os.path.join(agents, name), "w") as handle:
                handle.write(text)
        fleet = os.path.join(root, "fleet.toml")
        with open(fleet, "w") as handle:
            handle.write(FLEET_TOML)
        badfleet = os.path.join(root, "badfleet.toml")
        with open(badfleet, "w") as handle:
            handle.write("[fleet]\nstart = \"cap://zzz\"\n[[sessions]]\nname = \"a\"\nagent = \"fake\"\n")
        nofleet = os.path.join(root, "nofleet.toml")
        with open(nofleet, "w") as handle:
            handle.write("[fleet]\n")
        garbage = os.path.join(root, "garbage.toml")
        with open(garbage, "w") as handle:
            handle.write("a = = 1")
        env = dict(os.environ, HOME=home, TAI_HOME=os.path.join(root, "taihome"), COLUMNS="80", PYTHONDONTWRITEBYTECODE="1")
        env.pop("TAI_PASSPHRASE", None)
        cases = [
            ["--cap"], ["--cap", "help"], ["--cap", "-h"], ["--cap", "--help"], ["--cap", "bogus"], ["--cap", "resolve"],
            ["--cap", "resolve", "-h"], ["--cap", "resolve", "nosuchagent"], ["--cap", "resolve", "fake"],
            ["--cap", "resolve", "acpfake"], ["--cap", "resolve", "ptyfake"], ["--cap", "resolve", "a2afake"],
            ["--cap", "resolve", "broken"], ["--cap", "resolve", "wrongname"], ["--cap", "resolve", "badre"],
            ["--cap", "resolve", "notoml"], ["--cap", "resolve", "x", "--binary"], ["--cap", "resolve", "x", "y"],
            ["--cap", "chat"], ["--cap", "chat", "-h"], ["--cap", "chat", "--driver", "fake", "--max-risk", "bad"],
            ["--cap", "chat", "--driver", "nosuch"], ["--cap", "chat", "--driver", "fake", "--a", "x"],
            ["--cap", "run"], ["--cap", "run", "-h"], ["--cap", "run", os.path.join(root, "missing.toml")],
            ["--cap", "run", garbage], ["--cap", "run", nofleet], ["--cap", "run", badfleet],
            ["--cap", "run", fleet, "--task", "   "],
            ["--cap", "chat", "--driver", "fake", "--task", "hello world"],
            ["--cap", "chat", "--driver", "ptyfake", "--task", "hello pty"],
            ["--cap", "chat", "--driver", "acpfake", "--argv", "python3 %s" % os.path.join(root, "fake_acp.py"), "--task", "hi acp"],
            ["--cap", "chat", "--driver", "acpfake", "--task", "no argv"],
            ["--cap", "chat", "--driver", "acpfake", "--argv", "python3 'unterminated", "--task", "x"],
            ["--cap", "chat", "--driver", "a2afake", "--task", "hi a2a"],
            ["--cap", "run", fleet],
            ["--cap", "run", fleet, "--task", "override task"],
        ]
        for columns in ("40", "60", "100", "132"):
            for sub in ("resolve", "chat", "run"):
                cases.append(("COLUMNS=" + columns, ["--cap", sub, "-h"]))
                cases.append(("COLUMNS=" + columns, ["--cap", sub, "--bogus-flag"]))
        failures = 0
        total = 0
        for case in cases:
            extra = {}
            if isinstance(case, tuple):
                key, value = case[0].split("=", 1)
                extra[key] = value
                case = case[1]
            if only and only not in " ".join(case):
                continue
            total += 1
            cenv = dict(env, **extra)
            want = run([sys.executable, "-I", tai] + case, cenv, work)
            got = run([cpp] + case, cenv, work)
            want = (want[0], normalize(want[1], root), normalize(want[2], root))
            got = (got[0], normalize(got[1], root), normalize(got[2], root))
            if "--verbose" in args:
                print("CASE %s -> rc %s stdout %r stderr %r" % (" ".join(case).replace(root, "<ROOT>"), want[0], want[1][:200], want[2][:160]))
            if want != got:
                failures += 1
                print("MISMATCH %s" % " ".join(case).replace(root, "<ROOT>"))
                for label, a, b in (("rc", want[0], got[0]), ("stdout", want[1], got[1]), ("stderr", want[2], got[2])):
                    if a != b:
                        print("  %s py : %r" % (label, a if label == "rc" else a[:600]))
                        print("  %s cpp: %r" % (label, b if label == "rc" else b[:600]))
            # worktree dirs created by coding profiles must not leak between cases
            shutil.rmtree(os.path.join(work, ".cap-worktrees"), ignore_errors=True)
        want = run([sys.executable, "-I", tai, "--cap-manifest"], env, work)
        got = run([cpp, "--cap-manifest"], env, work)
        total += 1
        if want != got:
            failures += 1
            print("MISMATCH --cap-manifest")
        print("%d cases, %d mismatches" % (total, failures))
        return 1 if failures else 0
    finally:
        server.shutdown()
        shutil.rmtree(root, ignore_errors=True)


if __name__ == "__main__":
    sys.exit(main())