tai/tai_pkg/plexer.py

"""Terminal multiplexer controllers and wrapping setup."""

import json
import os
import signal
import socket
import struct
import subprocess
import sys
import time
from datetime import datetime, timezone

from tai_pkg.compat import fcntl, pty, select, termios, tty
from tai_pkg.consts import Consts
from tai_pkg.text import Text

__all__ = ["PlexerError", "PlexerController", "NativePlexerController", "TmuxPlexerController", "PlexerControllerFactory", "PlexerBackend", "NativePlexerBackend", "PlexerSetup"]


class PlexerError(Exception):
    """Raised by a PlexerController when the backend can't be reached or refuses a command."""


# --------------------------------------------------------------------------- #
# Controller side: what Tools.run_plexer_* actually calls. One small
# interface, two backends, chosen by which environment variable is set —
# TAI_PLEXER_SOCKET (our own NativePlexerBackend already wrapping this shell)
# or TMUX_PANE (tmux is already the active manager of this shell). Neither
# set means no plexer session is active and the plexer_* tools stay hidden
# (see TOOL_ENV_GATES below) rather than erroring.
# --------------------------------------------------------------------------- #

class PlexerController:
    kind = "abstract"

    def send_keys(self, text, enter=False):
        raise NotImplementedError

    def read_pane(self, lines=200):
        raise NotImplementedError

    def type_text(self, text, delay=0.03, enter=True):
        for ch in text:
            self.send_keys(ch)
            time.sleep(delay)
        if enter:
            self.send_keys("", enter=True)

    def status(self):
        raise NotImplementedError


class NativePlexerController(PlexerController):
    kind = "native"

    def __init__(self, socket_path):
        self.socket_path = socket_path

    def _call(self, payload):
        try:
            conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            conn.settimeout(10)
            conn.connect(self.socket_path)
        except OSError as exc:
            raise PlexerError("cannot reach native plexer backend: " + Text.short_error(exc))
        try:
            conn.sendall((json.dumps(payload) + "\n").encode("utf-8"))
            chunks = []
            while True:
                chunk = conn.recv(65536)
                if not chunk:
                    break
                chunks.append(chunk)
                if chunk.endswith(b"\n"):
                    break
        except OSError as exc:
            raise PlexerError("native plexer backend read/write failed: " + Text.short_error(exc))
        finally:
            conn.close()
        raw = b"".join(chunks).decode("utf-8", "replace").strip()
        try:
            reply = json.loads(raw) if raw else {}
        except ValueError:
            raise PlexerError("native plexer backend sent an unparseable reply")
        if not reply.get("ok", False):
            raise PlexerError(reply.get("error") or "native plexer backend refused the command")
        return reply.get("data")

    def send_keys(self, text, enter=False):
        return self._call({"cmd": "send", "text": text, "enter": bool(enter)})

    def read_pane(self, lines=200):
        return self._call({"cmd": "read", "lines": int(lines)})

    def status(self):
        return self._call({"cmd": "status"})


class TmuxPlexerController(PlexerController):
    kind = "tmux"

    def __init__(self, pane):
        self.pane = pane

    def _tmux(self, *args):
        try:
            done = subprocess.run(["tmux"] + list(args), stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=15)
        except (OSError, subprocess.SubprocessError) as exc:
            raise PlexerError("tmux call failed: " + Text.short_error(exc))
        if done.returncode != 0:
            raise PlexerError("tmux error: " + (done.stderr or "").strip()[:300])
        return done.stdout

    def send_keys(self, text, enter=False):
        if text:
            self._tmux("send-keys", "-t", self.pane, "-l", text)
        if enter:
            self._tmux("send-keys", "-t", self.pane, "Enter")
        return True

    def read_pane(self, lines=200):
        return self._tmux("capture-pane", "-p", "-t", self.pane, "-S", "-%d" % max(1, int(lines)))

    def status(self):
        # list-panes, not display-message -p: display-message -p silently
        # falls back to tmux's default/current pane when -t names a target
        # that doesn't exist (confirmed empirically -- it still exits 0 with
        # blank format substitutions), while list-panes -t actually validates
        # the target and errors like send-keys/capture-pane already do.
        info = self._tmux("list-panes", "-t", self.pane, "-F", "#{session_name}:#{window_index}.#{pane_index} #{pane_current_command}")
        return {"backend": "tmux", "pane": self.pane, "info": info.strip()}


class PlexerControllerFactory:
    """The whole detection flow lives in one place: which env var is set
    decides which backend is already managing this shell. Nothing here ever
    spawns a backend — spawning only happens via `tai.py --plexer-wrap`,
    invoked by the bashrc hook, and only when neither var is already set."""

    @staticmethod
    def detect():
        socket_path = os.environ.get(Consts.PLEXER_SOCKET_ENV)
        if socket_path:
            return NativePlexerController(socket_path)
        pane = os.environ.get("TMUX_PANE")
        if pane:
            return TmuxPlexerController(pane)
        return None


# --------------------------------------------------------------------------- #
# Backend side: only NativePlexerBackend actually exists, because tmux
# already IS a backend when it's active — there is nothing for tai to wrap
# in that case (see PlexerControllerFactory). PlexerBackend is kept as the
# extension point for a future alternative (e.g. a screen-based backend)
# without touching the controller side or the detection flow above.
# --------------------------------------------------------------------------- #

class PlexerBackend:
    def run(self):
        raise NotImplementedError


class NativePlexerBackend(PlexerBackend):
    """Owns the pty directly: forks the user's real shell onto one, relays
    bytes between the real terminal and that pty in both directions, forwards
    SIGWINCH as a resize, tees everything to a log file, and serves control
    commands (send/read/status) over a Unix socket so a separate tai process
    (the agent the user is actually talking to) can drive this session.

    One prefix key (Ctrl+A, PLEXER_PREFIX_KEY) is reserved for tai's own
    shortcuts in the raw keystroke stream. There is deliberately no
    tmux-compatible keymap here: when tmux is the active backend, tai never
    owns the raw keystroke stream at all (see TmuxPlexerController), so there
    is nothing to avoid colliding with.

    A characteristic worth knowing, not a bug: this relay loop is single-
    threaded and synchronous, so its own os.write(stdout_fd, ...) applies
    normal pty backpressure -- if whatever is reading the OTHER end of the
    outer pty (normally a real terminal emulator, which always reads
    continuously) ever stops draining for a stretch, a large enough burst of
    child output can block that write and stall the whole loop until the
    reader catches up. This is standard for any synchronous pty relay
    (tmux, ssh, `script` all have the same property) and is why writes
    headed INTO the child are queued and drained cooperatively rather than
    written directly (see _pending_to_child) -- that fix was for a
    self-inflicted deadlock on the inner side; this is normal, external
    backpressure on the outer side, not something to engineer around.
    """

    def __init__(self, home, profile, shell_cmd=None):
        self.home = home
        self.profile = profile
        self.shell_cmd = shell_cmd or [os.environ.get("SHELL") or "/bin/bash"]
        self.session_id = "%s-%d" % (datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S"), os.getpid())
        self.socket_path = os.path.join(PlexerSetup.plexer_dir(home), self.session_id + ".sock")
        self.log_path = os.path.join(PlexerSetup.plexer_dir(home), self.session_id + ".log")
        self.pid = None
        self.master_fd = None
        self._log_fh = None
        self._server = None
        self._saved_termios = None
        self._resize_pending = False
        self._running = False
        self._pending_to_child = bytearray()

    def _set_raw(self):
        if termios is None:
            return
        fd = sys.stdin.fileno()
        try:
            self._saved_termios = termios.tcgetattr(fd)
            tty.setraw(fd)
        except (termios.error, OSError, ValueError):
            self._saved_termios = None

    def _restore_tty(self):
        if termios is None or self._saved_termios is None:
            return
        try:
            termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN, self._saved_termios)
        except (termios.error, OSError):
            pass

    def _get_winsize(self, fd):
        try:
            packed = fcntl.ioctl(fd, termios.TIOCGWINSZ, b"\0" * 8)
        except OSError:
            return None
        rows, cols = struct.unpack("HH", packed[:4])
        return rows, cols

    def _apply_winsize(self):
        size = self._get_winsize(sys.stdin.fileno())
        if size is None or self.master_fd is None:
            return
        rows, cols = size
        try:
            fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, cols, 0, 0))
        except OSError:
            pass

    def _on_sigwinch(self, _signum, _frame):
        self._resize_pending = True

    def _spawn_child(self):
        pid, master_fd = pty.fork()
        if pid == 0:
            env = dict(os.environ)
            env[Consts.PLEXER_ACTIVE_ENV] = "1"
            env[Consts.PLEXER_SOCKET_ENV] = self.socket_path
            try:
                os.execvpe(self.shell_cmd[0], self.shell_cmd, env)
            except OSError as exc:
                sys.stderr.write("tai plexer: failed to exec %s: %s\n" % (self.shell_cmd[0], exc))
                os._exit(127)
        return pid, master_fd

    def _open_server(self):
        try:
            os.unlink(self.socket_path)
        except OSError:
            pass
        server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        server.bind(self.socket_path)
        # Same reasoning as the log file: a plain AF_UNIX bind() creates the
        # socket file with umask-default permissions, often world-connectable
        # on a multi-user machine -- which would let another local user send
        # keystrokes into this terminal session. Owner-only, no exceptions.
        try:
            os.chmod(self.socket_path, 0o600)
        except OSError:
            pass
        server.listen(8)
        server.setblocking(False)
        return server

    def _recv_request(self, conn):
        """Mirrors NativePlexerController._call()'s own framing exactly: loop
        until a newline-terminated chunk arrives, not a single recv() call.
        json.dumps always escapes a real newline inside a string as the two
        characters '\\n', so the one literal 0x0A byte in the whole message
        is always the trailing terminator added below -- a single recv() would
        silently truncate any request too large for one read (e.g. a long
        plexer_type payload), which is exactly the kind of bug that only shows
        up under real load, not in a quick manual check."""
        chunks = []
        total = 0
        while True:
            chunk = conn.recv(65536)
            if not chunk:
                break
            chunks.append(chunk)
            total += len(chunk)
            if chunk.endswith(b"\n") or total >= Consts.PLEXER_RECV_MAX:
                break
        return b"".join(chunks)

    def _handle_control(self, conn):
        try:
            conn.settimeout(2)
            raw = self._recv_request(conn)
            request = json.loads(raw.decode("utf-8", "replace")) if raw.strip() else {}
            if not isinstance(request, dict):
                raise ValueError("request must be a JSON object")
            reply = self._dispatch(request)
        except Exception as exc:
            # Deliberately broad: a malformed or unexpected request must never
            # be able to crash this relay loop and take the user's whole live
            # terminal session down with it. This is the one place in the
            # whole backend where "catch everything" is the correct choice,
            # not laziness -- everything else in this class keeps narrow,
            # specific except clauses on purpose.
            reply = {"ok": False, "error": Text.short_error(exc)}
        try:
            conn.sendall((json.dumps(reply) + "\n").encode("utf-8"))
        except OSError:
            pass
        finally:
            conn.close()

    def _dispatch(self, request):
        cmd = request.get("cmd")
        if cmd == "send":
            text = str(request.get("text") or "")
            if text:
                self._pending_to_child += text.encode("utf-8", "replace")
            if request.get("enter"):
                self._pending_to_child += b"\r"
            return {"ok": True, "data": True}
        if cmd == "read":
            limit = max(1, min(int(request.get("lines") or 200), 5000))
            return {"ok": True, "data": self._tail_log(limit)}
        if cmd == "status":
            return {"ok": True, "data": {"backend": "native", "pid": self.pid, "log": self.log_path, "socket": self.socket_path}}
        return {"ok": False, "error": "unknown command %r" % cmd}

    def _tail_log(self, lines, chunk_size=65536, max_bytes=1 << 24):
        """Reads only the tail, in growing byte chunks from the end, instead
        of loading the whole file -- a session log is exactly the kind of
        file this feature encourages to grow large over a long-running
        terminal session, and read_pane()/plexer_read get called repeatedly,
        so re-reading the entire file on every call would make the backend
        slower the longer the very session it's meant to support stays open."""
        try:
            size = os.path.getsize(self.log_path)
        except OSError:
            return ""
        with open(self.log_path, "rb") as handle:
            read_from = size
            data = b""
            while True:
                read_from = max(0, read_from - chunk_size)
                handle.seek(read_from)
                data = handle.read(size - read_from)
                if data.count(b"\n") >= lines or read_from == 0 or len(data) >= max_bytes:
                    break
            text = data.decode("utf-8", "replace")
            return "\n".join(text.splitlines()[-lines:])

    def run(self):
        if pty is None or termios is None or select is None:
            return "error: pty/termios/select not available on this platform"
        self.pid, self.master_fd = self._spawn_child()
        # The log captures the entire terminal session verbatim -- anything typed,
        # including passwords and tokens -- so it must never be created world- or
        # group-readable. os.open's mode is applied atomically at creation (no
        # window where a looser default permission briefly exists), unlike a
        # plain open() + a later os.chmod.
        log_fd = os.open(self.log_path, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
        self._log_fh = os.fdopen(log_fd, "ab", buffering=0)
        self._server = self._open_server()
        self._set_raw()
        self._apply_winsize()
        try:
            old_handler = signal.signal(signal.SIGWINCH, self._on_sigwinch) if hasattr(signal, "SIGWINCH") else None
        except (ValueError, OSError):
            old_handler = None
        self._running = True
        crash = None
        try:
            self._relay_loop()
        except Exception as exc:  # noqa: BLE001 -- see comment below
            # Deliberately broad, and deliberately not re-raised bare: letting
            # an unexpected exception propagate straight out of run() would
            # skip the waitpid() below entirely (a bare `finally` only runs
            # cleanup, it does not stop the exception from continuing past
            # it), orphaning the wrapped shell with no supervisor left to
            # reap it. The user's actual terminal child process must not be
            # allowed to leak just because something in the relay loop broke.
            crash = exc
        finally:
            self._running = False
            if old_handler is not None:
                try:
                    signal.signal(signal.SIGWINCH, old_handler)
                except (ValueError, OSError):
                    pass
            self._restore_tty()
            try:
                self._log_fh.close()
            except OSError:
                pass
            try:
                self._server.close()
            except OSError:
                pass
            try:
                os.unlink(self.socket_path)
            except OSError:
                pass
        if crash is not None:
            try:
                os.killpg(self.pid, signal.SIGKILL)
            except OSError:
                try:
                    os.kill(self.pid, signal.SIGKILL)
                except OSError:
                    pass
            sys.stderr.write("tai plexer: wrapper crashed (%s), terminating the session\n" % Text.short_error(crash))
        try:
            _, status = os.waitpid(self.pid, 0)
            return os.waitstatus_to_exitcode(status) if hasattr(os, "waitstatus_to_exitcode") else (1 if crash else 0)
        except OSError:
            return 1 if crash else 0

    def _relay_loop(self):
        """Every byte headed INTO the child goes through self._pending_to_child,
        drained here in bounded chunks only when select() says master_fd is
        writable -- never a direct, unbounded os.write(self.master_fd, ...)
        anywhere, including from the keystroke path below. A real deadlock
        was found and fixed here: a large plexer_send payload written
        directly and synchronously into the pty (echo-enabled, so the write
        itself generates output) filled the pty's buffer while the only
        thing that could ever drain that output -- this very loop -- was
        itself blocked inside that same write call. Queuing and draining
        cooperatively, alongside reading master_fd's output on the same
        select() pass, is what actually avoids that, not a bigger buffer."""
        stdin_fd = sys.stdin.fileno()
        stdout_fd = sys.stdout.fileno()
        pending_prefix = False
        while True:
            if self._resize_pending:
                self._resize_pending = False
                self._apply_winsize()
            watch_read = [stdin_fd, self.master_fd, self._server.fileno()]
            watch_write = [self.master_fd] if self._pending_to_child else []
            try:
                readable, writable, _ = select.select(watch_read, watch_write, [], 1.0)
            except InterruptedError:
                continue
            if not readable and not writable and not self._process_alive():
                break
            if self._server.fileno() in readable:
                try:
                    conn, _addr = self._server.accept()
                    self._handle_control(conn)
                except OSError:
                    pass
            if self.master_fd in writable and self._pending_to_child:
                try:
                    written = os.write(self.master_fd, bytes(self._pending_to_child[:4096]))
                    del self._pending_to_child[:written]
                except OSError:
                    self._pending_to_child = bytearray()
            if self.master_fd in readable:
                try:
                    data = os.read(self.master_fd, 65536)
                except OSError:
                    data = b""
                if not data:
                    break
                self._log_fh.write(data)
                os.write(stdout_fd, data)
            if stdin_fd in readable:
                try:
                    data = os.read(stdin_fd, 4096)
                except OSError:
                    data = b""
                if not data:
                    break
                data, pending_prefix = self._strip_prefix_commands(data, pending_prefix)
                if data:
                    self._pending_to_child += data
            if not self._process_alive():
                break

    def _strip_prefix_commands(self, data, pending_prefix):
        """Intercepts tai's own Ctrl+A shortcuts; everything else passes through untouched.

        'q' quits the wrapped session. Verified empirically (see
        tests/plexer.py) that in this environment an interactive bash running
        as a pty.fork() session leader does not die from an external SIGHUP
        OR SIGTERM -- confirmed independent of this class with a bare
        pty.fork()+bash+os.kill script, so this is not something tai is
        doing wrong, it is this environment's actual bash behavior. SIGKILL
        is the one signal no process can ignore, so that's what's used here.
        This is a clean kill, not a true tmux-style detach-and-reattach --
        there is no session persistence to reattach to yet, so it is named
        and documented as "quit," not "detach", to not promise what isn't built."""
        out = bytearray()
        i = 0
        while i < len(data):
            byte = data[i:i + 1]
            if pending_prefix:
                pending_prefix = False
                if byte == b"q":
                    # killpg, not kill: pty.fork()'s child is its own session
                    # and process-group leader (pid == pgid), so this reaches
                    # a plain foreground child the bare shell pid alone would
                    # miss. It does NOT reach every descendant, and verified
                    # empirically rather than assumed: an interactive bash's
                    # own job control puts each backgrounded job (`sleep &`)
                    # in a NEW process group of its own, specifically so
                    # signals to the shell's group -- Ctrl+C, this killpg,
                    # closing a real terminal -- don't also take down its
                    # background jobs. A backgrounded job surviving `quit` is
                    # therefore expected, standard shell behavior, identical
                    # to what closing a real terminal does, not a gap to
                    # close here.
                    try:
                        os.killpg(self.pid, signal.SIGKILL)
                    except OSError:
                        os.kill(self.pid, signal.SIGKILL)
                elif byte == Consts.PLEXER_PREFIX_KEY:
                    out += Consts.PLEXER_PREFIX_KEY
                i += 1
                continue
            if byte == Consts.PLEXER_PREFIX_KEY:
                pending_prefix = True
                i += 1
                continue
            out += byte
            i += 1
        return bytes(out), pending_prefix

    def _process_alive(self):
        try:
            pid, _status = os.waitpid(self.pid, os.WNOHANG)
        except OSError:
            return False
        return pid == 0


class PlexerSetup:
    """PlexerSetup."""

    @staticmethod
    def plexer_dir(home):
        path = os.path.join(home, Consts.PLEXER_DIR_NAME)
        os.makedirs(path, exist_ok=True)
        try:
            os.chmod(path, 0o700)
        except OSError:
            pass
        return path

    @staticmethod
    def plexer_should_wrap(env):
        """The Python-side half of the detection flow (see research/shell-multiplexer-auto-attach/).
    The cheap, fast checks ($- interactive flag, isatty) already happened in the bashrc
    block itself before tai.py was even invoked -- this only covers what needs Python:
    never wrap if a plexer is already active, and never wrap if tmux already is."""
        if env.get(Consts.PLEXER_ACTIVE_ENV):
            return False
        if env.get("TMUX"):
            return False
        if os.path.exists(os.path.join(os.path.expanduser("~"), ".tai-no-plexer")):
            return False
        return True