GopherGate/src/dashboard/websocket.rs

use axum::{
    extract::{
        State,
        ws::{Message, WebSocket, WebSocketUpgrade},
    },
    response::IntoResponse,
};
use serde_json;
use tracing::info;

use super::DashboardState;

// WebSocket handler
pub(super) async fn handle_websocket(ws: WebSocketUpgrade, State(state): State<DashboardState>) -> impl IntoResponse {
    ws.on_upgrade(|socket| handle_websocket_connection(socket, state))
}

pub(super) async fn handle_websocket_connection(mut socket: WebSocket, state: DashboardState) {
    info!("WebSocket connection established");

    // Subscribe to events from the global bus
    let mut rx = state.app_state.dashboard_tx.subscribe();

    // Send initial connection message
    let _ = socket
        .send(Message::Text(
            serde_json::json!({
                "type": "connected",
                "message": "Connected to LLM Proxy Dashboard"
            })
            .to_string()
            .into(),
        ))
        .await;

    // Handle incoming messages and broadcast events
    loop {
        tokio::select! {
            // Receive broadcast events
            Ok(event) = rx.recv() => {
                let Ok(json_str) = serde_json::to_string(&event) else {
                    continue;
                };
                let message = Message::Text(json_str.into());
                if socket.send(message).await.is_err() {
                    break;
                }
            }

            // Receive WebSocket messages
            result = socket.recv() => {
                match result {
                    Some(Ok(Message::Text(text))) => {
                        handle_websocket_message(&text, &state).await;
                    }
                    _ => break,
                }
            }
        }
    }

    info!("WebSocket connection closed");
}

pub(super) async fn handle_websocket_message(text: &str, state: &DashboardState) {
    // Parse and handle WebSocket messages
    if let Ok(data) = serde_json::from_str::<serde_json::Value>(text)
        && data.get("type").and_then(|v| v.as_str()) == Some("ping")
    {
        let _ = state.app_state.dashboard_tx.send(serde_json::json!({
            "type": "pong",
            "payload": {}
        }));
    }
}