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

"""High-performance terminal session with asyncio-native I/O.

Replaces blocking select.select() with event-driven asyncio.add_reader()
for sub-frame latency. Includes output batching and flow control.
"""

import asyncio
import logging
import os
import pty
import signal
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 with event-driven PTY I/O.

    Uses asyncio.add_reader() instead of polling for near-zero read latency.
    Output is batched (2ms window) and sent as binary WebSocket frames.
    Flow control prevents memory bloat on fast output.
    """

    # Circular buffer for replay (10KB)
    BUFFER_SIZE = 10 * 1024

    # Idle timeout in seconds (30 minutes)
    IDLE_TIMEOUT = 30 * 60

    # Output batching window in seconds
    BATCH_WINDOW_S = 0.002  # 2ms

    # Flow control: pause PTY reads when unacknowledged bytes exceed this
    FLOW_CONTROL_PAUSE = 64 * 1024

    # Flow control: resume PTY reads when unacknowledged bytes drop below this
    FLOW_CONTROL_RESUME = 32 * 1024

    # Max WebSocket frame size
    MAX_FRAME_SIZE = 64 * 1024

    # Session number counter per instance_id for auto-naming
    _instance_counters: dict[str, int] = {}

    def __init__(
        self,
        session_id: str,
        instance_id: uuid.UUID,
        container_id: str,
        startup_command: str | None = None,
        name: str | None = None,
    ) -> None:
        self.session_id = session_id
        self.instance_id = instance_id
        self.container_id = container_id
        self.startup_command = startup_command
        self.process: asyncio.subprocess.Process | None = None
        self._closed = False
        self._master_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

        # Session metadata
        self.name = name or self._generate_name(str(instance_id))
        self.status: str = "active"

        # Output batching
        self._batch_buffer = bytearray()
        self._batch_timer: asyncio.TimerHandle | None = None
        self._batch_lock = asyncio.Lock()

        # Flow control
        self._unacknowledged_bytes = 0
        self._paused = False
        self._read_handler_set = False
        self._flow_control_lock = asyncio.Lock()

        # Ack timeout fallback
        self._ack_timeout_handle: asyncio.TimerHandle | None = None

    @classmethod
    def _generate_name(cls, instance_id: str) -> str:
        """Generate an auto-incremented session name for the instance."""
        count = cls._instance_counters.get(instance_id, 0) + 1
        cls._instance_counters[instance_id] = count
        return f"Session {count}"

    async def start(self, startup_command: str | None = None) -> None:
        """Start the docker exec process with a shell using a PTY."""
        # Create a pseudo-terminal on the host
        self._master_fd, slave_fd = pty.openpty()

        # Set the terminal size initially
        self._set_terminal_size(self._cols, self._rows)
        logger.debug(
            "Starting terminal session %s for container %s with initial size %sx%s",
            self.session_id,
            self.container_id,
            self._cols,
            self._rows,
        )

        # Build the shell command
        cmd = startup_command or self.startup_command
        if cmd:
            shell_cmd = f'bash -c "{cmd}" || true; exec bash -il'
            logger.debug(
                "Using startup command for session %s: %s",
                self.session_id,
                cmd,
            )
        else:
            shell_cmd = "bash -il"

        # Start docker exec with the slave fd as stdin/stdout/stderr
        self.process = await asyncio.create_subprocess_exec(
            "docker",
            "exec",
            "-it",
            "-e",
            "TERM=xterm-256color",
            self.container_id,
            "bash",
            "-c",
            shell_cmd,
            stdin=slave_fd,
            stdout=slave_fd,
            stderr=slave_fd,
        )

        # Close slave fd in parent process
        os.close(slave_fd)

        self.last_activity = time.time()

        # Start event-driven reading
        self._start_reading()

    def _start_reading(self) -> None:
        """Register PTY master fd with asyncio event loop for event-driven reads."""
        if self._read_handler_set or self._master_fd is None or self._closed:
            return
        try:
            loop = asyncio.get_event_loop()
            loop.add_reader(self._master_fd, self._on_fd_readable)
            self._read_handler_set = True
            logger.debug("Started event-driven reading for session %s", self.session_id)
        except Exception as exc:
            logger.error(
                "Failed to start reading for session %s: %s", self.session_id, exc
            )

    def _stop_reading(self) -> None:
        """Unregister PTY master fd from asyncio event loop."""
        if not self._read_handler_set or self._master_fd is None:
            return
        try:
            loop = asyncio.get_event_loop()
            loop.remove_reader(self._master_fd)
            self._read_handler_set = False
        except Exception:
            pass

    def _on_fd_readable(self) -> None:
        """Callback when PTY master fd has data available (called by event loop)."""
        if self._master_fd is None or self._closed:
            return

        try:
            data = os.read(self._master_fd, 4096)
        except (OSError, IOError) as exc:
            logger.debug("PTY read error for session %s: %s", self.session_id, exc)
            self._handle_eof()
            return

        if not data:
            # EOF: docker exec process exited
            logger.debug("PTY EOF for session %s", self.session_id)
            self._handle_eof()
            return

        self._add_to_buffer(data)
        self.last_activity = time.time()

        # Queue for batching + flow control
        self._queue_output(data)

    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)
        while self._buffer_size > self.BUFFER_SIZE and self._output_buffer:
            removed = self._output_buffer.popleft()
            self._buffer_size -= len(removed)

    def _queue_output(self, data: bytes) -> None:
        """Add output to batch buffer and schedule flush."""
        self._batch_buffer.extend(data)
        self._unacknowledged_bytes += len(data)

        # Check flow control
        if self._unacknowledged_bytes > self.FLOW_CONTROL_PAUSE and not self._paused:
            self._pause_output()

        # Schedule batch flush if not already scheduled
        if self._batch_timer is None:
            loop = asyncio.get_event_loop()
            self._batch_timer = loop.call_later(
                self.BATCH_WINDOW_S,
                self._flush_batch_sync,
            )

    def _flush_batch_sync(self) -> None:
        """Synchronous entry point for batch flush (called from event loop)."""
        self._batch_timer = None
        if not self._batch_buffer or not self._websockets:
            self._batch_buffer.clear()
            return

        payload = bytes(self._batch_buffer)
        self._batch_buffer.clear()

        # Send to all websockets (asyncio.create_task for async send)
        dead_sockets = set()
        for ws in list(self._websockets):
            try:
                asyncio.create_task(self._send_bytes(ws, payload))
            except Exception:
                dead_sockets.add(ws)

        if dead_sockets:
            self._websockets -= dead_sockets

    async def _send_bytes(self, ws: Any, payload: bytes) -> None:
        """Send bytes to a single websocket, catching errors."""
        try:
            await ws.send_bytes(payload)
        except Exception:
            self._websockets.discard(ws)

    def acknowledge_data(self, char_count: int) -> None:
        """Client acknowledges processing char_count bytes.

        Called from the WebSocket handler when the client sends an 'ack' message.
        """
        self._unacknowledged_bytes = max(0, self._unacknowledged_bytes - char_count)

        if self._paused and self._unacknowledged_bytes < self.FLOW_CONTROL_RESUME:
            self._resume_output()

        # Reset ack timeout
        if self._ack_timeout_handle:
            self._ack_timeout_handle.cancel()
        loop = asyncio.get_event_loop()
        self._ack_timeout_handle = loop.call_later(5.0, self._ack_timeout_fallback)

    def _ack_timeout_fallback(self) -> None:
        """If no ack received for 5s, assume client is dead and resume."""
        logger.warning(
            "Flow control ack timeout for session %s, resuming output",
            self.session_id,
        )
        self._unacknowledged_bytes = 0
        if self._paused:
            self._resume_output()

    def _pause_output(self) -> None:
        """Pause reading from PTY due to flow control."""
        self._paused = True
        self._stop_reading()
        logger.debug(
            "Paused output for session %s (%d unacked)",
            self.session_id,
            self._unacknowledged_bytes,
        )

    def _resume_output(self) -> None:
        """Resume reading from PTY."""
        self._paused = False
        self._start_reading()
        logger.debug("Resumed output for session %s", self.session_id)

    def get_buffer(self) -> bytes:
        """Get buffered output for replay."""
        return b"".join(self._output_buffer)

    def _handle_eof(self) -> None:
        """Handle PTY EOF: process died, close websockets to force reconnect."""
        self._stop_reading()
        # Mark process as done so is_alive() returns False
        if self.process is not None and self.process.returncode is None:
            # Force returncode to a non-None value since the process is dead
            # but asyncio.subprocess may not have set it yet
            try:
                self.process._transport.close()  # type: ignore[attr-defined]
            except Exception:
                pass
        # Close all websockets to force frontend reconnection
        dead_sockets = set(self._websockets)
        self._websockets.clear()
        for ws in dead_sockets:
            try:
                asyncio.create_task(
                    ws.close(code=4001, reason="Session process exited")
                )
            except Exception:
                pass
        logger.info("Session %s EOF handled, websockets closed", self.session_id)

    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) as exc:
            logger.debug("PTY write error for session %s: %s", self.session_id, exc)
            self._handle_eof()

    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
        size = struct.pack("HHHH", rows, cols, 0, 0)
        try:
            fcntl.ioctl(self._master_fd, TIOCSWINSZ, size)
            logger.debug("Resized PTY to %sx%s (fd=%s)", cols, rows, self._master_fd)
        except (OSError, IOError) as e:
            logger.error("Failed to resize PTY: %s", e)

    async def resize(self, cols: int, rows: int) -> None:
        """Resize the terminal."""
        if self._closed:
            logger.warning("Cannot resize: session is closed")
            return

        if cols == self._cols and rows == self._rows:
            return

        self._cols = cols
        self._rows = rows
        logger.debug(
            "resize() called for session %s: %sx%s", self.session_id, cols, rows
        )
        self._set_terminal_size(cols, rows)

        # Send SIGWINCH to docker exec process
        if self.process and self.process.pid:
            try:
                os.kill(self.process.pid, signal.SIGWINCH)
            except ProcessLookupError:
                logger.warning("docker exec process %s not found", self.process.pid)
            except Exception as e:
                logger.warning("Failed to send SIGWINCH: %s", e)

    async def reset(self) -> None:
        """Reset the session by killing the process and clearing state."""
        self.status = "resetting"
        await self.close()
        self._closed = False
        self._output_buffer.clear()
        self._buffer_size = 0
        self._websockets.clear()
        self._batch_buffer.clear()
        self._batch_timer = None
        self._unacknowledged_bytes = 0
        self._paused = False
        self._read_handler_set = False
        self.process = None
        self._master_fd = None
        self.status = "active"

    async def close(self) -> None:
        """Close the session and cleanup."""
        if self._closed:
            return
        self._closed = True
        self.status = "closed"

        self._stop_reading()

        if self._batch_timer:
            self._batch_timer.cancel()
            self._batch_timer = None

        if self._ack_timeout_handle:
            self._ack_timeout_handle.cancel()
            self._ack_timeout_handle = None

        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 (used for control messages)."""
        dead_sockets = set()
        for ws in self._websockets:
            try:
                await ws.send_bytes(data)
            except Exception:
                dead_sockets.add(ws)
        for ws in dead_sockets:
            self._websockets.discard(ws)

    async def read_output(self) -> bytes:
        """Legacy method: read output synchronously.

        With event-driven I/O, output is automatically sent to websockets.
        This method returns any buffered data for callers that poll.
        """
        return b""