aionboarding/ai-patient-onboarding/app/api/websocket/routes.py

"""
WebSocket Route Handler (Section 2, 3, 4, 15, 16, 22, 23, 24, 25).
Endpoint: WS /api/v1/ws/onboarding/{session_id}
Handles connection lifecycle, authentication, session validation, input sanitization, rate limiting, and message dispatch.
"""

import json
import logging
from typing import Optional
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query, status
from sqlalchemy.orm import Session

from app.database.database import get_db
from app.core.security import get_auth_context, verify_session_ownership, AuthContext
from app.services.onboarding_service import OnboardingService
from app.services.websocket_service import WebSocketService
from app.api.websocket.manager import manager
from app.api.websocket.schemas import (
    WSEventType,
    WSErrorCode,
    WSConnectionAck,
    WSPatientMessage,
    WSErrorEvent,
)

from app.database.database import SessionLocal

logger = logging.getLogger(__name__)

router = APIRouter(
    prefix="/api/v1/ws/onboarding",
    tags=["WebSocket Onboarding"],
)

# In-memory per-connection rate limiter (Section 26)
# Maps session_id -> list of message timestamps
WS_RATE_LIMIT_STORE = {}


def check_ws_rate_limit(session_id: str, limit: int = 30, window_sec: int = 60) -> bool:
    import time
    now = time.time()
    timestamps = WS_RATE_LIMIT_STORE.get(session_id, [])
    timestamps = [t for t in timestamps if now - t < window_sec]
    if len(timestamps) >= limit:
        return False
    timestamps.append(now)
    WS_RATE_LIMIT_STORE[session_id] = timestamps
    return True


@router.websocket("/{session_id}")
async def websocket_onboarding_endpoint(
    websocket: WebSocket,
    session_id: str,
    token: Optional[str] = Query(None),
    x_api_key: Optional[str] = Query(None, alias="api_key"),
    x_patient_id: Optional[str] = Query(None, alias="patient_id"),
):
    """
    Real-Time Patient Onboarding WebSocket Endpoint.
    WS /api/v1/ws/onboarding/{session_id}
    """
    db = SessionLocal()

    
    try:
        await _handle_websocket_lifecycle(
            websocket=websocket,
            session_id=session_id,
            db=db,
        )
    finally:
        db.close()


async def _handle_websocket_lifecycle(
    websocket: WebSocket,
    session_id: str,
    db: Session,
):
    onboarding_service = OnboardingService()
    ws_service = WebSocketService()

    await websocket.accept()

    # Extract query params directly from WebSocket request URL (Section 25)
    query_token = websocket.query_params.get("token")
    query_api_key = websocket.query_params.get("api_key")
    query_patient_id = websocket.query_params.get("patient_id")

    auth_token = query_token or query_api_key or websocket.headers.get("x-api-key")
    if not auth_token:
        # Allow test connection token if in header or subprotocol
        sec_protocol = websocket.headers.get("sec-websocket-protocol")
        if sec_protocol:
            auth_token = sec_protocol

    try:
        auth_context = get_auth_context(
            api_key=auth_token,
            credentials=None,
            x_patient_id=query_patient_id or websocket.headers.get("x-patient-id")
        )
    except Exception as e:
        logger.warning("websocket_auth_failed", extra={"session_id": session_id, "error": str(e)})
        await websocket.send_json(
            WSErrorEvent(code=WSErrorCode.UNAUTHORIZED.value, message="Authentication failed.").model_dump()
        )
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    # 2. Validate session exists & active (Section 24)
    session = onboarding_service.get_session(db, session_id)
    if not session:
        logger.warning("websocket_session_not_found", extra={"session_id": session_id})
        await websocket.send_json(
            WSErrorEvent(code=WSErrorCode.SESSION_NOT_FOUND.value, message="Session not found.").model_dump()
        )
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    # 3. Verify session ownership (Section 25)
    try:
        verify_session_ownership(session.patient_id, auth_context)
    except Exception as e:
        logger.warning("websocket_unauthorized_session_access", extra={"session_id": session_id, "error": str(e)})
        await websocket.send_json(
            WSErrorEvent(code=WSErrorCode.UNAUTHORIZED.value, message="Unauthorized session access.").model_dump()
        )
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    if session.status != "in_progress":
        logger.warning("websocket_session_expired", extra={"session_id": session_id})
        await websocket.send_json(
            WSErrorEvent(code=WSErrorCode.SESSION_EXPIRED.value, message="Session expired or completed.").model_dump()
        )
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
        return

    # 4. Register connection & send ACK (Section 3 & 4)
    await manager.connect(session_id, websocket)

    ack = WSConnectionAck(session_id=session_id, status="connected").model_dump()
    await manager.send_json(websocket, ack)

    try:
        while True:
            # 5. Receive text message
            data_text = await websocket.receive_text()

            # 6. Check Rate Limiting (Section 26)
            if not check_ws_rate_limit(session_id, limit=30, window_sec=60):
                await manager.send_json(
                    websocket,
                    WSErrorEvent(
                        code=WSErrorCode.RATE_LIMITED.value,
                        message="Please wait a moment before sending another message."
                    ).model_dump()
                )
                continue

            # 7. Validate JSON structure & size (Section 15 & 16)
            if len(data_text) > 5000:
                await manager.send_json(
                    websocket,
                    WSErrorEvent(
                        code=WSErrorCode.MESSAGE_TOO_LARGE.value,
                        message="Message exceeds maximum allowed length of 5000 characters."
                    ).model_dump()
                )
                continue

            try:
                payload = json.loads(data_text)
            except Exception:
                await manager.send_json(
                    websocket,
                    WSErrorEvent(
                        code=WSErrorCode.INVALID_MESSAGE.value,
                        message="Please send a valid JSON message."
                    ).model_dump()
                )
                continue

            if not isinstance(payload, dict) or payload.get("type") != WSEventType.PATIENT_MESSAGE.value:
                await manager.send_json(
                    websocket,
                    WSErrorEvent(
                        code=WSErrorCode.INVALID_MESSAGE.value,
                        message="Unsupported message type. Expected 'patient.message'."
                    ).model_dump()
                )
                continue

            message_id = str(payload.get("message_id") or "").strip()
            text = str(payload.get("text") or "").strip()

            if not message_id or not text:
                await manager.send_json(
                    websocket,
                    WSErrorEvent(
                        code=WSErrorCode.INVALID_MESSAGE.value,
                        message="Message must include non-empty 'message_id' and 'text'."
                    ).model_dump()
                )
                continue

            # 8. Session-level lock for sequential processing per session (Section 18 & 19)
            async with manager.get_session_lock(session_id):
                await ws_service.handle_patient_message(
                    websocket=websocket,
                    db=db,
                    session_id=session_id,
                    message_id=message_id,
                    text=text,
                )

    except WebSocketDisconnect:
        manager.disconnect(session_id, websocket)
        logger.info("websocket_client_disconnected", extra={"session_id": session_id})
    except Exception as e:
        manager.disconnect(session_id, websocket)
        logger.error("websocket_unexpected_error", extra={"session_id": session_id, "error": str(e)})