add chat function

#

src/handlers/ws.rs

@@ -0,0 +1,74 @@
+use axum::{
+    extract::{
+        ws::{Message, WebSocket, WebSocketUpgrade},
+        Path, State,
+    },
+    response::IntoResponse,
+};
+use std::sync::Arc;
+
+use crate::state::AppState;
+
+/// WebSocket for customer — pushes notifications when agent replies
+pub async fn customer_ws(
+    ws: WebSocketUpgrade,
+    Path(session_id): Path<String>,
+    State(state): State<Arc<AppState>>,
+) -> impl IntoResponse {
+    ws.on_upgrade(move |socket| handle_customer_socket(socket, session_id, state))
+}
+
+async fn handle_customer_socket(mut socket: WebSocket, session_id: String, state: Arc<AppState>) {
+    let tx = state.get_or_create_notifier(&session_id);
+    let mut rx = tx.subscribe();
+
+    loop {
+        tokio::select! {
+            // Forward broadcast events to the customer browser
+            Ok(sid) = rx.recv() => {
+                if sid == session_id {
+                    // Signal HTMX to re-poll messages
+                    if socket.send(Message::Text("refresh".into())).await.is_err() {
+                        break;
+                    }
+                }
+            }
+            // Handle ping/close from client
+            msg = socket.recv() => {
+                match msg {
+                    Some(Ok(Message::Close(_))) | None => break,
+                    _ => {}
+                }
+            }
+        }
+    }
+}
+
+/// WebSocket for agent dashboard — pushes notifications when any session updates
+pub async fn agent_ws(
+    ws: WebSocketUpgrade,
+    Path(_agent_id): Path<String>,
+    State(state): State<Arc<AppState>>,
+) -> impl IntoResponse {
+    ws.on_upgrade(move |socket| handle_agent_socket(socket, state))
+}
+
+async fn handle_agent_socket(mut socket: WebSocket, state: Arc<AppState>) {
+    let mut rx = state.agent_notifier.subscribe();
+
+    loop {
+        tokio::select! {
+            Ok(session_id) = rx.recv() => {
+                if socket.send(Message::Text(session_id.into())).await.is_err() {
+                    break;
+                }
+            }
+            msg = socket.recv() => {
+                match msg {
+                    Some(Ok(Message::Close(_))) | None => break,
+                    _ => {}
+                }
+            }
+        }
+    }
+}