alex
b55300ff6f
- Add TerminalSessionModel DB table with instance_id FK, name, status, created_at, last_activity_at, closed_at columns - Add Alembic migration for terminal_sessions table - Refactor TerminalManager to use composite key (instance_id, session_id) supporting up to 5 concurrent sessions per instance - Add create_session, get_session, get_sessions_for_instance, close_session - Preserve get_or_create_session for backward compatibility (default session) - Fix attach_websocket to only close sockets within same session - Add name (auto-generated 'Session N') and status tracking to TerminalSession - Add 7 unit tests for multi-session logic Quality gates: pytest (7 new passed, 174 total passed, 51 pre-existing failures)
279 lines
9.3 KiB
Python
279 lines
9.3 KiB
Python
"""Terminal session management for tool instances."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import pty
|
|
import select
|
|
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 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
|
|
|
|
# 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
|
|
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
|
|
|
|
# Session metadata
|
|
self.name = name or self._generate_name(str(instance_id))
|
|
self.status: str = "active"
|
|
|
|
@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, self._slave_fd = pty.openpty()
|
|
|
|
# Set the terminal size initially
|
|
self._set_terminal_size(self._cols, self._rows)
|
|
logger.debug(
|
|
f"Starting terminal session {self.session_id} for container {self.container_id} with initial size {self._cols}x{self._rows}"
|
|
)
|
|
|
|
# Build the shell command
|
|
if startup_command:
|
|
shell_cmd = f'bash -c "{startup_command}" || true; exec bash -il'
|
|
logger.debug(
|
|
f"Using startup command for session {self.session_id}: {startup_command}"
|
|
)
|
|
else:
|
|
shell_cmd = "bash -il"
|
|
|
|
# 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",
|
|
"-c",
|
|
shell_cmd,
|
|
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.debug(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.debug(f"resize() called for session {self.session_id}: {cols}x{rows}")
|
|
self._set_terminal_size(cols, rows)
|
|
|
|
# Docker exec -it creates its own PTY inside the container,
|
|
# so host PTY resize doesn't propagate to the container shell.
|
|
# Send SIGWINCH to the docker exec process on the host.
|
|
# Docker exec forwards signals to the container process, which should
|
|
# cause the container's shell to re-read its terminal size.
|
|
if self.process and self.process.pid:
|
|
try:
|
|
os.kill(self.process.pid, signal.SIGWINCH)
|
|
logger.debug(
|
|
f"Sent SIGWINCH to docker exec process {self.process.pid} for session {self.session_id}"
|
|
)
|
|
except ProcessLookupError:
|
|
logger.warning(
|
|
f"docker exec process {self.process.pid} not found for session {self.session_id}"
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to send SIGWINCH: {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.process = None
|
|
self._master_fd = None
|
|
self._slave_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"
|
|
|
|
if self._master_fd is not None:
|
|
try:
|
|
os.close(self._master_fd)
|
|
except OSError:
|
|
pass # noqa: S110
|
|
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)
|