alex/headquarter
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: restore dev features lost in main→dev merge

Browse Source 
Merge b6f89f9 ('integrate main restructuring into dev') took main's
smaller refactored versions and overwrote dev-specific features.

Restored from dev commit 51a399c and ancestors:

Backend terminal (c754984, 37134b8):
- terminal_session.py: asyncio.add_reader event-driven I/O, flow control
  (64KB pause/32KB resume), 2ms output batching, binary WebSocket frames,
  EOF detection, circular replay buffer
- terminal_manager.py: named sessions, MaxSessionsExceededError,
  TerminalSessionModel persistence, idle timeout cleanup
- terminal.py: dual WebSocket routes (/terminal and /terminal/{session_id}),
  SessionRef mutable reference for session resets

SSH key mounting (c6b804b):
- instance_lifecycle.py: _mount_ssh_keys helper that prepares multiple
  SSH keys with unique filenames (id_ed25519_<name>), writes combined
  SSH config, mounts single ~/.ssh directory
- tool_instance.py: ssh_key_ids JSON column
- schemas/tool_instance.py: ssh_key_ids field in CreateInstanceRequest
- tool_instances.py: pass ssh_key_ids through create endpoint and response

Quality gates: py_compile passed, ruff passed on all edited files
This commit is contained in:
Alex
2026-06-03 14:06:04 +02:00
parent 8ef98eab94
commit 7ce2af18c9
7 changed files with 1520 additions and 276 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/src/services/terminal_session.py
+389 -84
View File
@@ -1,44 +1,136 @@
"""Terminal session management for tool instances."""
"""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 contextlib
import fcntl
import logging
import os
import pty
import select
import signal
import struct
import termios
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."""
"""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:
"""Initialize a terminal session."""
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
self._echo_enabled = True
self._exit_reason: str | None = None
async def start(self) -> 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."""
self._master_fd, self._slave_fd = pty.openpty()
self._set_terminal_size(80, 24)
# 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",
@@ -47,112 +139,287 @@ class TerminalSession:
"TERM=xterm-256color",
self.container_id,
"bash",
"-il",
stdin=self._slave_fd,
stdout=self._slave_fd,
stderr=self._slave_fd,
"-c",
shell_cmd,
stdin=slave_fd,
stdout=slave_fd,
stderr=slave_fd,
)
os.close(self._slave_fd)
self._slave_fd = None
self._echo_enabled = self._detect_echo_state()
# Close slave fd in parent process
os.close(slave_fd)
def _set_terminal_size(self, cols: int, rows: int) -> None:
"""Set the terminal size using TIOCSWINSZ."""
if self._master_fd is None:
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
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,
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
)
if readable:
return os.read(self._master_fd, 8192)
return b""
except (OSError, ValueError):
return b""
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
with contextlib.suppress(OSError):
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)
def get_exit_reason(self) -> str | None:
"""Return the reason the session ended, if known."""
return self._exit_reason
# 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"
# 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"
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:
with contextlib.suppress(OSError):
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 (TimeoutError, ProcessLookupError):
except (asyncio.TimeoutError, ProcessLookupError):
pass
def is_alive(self) -> bool:
@@ -160,3 +427,41 @@ class TerminalSession:
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""