mentedb_server/
websocket.rs

1//! WebSocket handler for live cognition streaming.
2
3use std::sync::Arc;
4
5use axum::extract::ws::{Message, WebSocket};
6use axum::extract::{State, WebSocketUpgrade};
7use axum::response::IntoResponse;
8use serde::{Deserialize, Serialize};
9use serde_json::json;
10use tracing::{info, warn};
11
12use crate::state::AppState;
13
14/// Inbound message from a WebSocket client.
15#[derive(Debug, Deserialize)]
16struct WsInbound {
17    #[serde(rename = "type")]
18    msg_type: String,
19    data: serde_json::Value,
20}
21
22/// Outbound message pushed to a WebSocket client.
23#[derive(Debug, Serialize)]
24struct WsOutbound {
25    #[serde(rename = "type")]
26    msg_type: String,
27    #[serde(skip_serializing_if = "Option::is_none")]
28    alert_type: Option<String>,
29    data: serde_json::Value,
30}
31
32/// Axum handler that upgrades an HTTP request to a WebSocket connection.
33pub async fn ws_handler(
34    ws: WebSocketUpgrade,
35    State(state): State<Arc<AppState>>,
36) -> impl IntoResponse {
37    ws.on_upgrade(move |socket| handle_socket(socket, state))
38}
39
40async fn handle_socket(mut socket: WebSocket, state: Arc<AppState>) {
41    info!("websocket client connected");
42
43    while let Some(msg) = socket.recv().await {
44        let msg = match msg {
45            Ok(m) => m,
46            Err(e) => {
47                warn!("websocket receive error: {e}");
48                break;
49            }
50        };
51
52        match msg {
53            Message::Text(text) => {
54                let inbound: WsInbound = match serde_json::from_str(&text) {
55                    Ok(v) => v,
56                    Err(e) => {
57                        let err = json!({ "type": "error", "data": format!("invalid JSON: {e}") });
58                        if socket
59                            .send(Message::Text(err.to_string().into()))
60                            .await
61                            .is_err()
62                        {
63                            break;
64                        }
65                        continue;
66                    }
67                };
68
69                let response = process_message(&inbound, &state).await;
70                let payload = serde_json::to_string(&response).unwrap_or_default();
71                if socket.send(Message::Text(payload.into())).await.is_err() {
72                    break;
73                }
74            }
75            Message::Close(_) => break,
76            _ => {}
77        }
78    }
79
80    info!("websocket client disconnected");
81}
82
83async fn process_message(inbound: &WsInbound, state: &AppState) -> WsOutbound {
84    match inbound.msg_type.as_str() {
85        "query" => {
86            let query = inbound.data.as_str().unwrap_or("");
87            let db = &*state.db;
88            match db.recall(query) {
89                Ok(window) => {
90                    let memory_count: usize = window.blocks.iter().map(|b| b.memories.len()).sum();
91                    WsOutbound {
92                        msg_type: "result".into(),
93                        alert_type: None,
94                        data: json!({
95                            "context": window.format,
96                            "total_tokens": window.total_tokens,
97                            "memory_count": memory_count,
98                        }),
99                    }
100                }
101                Err(e) => WsOutbound {
102                    msg_type: "error".into(),
103                    alert_type: None,
104                    data: json!(format!("recall failed: {e}")),
105                },
106            }
107        }
108        "token" => {
109            // Echo acknowledgment for token-type messages.
110            WsOutbound {
111                msg_type: "ack".into(),
112                alert_type: None,
113                data: inbound.data.clone(),
114            }
115        }
116        _ => WsOutbound {
117            msg_type: "error".into(),
118            alert_type: None,
119            data: json!(format!("unknown message type: {}", inbound.msg_type)),
120        },
121    }
122}
mentedb_server/websocket.rs

mentedb_server/
websocket.rs
