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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user