headquarter/apps/api/src/services/terminal_session.py

"""Terminal session management for tool instances."""

import asyncio
import os
import pty
import select
import struct
import fcntl
import time
import uuid
from collections import deque
from typing import Any


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)
        
        # 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:
            return
        # TIOCSWINSZ = 0x5414 on Linux
        TIOCSWINSZ = 0x5414
        size = struct.pack('HHHH', rows, cols, 0, 0)
        try:
            fcntl.ioctl(self._master_fd, TIOCSWINSZ, size)
        except (OSError, IOError):
            pass

    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:
            return
        self._cols = cols
        self._rows = rows
        self._set_terminal_size(cols, 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)