tai/tests/plexer.py

#!/usr/bin/env python3
"""tests/plexer.py - integration tests for tai's Plexer subsystem (the mplex backend and its control socket).

Deterministic throughout: this drives tai.py --plexer-wrap, a plumbing mode
that starts the mplex server and attaches a real shell on a real pty with no LLM
involved at all, so "test with the real AI" (tests/README.md) doesn't apply here
the way it does to tests/rag.py - there is no model decision anywhere in this
path to test. What's real instead is the pty/signal/socket mechanics themselves:
the whole point of this file is that it exercises the actual
`tai.py --plexer-wrap` binary under a real pty (via tai_test.py's own pty mode -
testing a pty wrapper by driving it through another pty), not mocked
termios/signal calls. The multiplexer's shortcuts and config live in tests/mplex.py;
this file covers what the agent tools rely on: the control socket, per-pane
targeting, file permissions and the session log.

A real bug was caught by actually running the original version of this test, not
by reading the code: an external SIGTERM (and SIGHUP) does not terminate an
interactive bash running as a pty.fork() session leader in this environment -
verified with a bare pty.fork()+bash+os.kill script, independent of any tai.py
code. The server therefore ends a pane with SIGHUP to its group, closes the pty,
and falls back to SIGKILL after MPLEX_REAP_KILL_AFTER seconds.
"""
import glob
import os
import shutil
import struct
import sys
import tempfile
import asyncio
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from tai_test import TestCase, TestAssertionFailed, main_cli  # noqa: E402
import tai  # noqa: E402


async def _drain_forever(proc):
    """A real terminal emulator is always reading, continuously, for as long
    as it's open -- that's what keeps a pty's output buffer from ever
    backing up. This test harness has to replicate that explicitly: without
    a background task doing this, NativePlexerBackend's relay loop can block
    on its own os.write(stdout_fd, ...) the moment a command's echoed output
    is large enough to fill the (unread) outer pty buffer, which stalls the
    *entire* relay loop -- including replying to later control-socket calls,
    which is what made this look like a control-protocol deadlock the first
    time this was debugged, when the actual missing piece was this drain.

    A second real bug, found immediately after fixing the first: once the
    wrapped shell exits, the stream hits EOF and read() returns b"" forever
    without blocking -- so a plain `while True: await proc._read_chunk()`
    busy-spins at that point instead of idling, and spins hard enough to
    starve the event loop so badly that even asyncio.wait_for's own timeout
    machinery stops getting a fair turn. Hence the explicit EOF check."""
    try:
        while True:
            chunk = await proc._read_chunk()
            if not chunk:
                await asyncio.sleep(0.05)
    except asyncio.CancelledError:
        pass


def _mplex_cli(test, *words):
    """Run `tai.py --mplex <words>` against the test's throwaway TAI_HOME (no tty, no attach)."""
    import subprocess
    env = dict(os.environ)
    env.update(test.env_overrides)
    for name in (tai.PLEXER_ACTIVE_ENV, tai.PLEXER_SOCKET_ENV, tai.PLEXER_PANE_ENV, "TMUX"):
        env.pop(name, None)
    done = subprocess.run([sys.executable, str(test.tai_path), "--mplex"] + list(words), env=env, stdin=subprocess.DEVNULL,
                          capture_output=True, text=True, timeout=30)
    return done.returncode, (done.stdout + done.stderr).strip()


def _pane_logs(home):
    return sorted(path for path in glob.glob(os.path.join(home, "plexer", "*-pane*.log")))


class NativePlexerBackendTest(TestCase):
    """Spawns `tai.py --plexer-wrap` under a real pty (tai_test's pty mode),
    then drives it as an external controller would: over its control socket,
    plus raw keystrokes written straight to the outer pty master to exercise
    the Ctrl+B prefix path exactly as a human's terminal would deliver it."""

    name = "plexer-native-backend"
    mode = "pty"
    timeout = 40.0

    async def setup(self):
        self.tmp_home = tempfile.mkdtemp(prefix="tai-test-home-")
        self.env_overrides = {"TAI_HOME": self.tmp_home, "TAI_PASSPHRASE": "", "TERM": "xterm-256color", "SHELL": "/bin/bash"}

    async def run(self):
        proc = await self.spawn(extra_args=["--plexer-wrap"])
        self._drain_task = asyncio.ensure_future(_drain_forever(proc))
        sock_path = os.path.join(self.tmp_home, "plexer", tai.MPLEX_SOCKET_NAME)
        for _ in range(50):
            if os.path.exists(sock_path):
                break
            await asyncio.sleep(0.2)
        else:
            raise TestAssertionFailed("the mplex server never created its control socket")

        controller = tai.NativePlexerController(sock_path)
        await asyncio.sleep(0.8)

        status = controller.status()
        child_pid = status.get("pid")
        if not child_pid or child_pid <= 0:
            raise TestAssertionFailed("status() did not report a valid child pid: %r" % status)

        controller.send_keys("echo plexer_test_marker", enter=True)
        await asyncio.sleep(0.4)
        after_send = controller.read_pane(50)
        if "plexer_test_marker" not in after_send:
            raise TestAssertionFailed("send_keys/read_pane round trip failed: %r" % after_send[-300:])

        controller.type_text("echo typed_marker_ok", delay=0.01, enter=True)
        await asyncio.sleep(0.4)
        after_type = controller.read_pane(50)
        if "typed_marker_ok" not in after_type:
            raise TestAssertionFailed("type_text did not reach the shell: %r" % after_type[-300:])

        import fcntl
        import termios
        fcntl.ioctl(proc._master_fd, termios.TIOCSWINSZ, struct.pack("HHHH", 17, 66, 0, 0))
        os.kill(proc.proc.pid, tai.signal.SIGWINCH)
        await asyncio.sleep(0.5)
        controller.send_keys("stty size", enter=True)
        await asyncio.sleep(0.4)
        after_resize = controller.read_pane(50)
        if "16 66" not in after_resize:
            raise TestAssertionFailed("SIGWINCH resize was not forwarded to the pane (17 rows minus the status row): %r" % after_resize[-300:])

        controller.send_keys("echo contains_a_d_character", enter=True)
        await asyncio.sleep(0.3)
        after_plain = controller.read_pane(50)
        if "contains_a_d_character" not in after_plain:
            raise TestAssertionFailed("a plain command containing the letter 'd' was mangled by prefix-key handling")

        os.write(proc._master_fd, tai.MplexKeys.to_bytes("C-b") + b"d")
        client_code = await proc.wait(timeout=10.0)
        await asyncio.sleep(0.3)
        socket_survived = os.path.exists(sock_path)
        import subprocess
        alive_after_detach = subprocess.run(["ps", "-p", str(child_pid)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0

        log_has_marker = False
        logs = _pane_logs(self.tmp_home)
        if logs:
            with open(logs[0], "rb") as handle:
                log_has_marker = b"plexer_test_marker" in handle.read()

        _mplex_cli(self, "kill-server")
        for _ in range(40):
            if not os.path.exists(sock_path):
                break
            await asyncio.sleep(0.1)
        await asyncio.sleep(0.5)
        socket_after_kill = os.path.exists(sock_path)
        alive_after_kill = subprocess.run(["ps", "-p", str(child_pid)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0

        return {
            "client_code": client_code,
            "socket_survived": socket_survived,
            "alive_after_detach": alive_after_detach,
            "log_has_marker": log_has_marker,
            "socket_after_kill": socket_after_kill,
            "alive_after_kill": alive_after_kill,
        }

    async def verify(self, result):
        if result["client_code"] != 0:
            raise TestAssertionFailed("detaching (C-b d) should end the client with exit code 0, got %r" % result["client_code"])
        if not result["socket_survived"] or not result["alive_after_detach"]:
            raise TestAssertionFailed("detach must leave the server and the shell running: %r" % result)
        if not result["log_has_marker"]:
            raise TestAssertionFailed("the continuous session log did not capture the session's actual output")
        if result["socket_after_kill"]:
            raise TestAssertionFailed("control socket file was not cleaned up after kill-server")
        if result["alive_after_kill"]:
            raise TestAssertionFailed("the pane's shell process was still alive after kill-server (leaked process)")

    async def teardown(self):
        self._drain_task.cancel()
        _mplex_cli(self, "kill-server")
        shutil.rmtree(self.tmp_home, ignore_errors=True)


class NativePlexerBackendRobustnessTest(TestCase):
    """Targets specific bugs found and fixed during a deliberate review pass after
    the happy-path test first went green - a happy-path test passing does not
    mean the code is reliable, and all of these would have sailed through it:

    1. File permissions: the control socket and the session log (which
       captures the ENTIRE terminal session verbatim -- anything typed,
       including passwords) were being created with umask-default
       permissions, often world-readable/connectable on a multi-user box.
    2. The control protocol's server side did a single recv() instead of
       looping like the client side already did, silently truncating any
       request too large for one read -- untestable with short demo strings,
       which is exactly why the first test never caught it.
    3. A malformed (non-JSON-object) request on the control socket was not
       defensively handled and would have propagated an uncaught exception
       out of the relay loop, crashing the user's entire live wrapped
       terminal session over one bad control message.
    4. A genuine deadlock: plexer_send used to write its text directly into
       the pty master from inside the control-request handler. For a large
       payload on an echo-enabled pty, that write can block waiting for
       buffer space that only frees up once something reads the echoed
       output back out -- which is the server loop's own job, and it was
       the one blocked. Fixed by queuing all writes-to-child and only
       ever draining them from the main select() loop, cooperatively with
       reading the child's output, never as a blocking write from inside
       request handling.
    """

    name = "plexer-native-backend-robustness"
    mode = "pty"
    timeout = 40.0

    async def setup(self):
        self.tmp_home = tempfile.mkdtemp(prefix="tai-test-home-")
        self.env_overrides = {"TAI_HOME": self.tmp_home, "TAI_PASSPHRASE": "", "TERM": "xterm-256color", "SHELL": "/bin/bash"}

    async def run(self):
        proc = await self.spawn(extra_args=["--plexer-wrap"])
        self._drain_task = asyncio.ensure_future(_drain_forever(proc))
        sock_path = os.path.join(self.tmp_home, "plexer", tai.MPLEX_SOCKET_NAME)
        for _ in range(50):
            if os.path.exists(sock_path):
                break
            await asyncio.sleep(0.2)
        else:
            raise TestAssertionFailed("the mplex server never created its control socket")
        controller = tai.NativePlexerController(sock_path)
        await asyncio.sleep(0.8)

        plexer_dir_path = os.path.dirname(sock_path)
        controller.status()
        logs = _pane_logs(self.tmp_home)
        if not logs:
            raise TestAssertionFailed("no pane log file found to check permissions on")
        log_path = logs[0]

        dir_mode = os.stat(plexer_dir_path).st_mode & 0o777
        sock_mode = os.stat(sock_path).st_mode & 0o777
        log_mode = os.stat(log_path).st_mode & 0o777

        # Large payload: bigger than the control protocol's 65536-byte recv
        # chunk size, to actually exercise the server-side framing loop
        # rather than the common case that fits in one read.
        needle_head = "LARGE_PAYLOAD_START_"
        needle_tail = "_LARGE_PAYLOAD_END"
        padding = "x" * 90000
        large_text = "echo %s%s%s" % (needle_head, padding, needle_tail)
        controller.send_keys(large_text, enter=True)
        await asyncio.sleep(0.8)
        large_result = controller.read_pane(2000)

        # Malformed control request: connect directly to the socket (bypassing
        # the well-formed NativePlexerController entirely) and send something
        # that is valid JSON but not an object, which the old code would have
        # handed straight to request.get("cmd") and crashed on.
        import socket as socket_module
        import json as json_module
        raw_sock = socket_module.socket(socket_module.AF_UNIX, socket_module.SOCK_STREAM)
        raw_sock.settimeout(5)
        raw_sock.connect(sock_path)
        raw_sock.sendall(b"[1, 2, 3]\n")
        malformed_reply_raw = raw_sock.recv(65536)
        raw_sock.close()
        malformed_reply = json_module.loads(malformed_reply_raw.decode("utf-8", "replace"))

        # Garbage that is neither a frame nor a line must not hang or kill the server either.
        junk = socket_module.socket(socket_module.AF_UNIX, socket_module.SOCK_STREAM)
        junk.connect(sock_path)
        junk.sendall(b"\xff\xff\xff not json at all")
        junk.close()

        # The server must still be alive and responsive after those requests -
        # this is the actual point of the test, not the error shape above.
        status_after = controller.status()

        controller.send_keys("exit", enter=True)
        await proc.wait(timeout=10.0)

        return {
            "dir_mode": dir_mode,
            "sock_mode": sock_mode,
            "log_mode": log_mode,
            "large_result": large_result,
            "malformed_reply": malformed_reply,
            "status_after_pid": status_after.get("pid"),
        }

    async def verify(self, result):
        if result["dir_mode"] != 0o700:
            raise TestAssertionFailed("plexer dir mode is %o, expected 0700 (owner-only)" % result["dir_mode"])
        if result["sock_mode"] != 0o600:
            raise TestAssertionFailed("control socket mode is %o, expected 0600 (owner-only)" % result["sock_mode"])
        if result["log_mode"] != 0o600:
            raise TestAssertionFailed("session log mode is %o, expected 0600 (owner-only) -- it captures everything typed" % result["log_mode"])

        large = result["large_result"]
        if "LARGE_PAYLOAD_START_" not in large or "_LARGE_PAYLOAD_END" not in large:
            raise TestAssertionFailed("large (>65536 byte) plexer_send payload was truncated by the control protocol")

        if result["malformed_reply"].get("ok") is not False:
            raise TestAssertionFailed("a malformed control request did not come back as a clean error: %r" % result["malformed_reply"])

        if not result["status_after_pid"] or result["status_after_pid"] <= 0:
            raise TestAssertionFailed("server did not survive the malformed requests -- a bad control message must never crash the whole session")

    async def teardown(self):
        self._drain_task.cancel()
        _mplex_cli(self, "kill-server")
        shutil.rmtree(self.tmp_home, ignore_errors=True)


class PlexerPaneTargetingTest(TestCase):
    """With several panes in one server, the agent tools must reach the pane they were started in:
    the controller sends the pane id from $TAI_PLEXER_PANE, and an unknown pane is an error, not a guess."""

    name = "plexer-pane-targeting"
    mode = "pty"
    timeout = 40.0

    async def setup(self):
        self.tmp_home = tempfile.mkdtemp(prefix="tai-test-home-")
        self.env_overrides = {"TAI_HOME": self.tmp_home, "TAI_PASSPHRASE": "", "TERM": "xterm-256color", "SHELL": "/bin/bash"}

    async def run(self):
        proc = await self.spawn(extra_args=["--plexer-wrap"])
        self._drain_task = asyncio.ensure_future(_drain_forever(proc))
        sock_path = os.path.join(self.tmp_home, "plexer", tai.MPLEX_SOCKET_NAME)
        for _ in range(50):
            if os.path.exists(sock_path):
                break
            await asyncio.sleep(0.2)
        await asyncio.sleep(0.8)
        code, out = _mplex_cli(self, "split-window", "-h", "-P", "-F", "#{pane_id}")
        await asyncio.sleep(0.8)
        first = tai.NativePlexerController(sock_path, "0")
        second = tai.NativePlexerController(sock_path, out.strip().lstrip("%"))
        first.send_keys("echo only-in-first", enter=True)
        second.send_keys("echo only-in-second", enter=True)
        await asyncio.sleep(0.6)
        first_text, second_text = first.read_pane(50), second.read_pane(50)
        env_seen = []
        for controller, label in ((first, "first"), (second, "second")):
            controller.send_keys("echo pane=$TAI_PLEXER_PANE active=$TAI_PLEXER_ACTIVE", enter=True)
        await asyncio.sleep(0.6)
        env_first, env_second = first.read_pane(50), second.read_pane(50)
        bad = None
        try:
            tai.NativePlexerController(sock_path, "999").status()
        except tai.PlexerError as exc:
            bad = str(exc)
        return {"first": first_text, "second": second_text, "env_first": env_first, "env_second": env_second, "bad": bad,
                "split_id": out.strip(), "split_code": code}

    async def verify(self, result):
        if result["split_code"] != 0 or not result["split_id"].startswith("%"):
            raise TestAssertionFailed("split-window -P did not report the new pane: %r" % result)
        if "only-in-first" not in result["first"] or "only-in-second" in result["first"].replace("echo only-in-second", ""):
            raise TestAssertionFailed("pane 0 output is wrong: %r" % result["first"][-300:])
        if "only-in-second" not in result["second"] or "only-in-first" in result["second"].replace("echo only-in-first", ""):
            raise TestAssertionFailed("second pane output is wrong: %r" % result["second"][-300:])
        if "pane=0 active=1" not in result["env_first"]:
            raise TestAssertionFailed("pane 0 does not see TAI_PLEXER_PANE=0: %r" % result["env_first"][-200:])
        if "pane=%s active=1" % result["split_id"].lstrip("%") not in result["env_second"]:
            raise TestAssertionFailed("second pane does not see its own TAI_PLEXER_PANE: %r" % result["env_second"][-200:])
        if not result["bad"] or "no such pane" not in result["bad"]:
            raise TestAssertionFailed("an unknown pane id must be an error, got %r" % result["bad"])

    async def teardown(self):
        self._drain_task.cancel()
        _mplex_cli(self, "kill-server")
        shutil.rmtree(self.tmp_home, ignore_errors=True)


ALL_CASES = [NativePlexerBackendTest(), NativePlexerBackendRobustnessTest(), PlexerPaneTargetingTest()]

if __name__ == "__main__":
    sys.exit(main_cli(ALL_CASES, default_log_dir=str(Path(__file__).resolve().parent / "logs")))