diff --git a/src-tauri/src/mcp.rs b/src-tauri/src/mcp.rs index fca8303..8104cba 100644 --- a/src-tauri/src/mcp.rs +++ b/src-tauri/src/mcp.rs @@ -62,6 +62,7 @@ .route("/sse", get(sse_handler)) .route("/messages", post(mcp_messages_handler)) .route("/llama_status", get(llama_status_handler)) + .route("/doc_count", get(doc_count_handler)) .layer(cors) .with_state(app_state); @@ -77,6 +78,15 @@ Json(serde_json::json!({ "status": status })) } +async fn doc_count_handler(State(state): State) -> impl IntoResponse { + let row = sqlx::query("SELECT COUNT(*) FROM items") + .fetch_one(&state.db_pool) + .await + .unwrap(); + let count: i64 = row.get(0); + Json(serde_json::json!({ "count": count })) +} + #[derive(Deserialize)] struct SseQuery { session_id: Option, @@ -101,17 +111,22 @@ let session_id_for_close = session_id.clone(); let sessions_for_close = state.sessions.clone(); + let mut global_rx = state.tx.subscribe(); let stream = futures::stream::unfold( - (rx, Some(endpoint_event), session_id_for_close, sessions_for_close), - |(mut rx, mut initial, sid, smap)| async move { + (rx, Some(endpoint_event), session_id_for_close, sessions_for_close, global_rx), + |(mut rx, mut initial, sid, smap, mut grx)| async move { if let Some(event) = initial.take() { - return Some((Ok(event), (rx, None, sid, smap))); + return Some((Ok(event), (rx, None, sid, smap, grx))); } tokio::select! { Some(msg) = rx.recv() => { - Some((Ok(Event::default().event("message").data(msg)), (rx, None, sid, smap))) + Some((Ok(Event::default().event("message").data(msg)), (rx, None, sid, smap, grx))) + } + Ok(msg) = grx.recv() => { + // Global notification (e.g. data update) + Some((Ok(Event::default().event("update").data(msg)), (rx, None, sid, smap, grx))) } else => { log::info!("MCP SSE Session Closed: {}", sid); @@ -265,6 +280,7 @@ .unwrap(); tx.commit().await.unwrap(); + let _ = state.tx.send("data_changed".to_string()); Some(serde_json::json!({ "content": [{ "type": "text", "text": format!("Successfully added item with ID: {}", id) }] })) } Err(e) => Some(serde_json::json!({ "error": format!("Embedding failed: {}", e) })) @@ -344,6 +360,7 @@ .unwrap(); tx.commit().await.unwrap(); + let _ = state.tx.send("data_changed".to_string()); Some(serde_json::json!({ "content": [{ "type": "text", "text": format!("Successfully updated item {}", id) }] })) } Err(e) => Some(serde_json::json!({ "error": format!("Embedding failed: {}", e) })) @@ -355,6 +372,7 @@ sqlx::query("DELETE FROM items WHERE id = ?").bind(id).execute(&mut *tx).await.unwrap(); sqlx::query("DELETE FROM vec_items WHERE id = ?").bind(id).execute(&mut *tx).await.unwrap(); tx.commit().await.unwrap(); + let _ = state.tx.send("data_changed".to_string()); Some(serde_json::json!({ "content": [{ "type": "text", "text": format!("Successfully deleted item {}", id) }] })) }, _ => Some(serde_json::json!({ "error": "Unknown tool" })), diff --git a/src/index.html b/src/index.html index 6d7282e..246eec7 100644 --- a/src/index.html +++ b/src/index.html @@ -341,19 +341,9 @@ async function updateDocCount() { try { - const res = await fetch(`${API_BASE}/messages`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - jsonrpc: "2.0", - method: "search_text", - params: { content: "", limit: 100000 }, - id: 2, - }), - }); + const res = await fetch(`${API_BASE}/doc_count`); const data = await res.json(); - const count = data.result?.content?.length || 0; - document.getElementById("doc-count-value").textContent = count; + document.getElementById("doc-count-value").textContent = data.count || "0"; } catch { document.getElementById("doc-count-value").textContent = "0"; } @@ -439,6 +429,16 @@ updateLlamaStatus(); updateDocCount(); + // Setup SSE for real-time updates + const eventSource = new EventSource(`${API_BASE}/sse`); + eventSource.addEventListener("update", (e) => { + console.log("Data update received:", e.data); + updateDocCount(); + }); + eventSource.onerror = () => { + console.warn("SSE connection lost. Reconnecting..."); + }; + // Enter key search document.getElementById("query").addEventListener("keypress", (e) => { if (e.key === "Enter") search();