feat: implement persistent terminal sessions
- Terminal sessions now persist across WebSocket disconnections - Added circular output buffer (10KB) for replay on reconnect - Added idle timeout cleanup (30 minutes) - Added reset functionality via WebSocket message and HTTP endpoint - Concurrent connections close old WebSocket when new one connects - Frontend: Added reset button with confirmation dialog - Frontend: Handle resetting status and reconnection Refs: persistent-terminal-sessions
This commit is contained in:
+154
-10
@@ -26,6 +26,7 @@ async def terminal_websocket(
|
||||
"""WebSocket endpoint for terminal access to a tool instance.
|
||||
|
||||
Provides an interactive terminal session inside a running tool instance container.
|
||||
Sessions persist across WebSocket disconnections.
|
||||
|
||||
Args:
|
||||
websocket: The WebSocket connection.
|
||||
@@ -70,32 +71,175 @@ 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)
|
||||
# Create terminal session
|
||||
# Get or create terminal session
|
||||
try:
|
||||
session = await terminal_manager.create_session(
|
||||
session = await terminal_manager.get_or_create_session(
|
||||
instance_uuid,
|
||||
instance.container_id,
|
||||
websocket,
|
||||
)
|
||||
logger.info("Terminal session created successfully for instance %s", instance_id)
|
||||
logger.info("Terminal session ready for instance %s (session_id=%s)", instance_id, session.session_id)
|
||||
|
||||
# Attach WebSocket to session
|
||||
await terminal_manager.attach_websocket(session, websocket)
|
||||
logger.info("WebSocket attached to session for instance %s", instance_id)
|
||||
|
||||
# Send connected status
|
||||
await websocket.send_json({"type": "status", "status": "connected"})
|
||||
|
||||
# Keep connection alive until session ends
|
||||
# The terminal_manager handles I/O loops, we just wait here
|
||||
while session.is_alive() and not session._closed:
|
||||
await asyncio.sleep(0.5)
|
||||
# Start I/O loops
|
||||
read_task = asyncio.create_task(_read_loop(session, websocket))
|
||||
write_task = asyncio.create_task(_write_loop(session, websocket))
|
||||
|
||||
# Wait for either task to complete (indicating disconnect or error)
|
||||
done, pending = await asyncio.wait(
|
||||
[read_task, write_task],
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
|
||||
# Cancel remaining tasks
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
|
||||
except Exception as exc:
|
||||
logger.error("Terminal session error for instance %s: %s", instance_id, str(exc), exc_info=True)
|
||||
await websocket.close(code=4000, reason=f"Error: {exc}")
|
||||
finally:
|
||||
# Cleanup will be handled by the session manager
|
||||
# Detach WebSocket, don't kill session
|
||||
try:
|
||||
if 'session' in locals():
|
||||
await terminal_manager.detach_websocket(session, websocket)
|
||||
logger.info("WebSocket detached from session for instance %s", instance_id)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def _read_loop(session, websocket) -> None:
|
||||
"""Read output from the container and send to WebSocket."""
|
||||
try:
|
||||
while session.is_alive() and not session._closed:
|
||||
data = await session.read_output()
|
||||
if data:
|
||||
try:
|
||||
await websocket.send_bytes(data)
|
||||
except Exception:
|
||||
break
|
||||
else:
|
||||
await asyncio.sleep(0.01)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
async def _write_loop(session, websocket) -> None:
|
||||
"""Read input from WebSocket and send to container."""
|
||||
try:
|
||||
while session.is_alive() and not session._closed:
|
||||
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)
|
||||
import json
|
||||
try:
|
||||
ctrl = json.loads(text)
|
||||
msg_type = ctrl.get("type")
|
||||
|
||||
if msg_type == "resize":
|
||||
await session.resize(
|
||||
ctrl.get("cols", 80),
|
||||
ctrl.get("rows", 24),
|
||||
)
|
||||
elif msg_type == "reset":
|
||||
# Reset terminal session
|
||||
logger.info("Resetting terminal session for instance %s", session.instance_id)
|
||||
await websocket.send_json({"type": "status", "status": "resetting"})
|
||||
|
||||
# Reset the session
|
||||
new_session = await terminal_manager.reset_session(
|
||||
session.instance_id,
|
||||
session.container_id,
|
||||
)
|
||||
|
||||
# Attach to new session
|
||||
await terminal_manager.attach_websocket(new_session, websocket)
|
||||
await websocket.send_json({"type": "status", "status": "connected"})
|
||||
|
||||
# Update session reference and restart loops
|
||||
# Note: This will cause the current loops to exit
|
||||
# The WebSocket handler will create new ones
|
||||
return
|
||||
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
else:
|
||||
await session.write_input(text.encode("utf-8"))
|
||||
elif message["type"] == "websocket.disconnect":
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@router.post(
|
||||
"/projects/{project_id}/repositories/{repo_id}/instances/{instance_id}/terminal/reset",
|
||||
summary="Reset terminal session",
|
||||
description="Reset the terminal session for a tool instance, killing the current shell and starting fresh.",
|
||||
)
|
||||
async def reset_terminal_session(
|
||||
project_id: uuid.UUID,
|
||||
repo_id: uuid.UUID,
|
||||
instance_id: uuid.UUID,
|
||||
db_session: AsyncSession = Depends(get_db_session),
|
||||
) -> dict:
|
||||
"""Reset the terminal session for an instance.
|
||||
|
||||
Args:
|
||||
project_id: UUID of the project.
|
||||
repo_id: UUID of the repository.
|
||||
instance_id: UUID of the tool instance.
|
||||
db_session: Database session.
|
||||
|
||||
Returns:
|
||||
Dictionary with status message.
|
||||
"""
|
||||
# Get instance and verify it exists and is 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.status != "running" or not instance.container_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Instance is not running"
|
||||
)
|
||||
|
||||
try:
|
||||
# Reset the session
|
||||
new_session = await terminal_manager.reset_session(
|
||||
instance_id,
|
||||
instance.container_id,
|
||||
)
|
||||
|
||||
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}"
|
||||
)
|
||||
|
||||
|
||||
async def _get_user_from_websocket(
|
||||
websocket: WebSocket,
|
||||
db_session: AsyncSession,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Terminal session manager for WebSocket connections."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
@@ -8,81 +9,132 @@ from fastapi import WebSocket
|
||||
|
||||
from src.services.terminal_session import TerminalSession
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TerminalManager:
|
||||
"""Manages active terminal sessions."""
|
||||
"""Manages active terminal sessions with persistence support."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
# Track sessions by instance_id for persistence
|
||||
self._sessions: dict[str, TerminalSession] = {}
|
||||
self._idle_check_task: asyncio.Task | None = None
|
||||
self._start_idle_check()
|
||||
|
||||
async def create_session(
|
||||
def _start_idle_check(self) -> None:
|
||||
"""Start the idle timeout background task."""
|
||||
if self._idle_check_task is None or self._idle_check_task.done():
|
||||
self._idle_check_task = asyncio.create_task(self._idle_check_loop())
|
||||
|
||||
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_sessions = []
|
||||
for instance_id, session in list(self._sessions.items()):
|
||||
if session.is_idle():
|
||||
idle_sessions.append(instance_id)
|
||||
|
||||
for instance_id in idle_sessions:
|
||||
logger.info("Cleaning up idle terminal session for instance %s", instance_id)
|
||||
session = self._sessions.pop(instance_id, None)
|
||||
if session:
|
||||
await session.close()
|
||||
|
||||
async def get_or_create_session(
|
||||
self,
|
||||
instance_id: uuid.UUID,
|
||||
container_id: str,
|
||||
websocket: WebSocket,
|
||||
) -> TerminalSession:
|
||||
"""Create a new terminal session."""
|
||||
"""Get existing session or create a new one."""
|
||||
instance_id_str = str(instance_id)
|
||||
|
||||
# Check for existing session
|
||||
if instance_id_str in self._sessions:
|
||||
session = self._sessions[instance_id_str]
|
||||
|
||||
# Check if session is still alive
|
||||
if session.is_alive():
|
||||
logger.info("Reattaching to existing terminal session for instance %s", instance_id)
|
||||
return session
|
||||
else:
|
||||
# Session died, clean it up
|
||||
logger.info("Existing session for instance %s is dead, cleaning up", instance_id)
|
||||
await session.close()
|
||||
del self._sessions[instance_id_str]
|
||||
|
||||
# Create new session
|
||||
logger.info("Creating new terminal session for instance %s", instance_id)
|
||||
session_id = str(uuid.uuid4())
|
||||
session = TerminalSession(session_id, instance_id, container_id)
|
||||
await session.start()
|
||||
self._sessions[session_id] = session
|
||||
|
||||
# Start background tasks for I/O streaming
|
||||
asyncio.create_task(self._read_loop(session, websocket))
|
||||
asyncio.create_task(self._write_loop(session, websocket))
|
||||
|
||||
self._sessions[instance_id_str] = session
|
||||
|
||||
return session
|
||||
|
||||
async def _read_loop(self, session: TerminalSession, websocket: WebSocket) -> None:
|
||||
"""Read output from the container and send to WebSocket."""
|
||||
try:
|
||||
while session.is_alive() and not session._closed:
|
||||
data = await session.read_output()
|
||||
if data:
|
||||
await websocket.send_bytes(data)
|
||||
else:
|
||||
await asyncio.sleep(0.01)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
await self._cleanup_session(session)
|
||||
async def attach_websocket(
|
||||
self,
|
||||
session: TerminalSession,
|
||||
websocket: WebSocket,
|
||||
) -> None:
|
||||
"""Attach a WebSocket to an existing session."""
|
||||
# Handle concurrent connections - close existing ones
|
||||
if session.has_websockets():
|
||||
logger.info("Closing existing WebSocket connections for instance %s", session.instance_id)
|
||||
for ws in list(session._websockets):
|
||||
try:
|
||||
await ws.close(code=4000, reason="New connection established")
|
||||
except Exception:
|
||||
pass
|
||||
session._websockets.clear()
|
||||
|
||||
# Attach new WebSocket
|
||||
session.attach_websocket(websocket)
|
||||
|
||||
# Replay buffer
|
||||
buffer = session.get_buffer()
|
||||
if buffer:
|
||||
try:
|
||||
await websocket.send_bytes(buffer)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _write_loop(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()
|
||||
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)
|
||||
import json
|
||||
try:
|
||||
ctrl = json.loads(text)
|
||||
if ctrl.get("type") == "resize":
|
||||
await session.resize(
|
||||
ctrl.get("cols", 80),
|
||||
ctrl.get("rows", 24),
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
else:
|
||||
await session.write_input(text.encode("utf-8"))
|
||||
elif message["type"] == "websocket.disconnect":
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
await self._cleanup_session(session)
|
||||
async def detach_websocket(
|
||||
self,
|
||||
session: TerminalSession,
|
||||
websocket: WebSocket,
|
||||
) -> None:
|
||||
"""Detach a WebSocket from a session."""
|
||||
session.detach_websocket(websocket)
|
||||
|
||||
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]
|
||||
await session.close()
|
||||
async def reset_session(
|
||||
self,
|
||||
instance_id: uuid.UUID,
|
||||
container_id: str,
|
||||
) -> TerminalSession:
|
||||
"""Reset a session by killing it and creating a new one."""
|
||||
instance_id_str = str(instance_id)
|
||||
|
||||
# Close existing session if any
|
||||
if instance_id_str in self._sessions:
|
||||
logger.info("Resetting terminal session for instance %s", instance_id)
|
||||
old_session = self._sessions.pop(instance_id_str)
|
||||
await old_session.close()
|
||||
|
||||
# Create new session
|
||||
session_id = str(uuid.uuid4())
|
||||
session = TerminalSession(session_id, instance_id, container_id)
|
||||
await session.start()
|
||||
self._sessions[instance_id_str] = session
|
||||
|
||||
return session
|
||||
|
||||
async def close_all(self) -> None:
|
||||
"""Close all active sessions."""
|
||||
@@ -90,6 +142,9 @@ class TerminalManager:
|
||||
self._sessions.clear()
|
||||
for session in sessions:
|
||||
await session.close()
|
||||
|
||||
if self._idle_check_task and not self._idle_check_task.done():
|
||||
self._idle_check_task.cancel()
|
||||
|
||||
|
||||
# Global terminal manager instance
|
||||
|
||||
@@ -6,12 +6,24 @@ import pty
|
||||
import select
|
||||
import struct
|
||||
import fcntl
|
||||
import time
|
||||
import uuid
|
||||
from collections import deque
|
||||
from typing import Any
|
||||
|
||||
|
||||
class TerminalSession:
|
||||
"""Manages a single terminal session connected to a docker container."""
|
||||
"""Manages a single terminal session connected to a docker container.
|
||||
|
||||
Supports persistent sessions that survive WebSocket disconnections.
|
||||
Multiple WebSocket connections can attach/detach from the same session.
|
||||
"""
|
||||
|
||||
# Circular buffer size (10KB)
|
||||
BUFFER_SIZE = 10 * 1024
|
||||
|
||||
# Idle timeout in seconds (30 minutes)
|
||||
IDLE_TIMEOUT = 30 * 60
|
||||
|
||||
def __init__(self, session_id: str, instance_id: uuid.UUID, container_id: str) -> None:
|
||||
self.session_id = session_id
|
||||
@@ -21,6 +33,20 @@ class TerminalSession:
|
||||
self._closed = False
|
||||
self._master_fd: int | None = None
|
||||
self._slave_fd: int | None = None
|
||||
|
||||
# Circular buffer for output replay
|
||||
self._output_buffer: deque[bytes] = deque(maxlen=self.BUFFER_SIZE)
|
||||
self._buffer_size = 0
|
||||
|
||||
# WebSocket connections
|
||||
self._websockets: set[Any] = set()
|
||||
|
||||
# Activity tracking
|
||||
self.last_activity = time.time()
|
||||
|
||||
# Terminal size
|
||||
self._cols = 80
|
||||
self._rows = 24
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the docker exec process with a shell using a PTY."""
|
||||
@@ -28,7 +54,7 @@ class TerminalSession:
|
||||
self._master_fd, self._slave_fd = pty.openpty()
|
||||
|
||||
# Set the terminal size initially
|
||||
self._set_terminal_size(80, 24)
|
||||
self._set_terminal_size(self._cols, self._rows)
|
||||
|
||||
# Start docker exec with the slave fd as stdin/stdout/stderr
|
||||
# Using -it because the slave fd IS a TTY
|
||||
@@ -49,6 +75,8 @@ class TerminalSession:
|
||||
# Close slave fd in parent process
|
||||
os.close(self._slave_fd)
|
||||
self._slave_fd = None
|
||||
|
||||
self.last_activity = time.time()
|
||||
|
||||
def _set_terminal_size(self, cols: int, rows: int) -> None:
|
||||
"""Set the terminal size using TIOCSWINSZ."""
|
||||
@@ -63,24 +91,43 @@ class TerminalSession:
|
||||
pass
|
||||
|
||||
async def read_output(self) -> bytes:
|
||||
"""Read output from the PTY master."""
|
||||
"""Read output from the PTY master and store in buffer."""
|
||||
if self._master_fd is None or self._closed:
|
||||
return b""
|
||||
try:
|
||||
# Use select to check if data is available
|
||||
readable, _, _ = select.select([self._master_fd], [], [], 0.1)
|
||||
if readable:
|
||||
return os.read(self._master_fd, 4096)
|
||||
data = os.read(self._master_fd, 4096)
|
||||
if data:
|
||||
self._add_to_buffer(data)
|
||||
self.last_activity = time.time()
|
||||
return data
|
||||
return b""
|
||||
except (OSError, IOError, ValueError):
|
||||
return b""
|
||||
|
||||
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)
|
||||
|
||||
# Trim if exceeds max size
|
||||
while self._buffer_size > self.BUFFER_SIZE and self._output_buffer:
|
||||
removed = self._output_buffer.popleft()
|
||||
self._buffer_size -= len(removed)
|
||||
|
||||
def get_buffer(self) -> bytes:
|
||||
"""Get buffered output for replay."""
|
||||
return b"".join(self._output_buffer)
|
||||
|
||||
async def write_input(self, data: bytes) -> None:
|
||||
"""Write input to the PTY master."""
|
||||
if self._master_fd is None or self._closed:
|
||||
return
|
||||
try:
|
||||
os.write(self._master_fd, data)
|
||||
self.last_activity = time.time()
|
||||
except (OSError, IOError):
|
||||
pass
|
||||
|
||||
@@ -88,8 +135,21 @@ class TerminalSession:
|
||||
"""Resize the terminal."""
|
||||
if self._closed:
|
||||
return
|
||||
self._cols = cols
|
||||
self._rows = rows
|
||||
self._set_terminal_size(cols, rows)
|
||||
|
||||
async def reset(self) -> None:
|
||||
"""Reset the session by killing the process and clearing state."""
|
||||
await self.close()
|
||||
self._closed = False
|
||||
self._output_buffer.clear()
|
||||
self._buffer_size = 0
|
||||
self._websockets.clear()
|
||||
self.process = None
|
||||
self._master_fd = None
|
||||
self._slave_fd = None
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the session and cleanup."""
|
||||
if self._closed:
|
||||
@@ -115,3 +175,35 @@ 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."""
|
||||
dead_sockets = set()
|
||||
for ws in self._websockets:
|
||||
try:
|
||||
await ws.send_bytes(data)
|
||||
except Exception:
|
||||
dead_sockets.add(ws)
|
||||
|
||||
# Clean up dead sockets
|
||||
for ws in dead_sockets:
|
||||
self._websockets.discard(ws)
|
||||
|
||||
Reference in New Issue
Block a user