tai/taicpp/tests/repl_persist.py

#!/usr/bin/env python3
"""Conversation history survives a restart in the same pane.

Starts the real REPL binary (build/repl_probe: Repl + Conversation + Store, a
scripted host in place of the model) repeatedly, each time as a child of a
long-lived shell that stands for one terminal pane (TMUX_PANE and
TAI_PLEXER_PANE set, a throwaway TAI_HOME, no TAI_SESSION). Every assertion
reads vault state (profiles/*.session.*.json, *.turn.*.json, memory.db), never
the program's output.

tai.py keys a conversation by the nearest shell ancestor (ShellSession.current),
and a pane is exactly one such shell, so:
  - same pane, restart      -> same session file, history resumes
  - different pane          -> its own session file, history independent
  - TAI_SESSION=<word>      -> panes share one pinned context
  - interrupted turn        -> --continue picks it up

Usage: python3 tests/repl_persist.py [--probe PATH]
"""
import argparse
import glob
import json
import os
import shlex
import shutil
import sqlite3
import subprocess
import sys
import tempfile
import time

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(os.path.dirname(HERE))


class Pane:
    """One long-lived shell = one terminal pane."""

    def __init__(self, home, probe, number, extra_env=None):
        self.probe = probe
        env = {key: value for key, value in os.environ.items() if not key.startswith(("TAI_", "TMUX"))}
        env.update({"TAI_HOME": home, "TAI_PASSPHRASE": "", "NO_COLOR": "1", "HOME": home,
                    "TMUX_PANE": "%%%d" % number, "TAI_PLEXER_PANE": str(number)})
        env.update(extra_env or {})
        self.shell = subprocess.Popen(["bash", "--norc", "--noprofile"], stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                                      stderr=subprocess.STDOUT, env=env, universal_newlines=True)

    def run(self, lines, *flags):
        """Run the probe in this pane with `lines` on stdin; returns its exit code."""
        feed = "".join(line + "\n" for line in lines)
        command = "printf %%s %s | %s %s; echo __RC=$?\n" % (shlex.quote(feed), shlex.quote(self.probe), " ".join(shlex.quote(f) for f in flags))
        self.shell.stdin.write(command)
        self.shell.stdin.flush()
        for raw in self.shell.stdout:
            if "__RC=" in raw:  # the prompt has no newline, so the marker can follow it on the same line
                return int(raw.strip().split("__RC=")[1])
        raise RuntimeError("pane shell died")

    def close(self):
        try:
            self.shell.stdin.close()
            self.shell.wait(timeout=10)
        except Exception:
            self.shell.kill()


def session_files(home, profile="default"):
    return sorted(glob.glob(os.path.join(home, "profiles", profile + ".session.*.json")))


def read_json(path):
    with open(path, "r", encoding="utf-8") as handle:
        return json.load(handle)


def contents(path):
    return [(item["role"], item["content"]) for item in read_json(path)]


def history(home):
    db = sqlite3.connect(os.path.join(home, "memory.db"))
    try:
        return [row[0] for row in db.execute("SELECT text FROM input_history WHERE profile = 'default' ORDER BY id").fetchall()]
    finally:
        db.close()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--probe", default=os.path.join(os.path.dirname(HERE), "build", "repl_probe"))
    opts = parser.parse_args()
    home = tempfile.mkdtemp(prefix="detz-repl-")
    checks = 0
    failures = []

    def check(label, ok, detail=""):
        nonlocal checks
        checks += 1
        if not ok:
            failures.append("%s %s" % (label, detail))
            print("FAIL %s %s" % (label, detail))

    pane_a = Pane(home, opts.probe, 11)
    pane_b = Pane(home, opts.probe, 12)
    try:
        # first start in pane A
        check("A1 exit", pane_a.run(["hello one", "hello two", "/quit"]) == 0)
        files = session_files(home)
        check("A1 one session file", len(files) == 1, str(files))
        sid_file = files[0]
        check("A1 file is keyed by a shell id", ".session.p" not in sid_file and "-p" in os.path.basename(sid_file), sid_file)
        expect = [("user", "hello one"), ("assistant", "echo: hello one"), ("user", "hello two"), ("assistant", "echo: hello two")]
        check("A1 messages saved", contents(sid_file) == expect, str(contents(sid_file)))
        check("A1 system message not stored", all(role != "system" for role, _ in contents(sid_file)))

        # quit and restart in the same pane: the same file resumes
        check("A2 exit", pane_a.run(["after restart", "/quit"]) == 0)
        files = session_files(home)
        check("A2 still one session file", files == [sid_file], str(files))
        check("A2 history resumed and extended", contents(sid_file) == expect + [("user", "after restart"), ("assistant", "echo: after restart")], str(contents(sid_file)))

        # another pane: independent conversation
        check("B1 exit", pane_b.run(["other pane"]) == 0)
        files = session_files(home)
        check("B1 second session file", len(files) == 2, str(files))
        other = [path for path in files if path != sid_file][0]
        check("B1 own history only", contents(other) == [("user", "other pane"), ("assistant", "echo: other pane")], str(contents(other)))
        check("B1 pane A untouched", len(contents(sid_file)) == 6)

        # pane A again, one-shot path also resumes and saves
        check("A3 oneshot exit", pane_a.run([], "--oneshot", "one shot q") == 0)
        check("A3 oneshot appended", contents(sid_file)[-2:] == [("user", "one shot q"), ("assistant", "echo: one shot q")] and len(contents(sid_file)) == 8, str(contents(sid_file)))
        check("A3 input history in the vault", history(home)[:3] == ["hello one", "hello two", "/quit"], str(history(home)))

        # shell_sessions rows registered, keyed like the files
        db = sqlite3.connect(os.path.join(home, "memory.db"))
        sids = {row[0] for row in db.execute("SELECT sid FROM shell_sessions").fetchall()}
        db.close()
        sid_a = os.path.basename(sid_file)[len("default.session."):-len(".json")]
        check("shell_sessions row for pane A", sid_a in sids, "%s not in %s" % (sid_a, sids))

        # python reads what C++ wrote (same files, same keys)
        sys.path.insert(0, ROOT)
        env_backup = dict(os.environ)
        os.environ["TAI_HOME"] = home
        import argparse as ap
        import tai
        config = tai.Config(ap.Namespace(profile="default", yes=False, no_opencode=False))
        store = tai.Store(config, tai.Seal(home, ""))
        loaded = store.load_session("default", ctx=sid_a)
        check("python loads the C++ session", [(m["role"], m["content"]) for m in loaded] == contents(sid_file))
        # and C++ resumes a session python wrote
        store.save_session("default", [{"role": "system", "content": "s"}, {"role": "user", "content": "from python"}, {"role": "assistant", "content": "ok"}], ctx=sid_a)
        store.close()
        os.environ.clear()
        os.environ.update(env_backup)
        check("A4 exit", pane_a.run(["and back"]) == 0)
        check("A4 C++ resumed the python-written session", contents(sid_file) == [("user", "from python"), ("assistant", "ok"), ("user", "and back"), ("assistant", "echo: and back")], str(contents(sid_file)))
    finally:
        pane_a.close()
        pane_b.close()

    # interrupted turn and --continue, in a fresh pane
    home2 = tempfile.mkdtemp(prefix="detz-repl-")
    pane_c = Pane(home2, opts.probe, 21)
    try:
        check("C1 crash exit code", pane_c.run(["CRASH"]) == 9)
        files = session_files(home2)
        check("C1 checkpoint saved the user message", len(files) == 1 and contents(files[0]) == [("user", "CRASH")], str(files))
        turns = glob.glob(os.path.join(home2, "profiles", "default.turn.*.json"))
        check("C1 turn left open", len(turns) == 1 and read_json(turns[0]).get("open") is True and read_json(turns[0]).get("goal") == "CRASH", str(turns))
        check("C2 continue exit", pane_c.run(["/quit"], "--continue") == 0)
        got = contents(files[0])
        check("C2 resume turn ran", got[0] == ("user", "CRASH") and got[1][0] == "user" and got[1][1].startswith("Continue where you left off: the interrupted request was: CRASH"), str(got))
        check("C2 turn closed", read_json(turns[0]).get("open") is False)
    finally:
        pane_c.close()

    # a pinned TAI_SESSION is shared by different panes
    home3 = tempfile.mkdtemp(prefix="detz-repl-")
    pane_d = Pane(home3, opts.probe, 31, {"TAI_SESSION": "pinned1"})
    pane_e = Pane(home3, opts.probe, 32, {"TAI_SESSION": "pinned1"})
    try:
        check("D1 exit", pane_d.run(["from d"]) == 0)
        check("D2 exit", pane_e.run(["from e"]) == 0)
        files = session_files(home3)
        check("D shared pinned context", [os.path.basename(f) for f in files] == ["default.session.pinned1.json"], str(files))
        check("D both panes wrote it", contents(files[0]) == [("user", "from d"), ("assistant", "echo: from d"), ("user", "from e"), ("assistant", "echo: from e")], str(contents(files[0])))
    finally:
        pane_d.close()
        pane_e.close()

    tmux_section(opts.probe, check)
    signal_section(opts.probe, check)

    print("%d checks, %d failures" % (checks, len(failures)))
    return 1 if failures else 0


def signal_section(probe, check):
    """Ctrl-C in the middle of a turn checkpoints it open; an uncaught exception reaches the debug log."""
    import signal
    home = tempfile.mkdtemp(prefix="detz-repl-sig-")
    env = {key: value for key, value in os.environ.items() if not key.startswith(("TAI_", "TMUX"))}
    env.update({"TAI_HOME": home, "TAI_PASSPHRASE": "", "NO_COLOR": "1", "HOME": home, "TAI_SESSION": "sigctx"})
    proc = subprocess.Popen([probe], env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
    proc.stdin.write("SLEEP\n")
    proc.stdin.flush()
    turn = os.path.join(home, "profiles", "default.turn.sigctx.json")
    end = time.time() + 15
    while time.time() < end and not os.path.exists(turn) and proc.poll() is None:
        time.sleep(0.05)
    time.sleep(0.4)
    proc.send_signal(signal.SIGINT)
    time.sleep(0.6)
    proc.stdin.write("after\n/quit\n")
    proc.stdin.flush()
    out, _ = proc.communicate(timeout=30)
    check("SIGINT: REPL survived and exited cleanly", proc.returncode == 0, str(proc.returncode))
    check("SIGINT: interrupted is reported", "interrupted" in out, repr(out[-300:]))
    sess = os.path.join(home, "profiles", "default.session.sigctx.json")
    got = contents(sess) if os.path.exists(sess) else []
    check("SIGINT: interrupted turn kept, next turn appended", got[:1] == [("user", "SLEEP")] and got[-2:] == [("user", "after"), ("assistant", "echo: after")], str(got))
    check("SIGINT: later turn closed the turn file", os.path.exists(turn) and read_json(turn).get("open") is False)

    home2 = tempfile.mkdtemp(prefix="detz-repl-thr-")
    env["TAI_HOME"] = home2
    done = subprocess.run([probe, "--throw"], env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=30)
    log = os.path.join(home2, "debug.log")
    text = open(log, encoding="utf-8").read() if os.path.exists(log) else ""
    check("terminate hook: process aborts", done.returncode != 0)
    check("terminate hook: ERROR app.exception logged", "ERROR" in text and "app.exception" in text and "kaboom" in text, text[:200])


def tmux_section(probe, check):
    """The same scenario in real tmux panes (a private tmux server, TMUX_PANE set by tmux itself)."""
    if shutil.which("tmux") is None:
        print("SKIP tmux panes: tmux is not installed")
        return
    root = tempfile.mkdtemp(prefix="detz-repl-tmux-")
    home = os.path.join(root, "home")
    os.makedirs(home)
    env = {key: value for key, value in os.environ.items() if not key.startswith("TMUX")}
    env["TMUX_TMPDIR"] = root

    def tmux(*args):
        return subprocess.run(["tmux", "-f", "/dev/null"] + list(args), env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)

    def screen(pane):
        return tmux("capture-pane", "-p", "-t", pane, "-S", "-200").stdout

    def wait_for(pane, needle, seconds=20):
        end = time.time() + seconds
        while time.time() < end:
            if needle in screen(pane):
                return True
            time.sleep(0.1)
        return False

    def type_line(pane, text):
        tmux("send-keys", "-t", pane, "-l", text)
        tmux("send-keys", "-t", pane, "Enter")

    started = tmux("new-session", "-d", "-s", "t", "-x", "200", "-y", "50", "-e", "TAI_HOME=" + home, "-e", "TAI_PASSPHRASE=", "-e", "NO_COLOR=1", "-e", "HOME=" + home,
                   "bash --norc --noprofile")
    check("tmux session started", started.returncode == 0, started.stderr)
    if started.returncode != 0:
        return
    try:
        pane_a = tmux("list-panes", "-t", "t", "-F", "#{pane_id}").stdout.strip()
        # first start, quit, restart in the same pane
        type_line(pane_a, "%s; echo __DONE_$((0+1))__" % shlex.quote(probe))
        check("tmux A1 prompt", wait_for(pane_a, "tai[default]"))
        type_line(pane_a, "hello tmux")
        type_line(pane_a, "/quit")
        check("tmux A1 exit", wait_for(pane_a, "__DONE_1__"))
        files = session_files(home)
        check("tmux A1 session file", len(files) == 1, str(files))
        type_line(pane_a, "clear; %s; echo __DONE_$((0+2))__" % shlex.quote(probe))
        check("tmux A2 prompt", wait_for(pane_a, "tai[default]"))
        type_line(pane_a, "back again")
        type_line(pane_a, "/quit")
        check("tmux A2 exit", wait_for(pane_a, "__DONE_2__"))
        files = session_files(home)
        check("tmux A2 same file after restart", len(files) == 1, str(files))
        check("tmux A2 history resumed", files and contents(files[0]) == [("user", "hello tmux"), ("assistant", "echo: hello tmux"), ("user", "back again"), ("assistant", "echo: back again")],
              str(contents(files[0])) if files else "")
        # a second pane of the same window is another shell
        split = tmux("split-window", "-d", "-P", "-F", "#{pane_id}", "-t", "t", "bash --norc --noprofile")
        pane_b = split.stdout.strip()
        check("tmux second pane", split.returncode == 0 and pane_b and pane_b != pane_a, split.stderr)
        type_line(pane_b, "%s; echo __DONE_$((0+3))__" % shlex.quote(probe))
        check("tmux B prompt", wait_for(pane_b, "tai[default]"))
        type_line(pane_b, "second pane")
        type_line(pane_b, "/quit")
        check("tmux B exit", wait_for(pane_b, "__DONE_3__"))
        files = session_files(home)
        check("tmux B has its own session", len(files) == 2, str(files))
        mine = sorted(contents(path) for path in files)
        check("tmux B history independent", [("user", "second pane"), ("assistant", "echo: second pane")] in mine and len(mine) == 2)
    finally:
        tmux("kill-server")


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