miyabi_a2a/http/
websocket.rs

1//! WebSocket support for real-time dashboard updates
2
3use axum::{
4    extract::{
5        ws::{Message, WebSocket, WebSocketUpgrade},
6        State,
7    },
8    response::Response,
9};
10use futures::{sink::SinkExt, stream::StreamExt};
11use serde_json;
12use std::sync::Arc;
13use tokio::sync::broadcast;
14use tokio::time::{interval, Duration};
15use tracing::{debug, error, info};
16
17use super::real_data::{fetch_real_agents, fetch_real_system_status};
18use super::server::AppState;
19
20/// Dashboard update message
21#[derive(Debug, Clone, serde::Serialize)]
22#[serde(tag = "type", rename_all = "lowercase")]
23pub enum DashboardUpdate {
24    Agents { agents: Vec<super::routes::Agent> },
25    SystemStatus { status: super::routes::SystemStatus },
26    Error { error: ErrorInfo },
27    TaskRetry { event: TaskRetryEvent },
28    TaskCancel { event: TaskCancelEvent },
29    Ping,
30}
31
32/// Task retry event
33#[derive(Debug, Clone, serde::Serialize)]
34pub struct TaskRetryEvent {
35    /// Task ID being retried
36    pub task_id: String,
37    /// Current retry attempt number (after increment)
38    pub retry_count: u32,
39    /// Reason for retry (if provided)
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub reason: Option<String>,
42    /// Next retry timestamp (exponential backoff)
43    #[serde(skip_serializing_if = "Option::is_none")]
44    pub next_retry_at: Option<chrono::DateTime<chrono::Utc>>,
45    /// Timestamp when retry was triggered
46    pub timestamp: chrono::DateTime<chrono::Utc>,
47}
48
49/// Task cancel event
50#[derive(Debug, Clone, serde::Serialize)]
51pub struct TaskCancelEvent {
52    /// Task ID being cancelled
53    pub task_id: String,
54    /// Reason for cancellation
55    pub reason: String,
56    /// Timestamp when cancellation was triggered
57    pub timestamp: chrono::DateTime<chrono::Utc>,
58}
59
60/// Error information for WebSocket broadcasting
61#[derive(Debug, Clone, serde::Serialize)]
62pub struct ErrorInfo {
63    /// Unique error ID
64    pub id: String,
65    /// Associated task ID (if any)
66    #[serde(skip_serializing_if = "Option::is_none")]
67    pub task_id: Option<String>,
68    /// Associated agent ID (if any)
69    #[serde(skip_serializing_if = "Option::is_none")]
70    pub agent_id: Option<String>,
71    /// Associated agent name (if any)
72    #[serde(skip_serializing_if = "Option::is_none")]
73    pub agent_name: Option<String>,
74    /// Error message
75    pub message: String,
76    /// Stack trace (if available)
77    #[serde(skip_serializing_if = "Option::is_none")]
78    pub stack_trace: Option<String>,
79    /// Timestamp when error occurred
80    pub timestamp: chrono::DateTime<chrono::Utc>,
81    /// Error severity level
82    pub severity: ErrorSeverity,
83    /// Whether this error can be retried
84    pub is_retryable: bool,
85}
86
87/// Error severity levels
88#[derive(Debug, Clone, Copy, serde::Serialize)]
89#[serde(rename_all = "lowercase")]
90pub enum ErrorSeverity {
91    /// Critical error - system failure
92    Critical,
93    /// High severity - major functionality broken
94    High,
95    /// Medium severity - some functionality impaired
96    Medium,
97    /// Low severity - minor issue
98    Low,
99}
100
101/// Shared state for WebSocket broadcasting
102#[derive(Clone)]
103pub struct WsState {
104    /// Tx
105    pub tx: broadcast::Sender<DashboardUpdate>,
106}
107
108impl Default for WsState {
109    fn default() -> Self {
110        Self::new()
111    }
112}
113
114impl WsState {
115    pub fn new() -> Self {
116        let (tx, _rx) = broadcast::channel(100);
117        Self { tx }
118    }
119}
120
121/// WebSocket upgrade handler
122pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<AppState>) -> Response {
123    ws.on_upgrade(move |socket| handle_socket(socket, state.ws_state))
124}
125
126/// Handle WebSocket connection
127async fn handle_socket(socket: WebSocket, state: Arc<WsState>) {
128    let (mut sender, mut receiver) = socket.split();
129    let mut rx = state.tx.subscribe();
130
131    info!("WebSocket client connected");
132
133    // Send initial data with timeout
134    let send_result =
135        tokio::time::timeout(Duration::from_secs(5), send_initial_data(&mut sender)).await;
136
137    match send_result {
138        Ok(Ok(())) => {
139            debug!("Initial data sent successfully");
140        },
141        Ok(Err(e)) => {
142            error!("Failed to send initial data: {}", e);
143            // Don't return - continue with broadcast updates
144        },
145        Err(_) => {
146            error!("Timeout while sending initial data");
147            // Don't return - continue with broadcast updates
148        },
149    }
150
151    // Spawn task to listen for broadcasts
152    let mut send_task = tokio::spawn(async move {
153        while let Ok(msg) = rx.recv().await {
154            let json = match serde_json::to_string(&msg) {
155                Ok(json) => json,
156                Err(e) => {
157                    error!("Failed to serialize message: {}", e);
158                    continue;
159                },
160            };
161
162            if sender.send(Message::Text(json)).await.is_err() {
163                break;
164            }
165        }
166    });
167
168    // Handle incoming messages (for ping/pong)
169    let mut recv_task = tokio::spawn(async move {
170        while let Some(Ok(msg)) = receiver.next().await {
171            match msg {
172                Message::Text(text) => {
173                    debug!("Received text message: {}", text);
174                },
175                Message::Close(_) => {
176                    info!("WebSocket client disconnected");
177                    break;
178                },
179                _ => {},
180            }
181        }
182    });
183
184    // Wait for either task to finish
185    tokio::select! {
186        _ = (&mut send_task) => {
187            recv_task.abort();
188        }
189        _ = (&mut recv_task) => {
190            send_task.abort();
191        }
192    }
193
194    info!("WebSocket connection closed");
195}
196
197/// Send initial data to newly connected client
198async fn send_initial_data<S>(sender: &mut S) -> Result<(), axum::Error>
199where
200    S: SinkExt<Message> + Unpin,
201    S::Error: std::error::Error + Send + Sync + 'static,
202{
203    // Send agents data with timeout
204    let agents_future = tokio::time::timeout(Duration::from_secs(3), fetch_real_agents());
205
206    match agents_future.await {
207        Ok(Ok(agents)) => {
208            info!("📊 Sending {} agents data", agents.len());
209            if !agents.is_empty() {
210                info!("📊 First agent: {:?}", agents[0]);
211                let task_counts: Vec<_> =
212                    agents.iter().map(|a| format!("{}:{}", a.name, a.tasks)).collect();
213                info!("📊 Agent task counts: {}", task_counts.join(", "));
214            }
215            let msg = DashboardUpdate::Agents { agents };
216            let json = serde_json::to_string(&msg).unwrap();
217            // Safely truncate at character boundary
218            let truncated = json.chars().take(200).collect::<String>();
219            info!("📤 WebSocket sending JSON (first 200 chars): {}...", truncated);
220            if let Err(e) = sender.send(Message::Text(json)).await {
221                debug!("Failed to send agents data (client may have disconnected): {}", e);
222                return Err(axum::Error::new(e));
223            }
224        },
225        Ok(Err(e)) => {
226            error!("Failed to fetch agents: {}", e);
227            // Continue to send other data
228        },
229        Err(_) => {
230            error!("Timeout fetching agents");
231            // Continue to send other data
232        },
233    }
234
235    // Send system status with timeout
236    let status_future = tokio::time::timeout(Duration::from_secs(3), fetch_real_system_status());
237
238    match status_future.await {
239        Ok(Ok(status)) => {
240            let msg = DashboardUpdate::SystemStatus { status };
241            let json = serde_json::to_string(&msg).unwrap();
242            if let Err(e) = sender.send(Message::Text(json)).await {
243                debug!("Failed to send system status (client may have disconnected): {}", e);
244                return Err(axum::Error::new(e));
245            }
246        },
247        Ok(Err(e)) => {
248            error!("Failed to fetch system status: {}", e);
249            // Continue anyway
250        },
251        Err(_) => {
252            error!("Timeout fetching system status");
253            // Continue anyway
254        },
255    }
256
257    Ok(())
258}
259
260/// Background task to periodically fetch data and broadcast updates
261pub async fn broadcast_updates(state: Arc<WsState>) {
262    let mut interval = interval(Duration::from_secs(10));
263
264    loop {
265        interval.tick().await;
266
267        // Fetch and broadcast agents
268        match fetch_real_agents().await {
269            Ok(agents) => {
270                let msg = DashboardUpdate::Agents { agents };
271                let _ = state.tx.send(msg);
272            },
273            Err(e) => {
274                error!("Failed to fetch agents for broadcast: {}", e);
275            },
276        }
277
278        // Fetch and broadcast system status
279        match fetch_real_system_status().await {
280            Ok(status) => {
281                let msg = DashboardUpdate::SystemStatus { status };
282                let _ = state.tx.send(msg);
283            },
284            Err(e) => {
285                error!("Failed to fetch system status for broadcast: {}", e);
286            },
287        }
288    }
289}
miyabi_a2a/http/websocket.rs

miyabi_a2a/http/
websocket.rs
