Adds assistants.

castle-api/src/castle_api/pty_session.py

@@ -0,0 +1,156 @@
+"""Run a command inside a pseudo-terminal and stream it over asyncio.
+
+Castle stays assistant-agnostic: it just launches a command in a real TTY and
+shunts raw bytes both ways. The command's own TUI (an agent CLI, a shell, …)
+renders exactly as it would in a terminal — Castle never parses its output.
+
+The PTY master is a raw OS fd, not an asyncio stream, so we watch it with
+``loop.add_reader`` (no extra thread, no busy-poll). Output is delivered to an
+``on_output`` callback (a fan-out layer above can broadcast it to any number of
+attached clients and buffer scrollback), so a session's lifetime is decoupled
+from any one WebSocket connection. Keystroke writes are tiny, so a plain
+``os.write`` is fine.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import fcntl
+import os
+import pty
+import signal
+import struct
+import termios
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+def _set_winsize(fd: int, rows: int, cols: int) -> None:
+    """Apply a terminal window size via TIOCSWINSZ (struct winsize)."""
+    winsize = struct.pack("HHHH", rows, cols, 0, 0)
+    fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
+
+
+@dataclass
+class PtySession:
+    """A child process attached to a pty master.
+
+    Output is pushed to ``on_output(bytes)`` as it arrives and ``on_exit()`` is
+    fired once when the child closes the tty. Both run in the event-loop thread.
+    """
+
+    proc: asyncio.subprocess.Process
+    master_fd: int
+    on_output: Callable[[bytes], None] | None = None
+    on_exit: Callable[[], None] | None = None
+    cols: int = 80
+    rows: int = 24
+    _closed: bool = field(default=False)
+
+    @classmethod
+    async def start(
+        cls,
+        command: str,
+        args: list[str] | None = None,
+        cwd: str | Path | None = None,
+        env: dict[str, str] | None = None,
+        cols: int = 80,
+        rows: int = 24,
+        on_output: Callable[[bytes], None] | None = None,
+        on_exit: Callable[[], None] | None = None,
+    ) -> PtySession:
+        """Spawn ``command args`` on a fresh pty, sized ``cols`` x ``rows``."""
+        master_fd, slave_fd = pty.openpty()
+        # Size the tty before exec so the first TUI render lays out correctly.
+        _set_winsize(master_fd, rows, cols)
+
+        proc_env = {**os.environ, "TERM": "xterm-256color"}
+        if env:
+            proc_env.update(env)
+
+        proc = await asyncio.create_subprocess_exec(
+            command,
+            *(args or []),
+            stdin=slave_fd,
+            stdout=slave_fd,
+            stderr=slave_fd,
+            cwd=str(cwd) if cwd else None,
+            env=proc_env,
+            start_new_session=True,  # own session/group → controlling tty + killpg
+            close_fds=True,
+        )
+        os.close(slave_fd)  # parent keeps only the master
+        os.set_blocking(master_fd, False)
+
+        session = cls(proc=proc, master_fd=master_fd, cols=cols, rows=rows)
+        if on_output is not None:
+            session.on_output = on_output
+        if on_exit is not None:
+            session.on_exit = on_exit
+        asyncio.get_running_loop().add_reader(master_fd, session._drain_master)
+        return session
+
+    def _drain_master(self) -> None:
+        """Reader callback: pull available bytes off the master and fan them out."""
+        try:
+            data = os.read(self.master_fd, 65536)
+        except (BlockingIOError, InterruptedError):
+            return
+        except OSError:
+            data = b""  # slave closed → child gone
+        if not data:
+            try:
+                asyncio.get_running_loop().remove_reader(self.master_fd)
+            except (ValueError, OSError):
+                pass
+            if self.on_exit is not None:
+                self.on_exit()
+            return
+        if self.on_output is not None:
+            self.on_output(data)
+
+    def write(self, data: bytes) -> None:
+        """Forward keystroke bytes to the child's tty."""
+        try:
+            os.write(self.master_fd, data)
+        except OSError:
+            pass
+
+    def resize(self, cols: int, rows: int) -> None:
+        """Resize the tty (browser terminal changed dimensions)."""
+        self.cols, self.rows = cols, rows
+        try:
+            _set_winsize(self.master_fd, rows, cols)
+        except OSError:
+            pass
+
+    @property
+    def running(self) -> bool:
+        return self.proc.returncode is None
+
+    async def close(self) -> None:
+        """Kill the whole process group and release the master fd (idempotent)."""
+        if self._closed:
+            return
+        self._closed = True
+        try:
+            asyncio.get_running_loop().remove_reader(self.master_fd)
+        except (ValueError, OSError):
+            pass
+        if self.proc.returncode is None:
+            try:
+                os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
+            except ProcessLookupError:
+                pass
+            try:
+                await asyncio.wait_for(self.proc.wait(), timeout=3)
+            except asyncio.TimeoutError:
+                try:
+                    os.killpg(os.getpgid(self.proc.pid), signal.SIGKILL)
+                except ProcessLookupError:
+                    pass
+        try:
+            os.close(self.master_fd)
+        except OSError:
+            pass