tai/tests/plexer.py

#!/usr/bin/env python3
"""tests/plexer.py — integration tests for tai's Plexer subsystem (NativePlexerBackend).

Deterministic throughout: this drives tai.py --plexer-wrap, a plumbing mode
that forks a real shell onto a real pty with no LLM involvement 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.

A real bug was caught by actually running this, 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.
NativePlexerBackend's Ctrl+A,q shortcut uses SIGKILL specifically because of
this, not out of caution -- SIGTERM was tried first and silently did nothing
for ~10s in a real run.
"""
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


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+A prefix-key path exactly as a human's terminal would deliver it."""

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

    async def setup(self):
        self.tmp_home = tempfile.mkdtemp(prefix="tai-test-home-")
        self.env_overrides = {"TAI_HOME": self.tmp_home, "TAI_PASSPHRASE": ""}

    async def run(self):
        proc = await self.spawn(extra_args=["--plexer-wrap"])
        self._drain_task = asyncio.ensure_future(_drain_forever(proc))
        sock_path = None
        for _ in range(50):
            matches = glob.glob(os.path.join(self.tmp_home, "plexer", "*.sock"))
            if matches:
                sock_path = matches[0]
                break
            await asyncio.sleep(0.2)
        if not sock_path:
            raise TestAssertionFailed("NativePlexerBackend never created its control socket")

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

        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.4)
        controller.send_keys("stty size", enter=True)
        await asyncio.sleep(0.4)
        after_resize = controller.read_pane(50)
        if "17 66" not in after_resize:
            raise TestAssertionFailed("SIGWINCH resize was not forwarded to the inner shell: %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.PLEXER_PREFIX_KEY)
        os.write(proc._master_fd, b"q")
        await proc.wait(timeout=10.0)
        await asyncio.sleep(0.2)
        socket_survived = os.path.exists(sock_path)

        import subprocess
        still_running = subprocess.run(["ps", "-p", str(child_pid)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0

        log_matches = glob.glob(os.path.join(self.tmp_home, "plexer", "*.log"))
        log_has_marker = False
        if log_matches:
            with open(log_matches[0], "rb") as handle:
                log_has_marker = b"plexer_test_marker" in handle.read()

        return {
            "socket_survived": socket_survived,
            "still_running": still_running,
            "log_has_marker": log_has_marker,
        }

    async def verify(self, result):
        if result["socket_survived"]:
            raise TestAssertionFailed("control socket file was not cleaned up after Ctrl+A,q")
        if result["still_running"]:
            raise TestAssertionFailed("the wrapped shell process was still alive after Ctrl+A,q (leaked process)")
        if not result["log_has_marker"]:
            raise TestAssertionFailed("the continuous session log did not capture the session's actual output")

    async def teardown(self):
        self._drain_task.cancel()
        shutil.rmtree(self.tmp_home, ignore_errors=True)


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

    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 this very relay 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 = 30.0

    async def setup(self):
        self.tmp_home = tempfile.mkdtemp(prefix="tai-test-home-")
        self.env_overrides = {"TAI_HOME": self.tmp_home, "TAI_PASSPHRASE": ""}

    async def run(self):
        proc = await self.spawn(extra_args=["--plexer-wrap"])
        self._drain_task = asyncio.ensure_future(_drain_forever(proc))
        sock_path = None
        for _ in range(50):
            matches = glob.glob(os.path.join(self.tmp_home, "plexer", "*.sock"))
            if matches:
                sock_path = matches[0]
                break
            await asyncio.sleep(0.2)
        if not sock_path:
            raise TestAssertionFailed("NativePlexerBackend never created its control socket")
        controller = tai.NativePlexerController(sock_path)
        await asyncio.sleep(0.3)

        plexer_dir_path = os.path.dirname(sock_path)
        log_matches = glob.glob(os.path.join(plexer_dir_path, "*.log"))
        if not log_matches:
            raise TestAssertionFailed("no log file found to check permissions on")
        log_path = log_matches[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.6)
        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"))

        # The backend must still be alive and responsive after that malformed
        # request -- 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("backend did not survive the malformed request -- a bad control message must never crash the whole wrapped session")

    async def teardown(self):
        self._drain_task.cancel()
        shutil.rmtree(self.tmp_home, ignore_errors=True)


ALL_CASES = [NativePlexerBackendTest(), NativePlexerBackendRobustnessTest()]

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