"""Terminal session management for tool instances.""" import asyncio import logging import os import pty import select import struct import fcntl import time import uuid from collections import deque from typing import Any logger = logging.getLogger(__name__) class TerminalSession: """Manages a single terminal session connected to a docker container. Supports persistent sessions that survive WebSocket disconnections. Multiple WebSocket connections can attach/detach from the same session. """ # Circular buffer size (10KB) BUFFER_SIZE = 10 * 1024 # Idle timeout in seconds (30 minutes) IDLE_TIMEOUT = 30 * 60 def __init__(self, session_id: str, instance_id: uuid.UUID, container_id: str) -> None: self.session_id = session_id self.instance_id = instance_id self.container_id = container_id self.process: asyncio.subprocess.Process | None = None self._closed = False self._master_fd: int | None = None self._slave_fd: int | None = None # Circular buffer for output replay self._output_buffer: deque[bytes] = deque(maxlen=self.BUFFER_SIZE) self._buffer_size = 0 # WebSocket connections self._websockets: set[Any] = set() # Activity tracking self.last_activity = time.time() # Terminal size self._cols = 80 self._rows = 24 async def start(self) -> None: """Start the docker exec process with a shell using a PTY.""" # Create a pseudo-terminal on the host self._master_fd, self._slave_fd = pty.openpty() # Set the terminal size initially self._set_terminal_size(self._cols, self._rows) logger.info(f"Starting terminal session {self.session_id} for container {self.container_id} with initial size {self._cols}x{self._rows}") # Start docker exec with the slave fd as stdin/stdout/stderr # Using -it because the slave fd IS a TTY self.process = await asyncio.create_subprocess_exec( "docker", "exec", "-it", "-e", "TERM=xterm", self.container_id, "bash", "-il", stdin=self._slave_fd, stdout=self._slave_fd, stderr=self._slave_fd, ) # Close slave fd in parent process os.close(self._slave_fd) self._slave_fd = None self.last_activity = time.time() def _set_terminal_size(self, cols: int, rows: int) -> None: """Set the terminal size using TIOCSWINSZ.""" if self._master_fd is None: logger.warning("Cannot resize: master_fd is None (session not started)") return # TIOCSWINSZ = 0x5414 on Linux TIOCSWINSZ = 0x5414 size = struct.pack('HHHH', rows, cols, 0, 0) try: fcntl.ioctl(self._master_fd, TIOCSWINSZ, size) logger.info(f"Resized PTY to {cols}x{rows} (fd={self._master_fd})") except (OSError, IOError) as e: logger.error(f"Failed to resize PTY: {e}") async def read_output(self) -> bytes: """Read output from the PTY master and store in buffer.""" if self._master_fd is None or self._closed: return b"" try: # Use select to check if data is available readable, _, _ = select.select([self._master_fd], [], [], 0.1) if readable: data = os.read(self._master_fd, 4096) if data: self._add_to_buffer(data) self.last_activity = time.time() return data return b"" except (OSError, IOError, ValueError): return b"" def _add_to_buffer(self, data: bytes) -> None: """Add data to circular buffer, maintaining size limit.""" self._output_buffer.append(data) self._buffer_size += len(data) # Trim if exceeds max size while self._buffer_size > self.BUFFER_SIZE and self._output_buffer: removed = self._output_buffer.popleft() self._buffer_size -= len(removed) def get_buffer(self) -> bytes: """Get buffered output for replay.""" return b"".join(self._output_buffer) async def write_input(self, data: bytes) -> None: """Write input to the PTY master.""" if self._master_fd is None or self._closed: return try: os.write(self._master_fd, data) self.last_activity = time.time() except (OSError, IOError): pass async def resize(self, cols: int, rows: int) -> None: """Resize the terminal.""" if self._closed: logger.warning("Cannot resize: session is closed") return # Only resize if dimensions actually changed if cols == self._cols and rows == self._rows: return self._cols = cols self._rows = rows logger.info(f"resize() called for session {self.session_id}: {cols}x{rows}") self._set_terminal_size(cols, rows) # Docker exec doesn't forward PTY resize to the container process, # so we need to explicitly set the size inside the container shell. # Send on every resize so the container shell always matches the frontend. # Use stty -echo to prevent the command from being visible, then clear the line. stty_cmd = ( f"stty -echo; stty cols {cols} rows {rows}; stty echo\n" f"\x1b[A\x1b[M" # Move up 1 line and delete it (clears the stty command) ).encode() await self.write_input(stty_cmd) logger.debug(f"Sent stty resize to container for session {self.session_id}: {cols}x{rows}") async def reset(self) -> None: """Reset the session by killing the process and clearing state.""" await self.close() self._closed = False self._output_buffer.clear() self._buffer_size = 0 self._websockets.clear() self.process = None self._master_fd = None self._slave_fd = None async def close(self) -> None: """Close the session and cleanup.""" if self._closed: return self._closed = True if self._master_fd is not None: try: os.close(self._master_fd) except OSError: pass self._master_fd = None if self.process is not None: try: self.process.kill() await asyncio.wait_for(self.process.wait(), timeout=2.0) except (asyncio.TimeoutError, ProcessLookupError): pass def is_alive(self) -> bool: """Check if the session process is still running.""" if self.process is None: return False return self.process.returncode is None def is_idle(self) -> bool: """Check if the session has been idle for too long.""" if self._websockets: return False return time.time() - self.last_activity > self.IDLE_TIMEOUT def attach_websocket(self, websocket: Any) -> None: """Attach a WebSocket to this session.""" self._websockets.add(websocket) self.last_activity = time.time() def detach_websocket(self, websocket: Any) -> None: """Detach a WebSocket from this session.""" self._websockets.discard(websocket) def has_websockets(self) -> bool: """Check if any WebSockets are attached.""" return len(self._websockets) > 0 async def send_to_all(self, data: bytes) -> None: """Send data to all attached WebSockets.""" dead_sockets = set() for ws in self._websockets: try: await ws.send_bytes(data) except Exception: dead_sockets.add(ws) # Clean up dead sockets for ws in dead_sockets: self._websockets.discard(ws)