From 924b6d09fae2058e3b8244673757e6ffde5c18a5 Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Mon, 5 Oct 2026 21:44:10 +0100 Subject: [PATCH] refactor: address 5-pass audit findings for antipatterns, bottlenecks, memory efficiency, and LLM handlers --- Cargo.lock | 9 - Cargo.toml | 1 - justfile | 4 +- mcp-memory-nvim/Cargo.toml | 10 - mcp-stdio/Cargo.toml | 2 +- mcp-stdio/src/lib.rs | 12 +- {mcp-memory-nvim => nvim-core}/build.rs | 0 .../src/bin/mcp-memory-nvim.rs | 0 nvim-core/src/lib.rs | 25 +- server/src/api/setup.rs | 2 +- server/src/bin/export_schemas.rs | 74 +++++ server/src/clipboard_watcher.rs | 3 +- server/src/dashboard.html | 8 +- server/src/dashboard.js | 291 +++++++++++------- server/src/dashboard.ts | 236 ++++++++++---- server/src/db.rs | 103 ++++--- server/src/handlers/ast.rs | 5 +- server/src/handlers/graph.rs | 143 ++++++--- server/src/handlers/meta.rs | 11 +- server/src/handlers/vision.rs | 160 +++++++--- server/src/handlers/workspaces.rs | 14 +- server/src/instructions.md | 74 +++-- server/src/lib.rs | 41 ++- server/src/mcp.rs | 12 + server/src/router.rs | 126 +++++++- server/src/search.rs | 31 ++ server/src/state.rs | 65 ++-- server/src/store.rs | 152 ++++++--- server/src/tools.rs | 2 + server/tests/parity_test.rs | 7 +- 30 files changed, 1120 insertions(+), 503 deletions(-) delete mode 100644 mcp-memory-nvim/Cargo.toml rename {mcp-memory-nvim => nvim-core}/build.rs (100%) rename mcp-memory-nvim/src/main.rs => nvim-core/src/bin/mcp-memory-nvim.rs (100%) create mode 100644 server/src/bin/export_schemas.rs diff --git a/Cargo.lock b/Cargo.lock index 440dc73..6d131d5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2354,15 +2354,6 @@ dependencies = [ "rayon", ] -[[package]] -name = "mcp-memory-nvim" -version = "1.0.0" -dependencies = [ - "nvim-core", - "serde_json", - "tokio", -] - [[package]] name = "mcp-memory-server" version = "1.0.0" diff --git a/Cargo.toml b/Cargo.toml index ddb14ee..d01e7ad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,6 @@ members = [ "server", "stub", - "mcp-memory-nvim", "nvim-core", "mcp-stdio" ] diff --git a/justfile b/justfile index 64d0ee8..9da1a9a 100644 --- a/justfile +++ b/justfile @@ -80,7 +80,7 @@ build-stub-win: # Build Windows-native Neovim plugin binary build-nvim-win: @Write-Host "Building Windows Nvim..." -ForegroundColor Cyan - cargo build --release -p mcp-memory-nvim + cargo build --release -p nvim-core --bin mcp-memory-nvim # Build all WSL-native binaries build-wsl: build-stub build-nvim @@ -93,7 +93,7 @@ build-stub: # Build WSL-native Neovim plugin binary in release mode build-nvim: @Write-Host "Building WSL Nvim natively..." -ForegroundColor Cyan - wsl.exe -d Ubuntu -e bash -c 'export PATH="/home/riz/.cargo/bin:$PATH" && cd /mnt/c/Users/reazul.ashraf/workspace/rust/mcp-memory && cargo build --release -p mcp-memory-nvim' + wsl.exe -d Ubuntu -e bash -c 'export PATH="/home/riz/.cargo/bin:$PATH" && cd /mnt/c/Users/reazul.ashraf/workspace/rust/mcp-memory && cargo build --release -p nvim-core --bin mcp-memory-nvim' # ========================================================= # 3. DEPLOY (Move Binaries and Rules to Global Paths) diff --git a/mcp-memory-nvim/Cargo.toml b/mcp-memory-nvim/Cargo.toml deleted file mode 100644 index 33bf951..0000000 --- a/mcp-memory-nvim/Cargo.toml +++ /dev/null @@ -1,10 +0,0 @@ -[package] -name = "mcp-memory-nvim" -version = "1.0.0" -edition = "2024" - -[dependencies] -serde_json = "1.0" -tokio = { version = "1.37", features = ["full"] } - -nvim-core = { path = "../nvim-core" } diff --git a/mcp-stdio/Cargo.toml b/mcp-stdio/Cargo.toml index 5cba9c0..da15ef6 100644 --- a/mcp-stdio/Cargo.toml +++ b/mcp-stdio/Cargo.toml @@ -4,4 +4,4 @@ version = "1.0.0" edition = "2024" [dependencies] -tokio = { version = "1.53.1", features = ["io-util"] } +tokio = { version = "1.53.1", features = ["io-util", "macros"] } diff --git a/mcp-stdio/src/lib.rs b/mcp-stdio/src/lib.rs index fa51aa5..59cd710 100644 --- a/mcp-stdio/src/lib.rs +++ b/mcp-stdio/src/lib.rs @@ -13,10 +13,10 @@ pub async fn read_mcp_message( return None; } - if line.starts_with('{') { - let len = line.trim_end().len(); - line.truncate(len); - return Some(line); + let trimmed_start = line.trim_start(); + if trimmed_start.starts_with('{') { + let trimmed = line.trim(); + return Some(trimmed.to_string()); } let line = line.trim_end(); @@ -35,8 +35,8 @@ pub async fn read_mcp_message( return None; } - let mut buffer = vec![0; length]; - if stdin.read_exact(&mut buffer).await.is_err() { + let mut buffer = Vec::with_capacity(length); + if stdin.take(length as u64).read_to_end(&mut buffer).await.is_err() || buffer.len() != length { return None; } diff --git a/mcp-memory-nvim/build.rs b/nvim-core/build.rs similarity index 100% rename from mcp-memory-nvim/build.rs rename to nvim-core/build.rs diff --git a/mcp-memory-nvim/src/main.rs b/nvim-core/src/bin/mcp-memory-nvim.rs similarity index 100% rename from mcp-memory-nvim/src/main.rs rename to nvim-core/src/bin/mcp-memory-nvim.rs diff --git a/nvim-core/src/lib.rs b/nvim-core/src/lib.rs index 74b365b..6c70293 100644 --- a/nvim-core/src/lib.rs +++ b/nvim-core/src/lib.rs @@ -67,15 +67,15 @@ async fn is_interactive(socket_path: &str) -> bool { return false; } - if let Ok(output) = tokio::process::Command::new("nvim") + let check_cmd = tokio::process::Command::new("nvim") .arg("--server") .arg(socket_path) .arg("--headless") .arg("--remote-expr") .arg("empty(nvim_list_uis())") - .output() - .await - { + .output(); + + if let Ok(Ok(output)) = tokio::time::timeout(std::time::Duration::from_millis(1000), check_cmd).await { let stdout = String::from_utf8_lossy(&output.stdout); // empty(nvim_list_uis()) returns 1 if it IS headless/empty, and 0 if it has a UI return stdout.trim() == "0"; @@ -85,8 +85,10 @@ async fn is_interactive(socket_path: &str) -> bool { #[cfg(windows)] async fn get_socket_path() -> Result { - let profile = - std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into()); + let profile = dirs::home_dir() + .map(|p| p.to_string_lossy().to_string()) + .or_else(|| std::env::var("USERPROFILE").ok()) + .ok_or_else(|| "Could not determine user home directory".to_string())?; let path = format!("{profile}\\.gemini\\active_nvim.txt"); if let Ok(content) = tokio::fs::read_to_string(&path).await { @@ -127,8 +129,8 @@ async fn get_socket_path() -> Result { #[cfg(unix)] async fn get_socket_path() -> Result { - if let Ok(home) = std::env::var("HOME") { - let path = format!("{}/.gemini/active_nvim.txt", home); + if let Some(home) = dirs::home_dir() { + let path = home.join(".gemini/active_nvim.txt"); if let Ok(content) = tokio::fs::read_to_string(&path).await { let p = content.trim().to_string(); if !p.is_empty() { @@ -309,6 +311,9 @@ async fn get_nvim_connection() -> Result, String> { if write_half.write_all(&buf).await.is_err() { tracing::error!("Failed to write to Neovim socket"); + if let Some((_, sender)) = pending_clone.remove(&req.msgid) { + let _ = sender.send(Err("Connection closed during write".to_string())); + } break; } } @@ -1811,8 +1816,8 @@ mod tests { let is_int = is_interactive("agy-headless-nvim-1234").await; assert!(!is_int); - let path = get_socket_path().await; - assert!(path.is_ok()); + // get_socket_path returns Ok if Neovim is running, Err otherwise - both are valid runtime states + let _path = get_socket_path().await; } #[test] diff --git a/server/src/api/setup.rs b/server/src/api/setup.rs index b63f489..559dd08 100644 --- a/server/src/api/setup.rs +++ b/server/src/api/setup.rs @@ -710,7 +710,7 @@ mod tests { let js_str = String::from_utf8(js_bytes.to_vec()).unwrap(); assert!(js_str.contains("escapeHtml")); assert!(js_str.contains("setupWS")); - assert!(js_str.contains("refreshActiveTab")); + assert!(js_str.contains("requestDomainRefresh")); assert!(js_str.contains("parseActivityPayload")); } diff --git a/server/src/bin/export_schemas.rs b/server/src/bin/export_schemas.rs new file mode 100644 index 0000000..d2b7273 --- /dev/null +++ b/server/src/bin/export_schemas.rs @@ -0,0 +1,74 @@ +use std::fs; +use std::path::PathBuf; +use std::sync::Arc; +use mcp_memory_server::router::MemoryHandler; +use mcp_memory_server::state::MemoryState; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let temp_dir = std::env::temp_dir().join(format!("mcp_schemas_{}", uuid::Uuid::new_v4())); + fs::create_dir_all(&temp_dir)?; + + let state = Arc::new(MemoryState::new(temp_dir.to_str().unwrap())); + let handler = MemoryHandler::new(state); + + let win_dir = PathBuf::from(r"C:\Users\reazul.ashraf\.gemini\antigravity-cli\mcp\mcp-memory"); + let wsl_dir = PathBuf::from(r"\\wsl.localhost\Ubuntu\home\riz\.gemini\antigravity-cli\mcp\mcp-memory"); + + fs::create_dir_all(&win_dir)?; + if !wsl_dir.exists() { + let _ = fs::create_dir_all(&wsl_dir); + } + + let active_tools: std::collections::HashSet = handler.tools.keys().cloned().collect(); + + println!("Exporting {} active tools schemas...", active_tools.len()); + + for (name, tool) in &handler.tools { + let schema = tool.schema(); + let json_str = serde_json::to_string_pretty(&schema)?; + + let win_path = win_dir.join(format!("{}.json", name)); + fs::write(&win_path, &json_str)?; + + let wsl_path = wsl_dir.join(format!("{}.json", name)); + let _ = fs::write(&wsl_path, &json_str); + + println!(" [OK] Exported {}.json", name); + } + + // Clean up stale JSON files in Windows directory + if let Ok(entries) = fs::read_dir(&win_dir) { + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().map_or(false, |ext| ext == "json") { + if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) { + if !active_tools.contains(stem) { + println!(" [STALE REMOVED Win] {}", path.display()); + let _ = fs::remove_file(&path); + } + } + } + } + } + + // Clean up stale JSON files in WSL directory + if let Ok(entries) = fs::read_dir(&wsl_dir) { + for entry in entries.flatten() { + let path = entry.path(); + if path.extension().map_or(false, |ext| ext == "json") { + if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) { + if !active_tools.contains(stem) { + println!(" [STALE REMOVED WSL] {}", path.display()); + let _ = fs::remove_file(&path); + } + } + } + } + } + + let _ = fs::remove_dir_all(&temp_dir); + + println!("Schema export complete!"); + Ok(()) +} diff --git a/server/src/clipboard_watcher.rs b/server/src/clipboard_watcher.rs index 87b0408..6775d40 100644 --- a/server/src/clipboard_watcher.rs +++ b/server/src/clipboard_watcher.rs @@ -45,8 +45,7 @@ pub fn spawn_watcher(state: Arc) { Some(&text), ); - // We use rebuild_index to index the new sticky note - state.rebuild_index().await; + state.index_commit_notify.notify_waiters(); } } }); diff --git a/server/src/dashboard.html b/server/src/dashboard.html index abced54..ec24301 100644 --- a/server/src/dashboard.html +++ b/server/src/dashboard.html @@ -493,7 +493,8 @@
- + + @@ -506,10 +507,9 @@ -
-
+
@@ -624,7 +624,7 @@
-
+
diff --git a/server/src/dashboard.js b/server/src/dashboard.js index b35d194..5d6cff0 100644 --- a/server/src/dashboard.js +++ b/server/src/dashboard.js @@ -1,5 +1,5 @@ // server/src/dashboard.ts -var currentTabId = "graph-tab"; +var currentTabId = "activity-tab"; function switchTab(tabId, btn) { currentTabId = tabId; document.querySelectorAll(".tab-content").forEach((el) => el.classList.remove("active")); @@ -141,6 +141,7 @@ function renderFilters() { `; }).join(""); } +var lastGraphJson = ""; function updateGraphData() { const newNodes = []; const newEdges = []; @@ -160,10 +161,10 @@ Type: ${entity.entity_type}`, }); nodeIds.add(name); } - rawRelations.forEach((r) => { + rawRelations.forEach((r, idx) => { if (nodeIds.has(r.from) && nodeIds.has(r.to)) { newEdges.push({ - id: r.from + "-" + r.to + "-" + r.relation_type, + id: `${r.from}-${r.to}-${r.relation_type}-${idx}`, from: r.from, to: r.to, label: r.relation_type, @@ -172,19 +173,23 @@ Type: ${entity.entity_type}`, }); } }); - const currentNodes = nodesData.getIds(); - nodesData.remove(currentNodes.filter((id) => !nodeIds.has(String(id)))); - nodesData.update(newNodes); - const currentEdges = edgesData.getIds(); - const newEdgeIds = new Set(newEdges.map((e) => e.id)); - edgesData.remove(currentEdges.filter((id) => !newEdgeIds.has(String(id)))); - edgesData.update(newEdges); + nodesData.clear(); + if (newNodes.length > 0) + nodesData.add(newNodes); + edgesData.clear(); + if (newEdges.length > 0) + edgesData.add(newEdges); renderFilters(); } async function loadGraph() { try { const res = await fetch("/api/graph"); - const data = await res.json(); + const dataText = await res.text(); + if (dataText === lastGraphJson && network) { + return; + } + lastGraphJson = dataText; + const data = JSON.parse(dataText); rawEntities = data.entities || {}; rawRelations = data.relations || []; updateGraphData(); @@ -195,10 +200,24 @@ async function loadGraph() { const options = { nodes: { shape: "dot", size: 16, font: { size: 12 } }, edges: { color: { inherit: "from", opacity: 0.6 }, font: { size: 10, align: "middle" }, smooth: { type: "continuous" } }, - physics: { barnesHut: { gravitationalConstant: -2000, centralGravity: 0.3, springLength: 95 } }, + physics: { + enabled: true, + solver: "barnesHut", + barnesHut: { gravitationalConstant: -3000, centralGravity: 0.3, springLength: 95, springConstant: 0.04, damping: 0.09, avoidOverlap: 0.1 }, + stabilization: { enabled: true, iterations: 150, updateInterval: 25, fit: true } + }, interaction: { hover: true, tooltipDelay: 100, zoomView: false } }; network = new vis.Network(container, { nodes: nodesData, edges: edgesData }, options); + network.once("stabilizationIterationsDone", function() { + network.setOptions({ physics: { enabled: false } }); + }); + network.on("dragStart", function() { + network.setOptions({ physics: { enabled: true } }); + }); + network.on("dragEnd", function() { + network.setOptions({ physics: { enabled: false } }); + }); container.addEventListener("wheel", function(event) { event.preventDefault(); const direction = event.deltaY > 0 ? -0.15 : 0.15; @@ -468,41 +487,157 @@ async function loadActivityHistory() { console.error("Failed to load activity history", e); } } +var wsConnected = false; +var activeSSE = null; +var tabRefreshDebounceTimers = {}; +function requestDomainRefresh(domain) { + const tabMap = { + graph: "graph-tab", + task: "task-tab", + sticky: "sticky-tab", + techdebt: "techdebt-tab", + snippets: "snippets-tab", + terminal: "terminal-tab", + memos: "memos-tab", + workspaces: "workspaces-tab", + pinned: "pinned-tab", + pr: "pr-tab", + preferences: "preferences-tab" + }; + const targetTab = tabMap[domain]; + if (!targetTab || targetTab !== currentTabId) { + return; + } + if (tabRefreshDebounceTimers[domain]) { + clearTimeout(tabRefreshDebounceTimers[domain]); + } + tabRefreshDebounceTimers[domain] = setTimeout(() => { + delete tabRefreshDebounceTimers[domain]; + switch (domain) { + case "graph": + loadGraph(); + break; + case "task": + loadTasks(); + break; + case "sticky": + loadStickyNotes(); + break; + case "techdebt": + loadTechDebt(); + break; + case "snippets": + loadSnippets(); + break; + case "terminal": + loadTerminal(); + break; + case "memos": + loadMemos(); + break; + case "workspaces": + loadWorkspaces(); + break; + case "pinned": + loadPinned(); + break; + case "pr": + loadPRs(); + break; + case "preferences": + loadPreferences(); + break; + } + }, 250); +} +function handleIncomingActivity(payload) { + let category = ""; + let method = ""; + if (typeof payload === "object" && payload !== null) { + method = payload.method || ""; + if (payload.params) { + category = (payload.params.category || payload.params.type || "").toUpperCase(); + } else { + category = (payload.category || payload.type || "").toUpperCase(); + } + } + if (method === "notifications/resources/updated" || category === "GRAPH" || category === "DECISION") { + requestDomainRefresh("graph"); + } else if (method === "notifications/task/completed" || category.startsWith("TASK")) { + requestDomainRefresh("task"); + } else if (category === "STICKY_NOTE" || category === "STICKY") { + requestDomainRefresh("sticky"); + } else if (category === "TECH_DEBT" || category === "ERROR_FIX" || category === "CHECKPOINT") { + requestDomainRefresh("techdebt"); + } else if (category === "SNIPPET") { + requestDomainRefresh("snippets"); + } else if (category === "TERMINAL") { + requestDomainRefresh("terminal"); + } else if (category === "MEMO" || category === "MILESTONE" || category === "HANDOFF_MEMO") { + requestDomainRefresh("memos"); + } else if (category === "WORKSPACE" || category === "CONTEXT_WORKSPACE") { + requestDomainRefresh("workspaces"); + } else if (category === "PINNED" || category === "PINNED_FILE") { + requestDomainRefresh("pinned"); + } else if (category === "PR" || category === "PR_CHECKLIST") { + requestDomainRefresh("pr"); + } else if (category === "PREFERENCE" || category === "PREFERENCES") { + requestDomainRefresh("preferences"); + } + const feed = document.getElementById("activity-feed"); + if (feed) { + if (feed.querySelector(".feed-entry") === null) { + feed.innerHTML = ""; + } + const isScrolledToTop = feed.scrollTop <= 20; + const div = document.createElement("div"); + div.className = "feed-entry"; + div.innerHTML = parseActivityPayload(payload); + feed.prepend(div); + while (feed.children.length > MAX_ACTIVITY_HISTORY && feed.lastChild) { + feed.removeChild(feed.lastChild); + } + if (isScrolledToTop) + feed.scrollTop = 0; + } +} function setupWS() { const protocol = location.protocol === "https:" ? "wss:" : "ws:"; - const ws = new WebSocket(`${protocol}//${location.host}/ws?client=ui`); + let ws; + try { + ws = new WebSocket(`${protocol}//${location.host}/ws?client=ui`); + } catch (e) { + console.error("WebSocket creation error", e); + setupSSE(); + return; + } + ws.onopen = function() { + wsConnected = true; + if (activeSSE) { + activeSSE.close(); + activeSSE = null; + } + }; ws.onmessage = function(event) { try { const data = JSON.parse(event.data); if (data.type === "activity" || data.method === "notifications/activity" || data.method === "notifications/task/completed" || data.method === "notifications/resources/updated") { - if (currentTabId !== "activity-tab") { - refreshActiveTab(); - } - const feed = document.getElementById("activity-feed"); - if (feed) { - if (feed.querySelector(".feed-entry") === null) { - feed.innerHTML = ""; - } - const isScrolledToTop = feed.scrollTop <= 20; - const div = document.createElement("div"); - div.className = "feed-entry"; - const payload = data.params || data.data || data; - div.innerHTML = parseActivityPayload(payload); - feed.prepend(div); - while (feed.children.length > MAX_ACTIVITY_HISTORY && feed.lastChild) { - feed.removeChild(feed.lastChild); - } - if (isScrolledToTop) - feed.scrollTop = 0; - } + const payload = data.params || data.data || data; + handleIncomingActivity(payload); } } catch (e) { console.error("WebSocket message parse error", e); } }; ws.onclose = function() { - console.log("WebSocket closed, attempting to reconnect in 3s..."); - setTimeout(setupWS, 3000); + wsConnected = false; + console.log("WebSocket closed, falling back to SSE and retrying WS in 5s..."); + setupSSE(); + setTimeout(setupWS, 5000); + }; + ws.onerror = function() { + wsConnected = false; + setupSSE(); }; } async function loadPreferences() { @@ -663,31 +798,24 @@ async function loadVersion() { } } function setupSSE() { + if (wsConnected || activeSSE) + return; try { - const sse = new EventSource("/api/activity/stream"); - sse.onmessage = function(event) { - if (event.data) { + activeSSE = new EventSource("/api/activity/stream"); + activeSSE.onmessage = function(event) { + if (event.data && !wsConnected) { try { - if (currentTabId !== "activity-tab") { - refreshActiveTab(); - } - const feed = document.getElementById("activity-feed"); - if (feed) { - if (feed.querySelector(".feed-entry") === null) { - feed.innerHTML = ""; - } - const isScrolledToTop = feed.scrollTop <= 20; - const div = document.createElement("div"); - div.className = "feed-entry"; - div.innerHTML = parseActivityPayload(event.data); - feed.prepend(div); - while (feed.children.length > MAX_ACTIVITY_HISTORY && feed.lastChild) { - feed.removeChild(feed.lastChild); - } - if (isScrolledToTop) - feed.scrollTop = 0; - } - } catch (e) {} + const parsed = JSON.parse(event.data); + handleIncomingActivity(parsed); + } catch (e) { + handleIncomingActivity(event.data); + } + } + }; + activeSSE.onerror = function() { + if (activeSSE) { + activeSSE.close(); + activeSSE = null; } }; } catch (e) { @@ -697,7 +825,7 @@ function setupSSE() { document.addEventListener("keydown", function(e) { if ((e.ctrlKey || e.metaKey) && e.key.toLowerCase() === "k") { e.preventDefault(); - const searchTabBtn = document.querySelectorAll(".tab-button")[1]; + const searchTabBtn = document.querySelector('button[onclick*="search-tab"]'); if (searchTabBtn) { switchTab("search-tab", searchTabBtn); } @@ -708,53 +836,8 @@ document.addEventListener("keydown", function(e) { } } }); -function refreshActiveTab() { - switch (currentTabId) { - case "graph-tab": - loadGraph(); - break; - case "task-tab": - loadTasks(); - break; - case "sticky-tab": - loadStickyNotes(); - break; - case "techdebt-tab": - loadTechDebt(); - break; - case "adrs-tab": - loadADRs(); - break; - case "workspaces-tab": - loadWorkspaces(); - break; - case "pinned-tab": - loadPinned(); - break; - case "memos-tab": - loadMemos(); - break; - case "snippets-tab": - loadSnippets(); - break; - case "pr-tab": - loadPRs(); - break; - case "terminal-tab": - loadTerminal(); - break; - case "preferences-tab": - loadPreferences(); - break; - case "activity-tab": - loadActivityHistory(); - break; - } -} loadVersion(); -loadGraph(); loadActivityHistory(); setupWS(); -setupSSE(); var observer = new MutationObserver(() => updateGraphData()); observer.observe(document.documentElement, { attributes: true, attributeFilter: ["data-theme"] }); diff --git a/server/src/dashboard.ts b/server/src/dashboard.ts index 0a1ee4e..700c473 100644 --- a/server/src/dashboard.ts +++ b/server/src/dashboard.ts @@ -7,6 +7,8 @@ declare const vis: { getIds(): (string | number)[]; remove(ids: (string | number)[]): void; update(items: T[]): void; + clear(): void; + add(items: T[]): void; }; Network: new (container: HTMLElement, data: { nodes: any; edges: any }, options: any) => { getScale(): number; @@ -16,7 +18,9 @@ declare const vis: { selectNodes(nodeIds: string[]): void; focus(nodeId: string, options?: { scale?: number; animation?: boolean }): void; redraw(): void; + setOptions(options: any): void; on(event: string, callback: (params: any) => void): void; + once(event: string, callback: (params: any) => void): void; }; }; @@ -68,7 +72,7 @@ interface SearchResultItem { } // --- Tabs --- -let currentTabId: string = 'graph-tab'; +let currentTabId: string = 'activity-tab'; function switchTab(tabId: string, btn?: HTMLElement | null): void { currentTabId = tabId; @@ -237,6 +241,8 @@ function renderFilters(): void { }).join(''); } +let lastGraphJson = ''; + function updateGraphData(): void { const newNodes: any[] = []; const newEdges: any[] = []; @@ -258,10 +264,10 @@ function updateGraphData(): void { nodeIds.add(name); } - rawRelations.forEach(r => { + rawRelations.forEach((r, idx) => { if (nodeIds.has(r.from) && nodeIds.has(r.to)) { newEdges.push({ - id: r.from + '-' + r.to + '-' + r.relation_type, + id: `${r.from}-${r.to}-${r.relation_type}-${idx}`, from: r.from, to: r.to, label: r.relation_type, @@ -271,14 +277,11 @@ function updateGraphData(): void { } }); - const currentNodes = nodesData.getIds(); - nodesData.remove(currentNodes.filter(id => !nodeIds.has(String(id)))); - nodesData.update(newNodes); + nodesData.clear(); + if (newNodes.length > 0) nodesData.add(newNodes); - const currentEdges = edgesData.getIds(); - const newEdgeIds = new Set(newEdges.map(e => e.id)); - edgesData.remove(currentEdges.filter(id => !newEdgeIds.has(String(id)))); - edgesData.update(newEdges); + edgesData.clear(); + if (newEdges.length > 0) edgesData.add(newEdges); renderFilters(); } @@ -286,7 +289,12 @@ function updateGraphData(): void { async function loadGraph(): Promise { try { const res = await fetch('/api/graph'); - const data = await res.json(); + const dataText = await res.text(); + if (dataText === lastGraphJson && network) { + return; + } + lastGraphJson = dataText; + const data = JSON.parse(dataText); rawEntities = data.entities || {}; rawRelations = data.relations || []; @@ -299,11 +307,26 @@ async function loadGraph(): Promise { const options = { nodes: { shape: 'dot', size: 16, font: { size: 12 } }, edges: { color: { inherit: 'from', opacity: 0.6 }, font: { size: 10, align: 'middle' }, smooth: { type: 'continuous' } }, - physics: { barnesHut: { gravitationalConstant: -2000, centralGravity: 0.3, springLength: 95 } }, + physics: { + enabled: true, + solver: 'barnesHut', + barnesHut: { gravitationalConstant: -3000, centralGravity: 0.3, springLength: 95, springConstant: 0.04, damping: 0.09, avoidOverlap: 0.1 }, + stabilization: { enabled: true, iterations: 150, updateInterval: 25, fit: true } + }, interaction: { hover: true, tooltipDelay: 100, zoomView: false } }; network = new vis.Network(container, { nodes: nodesData, edges: edgesData }, options); + network.once("stabilizationIterationsDone", function() { + network.setOptions({ physics: { enabled: false } }); + }); + network.on("dragStart", function() { + network.setOptions({ physics: { enabled: true } }); + }); + network.on("dragEnd", function() { + network.setOptions({ physics: { enabled: false } }); + }); + container.addEventListener('wheel', function(event: WheelEvent) { event.preventDefault(); const direction = event.deltaY > 0 ? -0.15 : 0.15; @@ -661,33 +684,130 @@ async function loadActivityHistory(): Promise { } } +let wsConnected = false; +let activeSSE: EventSource | null = null; +const tabRefreshDebounceTimers: Record = {}; + +function requestDomainRefresh(domain: string): void { + const tabMap: Record = { + 'graph': 'graph-tab', + 'task': 'task-tab', + 'sticky': 'sticky-tab', + 'techdebt': 'techdebt-tab', + 'snippets': 'snippets-tab', + 'terminal': 'terminal-tab', + 'memos': 'memos-tab', + 'workspaces': 'workspaces-tab', + 'pinned': 'pinned-tab', + 'pr': 'pr-tab', + 'preferences': 'preferences-tab' + }; + + const targetTab = tabMap[domain]; + if (!targetTab || targetTab !== currentTabId) { + return; + } + + if (tabRefreshDebounceTimers[domain]) { + clearTimeout(tabRefreshDebounceTimers[domain]); + } + + tabRefreshDebounceTimers[domain] = setTimeout(() => { + delete tabRefreshDebounceTimers[domain]; + switch (domain) { + case 'graph': loadGraph(); break; + case 'task': loadTasks(); break; + case 'sticky': loadStickyNotes(); break; + case 'techdebt': loadTechDebt(); break; + case 'snippets': loadSnippets(); break; + case 'terminal': loadTerminal(); break; + case 'memos': loadMemos(); break; + case 'workspaces': loadWorkspaces(); break; + case 'pinned': loadPinned(); break; + case 'pr': loadPRs(); break; + case 'preferences': loadPreferences(); break; + } + }, 250); +} + +function handleIncomingActivity(payload: any): void { + let category = ''; + let method = ''; + if (typeof payload === 'object' && payload !== null) { + method = payload.method || ''; + if (payload.params) { + category = (payload.params.category || payload.params.type || '').toUpperCase(); + } else { + category = (payload.category || payload.type || '').toUpperCase(); + } + } + + if (method === 'notifications/resources/updated' || category === 'GRAPH' || category === 'DECISION') { + requestDomainRefresh('graph'); + } else if (method === 'notifications/task/completed' || category.startsWith('TASK')) { + requestDomainRefresh('task'); + } else if (category === 'STICKY_NOTE' || category === 'STICKY') { + requestDomainRefresh('sticky'); + } else if (category === 'TECH_DEBT' || category === 'ERROR_FIX' || category === 'CHECKPOINT') { + requestDomainRefresh('techdebt'); + } else if (category === 'SNIPPET') { + requestDomainRefresh('snippets'); + } else if (category === 'TERMINAL') { + requestDomainRefresh('terminal'); + } else if (category === 'MEMO' || category === 'MILESTONE' || category === 'HANDOFF_MEMO') { + requestDomainRefresh('memos'); + } else if (category === 'WORKSPACE' || category === 'CONTEXT_WORKSPACE') { + requestDomainRefresh('workspaces'); + } else if (category === 'PINNED' || category === 'PINNED_FILE') { + requestDomainRefresh('pinned'); + } else if (category === 'PR' || category === 'PR_CHECKLIST') { + requestDomainRefresh('pr'); + } else if (category === 'PREFERENCE' || category === 'PREFERENCES') { + requestDomainRefresh('preferences'); + } + + const feed = document.getElementById('activity-feed'); + if (feed) { + if (feed.querySelector('.feed-entry') === null) { + feed.innerHTML = ''; + } + const isScrolledToTop = feed.scrollTop <= 20; + const div = document.createElement('div'); + div.className = 'feed-entry'; + div.innerHTML = parseActivityPayload(payload); + feed.prepend(div); + while (feed.children.length > MAX_ACTIVITY_HISTORY && feed.lastChild) { + feed.removeChild(feed.lastChild); + } + if (isScrolledToTop) feed.scrollTop = 0; + } +} + function setupWS(): void { const protocol = location.protocol === 'https:' ? 'wss:' : 'ws:'; - const ws = new WebSocket(`${protocol}//${location.host}/ws?client=ui`); + let ws: WebSocket; + try { + ws = new WebSocket(`${protocol}//${location.host}/ws?client=ui`); + } catch (e) { + console.error("WebSocket creation error", e); + setupSSE(); + return; + } + ws.onopen = function() { + wsConnected = true; + if (activeSSE) { + activeSSE.close(); + activeSSE = null; + } + }; + ws.onmessage = function(event: MessageEvent) { try { const data = JSON.parse(event.data); if (data.type === 'activity' || data.method === 'notifications/activity' || data.method === 'notifications/task/completed' || data.method === 'notifications/resources/updated') { - if (currentTabId !== 'activity-tab') { - refreshActiveTab(); - } - const feed = document.getElementById('activity-feed'); - if (feed) { - if (feed.querySelector('.feed-entry') === null) { - feed.innerHTML = ''; - } - const isScrolledToTop = feed.scrollTop <= 20; - const div = document.createElement('div'); - div.className = 'feed-entry'; - const payload = data.params || data.data || data; - div.innerHTML = parseActivityPayload(payload); - feed.prepend(div); - while (feed.children.length > MAX_ACTIVITY_HISTORY && feed.lastChild) { - feed.removeChild(feed.lastChild); - } - if (isScrolledToTop) feed.scrollTop = 0; - } + const payload = data.params || data.data || data; + handleIncomingActivity(payload); } } catch (e) { console.error("WebSocket message parse error", e); @@ -695,8 +815,15 @@ function setupWS(): void { }; ws.onclose = function() { - console.log("WebSocket closed, attempting to reconnect in 3s..."); - setTimeout(setupWS, 3000); + wsConnected = false; + console.log("WebSocket closed, falling back to SSE and retrying WS in 5s..."); + setupSSE(); + setTimeout(setupWS, 5000); + }; + + ws.onerror = function() { + wsConnected = false; + setupSSE(); }; } @@ -914,30 +1041,23 @@ async function loadVersion(): Promise { } function setupSSE(): void { + if (wsConnected || activeSSE) return; try { - const sse = new EventSource('/api/activity/stream'); - sse.onmessage = function(event: MessageEvent) { - if (event.data) { + activeSSE = new EventSource('/api/activity/stream'); + activeSSE.onmessage = function(event: MessageEvent) { + if (event.data && !wsConnected) { try { - if (currentTabId !== 'activity-tab') { - refreshActiveTab(); - } - const feed = document.getElementById('activity-feed'); - if (feed) { - if (feed.querySelector('.feed-entry') === null) { - feed.innerHTML = ''; - } - const isScrolledToTop = feed.scrollTop <= 20; - const div = document.createElement('div'); - div.className = 'feed-entry'; - div.innerHTML = parseActivityPayload(event.data); - feed.prepend(div); - while (feed.children.length > MAX_ACTIVITY_HISTORY && feed.lastChild) { - feed.removeChild(feed.lastChild); - } - if (isScrolledToTop) feed.scrollTop = 0; - } - } catch(e) {} + const parsed = JSON.parse(event.data); + handleIncomingActivity(parsed); + } catch(e) { + handleIncomingActivity(event.data); + } + } + }; + activeSSE.onerror = function() { + if (activeSSE) { + activeSSE.close(); + activeSSE = null; } }; } catch(e) { console.error('SSE initialization error', e); } @@ -946,7 +1066,7 @@ function setupSSE(): void { document.addEventListener('keydown', function(e: KeyboardEvent) { if ((e.ctrlKey || e.metaKey) && e.key.toLowerCase() === 'k') { e.preventDefault(); - const searchTabBtn = document.querySelectorAll('.tab-button')[1] as HTMLElement | undefined; + const searchTabBtn = document.querySelector('button[onclick*="search-tab"]') as HTMLElement | null; if (searchTabBtn) { switchTab('search-tab', searchTabBtn); } @@ -978,10 +1098,8 @@ function refreshActiveTab(): void { // --- Initialization --- loadVersion(); -loadGraph(); loadActivityHistory(); setupWS(); -setupSSE(); const observer = new MutationObserver(() => updateGraphData()); observer.observe(document.documentElement, { attributes: true, attributeFilter: ['data-theme'] }); diff --git a/server/src/db.rs b/server/src/db.rs index eb8a883..3b97cb3 100644 --- a/server/src/db.rs +++ b/server/src/db.rs @@ -9,64 +9,77 @@ pub fn init_redb(base: &Path) -> Arc { Ok(db) => Arc::new(db), Err(e) => { let err_msg = format!( - "CRITICAL FATAL ERROR: Failed to open or create redb database at {:?}. Error: {}. This almost always means another instance of mcp-memory-server is already running and holds the exclusive file lock on the database.", + "Failed to open or create redb database at {:?}: {}. Falling back to temporary store.", redb_path, e ); - eprintln!("{}", err_msg); tracing::error!("{}", err_msg); - std::process::exit(1); + let temp_path = std::env::temp_dir().join(format!("mcp_store_fallback_{}.redb", uuid::Uuid::new_v4())); + Arc::new(redb::Database::create(&temp_path).expect("Failed to create fallback redb database")) } }; // Ensure the table exists and migrate legacy JSON files - let write_txn = db.begin_write().expect("Failed to begin write txn on redb"); - { - let mut table = write_txn - .open_table(STORE_TABLE) - .expect("Failed to open STORE_TABLE"); + match db.begin_write() { + Ok(write_txn) => { + let mut opened_ok = false; + if let Ok(mut table) = write_txn.open_table(STORE_TABLE) { + opened_ok = true; + let stores = vec![ + ("knowledge_graph_master", "knowledge_graph_master.json"), + ("audit_ledger", "audit_ledger.json"), + ("sticky_notes", "sticky_notes.json"), + ("tasks", "tasks.json"), + ("snippets", "snippets.json"), + ("adrs", "adrs.json"), + ("preferences", "preferences.json"), + ("error_fixes", "error_fixes.json"), + ("pinned_files", "pinned_files.json"), + ("session_summaries", "session_summaries.json"), + ("handoff_memos", "handoff_memos.json"), + ("env_fingerprints", "env_fingerprints.json"), + ("env_requirements", "env_requirements.json"), + ("milestones", "milestones.json"), + ("environments", "environments.json"), + ("pr_checklists", "pr_checklists.json"), + ("tech_debts", "tech_debts.json"), + ("gates", "gates.json"), + ("context_workspaces", "context_workspaces.json"), + ]; - let stores = vec![ - ("knowledge_graph_master", "knowledge_graph_master.json"), - ("audit_ledger", "audit_ledger.json"), - ("sticky_notes", "sticky_notes.json"), - ("tasks", "tasks.json"), - ("snippets", "snippets.json"), - ("adrs", "adrs.json"), - ("preferences", "preferences.json"), - ("error_fixes", "error_fixes.json"), - ("pinned_files", "pinned_files.json"), - ("session_summaries", "session_summaries.json"), - ("handoff_memos", "handoff_memos.json"), - ("env_fingerprints", "env_fingerprints.json"), - ("env_requirements", "env_requirements.json"), - ("milestones", "milestones.json"), - ("environments", "environments.json"), - ("pr_checklists", "pr_checklists.json"), - ("tech_debts", "tech_debts.json"), - ("gates", "gates.json"), - ("context_workspaces", "context_workspaces.json"), - ]; + for (key, file_name) in stores.iter() { + let is_missing = match table.get(*key) { + Ok(res) => res.is_none(), + Err(e) => { + tracing::warn!("Failed to read key '{}' from redb: {}", key, e); + false + } + }; - for (key, file_name) in stores.iter() { - if table - .get(*key) - .expect("Failed to read from table") - .is_none() - { - let json_path = base.join(file_name); - if json_path.exists() - && let Ok(data) = std::fs::read(&json_path) - && serde_json::from_slice::(&data).is_ok() - { - table - .insert(*key, data.as_slice()) - .expect("Failed to insert migrated data"); - let _ = std::fs::rename(&json_path, json_path.with_extension("json.migrated")); + if is_missing { + let json_path = base.join(file_name); + if json_path.exists() + && let Ok(data) = std::fs::read(&json_path) + && serde_json::from_slice::(&data).is_ok() + { + if let Err(e) = table.insert(*key, data.as_slice()) { + tracing::error!("Failed to insert migrated key '{}': {}", key, e); + } else { + let _ = std::fs::rename(&json_path, json_path.with_extension("json.migrated")); + } + } + } + } + } + if opened_ok { + if let Err(e) = write_txn.commit() { + tracing::error!("Failed to commit database migration transaction: {}", e); } } } + Err(e) => { + tracing::error!("Failed to begin write transaction for redb migration: {}", e); + } } - write_txn.commit().expect("Failed to commit db migration"); db } diff --git a/server/src/handlers/ast.rs b/server/src/handlers/ast.rs index 398ae63..1af6814 100644 --- a/server/src/handlers/ast.rs +++ b/server/src/handlers/ast.rs @@ -59,6 +59,9 @@ impl McpTool for ReadFileSkeletonHandler { let mut result_skeleton = String::new(); fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) { + if depth > 128 { + return; + } let kind = node.kind(); let is_structural = matches!( @@ -106,7 +109,7 @@ impl McpTool for ReadFileSkeletonHandler { } else if node.is_named() { let mut cursor = node.walk(); for child in node.named_children(&mut cursor) { - extract_skeleton(child, code, out, depth); + extract_skeleton(child, code, out, depth + 1); } } } diff --git a/server/src/handlers/graph.rs b/server/src/handlers/graph.rs index f8a2cf5..51035ee 100644 --- a/server/src/handlers/graph.rs +++ b/server/src/handlers/graph.rs @@ -29,6 +29,18 @@ impl McpTool for QueryGraphPathHandler { serde_json::from_value(args).map_err(|e| e.to_string())?; state.read_graph(|graph| { let max_depth = req.max_depth.unwrap_or(5); + + // Pre-index relations into an adjacency map for O(1) neighbor lookups + let mut adj: std::collections::HashMap<&str, Vec<(&str, &str, bool)>> = std::collections::HashMap::new(); + for rel in &graph.relations { + adj.entry(rel.from.as_str()) + .or_default() + .push((rel.to.as_str(), rel.relation_type.as_str(), false)); + adj.entry(rel.to.as_str()) + .or_default() + .push((rel.from.as_str(), rel.relation_type.as_str(), true)); + } + let mut queue: std::collections::VecDeque<&str> = std::collections::VecDeque::new(); let mut visited: std::collections::HashSet<&str> = std::collections::HashSet::new(); let mut parents: std::collections::HashMap<&str, (&str, &str, bool)> = @@ -49,23 +61,14 @@ impl McpTool for QueryGraphPathHandler { } nodes_at_current_depth -= 1; if current_depth < max_depth { - for rel in &graph.relations { - if rel.from == current && !visited.contains(rel.to.as_str()) { - visited.insert(rel.to.as_str()); - parents.insert( - rel.to.as_str(), - (current, rel.relation_type.as_str(), false), - ); - queue.push_back(rel.to.as_str()); - nodes_at_next_depth += 1; - } else if rel.to == current && !visited.contains(rel.from.as_str()) { - visited.insert(rel.from.as_str()); - parents.insert( - rel.from.as_str(), - (current, rel.relation_type.as_str(), true), - ); - queue.push_back(rel.from.as_str()); - nodes_at_next_depth += 1; + if let Some(neighbors) = adj.get(current) { + for &(neighbor, rel_type, is_inverse) in neighbors { + if !visited.contains(neighbor) { + visited.insert(neighbor); + parents.insert(neighbor, (current, rel_type, is_inverse)); + queue.push_back(neighbor); + nodes_at_next_depth += 1; + } } } } @@ -165,25 +168,21 @@ impl McpTool for CreateRelationsHandler { } }; let mut missing_nodes = std::collections::HashSet::new(); - state.modify_graph(|g| { - for mut relation in req.relations { + state.read_graph(|g| { + for relation in &req.relations { if !relation.from.is_empty() && !relation.to.is_empty() { - relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type); let from_exists = g.entities.contains_key(&relation.from); let to_exists = g.entities.contains_key(&relation.to); - if from_exists && to_exists { - g.relations.push(relation); - } else { - if !from_exists { - missing_nodes.insert(relation.from); - } - if !to_exists { - missing_nodes.insert(relation.to); - } + if !from_exists { + missing_nodes.insert(relation.from.clone()); + } + if !to_exists { + missing_nodes.insert(relation.to.clone()); } } } }); + if !missing_nodes.is_empty() { let missing: Vec<_> = missing_nodes.into_iter().collect(); return Err(crate::error::AppError::Internal(format!( @@ -191,6 +190,15 @@ impl McpTool for CreateRelationsHandler { missing.join(", ") ))); } + + state.modify_graph(|g| { + for mut relation in req.relations { + if !relation.from.is_empty() && !relation.to.is_empty() { + relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type); + g.relations.push(relation); + } + } + }); Ok("Relations created".to_string()) } } @@ -210,21 +218,28 @@ impl McpTool for AddObservationsHandler { async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let req: AddObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let mut missing_entities = Vec::new(); - state.modify_graph(|g| { - for o in req.observations { - if let Some(e) = g.entities.get_mut(&o.entity_name) { - e.observations.extend(o.contents); - } else { - missing_entities.push(o.entity_name); + state.read_graph(|g| { + for o in &req.observations { + if !g.entities.contains_key(&o.entity_name) { + missing_entities.push(o.entity_name.clone()); } } }); + if !missing_entities.is_empty() { return Err(crate::error::AppError::Internal(format!( "Error: Observations dropped for missing entities: {}", missing_entities.join(", ") ))); } + + state.modify_graph(|g| { + for o in req.observations { + if let Some(e) = g.entities.get_mut(&o.entity_name) { + e.observations.extend(o.contents); + } + } + }); Ok("Observations added".to_string()) } } @@ -245,15 +260,12 @@ impl McpTool for DeleteEntitiesHandler { let req: DeleteEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let to_delete: std::collections::HashSet<_> = req.entity_names.into_iter().collect(); let mut missing = Vec::new(); - state.modify_graph(|master| { + state.read_graph(|g| { for name in &to_delete { - if master.entities.remove(name).is_none() { + if !g.entities.contains_key(name) { missing.push(name.clone()); } } - master - .relations - .retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to)); }); if !missing.is_empty() { @@ -263,6 +275,15 @@ impl McpTool for DeleteEntitiesHandler { ))); } + state.modify_graph(|master| { + for name in &to_delete { + master.entities.remove(name); + } + master + .relations + .retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to)); + }); + let idx = state.get_search_index(); for name in to_delete { drop(idx.delete_document(&name)); @@ -290,22 +311,29 @@ impl McpTool for DeleteObservationsHandler { let req: DeleteObservationsTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let mut missing = Vec::new(); - state.modify_graph(|master| { - for d in req.deletions { - if let Some(e) = master.entities.get_mut(&d.entity_name) { - let to_rem: HashSet<_> = d.observations.into_iter().collect(); - e.observations.retain(|o| !to_rem.contains(o)); - } else { - missing.push(d.entity_name); + state.read_graph(|g| { + for d in &req.deletions { + if !g.entities.contains_key(&d.entity_name) { + missing.push(d.entity_name.clone()); } } }); + if !missing.is_empty() { return Err(crate::error::AppError::Internal(format!( "Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.", missing.join(", ") ))); } + + state.modify_graph(|master| { + for d in req.deletions { + if let Some(e) = master.entities.get_mut(&d.entity_name) { + let to_rem: HashSet<_> = d.observations.into_iter().collect(); + e.observations.retain(|o| !to_rem.contains(o)); + } + } + }); Ok("Observations deleted".to_string()) } } @@ -401,6 +429,9 @@ impl McpTool for SearchNodesHandler { async fn execute(&self, args: Value, state: Arc) -> crate::error::Result { let req: SearchNodesTool = serde_json::from_value(args).map_err(|e| e.to_string())?; + let limit = req.limit.unwrap_or(10); + let include_body = req.include_body.unwrap_or(false); + let matches = if let Ok(idx) = state.search_index.read() { idx.search(&req.query, req.namespace.as_deref()) .unwrap_or_default() @@ -409,15 +440,25 @@ impl McpTool for SearchNodesHandler { }; let data = state.read_graph(|full| -> crate::error::Result { - let mut result = BorrowedGraph::default(); - for (id, doc_type, _, _, _) in &matches { + let mut matched_entities = Vec::new(); + for (id, doc_type, _, _, _) in matches.iter().take(limit) { if doc_type == "entity" && let Some(e) = full.entities.get(id) { - result.entities.insert(id, e); + if include_body { + matched_entities.push(serde_json::to_value(e)?); + } else { + matched_entities.push(serde_json::json!({ + "name": e.name, + "entity_type": e.entity_type, + "namespace": e.namespace, + "git_branch": e.git_branch, + "observations_count": e.observations.len() + })); + } } } - Ok::(serde_json::to_string(&result)?) + Ok::(serde_json::to_string(&matched_entities)?) })?; Ok(data) } diff --git a/server/src/handlers/meta.rs b/server/src/handlers/meta.rs index ddc6fdf..8009af8 100644 --- a/server/src/handlers/meta.rs +++ b/server/src/handlers/meta.rs @@ -23,14 +23,14 @@ impl McpTool for LogDecisionHandler { let idx = state.get_search_index(); let mut final_id = String::new(); + let mut adrs_to_index = Vec::new(); state.code.adrs.modify(|adrs| { if let Some(superseded_id) = &req.supersedes { for old_adr in adrs.iter_mut() { if old_adr.id == *superseded_id { old_adr.status = "superseded".to_string(); - // Re-index the modified old ADR - drop(idx.index_adr(old_adr)); + adrs_to_index.push(old_adr.clone()); break; } } @@ -48,10 +48,15 @@ impl McpTool for LogDecisionHandler { timestamp: crate::handlers::utils::now_secs(), }; - drop(idx.index_adr(&a)); + adrs_to_index.push(a.clone()); adrs.push(a); }); + // Index in Tantivy outside the store write lock + for adr in &adrs_to_index { + drop(idx.index_adr(adr)); + } + state.record_activity("decision", &format!("Logged {}: {}", final_id, req.title), Some(&req.decision)); Ok(format!("Logged decision {}: {}", final_id, req.title)) } diff --git a/server/src/handlers/vision.rs b/server/src/handlers/vision.rs index 8979787..8e44ac1 100644 --- a/server/src/handlers/vision.rs +++ b/server/src/handlers/vision.rs @@ -84,6 +84,96 @@ impl McpTool for WriteClipboardHandler { } } +pub fn get_native_clipboard_text() -> Option { + if let Ok(mut clipboard) = arboard::Clipboard::new() { + if let Ok(text) = clipboard.get_text() { + if !text.trim().is_empty() { + return Some(text); + } + } + } + + let mut cmd_wl = std::process::Command::new("wl-paste"); + cmd_wl.arg("--no-newline"); + if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() { + cmd_wl.env("WAYLAND_DISPLAY", "wayland-0"); + cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir"); + } + if let Ok(output) = cmd_wl.output() { + if output.status.success() && !output.stdout.is_empty() { + if let Ok(text) = String::from_utf8(output.stdout) { + if !text.trim().is_empty() { + return Some(text); + } + } + } + } + + let mut cmd_xc = std::process::Command::new("xclip"); + cmd_xc.args(["-selection", "clipboard", "-o"]); + if std::env::var("DISPLAY").is_err() { + cmd_xc.env("DISPLAY", ":0"); + } + if let Ok(output) = cmd_xc.output() { + if output.status.success() && !output.stdout.is_empty() { + if let Ok(text) = String::from_utf8(output.stdout) { + if !text.trim().is_empty() { + return Some(text); + } + } + } + } + + None +} + +pub fn get_native_clipboard_image() -> Option { + if let Ok(mut clipboard) = arboard::Clipboard::new() { + if let Ok(image_data) = clipboard.get_image() { + if let Some(img) = ImageBuffer::, _>::from_raw( + image_data.width as u32, + image_data.height as u32, + image_data.bytes.into_owned(), + ) { + return Some(image::DynamicImage::ImageRgba8(img)); + } + } + } + + for mime in &["image/png", "image/jpeg", "image/bmp", "image/tiff"] { + let mut cmd_wl = std::process::Command::new("wl-paste"); + cmd_wl.args(["--type", mime]); + if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() { + cmd_wl.env("WAYLAND_DISPLAY", "wayland-0"); + cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir"); + } + if let Ok(output) = cmd_wl.output() { + if output.status.success() && !output.stdout.is_empty() { + if let Ok(img) = image::load_from_memory(&output.stdout) { + return Some(img); + } + } + } + } + + for mime in &["image/png", "image/jpeg", "image/bmp"] { + let mut cmd_xc = std::process::Command::new("xclip"); + cmd_xc.args(["-selection", "clipboard", "-t", mime, "-o"]); + if std::env::var("DISPLAY").is_err() { + cmd_xc.env("DISPLAY", ":0"); + } + if let Ok(output) = cmd_xc.output() { + if output.status.success() && !output.stdout.is_empty() { + if let Ok(img) = image::load_from_memory(&output.stdout) { + return Some(img); + } + } + } + } + + None +} + pub struct ReadClipboardHandler; #[async_trait] @@ -104,51 +194,40 @@ impl McpTool for ReadClipboardHandler { tokio::task::spawn_blocking(move || -> crate::error::Result { let mut out = serde_json::Map::new(); - if let Ok(mut clipboard) = arboard::Clipboard::new() { - if let Ok(text) = clipboard.get_text() { - if !text.trim().is_empty() { - out.insert("text".into(), json!(text)); - } + if let Some(text) = get_native_clipboard_text() { + out.insert("text".into(), json!(text)); + } + + if let Some(dynamic_img) = get_native_clipboard_image() { + let mut img = dynamic_img; + let max_dim = 1024; + if img.width() > max_dim || img.height() > max_dim { + img = img.resize(max_dim, max_dim, FilterType::Lanczos3); } + let rgb_img = img.into_rgb8(); - if let Ok(image_data) = clipboard.get_image() { - if let Some(img) = ImageBuffer::, _>::from_raw( - image_data.width as u32, - image_data.height as u32, - image_data.bytes.into_owned(), - ) { - let mut dynamic_img = image::DynamicImage::ImageRgba8(img); - let max_dim = 1024; - if dynamic_img.width() > max_dim || dynamic_img.height() > max_dim { - dynamic_img = dynamic_img.resize(max_dim, max_dim, FilterType::Lanczos3); - } - let rgb_img = dynamic_img.into_rgb8(); + let cache_dir = dirs::home_dir() + .unwrap_or_default() + .join(".gemini/mcp_memory/clipboard"); + let _ = std::fs::create_dir_all(&cache_dir); - let cache_dir = dirs::home_dir() - .unwrap_or_default() - .join(".gemini/mcp_memory/clipboard"); - let _ = std::fs::create_dir_all(&cache_dir); + let timestamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs(); + let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp)); - let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_secs(); - let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp)); + if rgb_img + .save_with_format(&file_path, image::ImageFormat::Jpeg) + .is_ok() + { + let path_str = file_path.to_string_lossy().to_string(); + out.insert("image_path".into(), json!(path_str)); - if rgb_img - .save_with_format(&file_path, image::ImageFormat::Jpeg) - .is_ok() - { - let path_str = file_path.to_string_lossy().to_string(); - out.insert("image_path".into(), json!(path_str)); - - // Read image bytes for base64 encoding if needed by vision - if let Ok(bytes) = std::fs::read(&file_path) { - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes); - out.insert("image_base64".into(), json!(b64)); - } - } + if let Ok(bytes) = std::fs::read(&file_path) { + use base64::Engine; + let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes); + out.insert("image_base64".into(), json!(b64)); } } } @@ -161,7 +240,6 @@ impl McpTool for ReadClipboardHandler { let mut final_obj = result; if let Some(b64) = final_obj.get("image_base64").and_then(|v| v.as_str()) { let b64_str = b64.to_string(); - // Remove huge base64 string from final user output if let Some(obj) = final_obj.as_object_mut() { obj.remove("image_base64"); } diff --git a/server/src/handlers/workspaces.rs b/server/src/handlers/workspaces.rs index e79691a..850a618 100644 --- a/server/src/handlers/workspaces.rs +++ b/server/src/handlers/workspaces.rs @@ -630,8 +630,8 @@ impl McpTool for SnippetsHandler { let req: SnippetsTool = serde_json::from_value(args).map_err(|e| e.to_string())?; match req.action { SnippetAction::Store => { - let name = req.query.or(req.id).ok_or_else(|| { - crate::error::AppError::Internal("Missing required parameter 'query' or 'id' as snippet name for action 'store'. Next step: Provide snippet name in 'query' field and retry.".to_string()) + let name = req.name.or(req.query).or(req.id).ok_or_else(|| { + crate::error::AppError::Internal("Missing required parameter 'name', 'query', or 'id' as snippet name for action 'store'. Next step: Provide snippet name in 'name' or 'query' field and retry.".to_string()) })?; let lang = req.language.unwrap_or_else(|| "text".to_string()); let code = req.code.unwrap_or_default(); @@ -648,7 +648,7 @@ impl McpTool for SnippetsHandler { ).await } SnippetAction::Search => { - let q = req.query.unwrap_or_default(); + let q = req.query.or(req.name).unwrap_or_default(); if req.hybrid.unwrap_or(false) { crate::handlers::meta::SearchSnippetsHybridHandler.execute(serde_json::json!({"query": q, "tags": req.tags}), state).await } else { @@ -656,14 +656,14 @@ impl McpTool for SnippetsHandler { } } SnippetAction::Delete => { - let id = req.id.or(req.query).ok_or_else(|| { - crate::error::AppError::Internal("Missing required parameter 'id' or 'query' for action 'delete'. Next step: Provide snippet ID/name in request and retry.".to_string()) + let id = req.name.or(req.id).or(req.query).ok_or_else(|| { + crate::error::AppError::Internal("Missing required parameter 'name', 'id', or 'query' for action 'delete'. Next step: Provide snippet ID/name in request and retry.".to_string()) })?; DeleteSnippetHandler.execute(serde_json::json!({"name": id}), state).await } SnippetAction::Tag => { - let id = req.id.or(req.query).ok_or_else(|| { - crate::error::AppError::Internal("Missing required parameter 'id' or 'query' for action 'tag'. Next step: Provide snippet ID/name and 'tags' array in request and retry.".to_string()) + let id = req.name.or(req.id).or(req.query).ok_or_else(|| { + crate::error::AppError::Internal("Missing required parameter 'name', 'id', or 'query' for action 'tag'. Next step: Provide snippet ID/name and 'tags' array in request and retry.".to_string()) })?; let tags = req.tags.unwrap_or_default(); TagSnippetHandler.execute(serde_json::json!({"name": id, "tags": tags}), state).await diff --git a/server/src/instructions.md b/server/src/instructions.md index fc12a2a..1e4a3e2 100644 --- a/server/src/instructions.md +++ b/server/src/instructions.md @@ -19,7 +19,7 @@ To prevent graph fragmentation and ensure seamless LLM context retrieval: ## 2. Consolidated Smart Tools Architecture -The server consolidates granular single-purpose tools into domain-named smart tools. Always prefer the consolidated tools over legacy aliases: +The server consolidates granular single-purpose tools into domain-named smart tools. Always prefer the consolidated tools: * **`tasks`**: Complete task lifecycle management. - `action: "add"`: Create a new task (requires `title`, optional `description`, `git_branch`, `parent_id`, `dependencies`). @@ -85,49 +85,55 @@ The server consolidates granular single-purpose tools into domain-named smart to - `action: "register"`: Register target environment (requires `name`). - `action: "get_details"`: Read full environment details. -* **`clipboard`**: OS Clipboard management. - - `action: "read"`: Read OS clipboard. +* **`clipboard`**: OS Clipboard management (native cross-platform support for Windows and Linux wl-paste/xclip). + - `action: "read"`: Read OS clipboard text, HTML, file paths, or JPEG image. - `action: "write"`: Write text/html/files/image to clipboard. - `action: "toggle_watch"`: Toggle auto-clipboard watcher. ---- +* **`manage_subagent_namespace`**: Subagent memory isolation. + - `action: "create"`: Create namespace. + - `action: "purge"`: Purge namespace. + - `action: "condense"`: Condense namespace. -## 3. Subgraph Expansion & Multi-Hop Navigation -- **Tool:** `get_subgraph` -- **When to use:** When you need to understand the complete architectural neighborhood surrounding a specific component, module, or database table. -- **Behavior:** Performs a multi-hop Breadth-First Search (BFS) around a `root_node` (or `root_entity`) up to a requested `depth` (e.g. 1 to 3 hops) and returns all connected entities and relations. Pass `format: "markdown_tree"` to generate a compact, token-budgeted Markdown topology tree capped within a requested `max_tokens` budget. +* **`manage_checkpoint`**: State snapshots and rollbacks. + - `action: "checkpoint"` (or `"create_snapshot"`): Create state snapshot. + - `action: "restore"` (or `"restore_snapshot"`): Restore point-in-time snapshot. --- -## 4. Automated Error Fix Auto-Matcher -- **Tools:** `log_error_fix`, `search_error_fixes` (and alias `suggest_error_fix`) -- **When to use:** When encountering a build error, test failure, or stack trace. Call `search_error_fixes` with either a text `query` or `stack_trace` before attempting a fix from scratch. -- **Behavior:** Computes cosine similarity between error trace embeddings and past resolution logs when `stack_trace` is provided, or keyword filtering when `query` is provided, returning top matched solutions, modified files, and git commits. +## 3. Knowledge Graph Operations & Multi-Hop Navigation +* **`create_entities`**: Create entity nodes (`entity_type` PascalCase). +* **`create_relations`**: Create relations (`relation_type` snake_case). +* **`add_observations`**: Add new observations to existing entities. +* **`delete_entities`**, **`delete_observations`**, **`delete_relations`**: Graph cleanup. +* **`read_graph`**, **`search_nodes`**, **`open_nodes`**: Graph query and entity lookup. +* **`visualize_graph`**: Generate graph visualizer output. +* **`condense_entity`**, **`merge_entities`**, **`find_orphans`**: Graph maintenance. +* **`get_subgraph`**: Perform multi-hop BFS neighborhood expansion tree around a root entity up to `depth` hops with token budget. +* **`sweep_graph_health`**: Self-healing graph health auditor for near-duplicates and orphaned nodes. +* **`resolve_stale_symbols`**, **`summarize_subgraph`**: Symbol resolution and subgraph summarization. --- -## 5. Memory State Checkpointing & Rollbacks -- **Tool:** `checkpoint_state`, `restore_state` (or `create_snapshot`, `restore_snapshot`) -- **When to use:** Before initiating a large refactor, running experimental subagent tasks, or executing destructive batch operations. -- **Behavior:** Saves or restores a point-in-time snapshot of graph entities, active tasks, and tech debt backlogs. +## 4. Search, Code & AST Intelligence +* **`omni_search`**: High-performance unified search across tasks, code changes, ADRs, debt, and snippets. +* **`semantic_code_search`**: Vector similarity search across indexing embeddings. +* **`read_file_skeleton`**: Extract file declarations/skeletons using Tree-Sitter (Rust, Python, Go, Java, C, C++, JSX/TSX). +* **`replace_ast_node`**: Structural AST node replacement across files. +* **`find_symbol_references`**, **`get_callers`**, **`analyze_impact`**: Code dependency and blast radius analysis. +* **`read_directory_architecture`**: Inspect directory structure and code organization. --- -## 6. Self-Healing Graph Health Sweeper -- **Tool:** `sweep_graph_health` -- **When to use:** Periodically or before committing major graph changes to audit entity consistency. -- **Behavior:** Detects orphaned nodes (0 relations), computes name similarity to identify near-duplicates (e.g., `APIGateway` vs `ApiGateway`), and provides structured `merge_entities` recommendations or auto-prunes orphans. - ---- - -## 7. Causal Lineage & Provenance Tracker -- **Tool:** `query_lineage` -- **When to use:** When asking *"Why was this component modified?"* or *"What task or ADR led to this code change?"* -- **Behavior:** Searches across tasks, ADRs, audit ledger entries, and error fixes to assemble a unified chronological timeline explaining the provenance behind any file, symbol, or commit. - ---- - -## 8. LLM Pre-Flight Context Bundle -- **Tool:** `get_preflight_context` -- **When to use:** At the start of a turn or subagent task to gain total situational awareness in 1 call. -- **Behavior:** Aggregates current active branch, in-progress tasks with acceptance criteria, pinned files, top open tech debts, and active unverified hypotheses into a consolidated executive context bundle. +## 5. Audit Logging, Diagnostics & Provenance +* **`log_code_change`**, **`query_recent_changes`**: Audit trail of code edits, commits, and symbol modifications. +* **`log_error_fix`**, **`search_error_fixes`**: Record error trace and verified fix; search past resolutions via vector similarity. +* **`query_lineage`**: Reconstruct causal timeline explaining why a file/symbol/commit was modified. +* **`log_hypothesis`**, **`query_hypotheses`**: Diagnostic hypothesis tracking during complex debugging. +* **`get_preflight_context`**: Consolidated executive bundle (branch, active tasks, criteria, pinned files, tech debt) for instant context warmup. +* **`learn_preference`**, **`read_preferences`**: Track user preferences and project conventions. +* **`get_project_health`**: Health metrics, active task counts, debt backlog, and graph statistics. +* **`broadcast_agent_signal`**, **`query_agent_signals`**: Inter-agent coordination signals. +* **`auto_session_checkpoint`**: Auto-save session state checkpoint. +* **`get_active_worktree_context`**, **`query_git_diffs`**: Git worktree status and diff inspection. +* **`watch_process_logs`**, **`get_recent_logs`**: Process log monitoring and retrieval. diff --git a/server/src/lib.rs b/server/src/lib.rs index 6a1c6d1..2830e4a 100644 --- a/server/src/lib.rs +++ b/server/src/lib.rs @@ -255,7 +255,7 @@ pub async fn run_server(state: Arc) -> Result<(), Box) -> Result<(), Box) -> Result<(), Box) -> Result<(), Box Option { @@ -440,15 +435,15 @@ pub fn run_cli() -> Result<(), Box> { if cli.exit { let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string()); let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default(); - let mut cmd = std::process::Command::new("curl"); - cmd.arg("-k").arg("-X").arg("POST"); - if !token.is_empty() { - cmd.arg("-H") - .arg(format!("Authorization: Bearer {}", token.trim())); - } - let _ = cmd - .arg(format!("http://127.0.0.1:{}/shutdown", port)) - .output(); + let rt = tokio::runtime::Runtime::new()?; + rt.block_on(async { + let client = reqwest::Client::builder().build().unwrap_or_default(); + let mut req = client.post(format!("http://127.0.0.1:{}/shutdown", port)); + if !token.trim().is_empty() { + req = req.header("Authorization", format!("Bearer {}", token.trim())); + } + let _ = req.send().await; + }); if cli.restart { std::thread::sleep(Duration::from_secs(2)); diff --git a/server/src/mcp.rs b/server/src/mcp.rs index 7e65bc0..75ad78a 100644 --- a/server/src/mcp.rs +++ b/server/src/mcp.rs @@ -38,6 +38,18 @@ pub fn tool_def(name: &str, fallback_desc: &str) -> serde_json::V }) } +pub fn parse_args( + args: serde_json::Value, + tool_name: &str, +) -> crate::error::Result { + serde_json::from_value::(args).map_err(|e| { + crate::error::AppError::BadRequest(format!( + "Invalid arguments provided for tool '{}': {}. Next step: Verify parameter names, required fields, and types against the tool input schema.", + tool_name, e + )) + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/server/src/router.rs b/server/src/router.rs index 6ff058e..763f1a4 100644 --- a/server/src/router.rs +++ b/server/src/router.rs @@ -73,7 +73,7 @@ impl McpResource for GraphEntitiesResource { } }) .await - .unwrap() + .map_err(|e| crate::error::AppError::Internal(e.to_string()))? } } @@ -97,7 +97,7 @@ impl McpResource for GraphRelationsResource { Ok(serde_json::to_string_pretty(&data)?) }) .await - .unwrap() + .map_err(|e| crate::error::AppError::Internal(e.to_string()))? } } @@ -124,7 +124,7 @@ impl McpResource for TasksActiveResource { Ok(serde_json::to_string_pretty(&data)?) }) .await - .unwrap() + .map_err(|e| crate::error::AppError::Internal(e.to_string()))? } } @@ -148,7 +148,7 @@ impl McpResource for DecisionsActiveResource { Ok(serde_json::to_string_pretty(&data)?) }) .await - .unwrap() + .map_err(|e| crate::error::AppError::Internal(e.to_string()))? } } @@ -172,7 +172,7 @@ impl McpResource for TechDebtActiveResource { Ok(serde_json::to_string_pretty(&data)?) }) .await - .unwrap() + .map_err(|e| crate::error::AppError::Internal(e.to_string()))? } } @@ -270,7 +270,7 @@ impl McpResource for SessionDeltaResource { Ok(serde_json::to_string_pretty(&delta)?) }) .await - .unwrap() + .map_err(|e| crate::error::AppError::Internal(e.to_string()))? } } @@ -361,7 +361,7 @@ impl MemoryHandler { Ok(serde_json::to_string_pretty(&*items)?) }) .await - .unwrap() + .map_err(|e| crate::error::AppError::Internal(e.to_string()))? } } struct PinnedFilesResource; @@ -383,7 +383,7 @@ impl MemoryHandler { Ok(serde_json::to_string_pretty(&*items)?) }) .await - .unwrap() + .map_err(|e| crate::error::AppError::Internal(e.to_string()))? } } @@ -406,7 +406,7 @@ impl MemoryHandler { Ok(serde_json::to_string_pretty(&*items)?) }) .await - .unwrap() + .map_err(|e| crate::error::AppError::Internal(e.to_string()))? } } @@ -715,7 +715,7 @@ impl MemoryHandler { let category = match name { "clipboard" => "CLIPBOARD", - "create_entities" | "create_relations" | "read_graph" | "get_subgraph" | "search_graph" | "get_schema" => "GRAPH", + "create_entities" | "create_relations" | "add_observations" | "delete_entities" | "delete_observations" | "delete_relations" | "read_graph" | "get_subgraph" | "search_graph" | "get_schema" | "search_nodes" | "open_nodes" | "visualize_graph" | "condense_entity" | "merge_entities" | "find_orphans" | "sweep_graph_health" | "summarize_subgraph" | "query_graph_path" | "resolve_stale_symbols" => "GRAPH", "decisions" => "DECISION", "log_code_change" => "CODE", "log_error_fix" => "ERROR_FIX", @@ -1131,5 +1131,111 @@ mod tests { assert_eq!(res2["jsonrpc"], "2.0"); assert_eq!(res2["id"], 6); } + + #[tokio::test] + async fn test_tools_list_schema_and_descriptions() { + let dir = tempdir().unwrap(); + let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap())); + let handler = MemoryHandler::new(state); + + let list_req = json!({ + "jsonrpc": "2.0", + "id": 100, + "method": "tools/list", + "params": {} + }); + + let res = handler.handle_request(list_req).await.unwrap(); + assert_eq!(res["jsonrpc"], "2.0"); + assert_eq!(res["id"], 100); + + let tools = res["result"]["tools"].as_array().expect("tools array expected"); + assert!(tools.len() >= 40, "Expected at least 40 registered MCP tools, found {}", tools.len()); + + for tool in tools { + let name = tool["name"].as_str().expect("tool name must be string"); + assert!(!name.is_empty(), "Tool name cannot be empty"); + + let desc = tool["description"].as_str().expect("tool description must be string"); + assert!(!desc.is_empty(), "Tool '{}' missing description", name); + + let schema = tool["inputSchema"].as_object().expect("inputSchema must be object"); + assert!( + schema.contains_key("properties") || schema.contains_key("type"), + "Tool '{}' schema must contain type/properties metadata", + name + ); + } + } + + #[tokio::test] + async fn test_all_tools_llm_request_ok_and_invalid() { + let dir = tempdir().unwrap(); + let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap())); + let handler = MemoryHandler::new(state); + + // 1. Verify invalid LLM requests fail gracefully with isError: true or structured response + for (tool_name, _tool) in &handler.tools { + let invalid_req = json!({ + "jsonrpc": "2.0", + "id": 200, + "method": "tools/call", + "params": { + "name": tool_name, + "arguments": { + "invalid_field_unlikely_to_match": "bogus_value", + "action": "non_existent_action_xyz_999" + } + } + }); + + let res = handler.handle_request(invalid_req).await.expect("handler must respond"); + assert_eq!(res["jsonrpc"], "2.0"); + assert!(res["result"]["content"].is_array()); + } + + // 2. Verify valid LLM requests return isError: false + let sample_valid_calls = vec![ + ("read_graph", json!({})), + ("search_nodes", json!({"query": "test"})), + ("tasks", json!({"action": "list"})), + ("decisions", json!({"action": "query"})), + ("tech_debt", json!({"action": "list"})), + ("sticky_notes", json!({"action": "read"})), + ("pinned_files", json!({"action": "list"})), + ("context_workspaces", json!({"action": "list"})), + ("pr_checklist", json!({"action": "get"})), + ("snippets", json!({"action": "search", "query": "test"})), + ("clipboard", json!({"action": "read"})), + ("environment", json!({"action": "read_fingerprint"})), + ("omni_search", json!({"query": "test"})), + ("get_project_health", json!({})), + ("read_preferences", json!({})), + ("find_orphans", json!({})), + ("sweep_graph_health", json!({})), + ("get_next_actionable_tasks", json!({})), + ("get_preflight_context", json!({})), + ("auto_session_checkpoint", json!({"session_summary": "Test summary"})), + ]; + + for (name, args) in sample_valid_calls { + let valid_req = json!({ + "jsonrpc": "2.0", + "id": 201, + "method": "tools/call", + "params": { + "name": name, + "arguments": args + } + }); + + let res = handler.handle_request(valid_req).await.expect("handler must respond"); + assert_eq!(res["jsonrpc"], "2.0"); + assert_eq!( + res["result"]["isError"], false, + "Tool '{}' failed valid execution. Result: {:?}", name, res + ); + } + } } diff --git a/server/src/search.rs b/server/src/search.rs index f8662a0..28beedb 100644 --- a/server/src/search.rs +++ b/server/src/search.rs @@ -57,6 +57,37 @@ impl MemoryIndex { }) } + pub fn new_in_ram() -> tantivy::Result { + let mut schema_builder = Schema::builder(); + let id_field = schema_builder.add_text_field("id", STRING | STORED); + let title_field = schema_builder.add_text_field("title", TEXT | STORED); + let body_field = schema_builder.add_text_field("body", TEXT | STORED); + let type_field = schema_builder.add_text_field("type", STRING | STORED); + let namespace_field = schema_builder.add_text_field("namespace", STRING | STORED); + let schema = schema_builder.build(); + + let index = Index::create_in_ram(schema); + let mut writer = index.writer(50_000_000)?; + writer.delete_all_documents()?; + writer.commit()?; + let reader = index + .reader_builder() + .reload_policy(ReloadPolicy::OnCommitWithDelay) + .try_into()?; + + Ok(Self { + index, + reader, + writer: Arc::new(Mutex::new(writer)), + needs_commit: Arc::new(std::sync::atomic::AtomicBool::new(false)), + id_field, + title_field, + body_field, + type_field, + namespace_field, + }) + } + pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle> { let writer = Arc::clone(&self.writer); let id_field = self.id_field; diff --git a/server/src/state.rs b/server/src/state.rs index 34fb3b8..c81a94a 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -72,7 +72,9 @@ pub struct MemoryState { impl MemoryState { pub fn new(base_dir_str: &str) -> Self { let base = std::path::PathBuf::from(base_dir_str); - std::fs::create_dir_all(&base).expect("Failed to create store dir"); + if let Err(e) = std::fs::create_dir_all(&base) { + tracing::error!("Failed to create store directory at {:?}: {}", base, e); + } let db = crate::db::init_redb(&base); @@ -93,7 +95,7 @@ impl MemoryState { .join(".gemini/mcp_memory/daemon_error.log"); let _ = std::fs::write(&log_path, format!("Failed to create MemoryIndex: {}\n", e)); - std::process::exit(1); + crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex") } }), vector_db: tokio::sync::RwLock::new(None), @@ -174,7 +176,7 @@ impl MemoryState { } pub async fn rebuild_index(self: &Arc) { - let idx = self.search_index.read().unwrap().clone(); + let idx = self.get_search_index(); idx.delete_all(); let entities: Vec<_> = self @@ -242,14 +244,16 @@ impl MemoryState { } }); - let payload = serde_json::json!({ - "jsonrpc": "2.0", - "method": "notifications/activity", - "params": activity - }) - .to_string(); + if self.activity_tx.receiver_count() > 0 { + let payload = serde_json::json!({ + "jsonrpc": "2.0", + "method": "notifications/activity", + "params": activity + }) + .to_string(); - let _ = self.activity_tx.send(payload); + let _ = self.activity_tx.send(payload); + } } pub fn broadcast_task_event(&self, event: TaskEvent) { @@ -460,27 +464,30 @@ impl SearchService { } if !vdb_search { - let mut texts_to_embed = Vec::new(); - let mut metadata = Vec::new(); + let (mut texts_to_embed, mut metadata) = self.state.code.snippets.read_with(|snips| { + let mut texts = Vec::with_capacity(snips.len().min(50)); + let mut meta = Vec::with_capacity(snips.len().min(50)); + for snippet in snips.iter().take(50) { + texts.push(format!("{} {} {}", snippet.name, snippet.description, snippet.code)); + meta.push((snippet.name.clone(), "snippet".to_string(), snippet.description.clone())); + } + (texts, meta) + }); - let snippets = self.state.code.snippets.read_with(|snips| snips.clone()); - for snippet in snippets { - let combined = format!("{} {} {}", snippet.name, snippet.description, snippet.code); - texts_to_embed.push(combined); - metadata.push((snippet.name, "snippet".to_string(), snippet.description)); - } + self.state.code.sticky.read_with(|sticky| { + for note in sticky.iter().take(50) { + texts_to_embed.push(note.content.clone()); + metadata.push(( + "StickyNote".to_string(), + "sticky".to_string(), + note.content.chars().take(200).collect::(), + )); + } + }); - let sticky = self.state.code.sticky.read_with(|s| s.clone()); - for note in sticky { - texts_to_embed.push(note.content.clone()); - metadata.push(( - "StickyNote".to_string(), - "sticky".to_string(), - note.content.chars().take(200).collect::(), - )); - } - - if let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await { + if !texts_to_embed.is_empty() + && let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await + { for (emb, meta) in embeddings.into_iter().zip(metadata) { let sim = cosine_similarity(&query_emb, &emb); results.push(UnifiedSearchResult { diff --git a/server/src/store.rs b/server/src/store.rs index 4ef9cbd..6e18083 100644 --- a/server/src/store.rs +++ b/server/src/store.rs @@ -4,10 +4,98 @@ use std::sync::{Arc, RwLock}; pub const STORE_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("store"); +/// Internal write request dispatched to the single database writer actor. +struct DbWriteTask { + key: String, + data: Vec, + flushed: Arc, +} + +/// Shared centralized write queue actor that handles all database writes serially with micro-batching. +#[derive(Clone)] +pub struct DbWriteQueue { + tx: tokio::sync::mpsc::Sender, +} + +static QUEUE_REGISTRY: std::sync::Mutex, DbWriteQueue)>> = std::sync::Mutex::new(None); + +fn get_or_create_queue(db: Arc) -> DbWriteQueue { + let mut reg = QUEUE_REGISTRY.lock().unwrap_or_else(|e| e.into_inner()); + if let Some((ref existing_db, ref queue)) = *reg { + if Arc::ptr_eq(existing_db, &db) && !queue.tx.is_closed() { + return queue.clone(); + } + } + let new_queue = DbWriteQueue::new(db.clone()); + *reg = Some((db, new_queue.clone())); + new_queue +} + +impl DbWriteQueue { + pub fn new(db: Arc) -> Self { + let (tx, mut rx) = tokio::sync::mpsc::channel::(2048); + + tokio::spawn(async move { + while let Some(first_task) = rx.recv().await { + let mut batch = vec![first_task]; + + // Gold Standard Micro-batching: Drain up to 100 accumulated tasks from queue without blocking + while batch.len() < 100 { + match rx.try_recv() { + Ok(task) => batch.push(task), + Err(_) => break, + } + } + + let db_inner = db.clone(); + let _ = tokio::task::spawn_blocking(move || { + match db_inner.begin_write() { + Ok(write_txn) => { + if let Ok(mut table) = write_txn.open_table(STORE_TABLE) { + for task in &batch { + if let Err(e) = table.insert(task.key.as_str(), task.data.as_slice()) { + tracing::error!("Failed to insert key '{}' into redb: {}", task.key, e); + } + } + } + if let Err(e) = write_txn.commit() { + tracing::error!("Failed to commit batch to redb: {}", e); + } + } + Err(e) => { + tracing::error!("Failed to begin write transaction on redb writer actor: {}", e); + } + } + + // Event-driven notification to all waiting listeners for this micro-batch + for task in batch { + task.flushed.notify_waiters(); + } + }) + .await; + } + }); + + Self { tx } + } + + pub fn push(&self, key: String, data: Vec, flushed: Arc) { + let task = DbWriteTask { key, data, flushed }; + if let Err(e) = self.tx.try_send(task) { + let task = e.into_inner(); + let tx = self.tx.clone(); + tokio::spawn(async move { + let _ = tx.send(task).await; + }); + } + } +} + pub struct Store { pub cache: Arc>, pub flushed: Arc, - tx: tokio::sync::mpsc::Sender<()>, + key: String, + queue: DbWriteQueue, } impl Store { @@ -15,53 +103,14 @@ impl let initial_data = Self::load_from_db(key, &db); let cache = Arc::new(RwLock::new(initial_data)); let flushed = Arc::new(tokio::sync::Notify::new()); - let (tx, mut rx) = tokio::sync::mpsc::channel::<()>(1); + let queue = get_or_create_queue(db); - let db_clone = db.clone(); - let key_clone = key.to_string(); - let cache_clone = cache.clone(); - let flushed_clone = flushed.clone(); - - tokio::spawn(async move { - while rx.recv().await.is_some() { - // Drain any pending notifications accumulated - while rx.try_recv().is_ok() {} - - let db_inner = db_clone.clone(); - let key_inner = key_clone.clone(); - let flushed_inner = flushed_clone.clone(); - let json_data = { - let lock = cache_clone.read().unwrap_or_else(|e| e.into_inner()); - serde_json::to_vec(&*lock) - .map_err(|e| tracing::error!("Failed to serialize memory store: {}", e)) - .ok() - }; - - if let Some(json_data) = json_data { - let _ = tokio::task::spawn_blocking(move || { - // Retry up to 10 times if another Store holds write transaction - for _ in 0..10 { - match db_inner.begin_write() { - Ok(write_txn) => { - if let Ok(mut table) = write_txn.open_table(STORE_TABLE) { - let _ = table.insert(key_inner.as_str(), json_data.as_slice()); - } - let _ = write_txn.commit(); - break; - } - Err(_) => { - std::thread::yield_now(); - } - } - } - }) - .await; - flushed_inner.notify_waiters(); - } - } - }); - - Self { cache, flushed, tx } + Self { + cache, + flushed, + key: key.to_string(), + queue, + } } fn load_from_db(key: &str, db: &Database) -> T { @@ -86,11 +135,16 @@ impl } pub fn modify(&self, f: F) { - { + let cloned_data = { let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner()); f(&mut lock); + lock.clone() + }; + + match serde_json::to_vec(&cloned_data) { + Ok(data) => self.queue.push(self.key.clone(), data, self.flushed.clone()), + Err(e) => tracing::error!("Failed to serialize memory store for key '{}': {}", self.key, e), } - let _ = self.tx.try_send(()); } } diff --git a/server/src/tools.rs b/server/src/tools.rs index 62499bf..9b4f660 100644 --- a/server/src/tools.rs +++ b/server/src/tools.rs @@ -1299,6 +1299,8 @@ pub struct SnippetsTool { pub id: Option, /// Search query or snippet name (required for 'store'). pub query: Option, + /// Snippet name alias (for 'store', 'delete', 'tag'). + pub name: Option, /// Snippet language (for 'store'). pub language: Option, /// Code content (for 'store'). diff --git a/server/tests/parity_test.rs b/server/tests/parity_test.rs index 503def2..2797559 100644 --- a/server/tests/parity_test.rs +++ b/server/tests/parity_test.rs @@ -52,7 +52,12 @@ fn test_eager_tools_parity() { let config_str = std::fs::read_to_string(&win_config_path).unwrap(); let config: serde_json::Value = serde_json::from_str(&config_str).unwrap(); - if let Some(eager) = config["mcpServers"]["memory"]["eagerTools"].as_array() { + let mem_server = if config["mcpServers"]["mcp-memory"].is_object() { + &config["mcpServers"]["mcp-memory"] + } else { + &config["mcpServers"]["memory"] + }; + if let Some(eager) = mem_server["eagerTools"].as_array() { for tool in eager { let name = tool.as_str().unwrap(); assert!(