tai/tai_test.py

#!/usr/bin/env python3
"""tai_test.py — the integration-test framework for driving tai.py as a real subprocess.

This is a library, not a test itself (it has no "test_" prefix on purpose — nothing
here is collected by unittest/pytest). Integration tests that use it live under
tests/ and import it; see tests/README.md for the conventions and tests/rag.py for
a worked example.

Design, and why, in one paragraph per decision (full research notes are archived
under research/ — see research/README.md for how this project keeps that archive):

- Two transport modes, one interface. "pipe" mode uses asyncio's native
  create_subprocess_exec with PIPE streams — the simplest, most deterministic
  choice, and what one-shot `tai.py <prompt>` invocations should use, since they
  just run once and exit. "pty" mode opens a real pseudo-terminal (the stdlib
  `pty` module) and gives the child its slave end as stdin/stdout/stderr, because
  tai.py's own code branches on `sys.stdin.isatty()` (approval prompts, readline)
  — a plain pipe makes it behave as if piped, not as if a human were typing. Both
  modes end up exposing the same asyncio.StreamReader-based interface (the pty
  master fd is wrapped into one via loop.connect_read_pipe), so callers never
  need to know which transport a given test is using.
- start_new_session=True, always, plus a kill-the-whole-group escalation
  (SIGTERM, wait, then SIGKILL, then always proc.wait() to reap). A plain
  proc.terminate() only signals the direct child; if tai.py ever forks a
  subagent or a shell command, a timeout needs to take the whole tree down, not
  just the leader, and failing to reap after killing leaves a zombie.
- Non-blocking everything, no polling. The child's output is consumed through
  asyncio.StreamReader (native backpressure, no busy loop). A human's live
  stdin, when a test opts into --interactive, is bridged into the loop with
  loop.add_reader + raw os.read on the fd (never the buffered sys.stdin, which
  blocks) rather than a background thread — matching how the child's own output
  is handled, and confirmed against multiple independent sources as the correct
  pattern. Full sourcing for every design decision in this module — pty vs.
  pipes, the connect_read_pipe trick, the killpg escalation, the TestCase
  shape itself — is archived at research/async-pty-process-testing/ per this
  project's standing research rule (see AGENTS.md).
- A known, accepted limitation: the pty mode does not explicitly call
  ioctl(TIOCSCTTY), so the child's session may not become the pty's controlling
  terminal in the job-control sense. isatty() on the fd is unaffected by this
  (it's a property of the fd, not of session leadership) so approval-prompt and
  readline behavior work correctly; what would NOT work is SIGWINCH-based resize
  forwarding or Ctrl-C-as-SIGINT generated by the tty line discipline. Tests
  that only script line-based input/output (the normal case) are unaffected.
"""
import argparse
import asyncio
import dataclasses
import datetime
import os
import pty
import re
import signal
import sys
import time
import traceback
from pathlib import Path

DEFAULT_TAI_PATH = Path(__file__).resolve().parent / "tai.py"
ANSI_RE = re.compile(r"\x1b\[[0-9;]*[A-Za-z]")
REPL_PROMPT_RE = re.compile(r"tai\[[^\]]*\]\xe2\x80\xba|tai\[[^\]]*\]›")


def strip_ansi(text):
    return ANSI_RE.sub("", text)


# --------------------------------------------------------------------------- #
# Exceptions. Every failure a test can hit is one of these, and every one of
# them carries the transcript collected so far, so a failing test is debuggable
# from its log file alone, without reproducing it.
# --------------------------------------------------------------------------- #

class TaiTestError(Exception):
    def __init__(self, message, transcript=None):
        super().__init__(message)
        self.transcript = list(transcript or [])


class TestTimeout(TaiTestError):
    """Raised when a step, or the test as a whole, exceeds its timeout."""


class TestCrashed(TaiTestError):
    """Raised when the child process exited (or died to a signal) unexpectedly."""

    def __init__(self, message, returncode, transcript=None):
        super().__init__(message, transcript)
        self.returncode = returncode


class TestAssertionFailed(TaiTestError):
    """Raised by a test's verify() when the captured behavior doesn't match expectations."""


# --------------------------------------------------------------------------- #
# Transcript + result data types.
# --------------------------------------------------------------------------- #

@dataclasses.dataclass
class OutputLine:
    stream: str   # "stdout" | "stdin"  (stderr is merged into stdout in pipe mode; pty merges all three)
    text: str
    ts: float


@dataclasses.dataclass
class TestResult:
    name: str
    passed: bool
    duration: float
    error: str = None
    transcript: list = dataclasses.field(default_factory=list)
    log_path: str = None


# --------------------------------------------------------------------------- #
# TeeLogger: every byte a test sees also goes to screen and to a log file,
# synchronously (plain stdlib file I/O, flushed per write) — no extra
# dependency, and fast enough that it never becomes the bottleneck next to a
# network-bound LLM call or a human typing.
# --------------------------------------------------------------------------- #

class TeeLogger:
    def __init__(self, log_path, echo=True):
        self.log_path = Path(log_path)
        self.log_path.parent.mkdir(parents=True, exist_ok=True)
        self._fh = open(self.log_path, "a", encoding="utf-8")
        self.echo = echo

    def write(self, text):
        self._fh.write(text)
        self._fh.flush()
        if self.echo:
            sys.stdout.write(text)
            sys.stdout.flush()

    def close(self):
        try:
            self._fh.close()
        except OSError:
            pass


# --------------------------------------------------------------------------- #
# StdinBridge: bridges the REAL terminal's stdin into an asyncio.Queue of
# complete lines, without blocking the event loop and without a thread.
# No-ops safely (returns False from start()) if real stdin isn't a tty, e.g.
# under CI — a test must not hang waiting for input nobody can type.
# --------------------------------------------------------------------------- #

class StdinBridge:
    def __init__(self):
        self.queue = asyncio.Queue()
        self._fd = None
        self._buf = b""
        self._active = False
        self._loop = None

    def start(self):
        if not sys.stdin.isatty():
            return False
        self._loop = asyncio.get_running_loop()
        self._fd = sys.stdin.fileno()
        self._loop.add_reader(self._fd, self._on_readable)
        self._active = True
        return True

    def _on_readable(self):
        try:
            chunk = os.read(self._fd, 4096)
        except OSError:
            chunk = b""
        if not chunk:
            self.stop()
            self.queue.put_nowait(None)
            return
        self._buf += chunk
        while b"\n" in self._buf:
            line, self._buf = self._buf.split(b"\n", 1)
            self.queue.put_nowait(line.decode(errors="replace"))

    def stop(self):
        if self._active:
            try:
                self._loop.remove_reader(self._fd)
            except (ValueError, OSError):
                pass
            self._active = False


# --------------------------------------------------------------------------- #
# ManagedProcess: spawns and safely controls one subprocess, in "pipe" or
# "pty" mode, with a single interface either way.
# --------------------------------------------------------------------------- #

class ManagedProcess:
    def __init__(self, command, mode="pipe", cwd=None, env=None, on_line=None):
        self.command = list(command)
        self.mode = mode
        self.cwd = cwd
        self.env = env
        self._on_line = on_line
        self.proc = None
        self._master_fd = None
        self._master_file = None
        self._reader = None
        self._lines = []

    async def start(self):
        loop = asyncio.get_running_loop()
        if self.mode == "pty":
            master_fd, slave_fd = pty.openpty()
            self.proc = await asyncio.create_subprocess_exec(
                *self.command,
                stdin=slave_fd, stdout=slave_fd, stderr=slave_fd,
                cwd=self.cwd, env=self.env,
                start_new_session=True,
            )
            os.close(slave_fd)
            self._master_fd = master_fd
            self._master_file = os.fdopen(master_fd, "rb", buffering=0)
            reader = asyncio.StreamReader(loop=loop)
            protocol = asyncio.StreamReaderProtocol(reader, loop=loop)
            await loop.connect_read_pipe(lambda: protocol, self._master_file)
            self._reader = reader
        else:
            self.proc = await asyncio.create_subprocess_exec(
                *self.command,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.STDOUT,
                cwd=self.cwd, env=self.env,
                start_new_session=True,
            )
            self._reader = self.proc.stdout
        return self

    async def send_line(self, text):
        data = (text + "\n").encode("utf-8", "replace")
        self._record("stdin", text + "\n")
        if self.mode == "pty":
            os.write(self._master_fd, data)
        else:
            self.proc.stdin.write(data)
            await self.proc.stdin.drain()

    def _record(self, stream, text):
        line = OutputLine(stream, text, time.time())
        self._lines.append(line)
        if self._on_line:
            self._on_line(line)

    async def _read_chunk(self, n=4096):
        try:
            chunk = await self._reader.read(n)
        except (OSError, ValueError, asyncio.IncompleteReadError):
            return b""
        return chunk

    async def wait_for_pattern(self, pattern, timeout=30.0):
        """Accumulates output until `pattern` (compiled regex) matches the
        ANSI-stripped buffer; returns everything read before the match."""
        buf = ""
        deadline = time.monotonic() + timeout
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TestTimeout("timed out waiting for pattern %r" % pattern.pattern, self._lines)
            try:
                chunk = await asyncio.wait_for(self._read_chunk(), timeout=remaining)
            except asyncio.TimeoutError:
                raise TestTimeout("timed out waiting for pattern %r" % pattern.pattern, self._lines)
            if not chunk:
                raise TestCrashed(
                    "process ended before pattern %r was seen" % pattern.pattern,
                    self.proc.returncode, self._lines,
                )
            text = chunk.decode("utf-8", "replace")
            self._record("stdout", text)
            buf += strip_ansi(text)
            match = pattern.search(buf)
            if match:
                return buf[:match.start()]

    async def drain_until_eof(self, timeout=120.0):
        """Reads everything until the process closes its output — the one-shot case."""
        deadline = time.monotonic() + timeout
        out = []
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise TestTimeout("timed out waiting for process output to end", self._lines)
            try:
                chunk = await asyncio.wait_for(self._read_chunk(), timeout=remaining)
            except asyncio.TimeoutError:
                raise TestTimeout("timed out waiting for process output to end", self._lines)
            if not chunk:
                break
            text = chunk.decode("utf-8", "replace")
            self._record("stdout", text)
            out.append(text)
        return "".join(out)

    async def wait(self, timeout=None):
        if timeout is None:
            return await self.proc.wait()
        try:
            return await asyncio.wait_for(self.proc.wait(), timeout)
        except asyncio.TimeoutError:
            raise TestTimeout("process did not exit within %.1fs" % timeout, self._lines)

    async def terminate(self, grace=5.0):
        """Escalating, group-wide, always-reaping shutdown. Safe to call on an
        already-exited process, and safe to call twice."""
        if self.proc is None or self.proc.returncode is not None:
            return self.proc.returncode if self.proc else None
        pgid = self.proc.pid
        for sig, wait_s in ((signal.SIGTERM, grace), (signal.SIGKILL, None)):
            try:
                os.killpg(pgid, sig)
            except ProcessLookupError:
                break
            except PermissionError:
                break
            if wait_s is None:
                break
            try:
                await asyncio.wait_for(self.proc.wait(), wait_s)
                break
            except asyncio.TimeoutError:
                continue
        try:
            return await asyncio.wait_for(self.proc.wait(), 5.0)
        except asyncio.TimeoutError:
            return None  # best-effort; we tried SIGKILL already
        finally:
            if self._master_file is not None:
                try:
                    self._master_file.close()
                except OSError:
                    pass

    @property
    def transcript(self):
        return list(self._lines)


# --------------------------------------------------------------------------- #
# TestCase: the thing you extend. setup()/run()/verify()/teardown() are the
# override points; execute() is the orchestration and is not meant to be
# overridden. run() can spawn() zero, one, or many ManagedProcess instances —
# the framework doesn't assume a 1:1 test-to-process relationship, since a
# one-shot-style test naturally spawns many short-lived processes while a
# REPL-style test spawns exactly one long-lived one.
# --------------------------------------------------------------------------- #

class TestCase:
    name = None
    timeout = 60.0
    mode = "pipe"          # default transport for spawn() when not overridden per call
    extra_args = ()        # args always passed to tai.py, before whatever spawn() adds

    def __init__(self, tai_path=None, interactive=False):
        self.name = self.name or type(self).__name__
        self.tai_path = Path(tai_path or DEFAULT_TAI_PATH)
        self.interactive = interactive
        self.logger = None
        self._processes = []
        self.env_overrides = {}

    # ---- override points --------------------------------------------------

    async def setup(self):
        """Optional. Prepare fixtures/state (e.g. an isolated TAI_HOME) before any process starts."""

    async def run(self):
        """Required. Drive the process(es) via self.spawn(); return whatever verify() needs."""
        raise NotImplementedError

    async def verify(self, result):
        """Required. Raise TestAssertionFailed (with a clear message) if `result` doesn't hold up."""
        raise NotImplementedError

    async def teardown(self):
        """Optional. Runs after every spawned process has been terminated, pass or fail."""

    # ---- helpers for run() --------------------------------------------------

    def build_env(self):
        env = dict(os.environ)
        env.update(self.env_overrides)
        return env

    def build_command(self, extra_args=()):
        return [sys.executable, str(self.tai_path)] + list(self.extra_args) + list(extra_args)

    async def spawn(self, extra_args=(), mode=None, env_overrides=None):
        env = self.build_env()
        env.update(env_overrides or {})
        proc = ManagedProcess(self.build_command(extra_args), mode=mode or self.mode, env=env, on_line=self._on_line)
        await proc.start()
        self._processes.append(proc)
        return proc

    def _on_line(self, line):
        if self.logger:
            self.logger.write(line.text)

    def all_transcript(self):
        out = []
        for proc in self._processes:
            out.extend(proc.transcript)
        return out

    # ---- orchestration (do not override) -----------------------------------

    async def execute(self, log_dir):
        start = time.monotonic()
        stamp = datetime.datetime.now().strftime("%Y%m%dT%H%M%S")
        log_path = Path(log_dir) / ("%s-%s.log" % (self.name, stamp))
        self.logger = TeeLogger(log_path, echo=True)
        self.logger.write("=== %s starting %s ===\n" % (self.name, stamp))
        bridge = StdinBridge() if self.interactive else None
        bridge_task = None
        try:
            await self.setup()
            if bridge is not None and bridge.start():
                bridge_task = asyncio.ensure_future(self._pump_bridge(bridge))
            payload = await asyncio.wait_for(self.run(), timeout=self.timeout)
            await self.verify(payload)
            return TestResult(self.name, True, time.monotonic() - start, transcript=self.all_transcript(), log_path=str(log_path))
        except asyncio.TimeoutError:
            err = TestTimeout("test exceeded its %.1fs timeout" % self.timeout, self.all_transcript())
            return self._failure(err, start, log_path)
        except TaiTestError as exc:
            return self._failure(exc, start, log_path)
        except Exception:
            err = TaiTestError(traceback.format_exc(), self.all_transcript())
            return self._failure(err, start, log_path)
        finally:
            if bridge is not None:
                bridge.stop()
            if bridge_task is not None:
                bridge_task.cancel()
            for proc in self._processes:
                await proc.terminate()
            await self.teardown()
            self.logger.write("=== %s finished ===\n" % self.name)
            self.logger.close()

    def _failure(self, exc, start, log_path):
        return TestResult(self.name, False, time.monotonic() - start, error=str(exc), transcript=exc.transcript, log_path=str(log_path))

    async def _pump_bridge(self, bridge):
        """Forwards live human typing (when --interactive) to whichever process
        was spawned most recently — the normal case of exactly one long-lived
        REPL-mode process."""
        while True:
            line = await bridge.queue.get()
            if line is None:
                break
            if self._processes:
                await self._processes[-1].send_line(line)


# --------------------------------------------------------------------------- #
# TestRunner + CLI: runs a list of TestCase instances sequentially (they
# share a real terminal and, in --interactive mode, real human attention, so
# parallel execution is deliberately not offered here), prints a summary.
# --------------------------------------------------------------------------- #

class TestRunner:
    def __init__(self, cases, log_dir="tests/logs"):
        self.cases = cases
        self.log_dir = log_dir

    async def run_all(self):
        results = []
        for case in self.cases:
            print("\n--- running %s ---" % case.name)
            result = await case.execute(self.log_dir)
            status = "PASS" if result.passed else "FAIL"
            print("--- %s: %s (%.1fs, log: %s) ---" % (case.name, status, result.duration, result.log_path))
            if not result.passed:
                print("    %s" % (result.error or "").splitlines()[-1] if result.error else "    (no error message)")
            results.append(result)
        self.print_summary(results)
        return results

    def print_summary(self, results):
        print("\n=== summary ===")
        for result in results:
            status = "PASS" if result.passed else "FAIL"
            print("%-6s %-40s %6.1fs  %s" % (status, result.name, result.duration, result.log_path))
        passed = sum(1 for r in results if r.passed)
        print("%d/%d passed" % (passed, len(results)))


def main_cli(cases, default_log_dir=None):
    """Standard entrypoint for a tests/*.py module: sys.exit(main_cli(ALL_CASES)).

    default_log_dir should normally be the calling module's own directory,
    e.g. str(Path(__file__).resolve().parent / "logs") — passed explicitly
    because __file__ inside this function would resolve to tai_test.py's own
    location, not the caller's."""
    parser = argparse.ArgumentParser(description="run tai integration tests")
    parser.add_argument("--interactive", action="store_true", help="let you type into the running process live, alongside the scripted input")
    parser.add_argument("--filter", default=None, help="only run tests whose name contains this substring")
    parser.add_argument("--timeout", type=float, default=None, help="override every selected test's timeout, in seconds")
    parser.add_argument("--log-dir", default=default_log_dir or str(Path(__file__).resolve().parent / "logs"))
    args = parser.parse_args()
    selected = [case for case in cases if not args.filter or args.filter in case.name]
    if not selected:
        print("no tests matched filter %r" % args.filter)
        return 1
    for case in selected:
        case.interactive = args.interactive
        if args.timeout:
            case.timeout = args.timeout
    runner = TestRunner(selected, log_dir=args.log_dir)
    results = asyncio.run(runner.run_all())
    return 0 if all(r.passed for r in results) else 1


if __name__ == "__main__":
    print(__doc__)
    sys.exit(0)