merge: restore dev features lost in main→dev merge
This commit is contained in:
+641
-50
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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""
|
||||
|
||||
Reference in New Issue
Block a user