"""WebSocket connection manager with run-level subscriptions."""

from __future__ import annotations

import asyncio

import logging

from typing import Any

from fastapi import WebSocket

from uipath.dev.models.data import (

ChatData,

InterruptData,

LogData,

StateData,

TraceData,

)

from uipath.dev.models.execution import ExecutionRun

from uipath.dev.server.serializers import (

serialize_chat,

serialize_interrupt,

serialize_log,

serialize_run,

serialize_state,

serialize_trace,

)

from uipath.dev.server.ws.protocol import ServerEvent, server_message

logger = logging.getLogger(__name__)

_SENTINEL: dict[str, Any] = {} # Unique object used to signal queue shutdown

class ConnectionManager:

"""Manages WebSocket connections and run-level subscriptions."""

def __init__(self) -> None:

"""Initialize the connection manager."""

self._connections: set[WebSocket] = set()

self._subscriptions: dict[str, set[WebSocket]] = {}

self._queues: dict[int, asyncio.Queue[dict[str, Any]]] = {}

self._send_tasks: dict[int, asyncio.Task[None]] = {}

self._loop: asyncio.AbstractEventLoop | None = None

def _get_loop(self) -> asyncio.AbstractEventLoop:

"""Get or cache the running event loop."""

if self._loop is None or self._loop.is_closed():

try:

self._loop = asyncio.get_running_loop()

except RuntimeError:

self._loop = asyncio.get_event_loop()

return self._loop

async def connect(self, websocket: WebSocket) -> None:

"""Accept a new WebSocket connection."""

await websocket.accept()

self._connections.add(websocket)

queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

ws_id = id(websocket)

self._queues[ws_id] = queue

self._send_tasks[ws_id] = asyncio.create_task(self._sender(websocket, queue))

async def _sender(

self, ws: WebSocket, queue: asyncio.Queue[dict[str, Any]]

) -> None:

"""Consume messages from queue and send them serially."""

while True:

message = await queue.get()

if message is _SENTINEL:

break

try:

await ws.send_json(message)

except Exception:

self.disconnect(ws)

break

def _enqueue(self, ws: WebSocket, message: dict[str, Any]) -> None:

"""Put a message on a WebSocket's send queue (non-blocking)."""

queue = self._queues.get(id(ws))

if queue is not None:

queue.put_nowait(message)

def disconnect(self, websocket: WebSocket) -> None:

"""Remove a WebSocket connection and all its subscriptions."""

self._connections.discard(websocket)

for run_id in list(self._subscriptions):

self._subscriptions[run_id].discard(websocket)

if not self._subscriptions[run_id]:

del self._subscriptions[run_id]

ws_id = id(websocket)

queue = self._queues.pop(ws_id, None)

if queue is not None:

try:

queue.put_nowait(_SENTINEL)

except Exception:

pass

task = self._send_tasks.pop(ws_id, None)

if task is not None and not task.done():

task.cancel()

async def disconnect_all(self) -> None:

"""Close all WebSocket connections gracefully."""

for ws in list(self._connections):

try:

await ws.close()

except Exception as e:

logger.debug(f"Error closing WebSocket: {e}")

self.disconnect(ws)

self._connections.clear()

self._subscriptions.clear()

def subscribe(self, websocket: WebSocket, run_id: str) -> None:

"""Subscribe a WebSocket to a run's events."""

if run_id not in self._subscriptions:

self._subscriptions[run_id] = set()

self._subscriptions[run_id].add(websocket)

def unsubscribe(self, websocket: WebSocket, run_id: str) -> None:

"""Unsubscribe a WebSocket from a run's events."""

if run_id in self._subscriptions:

self._subscriptions[run_id].discard(websocket)

if not self._subscriptions[run_id]:

del self._subscriptions[run_id]

def remove_run_subscriptions(self, run_id: str) -> None:

"""Remove all subscriptions for a run."""

self._subscriptions.pop(run_id, None)

def broadcast_run_updated(self, run: ExecutionRun) -> None:

"""Broadcast a run update to all subscribers (safe from sync context)."""

msg = server_message(ServerEvent.RUN_UPDATED, serialize_run(run))

self._schedule_broadcast(run.id, msg)

def broadcast_log(self, log_data: LogData) -> None:

"""Broadcast a log entry to run subscribers."""

msg = server_message(ServerEvent.LOG, serialize_log(log_data))

self._schedule_broadcast(log_data.run_id, msg)

def broadcast_trace(self, trace_data: TraceData) -> None:

"""Broadcast a trace span to run subscribers."""

msg = server_message(ServerEvent.TRACE, serialize_trace(trace_data))

self._schedule_broadcast(trace_data.run_id, msg)

def broadcast_chat(self, chat_data: ChatData) -> None:

"""Broadcast a chat message to run subscribers."""

msg = server_message(ServerEvent.CHAT, serialize_chat(chat_data))

self._schedule_broadcast(chat_data.run_id, msg)

def broadcast_interrupt(self, interrupt_data: InterruptData) -> None:

"""Broadcast a chat interrupt to run subscribers."""

msg = server_message(

ServerEvent.CHAT_INTERRUPT, serialize_interrupt(interrupt_data)

)

self._schedule_broadcast(interrupt_data.run_id, msg)

def broadcast_state(self, state_data: StateData) -> None:

"""Broadcast a state transition to run subscribers."""

msg = server_message(ServerEvent.STATE, serialize_state(state_data))

self._schedule_broadcast(state_data.run_id, msg)

def broadcast_reload(self, changed_files: list[str]) -> None:

"""Broadcast a reload event to all connected clients."""

msg = server_message(ServerEvent.RELOAD, {"files": changed_files})

for ws in self._connections:

self._enqueue(ws, msg)

def _schedule_broadcast(self, run_id: str, message: dict[str, Any]) -> None:

"""Enqueue a message for all subscribers of a run."""

subscribers = self._subscriptions.get(run_id)

if not subscribers:

return

for ws in subscribers:

self._enqueue(ws, message)