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

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

import asyncio
import contextlib
import fcntl
import logging
import os
import pty
import select
import struct
import termios
import uuid

logger = logging.getLogger(__name__)


class TerminalSession:
    """Manages a single terminal session connected to a docker container."""

    def __init__(
        self,
        session_id: str,
        instance_id: uuid.UUID,
        container_id: str,
    ) -> None:
        """Initialize a terminal session."""
        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
        self._echo_enabled = True
        self._exit_reason: str | None = None

    async def start(self) -> None:
        """Start the docker exec process with a shell using a PTY."""
        self._master_fd, self._slave_fd = pty.openpty()
        self._set_terminal_size(80, 24)

        self.process = await asyncio.create_subprocess_exec(
            "docker",
            "exec",
            "-it",
            "-e",
            "TERM=xterm-256color",
            self.container_id,
            "bash",
            "-il",
            stdin=self._slave_fd,
            stdout=self._slave_fd,
            stderr=self._slave_fd,
        )

        os.close(self._slave_fd)
        self._slave_fd = None
        self._echo_enabled = self._detect_echo_state()

    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
        size = struct.pack("HHHH", rows, cols, 0, 0)
        with contextlib.suppress(OSError):
            fcntl.ioctl(self._master_fd, tiocswinsz, size)

    def _detect_echo_state(self) -> bool:
        """Detect whether the PTY has echo enabled via termios."""
        if self._master_fd is None:
            return True
        try:
            attrs = termios.tcgetattr(self._master_fd)
            return bool(attrs[3] & termios.ECHO)
        except OSError:
            return True

    async def check_echo_state(self) -> bool | None:
        """Check if echo state changed. Returns new state if changed, None otherwise."""
        current = self._detect_echo_state()
        if current != self._echo_enabled:
            self._echo_enabled = current
            return current
        return None

    @property
    def echo_enabled(self) -> bool:
        """Return whether the PTY currently has echo enabled."""
        return self._echo_enabled

    @property
    def closed(self) -> bool:
        """Return whether the session has been closed."""
        return self._closed

    async def read_output(self, select_timeout: float = 0.1) -> bytes:
        """Read output from the PTY master."""
        if self._master_fd is None or self._closed:
            return b""
        try:
            readable, _, _ = select.select(
                [self._master_fd],
                [],
                [],
                select_timeout,
            )
            if readable:
                return os.read(self._master_fd, 8192)
            return b""
        except (OSError, ValueError):
            return b""

    async def write_input(self, data: bytes) -> None:
        """Write input to the PTY master."""
        if self._master_fd is None or self._closed:
            return
        with contextlib.suppress(OSError):
            os.write(self._master_fd, data)

    async def resize(self, cols: int, rows: int) -> None:
        """Resize the terminal."""
        if self._closed:
            return
        self._set_terminal_size(cols, rows)

    def get_exit_reason(self) -> str | None:
        """Return the reason the session ended, if known."""
        return self._exit_reason

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

        # Determine exit reason
        if self.process is not None and self.process.returncode is not None:
            if self.process.returncode == 0:
                self._exit_reason = "process_exit"
            else:
                self._exit_reason = "process_exit"
        else:
            self._exit_reason = "timeout"

        if self._master_fd is not None:
            with contextlib.suppress(OSError):
                os.close(self._master_fd)
            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 (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