diff --git a/apps/api/src/api/terminal.py b/apps/api/src/api/terminal.py index 0ee52e4..c2f2a8d 100644 --- a/apps/api/src/api/terminal.py +++ b/apps/api/src/api/terminal.py @@ -1,65 +1,102 @@ """WebSocket terminal endpoint for tool instances.""" import asyncio +import json import logging import uuid +from contextlib import suppress -from fastapi import APIRouter, Depends, WebSocket +from fastapi import APIRouter, Depends, HTTPException, WebSocket, status +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession +from starlette.websockets import WebSocketDisconnect -from src.auth.dependencies import get_db_session +from src.auth.dependencies import get_current_user_id, get_db_session +from src.models.terminal_session import TerminalSessionModel from src.models.tool_instance import ToolInstance -from src.services.terminal_manager import terminal_manager +from src.models.tool_type import ToolType +from src.services.terminal_manager import MaxSessionsExceededError, terminal_manager router = APIRouter() logger = logging.getLogger(__name__) +class SessionRef: + """Mutable reference to a terminal session, allowing updates during reset.""" + + def __init__(self, session, slot_session_id: str | None = None): + self.session = session + self.slot_session_id = slot_session_id or session.session_id + + @router.websocket( "/ws/tool-instances/{instance_id}/terminal", ) -async def terminal_websocket( +async def terminal_websocket_default( websocket: WebSocket, instance_id: str, db_session: AsyncSession = Depends(get_db_session), ) -> None: - """WebSocket endpoint for terminal access to a tool instance. + """WebSocket endpoint for terminal access (default session alias). - Provides an interactive terminal session inside a running tool instance container. - Supports: - - Auto-reconnection (client reconnects, server spawns new session) - - Heartbeat ping/pong - - Binary and text input frames - - Graceful session end notifications + Backward-compatible route that maps to the default session. + """ + await _handle_terminal_websocket(websocket, instance_id, None, db_session) + + +@router.websocket( + "/ws/tool-instances/{instance_id}/terminal/{session_id}", +) +async def terminal_websocket_specific( + websocket: WebSocket, + instance_id: str, + session_id: str, + db_session: AsyncSession = Depends(get_db_session), +) -> None: + """WebSocket endpoint for a specific terminal session.""" + await _handle_terminal_websocket(websocket, instance_id, session_id, db_session) + + +async def _handle_terminal_websocket( + websocket: WebSocket, + instance_id: str, + target_session_id: str | None, + db_session: AsyncSession, +) -> None: + """Shared WebSocket handler for terminal sessions. Args: websocket: The WebSocket connection. instance_id: UUID string of the tool instance. + target_session_id: Specific session ID (slot key). None means default session. db_session: Database session. - - Returns: - None. Communicates via WebSocket messages. - """ - logger.info("Terminal WebSocket connection attempt for instance %s", instance_id) + logger.debug( + "Terminal WebSocket connection attempt for instance %s (session=%s)", + instance_id, + target_session_id or "default", + ) await websocket.accept() + logger.debug("Terminal WebSocket accepted for instance %s", instance_id) try: + # Parse instance_id instance_uuid = uuid.UUID(instance_id) except ValueError: logger.error("Invalid instance ID: %s", instance_id) await websocket.close(code=4001, reason="Invalid instance ID") return + # Authenticate user from session cookie user_id = await _get_user_from_websocket(websocket, db_session) if user_id is None: logger.warning( - "Unauthorized terminal access attempt for instance %s", - instance_id, + "Unauthorized terminal access attempt for instance %s", instance_id ) await websocket.close(code=4003, reason="Unauthorized") return + # Get instance and verify ownership instance = await db_session.get(ToolInstance, instance_uuid) if instance is None: logger.warning("Instance %s not found", instance_id) @@ -85,50 +122,605 @@ async def terminal_websocket( await websocket.close(code=4004, reason="Instance not running") return - logger.info( - "Creating terminal session for instance %s (container_id=%s)", - instance_id, - instance.container_id, - ) - try: - session = await terminal_manager.create_session( - instance_uuid, + logger.debug("Terminal auth passed for instance %s, user %s", instance_id, user_id) + + # Verify the container actually exists (may have been removed/recreated) + from src.services.docker import get_container_status + + container_status = get_container_status(instance.container_id) + if container_status["status"] == "not_found": + logger.error( + "Container %s for instance %s not found (may have been removed)", instance.container_id, - websocket, - ) - logger.info( - "Terminal session created successfully for instance %s", instance_id, ) + await websocket.close( + code=4004, reason="Container not found — restart the tool instance" + ) + return + + # Fetch tool type to get startup_command + tool_type = await db_session.get(ToolType, instance.tool_type_id) + startup_command = tool_type.startup_command if tool_type else None + if startup_command: + logger.debug( + "Using startup command for instance %s: %s", + instance_id, + startup_command, + ) + + session = None + + # Get or create terminal session + try: + if target_session_id is None: + # Default session alias + session = await terminal_manager.get_or_create_session( + instance_uuid, + instance.container_id, + startup_command=startup_command, + ) + slot_session_id = "default" + else: + # Specific session + session = terminal_manager.get_session( + instance_id, + target_session_id, + ) + if session is None: + # Session not in memory — may have been lost on server restart. + # Try to restore from the DB row. + db_row = await db_session.get( + TerminalSessionModel, uuid.UUID(target_session_id) + ) + if ( + db_row is not None + and db_row.instance_id == instance_uuid + and db_row.status != "closed" + ): + logger.info( + "Restoring terminal session %s for instance %s from DB", + target_session_id, + instance_id, + ) + session = await terminal_manager.create_session( + instance_uuid, + instance.container_id, + startup_command=startup_command, + name=db_row.name, + session_id=target_session_id, + ) + else: + logger.warning( + "Session %s not found for instance %s", + target_session_id, + instance_id, + ) + await websocket.close(code=4004, reason="Session not found") + return + # Determine slot key for reset scoping + key = terminal_manager._find_key_by_internal_id( + instance_id, session.session_id + ) + slot_session_id = key[1] if key else target_session_id + + logger.debug( + "Terminal session ready for instance %s (session_id=%s, slot=%s)", + instance_id, + session.session_id, + slot_session_id, + ) + + # Attach WebSocket to session + await terminal_manager.attach_websocket(session, websocket) + logger.debug("WebSocket attached to session for instance %s", instance_id) # Send connected status await websocket.send_json({"type": "status", "status": "connected"}) + logger.debug("Sent connected status for instance %s", instance_id) - # Monitor session health and echo state - while session.is_alive() and not session.closed: - # Check echo state periodically - new_echo_state = await session.check_echo_state() - if new_echo_state is not None: - await websocket.send_json( - {"type": "set_echo_state", "enabled": new_echo_state}, - ) - await asyncio.sleep(1.0) + # Use mutable session reference so loops can survive reset + session_ref = SessionRef(session, slot_session_id) - # Session ended — determine reason and notify client - exit_reason = session.get_exit_reason() or "process_exit" - await websocket.send_json({"type": "session_ended", "reason": exit_reason}) - await websocket.close(code=1000, reason=f"Session ended: {exit_reason}") - - except Exception: - logger.exception( - "Terminal session error for instance %s", - instance_id, + # Start write loop and heartbeat (read is now event-driven in TerminalSession) + write_task = asyncio.create_task( + _write_loop(session_ref, websocket, instance_id) ) - await websocket.close(code=4000, reason="Terminal session error") + heartbeat_task = asyncio.create_task(_heartbeat_loop(websocket)) + logger.debug("Started terminal loops for instance %s", instance_id) + + # Wait for either task to complete (indicating disconnect or error) + done, pending = await asyncio.wait( + [write_task, heartbeat_task], + return_when=asyncio.FIRST_COMPLETED, + ) + + logger.debug( + "Terminal loop completed for instance %s, done=%s", + instance_id, + len(done), + ) + + # Cancel remaining tasks + for task in pending: + task.cancel() + + except WebSocketDisconnect: + logger.debug("WebSocket disconnected for instance %s", instance_id) + except Exception as exc: + logger.error( + "Terminal session error for instance %s: %s", + instance_id, + str(exc), + exc_info=True, + ) + with suppress(Exception): + await websocket.close(code=4000, reason=f"Error: {exc}") finally: + # Detach WebSocket, don't kill session + with suppress(Exception): + if session is not None: + await terminal_manager.detach_websocket(session, websocket) + logger.debug( + "WebSocket detached from session for instance %s", instance_id + ) + + +async def _write_loop(session_ref: SessionRef, websocket, instance_id: str) -> None: + """Read input from WebSocket and send to container.""" + try: + while True: + session = session_ref.session + if not session.is_alive() or session._closed: + await asyncio.sleep(0.1) + continue + message = await websocket.receive() + if message["type"] == "websocket.receive": + if "bytes" in message: + await session.write_input(message["bytes"]) + elif "text" in message: + text = message["text"] + if text.startswith("{"): + # Control message (JSON) + try: + ctrl = json.loads(text) + msg_type = ctrl.get("type") + + if msg_type == "resize": + cols = ctrl.get("cols", 80) + rows = ctrl.get("rows", 24) + logger.debug( + "Received resize message for instance %s: %sx%s", + instance_id, + cols, + rows, + ) + await session.resize(cols, rows) + elif msg_type == "ack": + char_count = ctrl.get("chars", 0) + if char_count > 0: + session.acknowledge_data(char_count) + elif msg_type == "reset": + # Reset terminal session (scoped to current slot) + logger.debug( + "Resetting terminal session for instance %s (slot=%s)", + session.instance_id, + session_ref.slot_session_id, + ) + await websocket.send_json( + {"type": "status", "status": "resetting"} + ) + + # Reset the session scoped to its slot + new_session = await terminal_manager.reset_session( + session.instance_id, + session.container_id, + startup_command=session.startup_command, + session_id=session_ref.slot_session_id, + name=session.name, + ) + + # Update the mutable session reference + session_ref.session = new_session + + # Attach to new session + await terminal_manager.attach_websocket( + new_session, websocket + ) + await websocket.send_json( + {"type": "status", "status": "connected"} + ) + + # Continue the loop with the new session + continue + + except json.JSONDecodeError: + # Not a valid JSON control message, treat as regular input + await session.write_input(text.encode("utf-8")) + else: + await session.write_input(text.encode("utf-8")) + elif message["type"] == "websocket.disconnect": + break + except Exception: pass +async def _heartbeat_loop(websocket: WebSocket) -> None: + """Send periodic ping messages to detect disconnections.""" + try: + while True: + await asyncio.sleep(30) # Ping every 30 seconds + try: + await websocket.send_json({"type": "ping"}) + except Exception: + # WebSocket is closed or broken + break + except Exception: + pass + + +async def _get_terminal_instance( + instance_id: uuid.UUID, + user_id: uuid.UUID, + db_session: AsyncSession, +) -> ToolInstance: + """Fetch instance and validate auth, ownership, and running status. + + Args: + instance_id: UUID of the tool instance. + user_id: ID of the authenticated user. + db_session: Database session. + + Returns: + The validated ToolInstance. + + Raises: + HTTPException: If instance not found, not owned, or not running. + """ + instance = await db_session.get(ToolInstance, instance_id) + if instance is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="Instance not found" + ) + + if instance.owner_id != user_id: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Not authorized to access this instance", + ) + + if instance.status != "running" or not instance.container_id: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail="Instance is not running" + ) + + return instance + + +@router.get( + "/instances/{instance_id}/terminal/sessions", + summary="List terminal sessions", + description="List terminal sessions for a tool instance with live WebSocket state.", +) +async def list_terminal_sessions( + instance_id: uuid.UUID, + user_id: uuid.UUID = Depends(get_current_user_id), + db_session: AsyncSession = Depends(get_db_session), +) -> dict: + """List terminal sessions for an instance. + + Args: + instance_id: UUID of the tool instance. + user_id: ID of the authenticated user. + db_session: Database session. + + Returns: + Dictionary with sessions list. + """ + await _get_terminal_instance(instance_id, user_id, db_session) + + # Query active DB rows for this instance + result = await db_session.execute( + select(TerminalSessionModel) + .where(TerminalSessionModel.instance_id == instance_id) + .where(TerminalSessionModel.status != "closed") + .order_by(TerminalSessionModel.created_at.asc()) + ) + db_rows = result.scalars().all() + + # Build response with live has_websockets flag. + # Include DB rows even without in-memory counterparts (e.g. after + # server restart) so the frontend can display tabs and reconnect. + sessions = [] + for row in db_rows: + live_session = terminal_manager.get_session(str(instance_id), str(row.id)) + sessions.append( + { + "id": str(row.id), + "name": row.name, + "status": row.status, + "has_websockets": live_session.has_websockets() + if live_session + else False, + "created_at": row.created_at.isoformat() if row.created_at else None, + "last_activity_at": row.last_activity_at.isoformat() + if row.last_activity_at + else None, + } + ) + + return {"sessions": sessions} + + +@router.post( + "/instances/{instance_id}/terminal/sessions", + summary="Create terminal session", + description="Create a new terminal session for a running tool instance.", + status_code=status.HTTP_201_CREATED, +) +async def create_terminal_session( + instance_id: uuid.UUID, + data: dict, + user_id: uuid.UUID = Depends(get_current_user_id), + db_session: AsyncSession = Depends(get_db_session), +) -> dict: + """Create a new terminal session. + + Args: + instance_id: UUID of the tool instance. + data: Request body with optional name. + user_id: ID of the authenticated user. + db_session: Database session. + + Returns: + Dictionary with new session details. + + Raises: + HTTPException: 409 if max sessions reached. + """ + instance = await _get_terminal_instance(instance_id, user_id, db_session) + assert instance.container_id is not None + + # Fetch tool type to get startup_command + tool_type = await db_session.get(ToolType, instance.tool_type_id) + startup_command = tool_type.startup_command if tool_type else None + + name = data.get("name") + + try: + session = await terminal_manager.create_session( + instance_id, + instance.container_id, + startup_command=startup_command, + name=name, + ) + except MaxSessionsExceededError: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail="Maximum of 5 terminal sessions reached for this instance", + ) from None + + return { + "id": session.session_id, + "name": session.name, + "status": session.status, + "created_at": session.last_activity, + } + + +@router.delete( + "/instances/{instance_id}/terminal/sessions/{session_id}", + summary="Close terminal session", + description="Close a specific terminal session.", +) +async def close_terminal_session( + instance_id: uuid.UUID, + session_id: str, + user_id: uuid.UUID = Depends(get_current_user_id), + db_session: AsyncSession = Depends(get_db_session), +) -> dict: + """Close a terminal session. + + Args: + instance_id: UUID of the tool instance. + session_id: ID of the session to close. + user_id: ID of the authenticated user. + db_session: Database session. + + Returns: + Dictionary with closure status. + """ + await _get_terminal_instance(instance_id, user_id, db_session) + + # Find the session by internal ID to determine its slot key + key = terminal_manager._find_key_by_internal_id(str(instance_id), session_id) + if ( + key is None + and terminal_manager.get_session(str(instance_id), session_id) is not None + ): + key = (str(instance_id), session_id) + + if key is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="Session not found" + ) + + await terminal_manager.close_session(key[0], key[1]) + + return {"status": "closed", "session_id": session_id} + + +@router.post( + "/instances/{instance_id}/terminal/sessions/{session_id}/reset", + summary="Reset terminal session", + description="Reset a specific terminal session, killing the current shell and starting fresh.", +) +async def reset_specific_terminal_session( + instance_id: uuid.UUID, + session_id: str, + user_id: uuid.UUID = Depends(get_current_user_id), + db_session: AsyncSession = Depends(get_db_session), +) -> dict: + """Reset a specific terminal session. + + Args: + instance_id: UUID of the tool instance. + session_id: ID of the session to reset. + user_id: ID of the authenticated user. + db_session: Database session. + + Returns: + Dictionary with reset session details. + """ + instance = await _get_terminal_instance(instance_id, user_id, db_session) + assert instance.container_id is not None + + # Determine slot key for reset + key = terminal_manager._find_key_by_internal_id(str(instance_id), session_id) + if ( + key is None + and terminal_manager.get_session(str(instance_id), session_id) is not None + ): + key = (str(instance_id), session_id) + + if key is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="Session not found" + ) + + # Fetch tool type to get startup_command + tool_type = await db_session.get(ToolType, instance.tool_type_id) + startup_command = tool_type.startup_command if tool_type else None + + # Preserve name if possible + live_session = terminal_manager.get_session(str(instance_id), session_id) + name = live_session.name if live_session else None + + new_session = await terminal_manager.reset_session( + instance_id, + instance.container_id, + startup_command=startup_command, + session_id=key[1], + name=name, + ) + + return { + "id": new_session.session_id, + "name": new_session.name, + "status": new_session.status, + } + + +@router.post( + "/instances/{instance_id}/terminal/sessions/{session_id}/rename", + summary="Rename terminal session", + description="Rename a specific terminal session.", +) +async def rename_terminal_session( + instance_id: uuid.UUID, + session_id: str, + data: dict, + user_id: uuid.UUID = Depends(get_current_user_id), + db_session: AsyncSession = Depends(get_db_session), +) -> dict: + """Rename a terminal session. + + Args: + instance_id: UUID of the tool instance. + session_id: ID of the session to rename. + data: Request body with new name. + user_id: ID of the authenticated user. + db_session: Database session. + + Returns: + Dictionary with updated session details. + """ + await _get_terminal_instance(instance_id, user_id, db_session) + + new_name = data.get("name") + if not new_name or not isinstance(new_name, str): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, detail="Name is required" + ) + + # Update in-memory session name if live + live_session = terminal_manager.get_session(str(instance_id), session_id) + if live_session: + live_session.name = new_name + + # Update DB row + db_row = await db_session.get(TerminalSessionModel, uuid.UUID(session_id)) + if db_row is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail="Session not found" + ) + + db_row.name = new_name + await db_session.commit() + + return {"id": str(db_row.id), "name": new_name} + + +@router.post( + "/instances/{instance_id}/terminal/reset", + summary="Reset terminal session (legacy alias)", + description="Reset the default terminal session for a tool instance. Preserved for backward compatibility.", +) +async def reset_terminal_session( + instance_id: uuid.UUID, + user_id: uuid.UUID = Depends(get_current_user_id), + db_session: AsyncSession = Depends(get_db_session), +) -> dict: + """Reset the default terminal session for an instance (legacy alias). + + Args: + instance_id: UUID of the tool instance. + user_id: ID of the authenticated user. + db_session: Database session. + + Returns: + Dictionary with status message. + """ + instance = await _get_terminal_instance(instance_id, user_id, db_session) + assert instance.container_id is not None + + # Fetch tool type to get startup_command + tool_type = await db_session.get(ToolType, instance.tool_type_id) + startup_command = tool_type.startup_command if tool_type else None + + try: + # Reset the default session + new_session = await terminal_manager.reset_session( + instance_id, + instance.container_id, + startup_command=startup_command, + ) + + logger.info( + "Terminal session reset for instance %s (new session_id=%s)", + instance_id, + new_session.session_id, + ) + + return { + "status": "success", + "message": "Terminal session reset successfully", + "instance_id": str(instance_id), + "session_id": new_session.session_id, + } + except Exception as exc: + logger.error( + "Failed to reset terminal session for instance %s: %s", + instance_id, + str(exc), + exc_info=True, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Failed to reset terminal session: {exc}", + ) from exc + + async def _get_user_from_websocket( websocket: WebSocket, db_session: AsyncSession, @@ -141,7 +733,6 @@ async def _get_user_from_websocket( Returns: The user's UUID if authenticated, None otherwise. - """ from src.auth.session import decode_session_cookie from src.config import Settings diff --git a/apps/api/src/api/tool_instances.py b/apps/api/src/api/tool_instances.py index 965c8ba..3c6a734 100644 --- a/apps/api/src/api/tool_instances.py +++ b/apps/api/src/api/tool_instances.py @@ -57,7 +57,8 @@ async def create_instance( raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="config profile does not belong to user") instance = await lifecycle.create_new_instance( - session, project, repo, tool_type, user, data.display_name, selected_profile + session, project, repo, tool_type, user, data.display_name, selected_profile, + ssh_key_ids=data.ssh_key_ids or None, ) return { "id": str(instance.id), @@ -66,6 +67,7 @@ async def create_instance( "tool_type_id": str(instance.tool_type_id), "status": instance.status, "config_profile_id": str(instance.selected_profile_id) if instance.selected_profile_id else None, + "ssh_key_ids": instance.ssh_key_ids or [], "created_at": instance.created_at.isoformat(), } @router.get("/{project_id}/repositories/{repo_id}/instances") diff --git a/apps/api/src/models/tool_instance.py b/apps/api/src/models/tool_instance.py index 4fa0a89..d4bffd8 100644 --- a/apps/api/src/models/tool_instance.py +++ b/apps/api/src/models/tool_instance.py @@ -2,7 +2,7 @@ import uuid from datetime import datetime from typing import TYPE_CHECKING -from sqlalchemy import DateTime, ForeignKey, Integer, String +from sqlalchemy import DateTime, ForeignKey, Integer, JSON, String from sqlalchemy import Uuid as UUID from sqlalchemy.orm import Mapped, mapped_column, relationship @@ -66,6 +66,9 @@ class ToolInstance(UUIDPrimaryKeyMixin, TimestampMixin, Base): selected_profile_id: Mapped[uuid.UUID | None] = mapped_column( UUID(), ForeignKey("config_profiles.id", ondelete="SET NULL"), nullable=True ) + ssh_key_ids: Mapped[list[str] | None] = mapped_column( + JSON, nullable=True + ) tool_type: Mapped["ToolType"] = relationship() repository: Mapped["GitRepository"] = relationship() diff --git a/apps/api/src/schemas/tool_instance.py b/apps/api/src/schemas/tool_instance.py index 09454d7..30a5a5d 100644 --- a/apps/api/src/schemas/tool_instance.py +++ b/apps/api/src/schemas/tool_instance.py @@ -15,6 +15,9 @@ class CreateInstanceRequest(BaseModel): config_profile_id: str | None = Field( default=None, description="Optional config profile ID to apply to the instance" ) + ssh_key_ids: list[str] = Field( + default_factory=list, description="SSH key IDs to mount into container ~/.ssh" + ) class SessionItemResponse(BaseModel): diff --git a/apps/api/src/services/instance_lifecycle.py b/apps/api/src/services/instance_lifecycle.py index 2b57b1e..ae913ab 100644 --- a/apps/api/src/services/instance_lifecycle.py +++ b/apps/api/src/services/instance_lifecycle.py @@ -7,6 +7,7 @@ to create, start, stop, restart, and delete tool instances. import logging import os import shutil +import uuid from datetime import datetime from typing import Any @@ -16,6 +17,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from src.models.config_profile import ConfigProfile from src.models.git_repository import GitRepository from src.models.project import Project +from src.models.ssh_key import SSHKey from src.models.tool_instance import ToolInstance from src.models.tool_type import ToolType from src.models.user import User @@ -37,6 +39,7 @@ async def create_new_instance( user: User, display_name: str | None, selected_profile: ConfigProfile | None, + ssh_key_ids: list[str] | None = None, ) -> ToolInstance: """Create a new tool instance record and its compose file.""" instance_name = await compose_svc._generate_instance_name( @@ -61,6 +64,7 @@ async def create_new_instance( compose_path=compose_path, port=tool_port, selected_profile_id=selected_profile.id if selected_profile else None, + ssh_key_ids=ssh_key_ids or None, ) session.add(instance) await session.commit() @@ -117,6 +121,12 @@ async def start_existing_instance( os.path.dirname(instance.compose_path), env_vars, config_files, extra_volumes ) + # Mount selected SSH keys into container ~/.ssh + if instance.ssh_key_ids: + extra_volumes = await _mount_ssh_keys( + session, instance, user, extra_volumes + ) + if port_override or start_command or working_directory or extra_volumes: compose_svc._modify_compose_file( instance.compose_path, @@ -227,6 +237,12 @@ async def restart_existing_instance( os.path.dirname(instance.compose_path), env_vars, config_files, extra_volumes ) + # Mount selected SSH keys into container ~/.ssh + if instance.ssh_key_ids: + extra_volumes = await _mount_ssh_keys( + session, instance, user, extra_volumes + ) + if port_override or start_command or working_directory or extra_volumes: compose_svc._modify_compose_file( instance.compose_path, @@ -374,6 +390,97 @@ async def _stage_configs( return env_file_path, extra_volumes +async def _mount_ssh_keys( + session: AsyncSession, + instance: ToolInstance, + user: User, + extra_volumes: list[dict], +) -> list[dict]: + """Prepare and mount SSH keys into the container.""" + from src.services.ssh_keys import ( + _sanitize_filename, + prepare_ssh_key_files, + write_ssh_config, + ) + + ssh_keys_to_mount = [] + for key_id in instance.ssh_key_ids or []: + try: + key_uuid = uuid.UUID(key_id) + except ValueError: + logger.warning( + "Invalid SSH key ID %s for instance %s", key_id, instance.id + ) + continue + ssh_key = await session.get(SSHKey, key_uuid) + if ssh_key and ssh_key.user_id == user.id: + ssh_keys_to_mount.append(ssh_key) + else: + logger.warning( + "SSH key %s not found or not authorized for user %s", + key_id, + user.id, + ) + + if not ssh_keys_to_mount: + return extra_volumes + + instance_dir = os.path.dirname(instance.compose_path or "") + ssh_dir = os.path.join(instance_dir, "mounts", "ssh", ".ssh") + os.makedirs(ssh_dir, exist_ok=True) + + key_filenames = [] + for ssh_key in ssh_keys_to_mount: + key_name = _sanitize_filename(ssh_key.name) + base_filename = f"id_ed25519_{key_name}" + filename = base_filename + counter = 1 + while filename in key_filenames: + filename = f"{base_filename}_{counter}" + counter += 1 + key_filenames.append(filename) + + try: + prepare_ssh_key_files( + instance_dir, + ssh_key, + subdir="mounts/ssh/.ssh", + key_filename=filename, + write_config=False, + ) + except Exception as exc: + logger.error( + "Failed to prepare SSH key %s for instance %s: %s", + ssh_key.id, + instance.id, + exc, + ) + + try: + write_ssh_config(ssh_dir, key_filenames) + except Exception as exc: + logger.error( + "Failed to write SSH config for instance %s: %s", instance.id, exc + ) + + ssh_target = "/root/.ssh" + extra_volumes.append( + { + "source": ssh_dir, + "target": ssh_target, + "type": "bind", + } + ) + logger.info( + "Mounted %d SSH key(s) for instance %s to %s", + len(ssh_keys_to_mount), + instance.id, + ssh_target, + ) + + return extra_volumes + + async def _start_tunnel_if_web(instance: ToolInstance, tool_type: ToolType) -> None: """Create Cloudflare tunnel for web-enabled tools.""" if tool_type.interface_type != "web" or not tool_type.default_port: diff --git a/apps/api/src/services/terminal_manager.py b/apps/api/src/services/terminal_manager.py index 6cb8f73..b9605db 100644 --- a/apps/api/src/services/terminal_manager.py +++ b/apps/api/src/services/terminal_manager.py @@ -1,193 +1,426 @@ """Terminal session manager for WebSocket connections.""" import asyncio -import contextlib -import json import logging -import time import uuid -from collections.abc import Coroutine -from typing import Any +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.terminal_session import TerminalSessionModel from src.services.terminal_session import TerminalSession logger = logging.getLogger(__name__) -_READ_BATCH_INTERVAL_S = 0.016 # 16ms max batching delay -_READ_POLL_TIMEOUT_S = 0.005 -_READ_POLL_SLEEP_S = 0.001 -_HEARTBEAT_INTERVAL_S = 15.0 -_IDLE_TIMEOUT_S = 60.0 + +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.""" + """Manages active terminal sessions with persistence support.""" + + # Maximum sessions per tool instance + MAX_SESSIONS_PER_INSTANCE = 5 def __init__(self) -> None: - """Initialise the terminal manager.""" - self._sessions: dict[str, TerminalSession] = {} - self._last_client_message: dict[str, float] = {} - self._background_tasks: set[asyncio.Task[Any]] = set() + # 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, - websocket: WebSocket, + startup_command: str | None = None, + name: str | None = None, + session_id: str | None = None, ) -> TerminalSession: - """Create a new terminal session.""" - session_id = str(uuid.uuid4()) - session = TerminalSession(session_id, instance_id, container_id) - await session.start() - self._sessions[session_id] = session - self._last_client_message[session_id] = time.monotonic() + """Create a new terminal session for an instance. - # Start background tasks for I/O streaming - self._start_task(self._read_loop(session, websocket)) - self._start_task(self._write_loop(session, websocket)) - self._start_task(self._heartbeat_loop(session, websocket)) + 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 _start_task(self, coro: Coroutine[Any, Any, None]) -> None: - """Start a background task and store a reference to prevent GC.""" - task = asyncio.create_task(coro) - self._background_tasks.add(task) - task.add_done_callback(self._background_tasks.discard) - - async def _read_loop( + def get_session( self, - session: TerminalSession, - websocket: WebSocket, - ) -> None: - """Read output from the container and send to WebSocket with batching.""" - try: - buffer = bytearray() - last_flush = time.monotonic() + 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 - while session.is_alive() and not session.closed: - data = await session.read_output(select_timeout=_READ_POLL_TIMEOUT_S) - if data: - buffer.extend(data) - - now = time.monotonic() - flush_due = buffer and ( - now - last_flush >= _READ_BATCH_INTERVAL_S or not data - ) - - if flush_due: - await websocket.send_bytes(bytes(buffer)) - buffer.clear() - last_flush = now - elif not data: - await asyncio.sleep(_READ_POLL_SLEEP_S) - - # Flush any remaining data - if buffer: - with contextlib.suppress(Exception): - await websocket.send_bytes(bytes(buffer)) - - except Exception: - logger.exception("Read loop error for session %s", session.session_id) - finally: - await self._cleanup_session(session) - - async def _write_loop( + def _find_key_by_internal_id( self, - session: TerminalSession, - websocket: WebSocket, - ) -> None: - """Read input from WebSocket and send to container.""" - try: - while session.is_alive() and not session.closed: - message = await websocket.receive() - self._last_client_message[session.session_id] = time.monotonic() + 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 - if message["type"] == "websocket.receive": - if "bytes" in message: - await session.write_input(message["bytes"]) - elif "text" in message: - text = message["text"] - if text.startswith("{"): - try: - ctrl = json.loads(text) - await self._handle_control_message( - session, - websocket, - ctrl, - ) - except json.JSONDecodeError: - logger.debug("Invalid JSON control message: %s", text) - else: - await session.write_input(text.encode("utf-8")) - elif message["type"] == "websocket.disconnect": - break - except Exception: - logger.exception("Write loop error for session %s", session.session_id) - finally: - await self._cleanup_session(session) - - async def _handle_control_message( + def get_sessions_for_instance( self, - session: TerminalSession, - websocket: WebSocket, - ctrl: dict[str, Any], + 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: - """Handle a JSON control message from the client.""" - msg_type = ctrl.get("type") - if msg_type == "resize": - await session.resize( - ctrl.get("cols", 80), - ctrl.get("rows", 24), - ) - elif msg_type == "ping": - await websocket.send_json( - {"type": "pong", "id": ctrl.get("id")}, + """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 _heartbeat_loop( + async def attach_websocket( self, session: TerminalSession, websocket: WebSocket, ) -> None: - """Monitor client activity and close idle connections.""" - try: - while session.is_alive() and not session.closed: - await asyncio.sleep(_HEARTBEAT_INTERVAL_S) - last_msg = self._last_client_message.get(session.session_id, 0) - if time.monotonic() - last_msg > _IDLE_TIMEOUT_S: - # Client has been silent for 60s — close connection - with contextlib.suppress(Exception): - await websocket.close( - code=1000, - reason="Idle timeout", - ) - break - except Exception: - logger.exception( - "Heartbeat loop error for session %s", + """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, ) - finally: - await self._cleanup_session(session) + for ws in list(session._websockets): + try: + await ws.close(code=4000, reason="New connection established") + except Exception: + pass # noqa: S110 + session._websockets.clear() - async def _cleanup_session(self, session: TerminalSession) -> None: - """Clean up a session.""" - if session.session_id in self._sessions: - del self._sessions[session.session_id] - self._last_client_message.pop(session.session_id, None) - await session.close() + # 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() - self._last_client_message.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() diff --git a/apps/api/src/services/terminal_session.py b/apps/api/src/services/terminal_session.py index 7f3a490..dff40cd 100644 --- a/apps/api/src/services/terminal_session.py +++ b/apps/api/src/services/terminal_session.py @@ -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""