Fusion
125 lines
4.2 KiB
Python
125 lines
4.2 KiB
Python
"""WebSocket terminal endpoint for tool instances."""
|
|
|
|
import asyncio
|
|
import logging
|
|
import uuid
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect, status
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from src.auth.dependencies import get_db_session
|
|
from src.models.tool_instance import ToolInstance
|
|
from src.services.terminal_manager import terminal_manager
|
|
|
|
router = APIRouter()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@router.websocket(
|
|
"/ws/tool-instances/{instance_id}/terminal",
|
|
)
|
|
async def terminal_websocket(
|
|
websocket: WebSocket,
|
|
instance_id: str,
|
|
db_session: AsyncSession = Depends(get_db_session),
|
|
) -> None:
|
|
"""WebSocket endpoint for terminal access to a tool instance.
|
|
|
|
Provides an interactive terminal session inside a running tool instance container.
|
|
|
|
Args:
|
|
websocket: The WebSocket connection.
|
|
instance_id: UUID string of the tool instance.
|
|
db_session: Database session.
|
|
|
|
Returns:
|
|
None. Communicates via WebSocket messages.
|
|
"""
|
|
logger.info("Terminal WebSocket connection attempt for instance %s", instance_id)
|
|
await websocket.accept()
|
|
|
|
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)
|
|
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)
|
|
await websocket.close(code=4004, reason="Instance not found")
|
|
return
|
|
|
|
if instance.owner_id != user_id:
|
|
logger.warning("Forbidden terminal access for instance %s by user %s", instance_id, user_id)
|
|
await websocket.close(code=4003, reason="Forbidden")
|
|
return
|
|
|
|
if instance.status != "running" or not instance.container_id:
|
|
logger.warning("Instance %s not running (status=%s, container_id=%s)", instance_id, instance.status, instance.container_id)
|
|
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
|
|
try:
|
|
session = await terminal_manager.create_session(
|
|
instance_uuid,
|
|
instance.container_id,
|
|
websocket,
|
|
)
|
|
logger.info("Terminal session created successfully 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)
|
|
|
|
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
|
|
pass
|
|
|
|
|
|
async def _get_user_from_websocket(
|
|
websocket: WebSocket,
|
|
db_session: AsyncSession,
|
|
) -> uuid.UUID | None:
|
|
"""Extract and validate user ID from session cookie in WebSocket.
|
|
|
|
Args:
|
|
websocket: The WebSocket connection.
|
|
db_session: Database session.
|
|
|
|
Returns:
|
|
The user's UUID if authenticated, None otherwise.
|
|
"""
|
|
from src.auth.session import decode_session_cookie
|
|
from src.config import Settings
|
|
|
|
session_cookie = websocket.cookies.get("session")
|
|
if not session_cookie:
|
|
return None
|
|
|
|
settings = Settings()
|
|
try:
|
|
payload = decode_session_cookie(settings=settings, cookie_value=session_cookie)
|
|
return uuid.UUID(str(payload["user_id"]))
|
|
except (ValueError, KeyError):
|
|
return None
|