chore: import upstream snapshot with attribution

libs/python/computer-server/computer_server/pty_manager.py

@@ -0,0 +1,203 @@
+"""Async PTY session manager for computer-server.
+
+Wraps :class:`cua_auto.terminal.Terminal` and adds asyncio-compatible
+output broadcasting via per-consumer :class:`asyncio.Queue` objects.
+
+Queue items are dicts:
+  - ``{"type": "output", "data": <bytes>}``  — terminal output chunk
+  - ``{"type": "exit", "code": <int>}``      — session terminated (sentinel)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import threading
+from typing import Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PtyManager:
+    """Manage PTY sessions with async broadcast queues.
+
+    Usage::
+
+        mgr = PtyManager()
+
+        # Create a session
+        info = await mgr.create(command="bash", cols=80, rows=24)
+        pid = info["pid"]
+
+        # Subscribe to output
+        q = mgr.subscribe(pid)
+        async for msg in _drain(q):
+            ...  # msg is {"type": "output", "data": b"..."} or {"type": "exit", "code": 0}
+
+        # Write to stdin
+        await mgr.send_stdin(pid, b"echo hello\\n")
+
+        # Resize
+        await mgr.resize(pid, 120, 40)
+
+        # Kill
+        await mgr.kill(pid)
+    """
+
+    def __init__(self) -> None:
+        # pid → list of subscriber queues
+        self._queues: Dict[int, List[asyncio.Queue]] = {}
+        # pid → {"cols": int, "rows": int}
+        self._sessions: Dict[int, dict] = {}
+        self._terminal = None
+
+    # ------------------------------------------------------------------
+    # Lazy terminal access
+    # ------------------------------------------------------------------
+
+    def _get_terminal(self):
+        if self._terminal is None:
+            from cua_auto.terminal import Terminal
+
+            self._terminal = Terminal()
+        return self._terminal
+
+    # ------------------------------------------------------------------
+    # Public async API
+    # ------------------------------------------------------------------
+
+    async def create(
+        self,
+        command: Optional[str] = None,
+        cols: int = 80,
+        rows: int = 24,
+        cwd: Optional[str] = None,
+        envs: Optional[dict] = None,
+    ) -> dict:
+        """Spawn a new PTY session.
+
+        Returns:
+            ``{"pid": int, "cols": int, "rows": int}``
+        """
+        loop = asyncio.get_running_loop()
+        terminal = self._get_terminal()
+
+        # We need pid to route output, but we only know it after create().
+        # Use a mutable cell so the callback can look up the pid after creation.
+        # early_buffer holds chunks that arrive before pid_cell[0] is set (the
+        # reader thread can fire before asyncio.to_thread returns the session).
+        pid_cell: List[Optional[int]] = [None]
+        early_buffer: List[bytes] = []
+
+        def _on_data(data: bytes) -> None:
+            pid = pid_cell[0]
+            if pid is None:
+                early_buffer.append(data)
+                return
+            msg = {"type": "output", "data": data}
+            for q in list(self._queues.get(pid, [])):
+                try:
+                    loop.call_soon_threadsafe(q.put_nowait, msg)
+                except Exception:
+                    pass
+
+        session = await asyncio.to_thread(
+            terminal.create,
+            command=command,
+            cols=cols,
+            rows=rows,
+            on_data=_on_data,
+            cwd=cwd,
+            envs=envs,
+        )
+
+        pid_cell[0] = session.pid
+        self._queues[session.pid] = []
+        self._sessions[session.pid] = {"cols": session.cols, "rows": session.rows}
+
+        # Flush any output that arrived before pid_cell[0] was set.
+        for chunk in early_buffer:
+            msg = {"type": "output", "data": chunk}
+            for q in list(self._queues[session.pid]):
+                try:
+                    loop.call_soon_threadsafe(q.put_nowait, msg)
+                except Exception:
+                    pass
+
+        # Watch for process exit in a daemon thread, then broadcast sentinel.
+        def _watch_exit() -> None:
+            exit_code = terminal.wait(session.pid) or 0
+            sentinel = {"type": "exit", "code": exit_code}
+            for q in list(self._queues.get(session.pid, [])):
+                try:
+                    loop.call_soon_threadsafe(q.put_nowait, sentinel)
+                except Exception:
+                    pass
+
+        threading.Thread(target=_watch_exit, daemon=True, name=f"pty-exit-{session.pid}").start()
+
+        logger.info(
+            "PTY session created: pid=%d cmd=%r cols=%d rows=%d", session.pid, command, cols, rows
+        )
+        return {"pid": session.pid, "cols": session.cols, "rows": session.rows}
+
+    async def send_stdin(self, pid: int, data: bytes) -> None:
+        """Write *data* to the stdin of session *pid*."""
+        terminal = self._get_terminal()
+        await asyncio.to_thread(terminal.send_stdin, pid, data)
+
+    async def resize(self, pid: int, cols: int, rows: int) -> None:
+        """Resize the terminal for session *pid*."""
+        terminal = self._get_terminal()
+        await asyncio.to_thread(terminal.resize, pid, cols, rows)
+        if pid in self._sessions:
+            self._sessions[pid]["cols"] = cols
+            self._sessions[pid]["rows"] = rows
+
+    def get_info(self, pid: int) -> Optional[dict]:
+        """Return ``{"pid": int, "cols": int, "rows": int}`` for *pid*, or ``None`` if unknown."""
+        info = self._sessions.get(pid)
+        if info is None:
+            return None
+        return {"pid": pid, "cols": info["cols"], "rows": info["rows"]}
+
+    async def kill(self, pid: int) -> bool:
+        """Kill session *pid*.
+
+        Also broadcasts the exit sentinel to all subscribers.
+        Returns:
+            ``True`` if the signal was delivered.
+        """
+        terminal = self._get_terminal()
+        result = await asyncio.to_thread(terminal.kill, pid)
+        # Broadcast sentinel immediately (the exit watcher will also fire,
+        # but duplicate sentinels are harmless — consumers stop after the first).
+        sentinel = {"type": "exit", "code": -1}
+        for q in list(self._queues.get(pid, [])):
+            try:
+                q.put_nowait(sentinel)
+            except Exception:
+                pass
+        return result
+
+    # ------------------------------------------------------------------
+    # Queue-based pub/sub
+    # ------------------------------------------------------------------
+
+    def subscribe(self, pid: int) -> asyncio.Queue:
+        """Return a new :class:`asyncio.Queue` that will receive output for *pid*.
+
+        The queue receives dicts with keys ``type`` (``"output"`` or ``"exit"``)
+        and either ``data`` (bytes) or ``code`` (int).
+        """
+        q: asyncio.Queue = asyncio.Queue()
+        self._queues.setdefault(pid, []).append(q)
+        return q
+
+    def unsubscribe(self, pid: int, queue: asyncio.Queue) -> None:
+        """Remove *queue* from the subscriber list for *pid*."""
+        qs = self._queues.get(pid, [])
+        try:
+            qs.remove(queue)
+        except ValueError:
+            pass