tai/taicpp/tests/orch_mplex_e2e.py

#!/usr/bin/env python3
"""End to end: the real multiplexer server plus the real OrchBoard/OrchMplexTransport.

    python3 -I tests/orch_mplex_e2e.py DETZ_BINARY ORCH_LIVE_BINARY

Checks layout geometry for 1, 2, 3 and 6 sub agents, streamed content, footers, the space limit,
linger and rebalance, the user closing a pane, the main pane exiting, and that a one-shot control
write reports panes that no longer exist. Everything runs in a throwaway TAI_HOME.
"""
import json, os, re, shutil, socket, subprocess, sys, tempfile, time

DETZ, LIVE = os.path.abspath(sys.argv[1]), os.path.abspath(sys.argv[2])
HOME = tempfile.mkdtemp(prefix="oe", dir="/tmp")
ENV = dict(os.environ, TAI_HOME=HOME, HOME=HOME)
for k in ("TAI_PLEXER_ACTIVE", "TAI_PLEXER_SOCKET", "TAI_PLEXER_PANE", "TMUX"):
    ENV.pop(k, None)
SOCK = os.path.join(HOME, "plexer", "server.sock")
fails = []


def check(cond, what):
    print(("ok   " if cond else "FAIL ") + what)
    if not cond:
        fails.append(what)


def mp(*args):
    return subprocess.run([DETZ, "--mplex", *args], env=ENV, capture_output=True, text=True, timeout=30)


def call(**kw):
    s = socket.socket(socket.AF_UNIX)
    s.connect(SOCK)
    s.sendall((json.dumps(kw) + "\n").encode())
    buf = b""
    while not buf.endswith(b"\n"):
        chunk = s.recv(65536)
        if not chunk:
            break
        buf += chunk
    s.close()
    return json.loads(buf)


def panes():
    out = mp("list-panes", "-F", "#{pane_id} #{pane_left} #{pane_top} #{pane_width} #{pane_height}").stdout.split("\n")
    rows = []
    for line in out:
        m = re.match(r"%(\d+) (\d+) (\d+) (\d+) (\d+)$", line.strip())
        if m:
            rows.append(tuple(int(x) for x in m.groups()))
    return rows


def run_live(n, tokens=40, hold=0, fail=False, wait=True):
    env = dict(ENV, TAI_PLEXER_ACTIVE="1", TAI_PLEXER_SOCKET=SOCK, TAI_PLEXER_PANE="0")
    cmd = [LIVE, str(n), str(tokens), str(hold)] + (["fail"] if fail else [])
    if wait:
        r = subprocess.run(cmd, env=env, capture_output=True, text=True, timeout=60)
        return r.stdout
    return subprocess.Popen(cmd, env=env, stdout=subprocess.PIPE, text=True)


def capture(pid):
    return mp("capture-pane", "-p", "-t", "%" + str(pid)).stdout


try:
    r = mp("new-session", "-d", "-s", "t", "-x", "120", "-y", "40")
    check(r.returncode == 0, "detached session starts")

    # ---- 1 sub agent: main left full height, one pane right
    p = run_live(1, hold=1500, wait=False)
    time.sleep(0.6)
    rows = panes()
    check(len(rows) == 2, "one sub agent: two panes")
    main = [r for r in rows if r[0] == 0][0]
    sub = [r for r in rows if r[0] != 0]
    check(main[1] == 0 and main[2] == 0 and main[4] == 40, "main pane on the left at full height")
    check(len(sub) == 1 and sub[0][1] > main[1] + main[3] - 1 and sub[0][4] == 40, "sub pane on the right at full height")
    check(abs(main[3] - 120 * 0.55) <= 2, "main pane is about 55 percent wide")
    out = p.communicate(timeout=30)[0]
    check("shown" in out, "board reports the panel shown")
    time.sleep(0.3)
    text = capture(sub[0][0])
    check("scripted task number 1" in text, "task echoed at the top")
    check("shell" in text and "line three" in text and "line four" not in text, "tool call and a three-line result")
    check("THE END 0" in text and "word0_39" in text, "streamed text arrived complete")
    check("done" in text and "all good" in text, "footer states the result")
    # linger then rebalance
    time.sleep(21.5)
    check(len(panes()) == 1, "panel leaves after the linger and the main pane takes the screen")

    # ---- 3 sub agents: stacked evenly on the right
    p = run_live(3, tokens=30, hold=1200, wait=False)
    time.sleep(0.8)
    rows = panes()
    subs = sorted([r for r in rows if r[0] != 0], key=lambda r: r[2])
    check(len(subs) == 3, "three sub agents: three panes on the right")
    check(len({r[1] for r in subs}) == 1 and subs[0][1] > 0, "all in one column")
    heights = [r[4] for r in subs]
    check(max(heights) - min(heights) <= 2, "even heights " + str(heights))
    check(sum(heights) + 2 == 40, "column fills the window")
    p.communicate(timeout=30)
    time.sleep(0.3)
    for r in subs:
        check("THE END" in capture(r[0]), "pane %%%d content" % r[0])
    time.sleep(21.5)
    check(len(panes()) == 1, "all three leave")

    # ---- space limit: the pane that does not fit runs without a panel
    p = run_live(8, tokens=5, hold=800, wait=False)
    time.sleep(1.2)
    rows = panes()
    check(len(rows) - 1 == 6, "only as many panels as fit (6 rows each): %d" % (len(rows) - 1))
    out = p.communicate(timeout=30)[0]
    check(out.count("shown") == 6 and "no space" in out, "the others report no space instead of failing")
    time.sleep(21.5)

    # ---- failure state and a manual close = cancel
    p = run_live(2, tokens=400, hold=3000, fail=True, wait=False)
    time.sleep(0.8)
    subs = sorted([r for r in panes() if r[0] != 0], key=lambda r: r[2])
    victim = subs[0][0]
    lst = call(cmd="vpane_list")["data"]
    check(len(lst) == 2 and all(x["state"] == "running" for x in lst), "list reports running panels")
    mp("kill-pane", "-t", "%" + str(victim))
    time.sleep(0.5)
    reply = call(cmd="vpane_write", ops=[{"pane": victim, "data": "x"}])
    check(victim in reply["data"]["closed"], "a write to a closed pane reports it closed")
    p.communicate(timeout=60)
    time.sleep(0.4)
    left = [r for r in panes() if r[0] != 0]
    check(len(left) == 1, "the other pane survived")
    check("scripted failure" in capture(left[0][0]) and "failed" in capture(left[0][0]), "failed footer")

    # ---- main pane exits: feed panes go with the window
    call(cmd="vpane_open", pane="0", label="orphan")
    check(len(panes()) >= 2, "a feed pane is open")
    mp("send-keys", "-t", "%0", "exit", "Enter")
    time.sleep(1.5)
    check(mp("list-sessions").stdout.strip() == "" or "no server" in (mp("list-sessions").stderr + mp("list-sessions").stdout).lower(),
          "session ends with its main pane")
finally:
    mp("kill-server")
    shutil.rmtree(HOME, ignore_errors=True)

print("\n%d failures" % len(fails))
sys.exit(1 if fails else 0)