tai/taicpp/tests/mplex_matrix_exec.py

"""Differential executor: runs every task of mplex_matrix_gen.py on tai.py --mplex and on the C++
binary, each in its own throwaway TAI_HOME, and compares step outputs and full state dumps.

usage: python3 -I mplex_matrix_exec.py <repo root> <cpp binary> <tasks.json> <out dir> [--jobs N] [--filter TEXT]

Normalized (only these): the TAI_HOME and work directory paths (replaced by <H> and <W>), times of the
forms "Mon Jan  1 12:00:00 2026" and "12:00 01-Jan-26" (tmux default formats), and the minute digits of the
clock mode screen. Everything else must be identical.
"""
import concurrent.futures
import difflib
import fcntl
import importlib.util
import json
import os
import pty
import re
import select
import shutil
import signal
import struct
import subprocess
import sys
import tempfile
import termios
import threading
import time

ROOT, CPP, TASKS, OUT = [os.path.abspath(p) for p in sys.argv[1:5]]
ARGS = sys.argv[5:]
JOBS = int(ARGS[ARGS.index("--jobs") + 1]) if "--jobs" in ARGS else 8
FILTER = ARGS[ARGS.index("--filter") + 1] if "--filter" in ARGS else ""

spec = importlib.util.spec_from_file_location("tai_mod", ROOT + "/tai.py")
tai = importlib.util.module_from_spec(spec)
sys.modules["tai_mod"] = tai
spec.loader.exec_module(tai)

CWD = "/tmp/mxcwd"
os.makedirs(CWD, exist_ok=True)
TIME_RES = [re.compile(r"[A-Z][a-z]{2} [A-Z][a-z]{2} [ \d]\d \d\d:\d\d:\d\d \d{4}"), re.compile(r"\d\d:\d\d \d\d-[A-Z][a-z]{2}-\d\d"), re.compile(r"\d\d:\d\d:\d\d"), re.compile(r"/dev/pts/\d+")]

BASE_CONFIG = '[options]\nautomatic-rename = false\nstatus-right = "R"\ndefault-command = "exec bash --norc --noprofile"\n'
P5_CONFIG = BASE_CONFIG + 'base-index = 1\npane-base-index = 1\nmode-keys = "vi"\nmain-pane-width = 40\nmain-pane-height = 12\n'

PRESETS = {
    "P0": [["new-session", "-d", "-s", "main", "-x", "80", "-y", "24"]],
}
P0 = PRESETS["P0"]
P1 = P0 + [["new-session", "-d", "-s", "other", "-x", "80", "-y", "24"], ["new-window", "-t", "other"]]
P2 = P0 + [["new-window"], ["split-window", "-h"], ["new-window"], ["split-window", "-v"], ["split-window", "-h"], ["split-window", "-v"],
           ["split-window", "-h"], ["select-layout", "tiled"], ["select-window", "-t", "main:0"]]
P3 = P2 + [["select-window", "-t", "main:2"], ["resize-pane", "-Z"], ["select-pane", "-t", "main:1.0"], ["select-pane", "-m"],
           ["set-option", "-g", "status-position", "top"]]
P4 = P2 + [["set-buffer", "-b", "b1", "hello"], ["set-buffer", "world"], ["set-option", "-g", "status-position", "top"],
           ["select-window", "-t", "main:1"], ["send-keys", "-t", "main:1.0", "seq 1 100", "Enter"]]
P5 = P0 + [["new-window"], ["split-window", "-h"], ["new-window"], ["split-window", "-v"], ["split-window", "-h"], ["select-layout", "main-vertical"],
           ["select-window", "-t", "main:1"], ["select-layout", "even-vertical"], ["set-buffer", "-b", "b1", "hello"],
           ["send-keys", "-t", "main:1.1", "seq 1 100", "Enter"]]
PRESET_STEPS = {"P0": P0, "P1": P1, "P2": P2, "P3": P3, "P4": P4, "P5": P5, "P4c": P4 + [["@sleep", "0.8"], ["copy-mode"]], "P5c": P5 + [["@sleep", "0.8"], ["copy-mode"]]}
PRESET_CONFIG = {"P5": P5_CONFIG, "P5c": P5_CONFIG}

LIST_FORMAT = ("#{session_name}:#{window_index}.#{pane_index} #{pane_id} #{window_id} #{pane_left},#{pane_top} #{pane_width}x#{pane_height} "
               "a=#{pane_active} l=#{pane_last} m=#{pane_marked} z=#{window_zoomed_flag} mode=#{pane_mode} hist=#{history_size} dead=#{pane_dead} "
               "sp=#{scroll_position} cc=#{copy_cursor_x},#{copy_cursor_y} sel=#{selection_present} s=#{pane_search_string} "
               "title=#{pane_title} syn=#{pane_synchronized} cmd=#{pane_current_command}")
WIN_FORMAT = ("#{session_name}:#{window_index} #{window_id} #{window_name} #{window_flags} #{window_panes} #{window_width}x#{window_height} "
              "act=#{window_active} last=#{window_last_flag} z=#{window_zoomed_flag} mk=#{window_marked_flag}")
SES_FORMAT = "#{session_name} #{session_id} w=#{session_windows} att=#{session_attached}"
MARK = ["display-message", "-p", "@@ %s"]


class Impl:
    def __init__(self, name, argv, config):
        self.name = name
        self.argv = argv
        self.tmp = tempfile.mkdtemp(prefix="mx", dir="/tmp")
        self.home = os.path.join(self.tmp, "h")
        self.work = os.path.join(self.tmp, "w")
        os.makedirs(os.path.join(self.home))
        os.makedirs(self.work)
        with open(os.path.join(self.home, "mplex.toml"), "w") as handle:
            handle.write(config)
        with open(os.path.join(self.work, "src.conf"), "w") as handle:
            handle.write("set -g status-left X\nbogus-command\nset -g status-right Y\n")
        with open(os.path.join(self.work, "buf.txt"), "w") as handle:
            handle.write("from file\n")
        self.env = dict(os.environ)
        for name_ in ("TAI_PLEXER_ACTIVE", "TAI_PLEXER_SOCKET", "TAI_PLEXER_PANE", "TMUX", "VISUAL", "EDITOR"):
            self.env.pop(name_, None)
        self.env.update({"TAI_HOME": self.home, "HOME": self.work, "SHELL": "/bin/bash", "TERM": "xterm-256color", "TAI_PASSPHRASE": "",
                         "LANG": "C.UTF-8", "PS1": "$ "})

    def norm(self, text):
        text = text.replace(self.home, "<H>").replace(self.work, "<W>").replace(self.tmp, "<T>")
        for regex in TIME_RES:
            text = regex.sub("<TIME>", text)
        return text

    def sub(self, words):
        return [w.replace("{W}", self.work) for w in words]

    def cli(self, *commands):
        """commands: lists of words; chained with ';' like the command line does."""
        words = []
        for index, command in enumerate(commands):
            if index:
                words.append(";")
            words.extend(self.sub(command))
        try:
            done = subprocess.run(self.argv + words, env=self.env, cwd=CWD, stdin=subprocess.DEVNULL, capture_output=True, text=True,
                                  errors="replace", timeout=30)
        except subprocess.TimeoutExpired:
            return 124, "TIMEOUT"
        return done.returncode, self.norm(done.stdout + "||" + done.stderr)

    STATE_COMMANDS = [["list-sessions", "-F", SES_FORMAT], MARK[:2] + ["@@ 1"], ["list-windows", "-a", "-F", WIN_FORMAT], MARK[:2] + ["@@ 2"],
                      ["list-panes", "-a", "-F", LIST_FORMAT], MARK[:2] + ["@@ 3"], ["show-options", "-g"], MARK[:2] + ["@@ 4"], ["show-options"],
                      MARK[:2] + ["@@ 5"], ["show-options", "-w"], MARK[:2] + ["@@ 6"], ["list-buffers"], MARK[:2] + ["@@ 7"], ["show-messages"],
                      MARK[:2] + ["@@ 8"], ["list-keys", "-T", "prefix"], MARK[:2] + ["@@ 9"], ["list-keys", "-T", "root"],
                      MARK[:2] + ["@@ 10"], ["list-clients", "-F", "#{client_session} #{client_width}x#{client_height} #{client_key_table}"]]

    def dump(self):
        """Poll until every pane shows a prompt, then return captures plus the full state listing (two calls)."""
        out = ""
        for _ in range(40):
            rc, ids = self.cli(["list-panes", "-a", "-F", "#{pane_id}"])
            if rc != 0:
                return "NOSERVER " + ids
            panes = ids.split("||")[0].split()
            commands = []
            for pane in panes:
                commands.append(["capture-pane", "-p", "-t", pane])
                commands.append(["display-message", "-p", "@@ end"])
            commands.append(["display-message", "-p", "@@ state"])
            rc, out = self.cli(*(commands + self.STATE_COMMANDS))
            head = out.split("@@ state")[0]
            parts = head.split("@@ end")
            if not panes or all(("$" in part or "bash" in part) for part in parts[:-1]):
                break
            time.sleep(0.1)
        return out

    def build(self, preset):
        steps = PRESET_STEPS[preset]
        for step in steps:
            if step[0] == "@sleep":
                time.sleep(float(step[1]))
                continue
            rc, out = self.cli(step)
            if rc != 0:
                return "BUILD FAIL %d %s %s" % (rc, step, out)
        if preset in ("P4", "P4c", "P5", "P5c"):
            time.sleep(0.4)
        return ""

    def kill(self):
        try:
            self.cli(["kill-server"])
        except Exception:
            pass
        time.sleep(0.1)
        shutil.rmtree(self.tmp, ignore_errors=True)


def delta(a, b):
    if isinstance(a, tuple) or isinstance(b, tuple):
        a, b = repr(a), repr(b)
    return "\n".join(list(difflib.unified_diff(str(a).split("\n"), str(b).split("\n"), "py", "cpp", lineterm="", n=1))[:60])


def make_impls(preset):
    config = PRESET_CONFIG.get(preset, BASE_CONFIG)
    py = Impl("py", [sys.executable, "-I", ROOT + "/tai.py", "--mplex"], config)
    cpp = Impl("cpp", [CPP, "--mplex"], config)
    return py, cpp


def both(fn_py, fn_cpp):
    results = [None, None]

    def run(index, fn):
        try:
            results[index] = fn()
        except Exception as exc:  # noqa: BLE001
            results[index] = "EXC %r" % (exc,)

    threads = [threading.Thread(target=run, args=(0, fn_py)), threading.Thread(target=run, args=(1, fn_cpp))]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    return results


def run_cmd_task(task):
    py, cpp = make_impls(task["preset"])
    log = []
    try:
        fails = both(lambda: py.build(task["preset"]), lambda: cpp.build(task["preset"]))
        if fails[0] != fails[1] or fails[0]:
            log.append({"step": "build", "diff": delta(fails[0], fails[1])})
            return log
        # state before any step
        a, b = both(py.dump, cpp.dump)
        if a != b:
            log.append({"step": "initial dump", "diff": delta(a, b)})
            return log
        for step in task["steps"]:
            ra, rb = both(lambda: py.cli(step), lambda: cpp.cli(step))
            if ra != rb:
                log.append({"step": step, "kind": "output", "diff": delta(ra, rb)})
            a, b = both(py.dump, cpp.dump)
            if a != b:
                log.append({"step": step, "kind": "state", "diff": delta(a, b)})
                # a diverged state would cascade: stop this chunk
                break
    finally:
        both(py.kill, cpp.kill)
    return log


class Client:
    def __init__(self, impl, words):
        self.impl = impl
        self.screen = tai.TuiVtScreen(24, 80)
        pid, fd = pty.fork()
        if pid == 0:
            os.chdir(CWD)
            os.execvpe(impl.argv[0], impl.argv + words, impl.env)
        self.pid, self.fd = pid, fd
        fcntl.ioctl(fd, termios.TIOCSWINSZ, struct.pack("HHHH", 24, 80, 0, 0))
        self.settle(1.0)

    def settle(self, quiet=0.3, limit=3.0):
        end, last = time.time() + limit, time.time()
        while time.time() < end and time.time() - last < quiet:
            r, _, _ = select.select([self.fd], [], [], 0.05)
            if r:
                try:
                    data = os.read(self.fd, 65536)
                except OSError:
                    return
                if not data:
                    return
                self.screen.feed(data)
                last = time.time()

    def send(self, data):
        try:
            os.write(self.fd, data)
        except OSError:
            pass
        self.settle()

    def text(self):
        text = self.screen.text()
        text = re.sub(r"/dev/pts/\d+", "/dev/pts/N", text).replace("^[", "")
        text = "\n".join(line.rstrip() for line in text.split("\n")).rstrip("\n")
        return re.sub(r"\d", "#", text) if "clock" in text.lower() else text

    def close(self):
        try:
            os.kill(self.pid, signal.SIGKILL)
        except OSError:
            pass
        try:
            os.waitpid(self.pid, 0)
        except OSError:
            pass
        try:
            os.close(self.fd)
        except OSError:
            pass


def key_bytes(name):
    return tai.MplexKeys.to_bytes(name)


def run_key_task(task):
    py, cpp = make_impls(task["preset"])
    log = []
    clients = []
    try:
        fails = both(lambda: py.build(task["preset"]), lambda: cpp.build(task["preset"]))
        if fails[0] != fails[1] or fails[0]:
            return [{"step": "build", "diff": delta(fails[0], fails[1])}]
        for impl in (py, cpp):
            clients.append(Client(impl, ["attach-session", "-t", "main"]))
        # copy-mode tables are entered by the preset (P4c/P5c); drop the attach screen noise first
        data = key_bytes(task["key"])
        if task["table"] == "prefix":
            data = b"\x02" + data
        for client in clients:
            client.send(data)
        time.sleep(0.8)
        for client in clients:
            client.settle(0.2)
        screens = [c.text() for c in clients]
        a, b = both(py.dump, cpp.dump)
        if screens[0] != screens[1]:
            log.append({"step": task["key"], "kind": "screen", "diff": delta(screens[0], screens[1])})
        if a != b:
            log.append({"step": task["key"], "kind": "state", "diff": delta(a, b)})
        # second key: dismiss anything open, then compare again
        for client in clients:
            client.send(b"\x1b")
        time.sleep(0.8)
        for client in clients:
            client.settle(0.2)
        screens = [c.text() for c in clients]
        a, b = both(py.dump, cpp.dump)
        if screens[0] != screens[1]:
            log.append({"step": task["key"] + " +Esc", "kind": "screen", "diff": delta(screens[0], screens[1])})
        if a != b:
            log.append({"step": task["key"] + " +Esc", "kind": "state", "diff": delta(a, b)})
    finally:
        for client in clients:
            client.close()
        both(py.kill, cpp.kill)
    return log


CUSTOM_CONFIG = ('[options]\nautomatic-rename = false\nstatus-right = "R"\ndefault-command = "exec bash --norc --noprofile"\n'
                 'prefix2 = "F12"\nescape-time = 50\nrepeat-time = 2500\n[keys.root]\n"M-x" = "display-message root-hit"\n"F7" = "split-window -h"\n'
                 '[keys.prefix]\nr = { command = "select-pane -R", repeat = true }\nl = { command = "select-pane -L", repeat = true }\n"C-b" = false\n'
                 'q = "command-prompt -p x { rename-window \\"%%\\" }"\n[keys.mine]\na = "rename-window mine-a"\n')


def run_custom_task(task):
    py, cpp = make_impls("P2")
    for impl in (py, cpp):
        with open(os.path.join(impl.home, "mplex.toml"), "w") as handle:
            handle.write(CUSTOM_CONFIG)
    log = []
    clients = []
    try:
        fails = both(lambda: py.build("P2"), lambda: cpp.build("P2"))
        if fails[0] != fails[1] or fails[0]:
            return [{"step": "build", "diff": delta(fails[0], fails[1])}]
        for impl in (py, cpp):
            clients.append(Client(impl, ["attach-session", "-t", "main"]))
        script = [b"\x1bx", b"\x1b[18~", b"\x1b[18~", b"\x02r", b"r", b"r", b"l", b"\x02", b"\x1b", b"\x1b[18~l", b"\x02q", b"hello\r", b"\x02" + b"b",
                  b"\x1b[18~", b"\x1b[18~" + b"x", b"\x1b[18~" + b"c", b"\x02l", b"x", b"\x1b[18~[", b"\x02[", b"q", b"\x1bOP", b"\x1b[A"]
        for number, data in enumerate(script):
            for client in clients:
                client.send(data)
            if number in (4, 9, 16):
                time.sleep(3.0)
                for client in clients:
                    client.settle(0.2)
            screens = [c.text() for c in clients]
            a, b = both(py.dump, cpp.dump)
            if screens[0] != screens[1]:
                log.append({"step": "script %d %r" % (number, data), "kind": "screen", "diff": delta(screens[0], screens[1])})
            if a != b:
                log.append({"step": "script %d %r" % (number, data), "kind": "state", "diff": delta(a, b)})
        # custom table through switch-client -T
        for impl in (py, cpp):
            impl.cli(["switch-client", "-T", "mine"])
        for client in clients:
            client.send(b"a")
        a, b = both(py.dump, cpp.dump)
        if a != b:
            log.append({"step": "table mine", "kind": "state", "diff": delta(a, b)})
    finally:
        for client in clients:
            client.close()
        both(py.kill, cpp.kill)
    return log


def run_task(task):
    try:
        if task["kind"] == "cmd":
            return run_cmd_task(task)
        if task["kind"] == "key":
            return run_key_task(task)
        return run_custom_task(task)
    except Exception as exc:  # noqa: BLE001
        return [{"step": "harness", "error": repr(exc)}]


def main():
    tasks = json.load(open(TASKS))
    if FILTER:
        tasks = [t for t in tasks if any(f in t["id"] for f in FILTER.split(","))]
    os.makedirs(OUT, exist_ok=True)
    started = time.time()
    done = mismatched = 0
    steps = sum(len(t.get("steps", [])) or 1 for t in tasks)
    report = open(os.path.join(OUT, "mismatches.jsonl"), "w")
    with concurrent.futures.ThreadPoolExecutor(JOBS) as pool:
        futures = {pool.submit(run_task, t): t for t in tasks}
        for future in concurrent.futures.as_completed(futures):
            task = futures[future]
            log = future.result()
            done += 1
            if log:
                mismatched += 1
                report.write(json.dumps({"task": task["id"], "log": log}) + "\n")
                report.flush()
            if done % 25 == 0 or log:
                print("[%d/%d] %s %s (%.0fs)" % (done, len(tasks), task["id"], "MISMATCH" if log else "ok", time.time() - started), flush=True)
    print("tasks %d, steps %d, mismatching tasks %d, %.0fs" % (len(tasks), steps, mismatched, time.time() - started))
    return 1 if mismatched else 0


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