from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status, Query from datetime import datetime, timezone import json from app.repositories.session_repository import ( get_session_by_id, update_session_status, update_session_connection_status ) from app.repositories.conversation_repository import get_conversation_by_session_id from app.database import get_database from app.services.presence import ( active_sessions, add_active_session, remove_active_session, add_agent_connection, remove_agent_connection, broadcast_to_agents, handle_agent_disconnect, cancel_grace_period_timer, start_grace_period_timer ) from app.services.video_signaling import ( handle_visitor_signaling, handle_agent_signaling, terminate_video_call ) router = APIRouter(tags=["websocket"]) @router.websocket("/ws/session/{session_id}") async def session_presence_websocket(websocket: WebSocket, session_id: str): db = get_database() # Step 1: Validate session existence in MongoDB session = await get_session_by_id(db, session_id) if not session: await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason=f"Session '{session_id}' not found") return # Step 2: Accept connection and add/update active_sessions await websocket.accept() site_id = session.get("siteId", "") user_data = None if session.get("userId") or session.get("userName") or session.get("userEmail"): user_data = { "id": session.get("userId"), "name": session.get("userName"), "email": session.get("userEmail") } ip_address = session.get("ipAddress") country = session.get("country", "Unknown") session_status = session.get("status", "RINGING") assigned_agent_id = session.get("assignedAgentId") assigned_at_val = session.get("assignedAt") assigned_at_str = assigned_at_val.isoformat() if hasattr(assigned_at_val, "isoformat") else (str(assigned_at_val) if assigned_at_val else None) conv = await get_conversation_by_session_id(db, session_id) if db is not None else None conv_id = conv.get("conversationId") if conv else None is_abandoned = (session_status == "ABANDONED") or (conv is not None and conv.get("status") == "ABANDONED") existing_in_mem = active_sessions.get(session_id) is_reconnect = existing_in_mem is not None and existing_in_mem.get("connectionStatus") == "disconnected" if is_abandoned: # Re-ring: ABANDONED -> RINGING session_status = "RINGING" assigned_agent_id = None assigned_at_str = None if db is not None: await update_session_status(db, session_id, "RINGING") await update_session_connection_status(db, session_id, "online", disconnected_at=None) await db["sessions"].update_one( {"sessionId": session_id}, {"$set": {"assignedAgentId": None, "assignedAt": None}} ) await db["conversations"].update_one( {"sessionId": session_id}, {"$set": {"status": "open"}, "$unset": {"assignedAgentId": ""}} ) cancel_grace_period_timer(session_id) active_info = add_active_session( session_id=session_id, site_id=site_id, websocket=websocket, user_data=user_data, ip_address=ip_address, country=country, status_val="RINGING", assigned_agent_id=None, assigned_at=None, connection_status="online", disconnected_at=None ) await broadcast_to_agents({ "type": "visitor_connected", "sessionId": session_id, "conversationId": conv_id, "siteId": site_id, "status": "RINGING", "connectionStatus": "online", "assignedAgentId": None, "user": user_data, "ipAddress": ip_address, "country": country, "connectedAt": active_info["connectedAt"], "isReturningVisitor": True }) await websocket.send_text(json.dumps({ "type": "status_changed", "status": "RINGING", "message": "Waiting for an agent to connect..." })) elif is_reconnect: cancel_grace_period_timer(session_id) existing_in_mem["websocket"] = websocket existing_in_mem["connectionStatus"] = "online" existing_in_mem["disconnectedAt"] = None if db is not None: await update_session_connection_status(db, session_id, "online", disconnected_at=None) await broadcast_to_agents({ "type": "visitor_reconnected", "sessionId": session_id, "conversationId": conv_id, "siteId": site_id, "status": session_status, "connectionStatus": "online", "assignedAgentId": assigned_agent_id, "user": user_data, "ipAddress": ip_address, "country": country, "connectedAt": existing_in_mem.get("connectedAt") }) else: cancel_grace_period_timer(session_id) active_info = add_active_session( session_id=session_id, site_id=site_id, websocket=websocket, user_data=user_data, ip_address=ip_address, country=country, status_val=session_status, assigned_agent_id=assigned_agent_id, assigned_at=assigned_at_str, connection_status="online", disconnected_at=None ) if db is not None: await update_session_connection_status(db, session_id, "online", disconnected_at=None) await broadcast_to_agents({ "type": "visitor_connected", "sessionId": session_id, "conversationId": conv_id, "siteId": site_id, "status": session_status, "connectionStatus": "online", "assignedAgentId": assigned_agent_id, "user": user_data, "ipAddress": ip_address, "country": country, "connectedAt": active_info["connectedAt"] }) # Step 4: Handle incoming signaling messages & keep connection alive try: while True: raw_msg = await websocket.receive_text() try: data = json.loads(raw_msg) if isinstance(data, dict): await handle_visitor_signaling(session_id, data) except json.JSONDecodeError: pass except WebSocketDisconnect: pass except Exception: pass finally: await terminate_video_call(session_id, "visitor_disconnected") current_mem = active_sessions.get(session_id) # Check if a replacement WebSocket connection has already taken over if current_mem and current_mem.get("websocket") is not None and current_mem.get("websocket") != websocket: return current_status = current_mem.get("status") if current_mem else session_status if current_status == "IN_CHAT": now_dt = datetime.now(timezone.utc) now_iso = now_dt.isoformat() if current_mem: current_mem["connectionStatus"] = "disconnected" current_mem["disconnectedAt"] = now_iso current_mem["websocket"] = None if db is not None: await update_session_connection_status(db, session_id, "disconnected", disconnected_at=now_dt) await broadcast_to_agents({ "type": "visitor_disconnected", "sessionId": session_id, "siteId": site_id, "connectionStatus": "disconnected", "disconnectedAt": now_iso }) await start_grace_period_timer(session_id, site_id, grace_seconds=30.0) else: remove_active_session(session_id) await broadcast_to_agents({ "type": "visitor_disconnected", "sessionId": session_id, "siteId": site_id }) @router.websocket("/ws/agent") async def agent_dashboard_websocket(websocket: WebSocket, agentId: str | None = Query(default=None)): """WebSocket endpoint for the Agent Dashboard live feed.""" db = get_database() await websocket.accept() agent_id = agentId or "agent_1" add_agent_connection(agent_id, websocket) try: while True: raw_msg = await websocket.receive_text() try: data = json.loads(raw_msg) if isinstance(data, dict): await handle_agent_signaling(agent_id, data) except json.JSONDecodeError: pass except WebSocketDisconnect: pass except Exception: pass finally: remove_agent_connection(agent_id, websocket) await handle_agent_disconnect(db, agent_id)