qt-chat-app

sse_client.rs (8068B)

---

 //! Server-Sent Events (SSE) client with auto-reconnection
 //!
 //! Handles real-time message subscriptions from the TrailBase API.
 
 use crate::types::{SseEventWrapper, BASE_URL, MAX_RECONNECT_DELAY_MS};
 use futures::StreamExt;
 use reqwest_eventsource::{Event, EventSource};
 use serde_json::Value;
 use std::collections::HashSet;
 use std::time::Duration;
 use tokio::sync::mpsc;
 
 const MAX_SEEN_IDS: usize = 100;
 
 /// Commands that can be sent to the SSE task
 #[derive(Debug)]
 pub enum SseCommand {
     /// Stop the SSE connection
     Stop,
 }
 
 /// Events emitted by the SSE client
 #[derive(Debug)]
 pub enum SseEvent {
     /// Successfully connected
     Connected { chat_id: i32 },
     /// Received a message
     Message { data: Value },
     /// Connection was disconnected
     Disconnected { chat_id: i32, reason: String },
     /// An error occurred
     Error { message: String },
 }
 
 /// Parse an SSE event and extract the message data
 fn parse_sse_message(data: &str, expected_chat_id: Option<i32>) -> Option<Value> {
     let json: Value = serde_json::from_str(data).ok()?;
 
     // TrailBase wraps data in "Insert" key
     let wrapper: SseEventWrapper = serde_json::from_value(json.clone()).ok()?;
     let message = wrapper.insert.or(Some(json))?;
 
     // Validate chat_id if we're subscribed to a specific chat
     if let Some(expected) = expected_chat_id {
         let msg_chat_id = message.get("chat_id")?.as_i64()? as i32;
         if msg_chat_id != expected {
             tracing::debug!(
                 "Message for wrong chat: got {}, expected {}",
                 msg_chat_id,
                 expected
             );
             return None;
         }
     }
 
     // Validate message has an ID
     if message.get("id").is_none() {
         tracing::warn!("SSE: Message without valid ID");
         return None;
     }
 
     Some(message)
 }
 
 /// Start an SSE subscription for a single chat
 pub async fn subscribe_to_chat(
     auth_token: String,
     chat_id: i32,
     event_tx: mpsc::Sender<SseEvent>,
     command_rx: mpsc::Receiver<SseCommand>,
 ) {
     let url = format!(
         "{}/api/records/v1/message/subscribe/*?filter[chat_id][$eq]={}",
         BASE_URL, chat_id
     );
 
     run_sse_loop(url, auth_token, Some(chat_id), event_tx, command_rx).await;
 }
 
 /// Start an SSE subscription for multiple chats
 pub async fn subscribe_to_chats(
     auth_token: String,
     chat_ids: Vec<i32>,
     event_tx: mpsc::Sender<SseEvent>,
     command_rx: mpsc::Receiver<SseCommand>,
 ) {
     if chat_ids.is_empty() {
         return;
     }
 
     let chat_id_list = chat_ids
         .iter()
         .map(|id| id.to_string())
         .collect::<Vec<_>>()
         .join(",");
 
     let url = format!(
         "{}/api/records/v1/message/subscribe/*?chat_id={}",
         BASE_URL, chat_id_list
     );
 
     // -1 indicates multi-chat subscription
     run_sse_loop(url, auth_token, None, event_tx, command_rx).await;
 }
 
 /// Main SSE event loop with reconnection logic and deduplication
 async fn run_sse_loop(
     url: String,
     auth_token: String,
     expected_chat_id: Option<i32>,
     event_tx: mpsc::Sender<SseEvent>,
     mut command_rx: mpsc::Receiver<SseCommand>,
 ) {
     let chat_id_for_signals = expected_chat_id.unwrap_or(-1);
     let mut reconnect_attempts = 0u32;
     let mut seen_ids: HashSet<i64> = HashSet::new();
 
     loop {
         // Check for stop command before connecting
         if let Ok(SseCommand::Stop) = command_rx.try_recv() {
             let _ = event_tx
                 .send(SseEvent::Disconnected {
                     chat_id: chat_id_for_signals,
                     reason: "User initiated".to_string(),
                 })
                 .await;
             return;
         }
 
         // Create EventSource with proper headers
         let request = reqwest::Client::new()
             .get(&url)
             .header("Accept", "text/event-stream")
             .header("Cache-Control", "no-cache")
             .bearer_auth(&auth_token);
 
         let mut es = EventSource::new(request).expect("Failed to create EventSource");
 
         // Notify connected
         let _ = event_tx
             .send(SseEvent::Connected {
                 chat_id: chat_id_for_signals,
             })
             .await;
 
         reconnect_attempts = 0; // Reset on successful connect
 
         // Process events
         loop {
             tokio::select! {
                 // Check for stop command
                 cmd = command_rx.recv() => {
                     if matches!(cmd, Some(SseCommand::Stop) | None) {
                         es.close();
                         let _ = event_tx.send(SseEvent::Disconnected {
                             chat_id: chat_id_for_signals,
                             reason: "User initiated".to_string(),
                         }).await;
                         return;
                     }
                 }
 
                 // Process SSE events
                 event = es.next() => {
                     match event {
                         Some(Ok(Event::Open)) => {
                             tracing::info!("SSE connection opened for chat {:?}", expected_chat_id);
                         }
                         Some(Ok(Event::Message(message))) => {
                             if let Some(data) = parse_sse_message(&message.data, expected_chat_id) {
                                 // Deduplicate by message ID
                                 if let Some(msg_id) = data.get("id").and_then(|v| v.as_i64()) {
                                     if seen_ids.contains(&msg_id) {
                                         continue;
                                     }
                                     seen_ids.insert(msg_id);
 
                                     // Keep bounded
                                     if seen_ids.len() > MAX_SEEN_IDS {
                                         // Remove oldest entries (arbitrary since HashSet is unordered)
                                         let to_remove: Vec<_> = seen_ids.iter().take(MAX_SEEN_IDS / 2).cloned().collect();
                                         for id in to_remove {
                                             seen_ids.remove(&id);
                                         }
                                     }
                                 }
 
                                 let _ = event_tx.send(SseEvent::Message { data }).await;
                             }
                         }
                         Some(Err(e)) => {
                             let error_msg = e.to_string();
                             tracing::warn!("SSE error: {}", error_msg);
                             let _ = event_tx.send(SseEvent::Error { message: error_msg }).await;
                             break; // Exit inner loop to reconnect
                         }
                         None => {
                             // Stream ended
                             tracing::info!("SSE stream ended");
                             break;
                         }
                     }
                 }
             }
         }
 
         // Connection lost - notify and schedule reconnect
         let _ = event_tx
             .send(SseEvent::Disconnected {
                 chat_id: chat_id_for_signals,
                 reason: "Connection closed by server".to_string(),
             })
             .await;
 
         // Exponential backoff: 1s, 2s, 4s, 8s, 16s, max 30s
         let delay = std::cmp::min(
             1000u64 * (1u64 << reconnect_attempts),
             MAX_RECONNECT_DELAY_MS,
         );
         reconnect_attempts += 1;
 
         tracing::info!(
             "Scheduling SSE reconnect in {}ms (attempt {})",
             delay,
             reconnect_attempts
         );
 
         // Wait for delay or stop command
         tokio::select! {
             _ = tokio::time::sleep(Duration::from_millis(delay)) => {
                 // Continue to reconnect
             }
             cmd = command_rx.recv() => {
                 if matches!(cmd, Some(SseCommand::Stop) | None) {
                     return;
                 }
             }
         }
     }
 }