refactor: organize API routers and services into subpackages
Service organization (19 files moved into 6 subpackages): - services/instance/ — event_bus, health_monitor, lifecycle_hooks - services/config/ — config_profile_resolver - services/git/ — clone, git_operations, git_service - services/build/ — docker_build, manifest_compiler - services/terminal/ — terminal_manager, terminal_session - services/shared/ — correlation, file_service, notification_service, permission_fixer, readiness_probe, ssh_keys, tunnel, workspace_manager API router organization (16 files moved into 6 subpackages): - api/tool/ — tool_instances, tool_types, tool_definitions, tool_types_validation, sessions (extracted from tool_instances) - api/config/ — config_profiles, user_config - api/workspace/ — workspaces, workspace_files, workspace_git, workspace_instances - api/user/ — users, auth, ssh_keys - api/project/ — projects, git_repositories - api/system/ — health, events, notifications, dashboard, terminal, instance_proxy Updated main.py imports and all __init__.py re-exports. Sessions router extracted from tool_instances.py into api/tool/sessions.py. Quality gates: py_compile passed, ruff passed.
This commit is contained in:
@@ -1 +1,9 @@
|
||||
"""Terminal module."""
|
||||
"""Terminal services module."""
|
||||
|
||||
from src.services.terminal.terminal_manager import (
|
||||
MaxSessionsExceededError,
|
||||
TerminalManager,
|
||||
)
|
||||
from src.services.terminal.terminal_session import TerminalSession
|
||||
|
||||
__all__ = ["MaxSessionsExceededError", "TerminalManager", "TerminalSession"]
|
||||
|
||||
@@ -0,0 +1,426 @@
|
||||
"""Terminal session manager for WebSocket connections."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from fastapi import WebSocket
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
|
||||
from src.database import SessionLocal
|
||||
from src.models import TerminalSessionModel
|
||||
from src.services.terminal.terminal_session import TerminalSession
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MaxSessionsExceededError(Exception):
|
||||
"""Raised when the maximum number of terminal sessions per instance is reached."""
|
||||
|
||||
def __init__(self, instance_id: str, max_sessions: int = 5) -> None:
|
||||
self.instance_id = instance_id
|
||||
self.max_sessions = max_sessions
|
||||
super().__init__(
|
||||
f"Maximum of {max_sessions} terminal sessions reached for instance {instance_id}"
|
||||
)
|
||||
|
||||
|
||||
class TerminalManager:
|
||||
"""Manages active terminal sessions with persistence support."""
|
||||
|
||||
# Maximum sessions per tool instance
|
||||
MAX_SESSIONS_PER_INSTANCE = 5
|
||||
|
||||
def __init__(self) -> None:
|
||||
# Track sessions by (instance_id, session_id) for multi-session support
|
||||
self._sessions: dict[tuple[str, str], TerminalSession] = {}
|
||||
self._idle_check_task: asyncio.Task | None = None
|
||||
self._start_idle_check()
|
||||
|
||||
def _start_idle_check(self) -> None:
|
||||
"""Start the idle timeout background task."""
|
||||
if self._idle_check_task is not None and not self._idle_check_task.done():
|
||||
return
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
self._idle_check_task = loop.create_task(self._idle_check_loop())
|
||||
except RuntimeError:
|
||||
# No event loop running yet, will be started lazily
|
||||
pass
|
||||
|
||||
async def _idle_check_loop(self) -> None:
|
||||
"""Periodically check for idle sessions and clean them up."""
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(60) # Check every minute
|
||||
await self._cleanup_idle_sessions()
|
||||
except Exception as exc:
|
||||
logger.error("Error in idle check loop: %s", exc)
|
||||
|
||||
async def _cleanup_idle_sessions(self) -> None:
|
||||
"""Clean up sessions that have been idle for too long."""
|
||||
idle_keys = []
|
||||
for (instance_id, session_id), session in list(self._sessions.items()):
|
||||
if session.is_idle():
|
||||
idle_keys.append((instance_id, session_id))
|
||||
|
||||
for key in idle_keys:
|
||||
instance_id, session_id = key
|
||||
logger.info(
|
||||
"Cleaning up idle terminal session %s for instance %s",
|
||||
session_id,
|
||||
instance_id,
|
||||
)
|
||||
session = self._sessions.pop(key, None)
|
||||
if session:
|
||||
await session.close()
|
||||
# Update DB status fire-and-forget
|
||||
asyncio.create_task(self._mark_closed_in_db(session_id))
|
||||
|
||||
async def _insert_db_session_row(
|
||||
self,
|
||||
session_id: str,
|
||||
instance_id: uuid.UUID,
|
||||
name: str,
|
||||
) -> None:
|
||||
"""Insert a TerminalSessionModel row into the database.
|
||||
|
||||
Uses ON CONFLICT DO NOTHING to handle races when a session is
|
||||
restored from DB and then re-inserted.
|
||||
"""
|
||||
try:
|
||||
async with SessionLocal() as db_session:
|
||||
stmt = (
|
||||
pg_insert(TerminalSessionModel)
|
||||
.values(
|
||||
id=uuid.UUID(session_id),
|
||||
instance_id=instance_id,
|
||||
name=name,
|
||||
status="active",
|
||||
created_at=datetime.now(timezone.utc),
|
||||
last_activity_at=datetime.now(timezone.utc),
|
||||
)
|
||||
.on_conflict_do_nothing(index_elements=["id"])
|
||||
)
|
||||
await db_session.execute(stmt)
|
||||
await db_session.commit()
|
||||
logger.debug(
|
||||
"Inserted terminal session row %s for instance %s",
|
||||
session_id,
|
||||
instance_id,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to insert terminal session row: %s", exc)
|
||||
|
||||
async def _mark_closed_in_db(self, session_id: str) -> None:
|
||||
"""Mark a terminal session as closed in the database."""
|
||||
try:
|
||||
async with SessionLocal() as db_session:
|
||||
db_row = await db_session.get(
|
||||
TerminalSessionModel, uuid.UUID(session_id)
|
||||
)
|
||||
if db_row:
|
||||
db_row.status = "closed"
|
||||
db_row.closed_at = datetime.now(timezone.utc)
|
||||
await db_session.commit()
|
||||
logger.debug(
|
||||
"Marked terminal session %s as closed in DB", session_id
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to mark terminal session as closed in DB: %s", exc)
|
||||
|
||||
def _count_sessions_for_instance(self, instance_id_str: str) -> int:
|
||||
"""Count active in-memory sessions for a given instance."""
|
||||
return sum(1 for (iid, _sid) in self._sessions if iid == instance_id_str)
|
||||
|
||||
async def create_session(
|
||||
self,
|
||||
instance_id: uuid.UUID,
|
||||
container_id: str,
|
||||
startup_command: str | None = None,
|
||||
name: str | None = None,
|
||||
session_id: str | None = None,
|
||||
) -> TerminalSession:
|
||||
"""Create a new terminal session for an instance.
|
||||
|
||||
Enforces a maximum of MAX_SESSIONS_PER_INSTANCE sessions per instance.
|
||||
Inserts a DB row fire-and-forget.
|
||||
|
||||
Args:
|
||||
instance_id: UUID of the tool instance.
|
||||
container_id: Docker container ID.
|
||||
startup_command: Optional startup command to run.
|
||||
name: Optional session name (auto-generated if omitted).
|
||||
|
||||
Returns:
|
||||
The newly created TerminalSession.
|
||||
|
||||
Raises:
|
||||
MaxSessionsExceededError: If the instance already has max sessions.
|
||||
"""
|
||||
instance_id_str = str(instance_id)
|
||||
|
||||
if (
|
||||
self._count_sessions_for_instance(instance_id_str)
|
||||
>= self.MAX_SESSIONS_PER_INSTANCE
|
||||
):
|
||||
raise MaxSessionsExceededError(
|
||||
instance_id_str, self.MAX_SESSIONS_PER_INSTANCE
|
||||
)
|
||||
|
||||
if session_id is None:
|
||||
session_id = str(uuid.uuid4())
|
||||
session = TerminalSession(
|
||||
session_id=session_id,
|
||||
instance_id=instance_id,
|
||||
container_id=container_id,
|
||||
startup_command=startup_command,
|
||||
name=name,
|
||||
)
|
||||
await session.start(startup_command=startup_command)
|
||||
|
||||
key = (instance_id_str, session_id)
|
||||
self._sessions[key] = session
|
||||
|
||||
# Fire-and-forget DB insert (skip if row already exists)
|
||||
asyncio.create_task(
|
||||
self._insert_db_session_row(session_id, instance_id, session.name)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Created terminal session %s for instance %s (name=%s)",
|
||||
session_id,
|
||||
instance_id,
|
||||
session.name,
|
||||
)
|
||||
return session
|
||||
|
||||
async def get_or_create_session(
|
||||
self,
|
||||
instance_id: uuid.UUID,
|
||||
container_id: str,
|
||||
startup_command: str | None = None,
|
||||
) -> TerminalSession:
|
||||
"""Get existing session or create a new one.
|
||||
|
||||
Backward-compatible alias that uses 'default' as the session_id.
|
||||
"""
|
||||
# Ensure idle check is running (lazy start)
|
||||
self._start_idle_check()
|
||||
|
||||
instance_id_str = str(instance_id)
|
||||
key = (instance_id_str, "default")
|
||||
|
||||
# Check for existing default session
|
||||
if key in self._sessions:
|
||||
session = self._sessions[key]
|
||||
|
||||
# Check if session is still alive
|
||||
if session.is_alive():
|
||||
logger.debug(
|
||||
"Reattaching to existing terminal session for instance %s",
|
||||
instance_id,
|
||||
)
|
||||
return session
|
||||
else:
|
||||
# Session died, clean it up
|
||||
logger.debug(
|
||||
"Existing session for instance %s is dead, cleaning up",
|
||||
instance_id,
|
||||
)
|
||||
await session.close()
|
||||
del self._sessions[key]
|
||||
|
||||
# Create new default session
|
||||
logger.info(
|
||||
"Creating new default terminal session for instance %s", instance_id
|
||||
)
|
||||
session_id = str(uuid.uuid4())
|
||||
session = TerminalSession(
|
||||
session_id=session_id,
|
||||
instance_id=instance_id,
|
||||
container_id=container_id,
|
||||
startup_command=startup_command,
|
||||
name="Session 1",
|
||||
)
|
||||
await session.start(startup_command=startup_command)
|
||||
self._sessions[key] = session
|
||||
|
||||
# Fire-and-forget DB insert
|
||||
asyncio.create_task(
|
||||
self._insert_db_session_row(session_id, instance_id, session.name)
|
||||
)
|
||||
|
||||
return session
|
||||
|
||||
def get_session(
|
||||
self,
|
||||
instance_id: str,
|
||||
session_id: str,
|
||||
) -> TerminalSession | None:
|
||||
"""Lookup a session by composite key, or by internal session_id."""
|
||||
session = self._sessions.get((instance_id, session_id))
|
||||
if session is not None:
|
||||
return session
|
||||
# Fallback: search by internal TerminalSession.session_id
|
||||
for (iid, _sid), sess in self._sessions.items():
|
||||
if iid == instance_id and sess.session_id == session_id:
|
||||
return sess
|
||||
return None
|
||||
|
||||
def _find_key_by_internal_id(
|
||||
self,
|
||||
instance_id: str,
|
||||
internal_session_id: str,
|
||||
) -> tuple[str, str] | None:
|
||||
"""Find the manager dict key for a session by its internal session_id."""
|
||||
for (iid, sid), session in self._sessions.items():
|
||||
if iid == instance_id and session.session_id == internal_session_id:
|
||||
return (iid, sid)
|
||||
return None
|
||||
|
||||
def get_sessions_for_instance(
|
||||
self,
|
||||
instance_id: str,
|
||||
) -> list[TerminalSession]:
|
||||
"""Return all in-memory sessions for a given instance."""
|
||||
return [
|
||||
session
|
||||
for (iid, _sid), session in self._sessions.items()
|
||||
if iid == instance_id
|
||||
]
|
||||
|
||||
async def close_session(
|
||||
self,
|
||||
instance_id: str,
|
||||
session_id: str,
|
||||
) -> None:
|
||||
"""Close a specific session and update its DB status."""
|
||||
key = (instance_id, session_id)
|
||||
session = self._sessions.pop(key, None)
|
||||
if session:
|
||||
await session.close()
|
||||
# Fire-and-forget DB update
|
||||
asyncio.create_task(self._mark_closed_in_db(session_id))
|
||||
logger.info(
|
||||
"Closed terminal session %s for instance %s",
|
||||
session_id,
|
||||
instance_id,
|
||||
)
|
||||
|
||||
async def attach_websocket(
|
||||
self,
|
||||
session: TerminalSession,
|
||||
websocket: WebSocket,
|
||||
) -> None:
|
||||
"""Attach a WebSocket to an existing session.
|
||||
|
||||
Closes existing WebSocket connections only for this specific session.
|
||||
"""
|
||||
# Handle concurrent connections - close existing ones within the same session
|
||||
if session.has_websockets():
|
||||
logger.debug(
|
||||
"Closing existing WebSocket connections for session %s (instance %s)",
|
||||
session.session_id,
|
||||
session.instance_id,
|
||||
)
|
||||
for ws in list(session._websockets):
|
||||
try:
|
||||
await ws.close(code=4000, reason="New connection established")
|
||||
except Exception:
|
||||
pass # noqa: S110
|
||||
session._websockets.clear()
|
||||
|
||||
# Attach new WebSocket
|
||||
session.attach_websocket(websocket)
|
||||
|
||||
# Replay buffer
|
||||
buffer = session.get_buffer()
|
||||
if buffer:
|
||||
try:
|
||||
await websocket.send_bytes(buffer)
|
||||
except Exception:
|
||||
pass # noqa: S110
|
||||
|
||||
async def detach_websocket(
|
||||
self,
|
||||
session: TerminalSession,
|
||||
websocket: WebSocket,
|
||||
) -> None:
|
||||
"""Detach a WebSocket from a session."""
|
||||
session.detach_websocket(websocket)
|
||||
|
||||
async def reset_session(
|
||||
self,
|
||||
instance_id: uuid.UUID,
|
||||
container_id: str,
|
||||
startup_command: str | None = None,
|
||||
session_id: str | None = None,
|
||||
name: str | None = None,
|
||||
) -> TerminalSession:
|
||||
"""Reset a session by killing it and creating a new one.
|
||||
|
||||
Args:
|
||||
instance_id: UUID of the tool instance.
|
||||
container_id: Docker container ID.
|
||||
startup_command: Optional startup command.
|
||||
session_id: Specific session to reset. If None, resets the default session.
|
||||
name: Optional name to preserve for the new session.
|
||||
|
||||
Returns:
|
||||
The newly created TerminalSession.
|
||||
"""
|
||||
instance_id_str = str(instance_id)
|
||||
target_session_id = session_id or "default"
|
||||
key = (instance_id_str, target_session_id)
|
||||
|
||||
# Preserve old name if not provided
|
||||
old_name = name
|
||||
if old_name is None and key in self._sessions:
|
||||
old_name = self._sessions[key].name
|
||||
|
||||
# Close existing session if any
|
||||
if key in self._sessions:
|
||||
logger.debug(
|
||||
"Resetting terminal session %s for instance %s",
|
||||
target_session_id,
|
||||
instance_id,
|
||||
)
|
||||
old_session = self._sessions.pop(key)
|
||||
await old_session.close()
|
||||
# Fire-and-forget DB update for old session
|
||||
asyncio.create_task(self._mark_closed_in_db(old_session.session_id))
|
||||
|
||||
# Create new session preserving the same session_id slot
|
||||
new_session_id = str(uuid.uuid4())
|
||||
new_session = TerminalSession(
|
||||
session_id=new_session_id,
|
||||
instance_id=instance_id,
|
||||
container_id=container_id,
|
||||
startup_command=startup_command,
|
||||
name=old_name or ("Session 1" if target_session_id == "default" else None),
|
||||
)
|
||||
await new_session.start(startup_command=startup_command)
|
||||
self._sessions[key] = new_session
|
||||
|
||||
# Fire-and-forget DB insert
|
||||
asyncio.create_task(
|
||||
self._insert_db_session_row(new_session_id, instance_id, new_session.name)
|
||||
)
|
||||
|
||||
return new_session
|
||||
|
||||
async def close_all(self) -> None:
|
||||
"""Close all active sessions."""
|
||||
sessions = list(self._sessions.values())
|
||||
self._sessions.clear()
|
||||
for session in sessions:
|
||||
await session.close()
|
||||
|
||||
if self._idle_check_task and not self._idle_check_task.done():
|
||||
self._idle_check_task.cancel()
|
||||
|
||||
|
||||
# Global terminal manager instance
|
||||
terminal_manager = TerminalManager()
|
||||
@@ -0,0 +1,467 @@
|
||||
"""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""
|
||||
Reference in New Issue
Block a user