feat: complete production hardening implementation, specialty canonicalization, google-genai migration, voice reliability, sqlite concurrency, and test suite consolidation
This commit is contained in:
@@ -0,0 +1,225 @@
|
||||
"""
|
||||
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)})
|
||||
Reference in New Issue
Block a user