Developer
49180e4c6d
- Make the host PTY master non-blocking and wait for write readiness. - Retry partial writes so a large bracketed paste always delivers its closing marker instead of leaving pi in paste mode. - Add regression coverage for partial PTY writes and resolve diagnostics. Quality gates: ruff, mypy, focused pytest (3 passed)
522 lines
18 KiB
Python
522 lines
18 KiB
Python
"""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 tty
|
|
import uuid
|
|
from collections import deque
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_session_counters_by_instance: dict[str, int] = {}
|
|
|
|
|
|
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
|
|
|
|
def __init__(
|
|
self,
|
|
session_id: str,
|
|
instance_id: uuid.UUID,
|
|
container_id: str,
|
|
startup_command: str | None = None,
|
|
name: str | None = None,
|
|
container_user: str | None = None,
|
|
) -> None:
|
|
self.session_id = session_id
|
|
self.instance_id = instance_id
|
|
self.container_id = container_id
|
|
self.container_user = container_user
|
|
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
|
|
|
|
@staticmethod
|
|
def _generate_name(instance_id: str) -> str:
|
|
"""Generate an auto-incremented session name for the instance."""
|
|
count = _session_counters_by_instance.get(instance_id, 0) + 1
|
|
_session_counters_by_instance[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."""
|
|
# Build the docker exec command. When the tool manifest declares a
|
|
# non-root container user, run the shell as that user so terminal
|
|
# sessions match the privileges of the main container process.
|
|
exec_cmd = [
|
|
"docker",
|
|
"exec",
|
|
"-it",
|
|
"-e",
|
|
"TERM=xterm-256color",
|
|
]
|
|
if self.container_user:
|
|
exec_cmd.extend(["--user", self.container_user])
|
|
|
|
# Create a pseudo-terminal on the host. The master must be
|
|
# non-blocking: a browser paste can be larger than the PTY input
|
|
# buffer, and write_input() drains it asynchronously without dropping
|
|
# the closing bracketed-paste marker.
|
|
self._master_fd, slave_fd = pty.openpty()
|
|
os.set_blocking(self._master_fd, False)
|
|
|
|
# Put the host PTY into raw mode so it behaves as a pass-through
|
|
# pipe. openpty() leaves the slave in canonical mode by default,
|
|
# which line-buffers input, splits multiline pastes at newlines,
|
|
# and mangles bracketed-paste markers before docker exec / the
|
|
# foreground app (e.g. pi) ever see them. Real terminal discipline
|
|
# (echo, canonical editing for readline) is provided by the
|
|
# in-container PTY that `docker exec -t` allocates.
|
|
tty.setraw(slave_fd)
|
|
|
|
# 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
|
|
exec_cmd.extend([self.container_id, "bash", "-c", shell_cmd])
|
|
self.process = await asyncio.create_subprocess_exec(
|
|
*exec_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 BlockingIOError:
|
|
# The readiness notification raced with another callback.
|
|
return
|
|
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 _wait_for_write_ready(self, fd: int) -> None:
|
|
"""Wait until a non-blocking PTY master can accept more input."""
|
|
loop = asyncio.get_running_loop()
|
|
writable = loop.create_future()
|
|
|
|
def mark_writable() -> None:
|
|
if not writable.done():
|
|
writable.set_result(None)
|
|
|
|
loop.add_writer(fd, mark_writable)
|
|
try:
|
|
await writable
|
|
finally:
|
|
loop.remove_writer(fd)
|
|
|
|
async def write_input(self, data: bytes) -> None:
|
|
"""Write all terminal input bytes to the PTY master in order."""
|
|
if self._master_fd is None or self._closed:
|
|
return
|
|
|
|
fd = self._master_fd
|
|
remaining = memoryview(data)
|
|
try:
|
|
while remaining and not self._closed and self._master_fd == fd:
|
|
try:
|
|
written = os.write(fd, remaining)
|
|
except BlockingIOError:
|
|
await self._wait_for_write_ready(fd)
|
|
continue
|
|
|
|
if written == 0:
|
|
await self._wait_for_write_ready(fd)
|
|
continue
|
|
|
|
remaining = remaining[written:]
|
|
|
|
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""
|