"""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