headquarter/apps/api/src/api/terminal.py

"""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__)


class SessionRef:
    """Mutable reference to a terminal session, allowing updates during reset."""

    def __init__(self, session):
        self.session = session


@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.
    Sessions persist across WebSocket disconnections.

    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

    # Get or create terminal session
    try:
        session = await terminal_manager.get_or_create_session(
            instance_uuid,
            instance.container_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"})

        # Use mutable session reference so loops can survive reset
        session_ref = SessionRef(session)

        # Start I/O loops and heartbeat
        read_task = asyncio.create_task(_read_loop(session_ref, websocket))
        write_task = asyncio.create_task(_write_loop(session_ref, websocket, instance_id))
        heartbeat_task = asyncio.create_task(_heartbeat_loop(websocket))

        # Wait for either task to complete (indicating disconnect or error)
        done, pending = await asyncio.wait(
            [read_task, write_task, heartbeat_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:
        # 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_ref: SessionRef, websocket) -> None:
    """Read output from the container and send to WebSocket."""
    try:
        while True:
            session = session_ref.session
            if not session.is_alive() or session._closed:
                await asyncio.sleep(0.1)
                continue
            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_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)
                        import 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.info(f"Received resize message for instance {instance_id}: {cols}x{rows}")
                                await session.resize(cols, rows)
                            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,
                                )
                                
                                # Update the mutable session reference so read_loop uses the new session
                                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


@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,
) -> 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