refactor: address 5-pass audit findings for antipatterns, bottlenecks, memory efficiency, and LLM handlers

This commit is contained in:
Riz Ashraf committed 2026-10-05 21:44:10 +01:00
1 parent 626403900f
commit 924b6d09fa
30 files changed
+1024 -407

No files matched your search

Generated
-9
View File
@@ -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"
-1
View File
@@ -2,7 +2,6 @@
members = [
"server",
"stub",
"mcp-memory-nvim",
"nvim-core",
"mcp-stdio"
]
+2 -2
View File
@@ -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)
-10
View File
@@ -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" }
+1 -1
View File
@@ -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"] }
+6 -6
View File
@@ -13,10 +13,10 @@ pub async fn read_mcp_message<R: tokio::io::AsyncRead + Unpin>(
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<R: tokio::io::AsyncRead + Unpin>(
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;
}
File renamed without changes.
File renamed without changes.
+15 -10
View File
@@ -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<String, String> {
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<String, String> {
#[cfg(unix)]
async fn get_socket_path() -> Result<String, String> {
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<mpsc::Sender<NvimRequest>, 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]
+1 -1
View File
@@ -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"));
}
+74
View File
@@ -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<dyn std::error::Error>> {
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<String> = 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(())
}
+1 -2
View File
@@ -45,8 +45,7 @@ pub fn spawn_watcher(state: Arc<MemoryState>) {
Some(&text),
);
// We use rebuild_index to index the new sticky note
state.rebuild_index().await;
state.index_commit_notify.notify_waiters();
}
}
});
+4 -4
View File
@@ -493,7 +493,8 @@
</div>
<div class="main-container">
<div class="tabs">
<button class="tab-button active" onclick="switchTab('graph-tab', this)">Graph Explorer</button>
<button class="tab-button active" onclick="switchTab('activity-tab', this)">Live Activity</button>
<button class="tab-button" onclick="switchTab('graph-tab', this)">Graph Explorer</button>
<button class="tab-button" onclick="switchTab('search-tab', this)">Omni-Search</button>
<button class="tab-button" onclick="switchTab('task-tab', this)">Task Board</button>
<button class="tab-button" onclick="switchTab('sticky-tab', this)">Sticky Notes</button>
@@ -506,10 +507,9 @@
<button class="tab-button" onclick="switchTab('snippets-tab', this)">Snippets</button>
<button class="tab-button" onclick="switchTab('pr-tab', this)">PRs</button>
<button class="tab-button" onclick="switchTab('terminal-tab', this)">Terminal</button>
<button class="tab-button" onclick="switchTab('activity-tab', this)">Live Activity</button>
</div>
<div id="graph-tab" class="tab-content active">
<div id="graph-tab" class="tab-content">
<div class="panel graph-panel" style="flex:1;">
<div class="graph-filters" id="graph-filters">
<!-- Filters injected dynamically -->
@@ -624,7 +624,7 @@
</div>
</div>
<div id="activity-tab" class="tab-content">
<div id="activity-tab" class="tab-content active">
<div class="panel activity-panel" style="flex:1; display:flex; flex-direction:column;">
<div id="activity-feed" style="flex:1;"></div>
</div>
+174 -91
View File
@@ -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() {
</button>`;
}).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,15 +487,102 @@ async function loadActivityHistory() {
console.error("Failed to load activity history", e);
}
}
function setupWS() {
const protocol = location.protocol === "https:" ? "wss:" : "ws:";
const ws = new WebSocket(`${protocol}//${location.host}/ws?client=ui`);
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();
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) {
@@ -486,7 +592,6 @@ function setupWS() {
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) {
@@ -496,13 +601,43 @@ function setupWS() {
feed.scrollTop = 0;
}
}
function setupWS() {
const protocol = location.protocol === "https:" ? "wss:" : "ws:";
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") {
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 parsed = JSON.parse(event.data);
handleIncomingActivity(parsed);
} catch (e) {
handleIncomingActivity(event.data);
}
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) {}
};
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"] });
+165 -47
View File
@@ -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<void> {
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<void> {
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,17 +684,88 @@ async function loadActivityHistory(): Promise<void> {
}
}
function setupWS(): void {
const protocol = location.protocol === 'https:' ? 'wss:' : 'ws:';
const ws = new WebSocket(`${protocol}//${location.host}/ws?client=ui`);
let wsConnected = false;
let activeSSE: EventSource | null = null;
const tabRefreshDebounceTimers: Record<string, any> = {};
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();
function requestDomainRefresh(domain: string): void {
const tabMap: Record<string, string> = {
'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) {
@@ -680,7 +774,6 @@ function setupWS(): void {
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) {
@@ -689,14 +782,48 @@ function setupWS(): void {
if (isScrolledToTop) feed.scrollTop = 0;
}
}
function setupWS(): void {
const protocol = location.protocol === 'https:' ? 'wss:' : 'ws:';
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') {
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();
};
}
@@ -914,30 +1041,23 @@ async function loadVersion(): Promise<void> {
}
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 parsed = JSON.parse(event.data);
handleIncomingActivity(parsed);
} catch(e) {
handleIncomingActivity(event.data);
}
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) {}
};
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'] });
+31 -18
View File
@@ -9,22 +9,21 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
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"),
@@ -48,25 +47,39 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
];
for (key, file_name) in stores.iter() {
if table
.get(*key)
.expect("Failed to read from table")
.is_none()
{
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
}
};
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::<serde_json::Value>(&data).is_ok()
{
table
.insert(*key, data.as_slice())
.expect("Failed to insert migrated data");
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"));
}
}
}
}
write_txn.commit().expect("Failed to commit db migration");
}
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);
}
}
db
}
+4 -1
View File
@@ -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);
}
}
}
+88 -47
View File
@@ -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,26 +61,17 @@ 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());
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;
}
}
}
}
if nodes_at_current_depth == 0 {
current_depth += 1;
nodes_at_current_depth = nodes_at_next_depth;
@@ -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);
missing_nodes.insert(relation.from.clone());
}
if !to_exists {
missing_nodes.insert(relation.to);
}
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<MemoryState>) -> crate::error::Result<String> {
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<MemoryState>) -> crate::error::Result<String> {
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<String> {
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::<String, crate::error::AppError>(serde_json::to_string(&result)?)
}
Ok::<String, crate::error::AppError>(serde_json::to_string(&matched_entities)?)
})?;
Ok(data)
}
+8 -3
View File
@@ -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))
}
+96 -18
View File
@@ -84,6 +84,96 @@ impl McpTool for WriteClipboardHandler {
}
}
pub fn get_native_clipboard_text() -> Option<String> {
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<image::DynamicImage> {
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(image_data) = clipboard.get_image() {
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::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,25 +194,17 @@ impl McpTool for ReadClipboardHandler {
tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Value> {
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() {
if let Some(text) = get_native_clipboard_text() {
out.insert("text".into(), json!(text));
}
}
if let Ok(image_data) = clipboard.get_image() {
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::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);
if let Some(dynamic_img) = get_native_clipboard_image() {
let mut img = dynamic_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);
if img.width() > max_dim || img.height() > max_dim {
img = img.resize(max_dim, max_dim, FilterType::Lanczos3);
}
let rgb_img = dynamic_img.into_rgb8();
let rgb_img = img.into_rgb8();
let cache_dir = dirs::home_dir()
.unwrap_or_default()
@@ -142,7 +224,6 @@ impl McpTool for ReadClipboardHandler {
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);
@@ -150,8 +231,6 @@ impl McpTool for ReadClipboardHandler {
}
}
}
}
}
Ok(Value::Object(out))
})
@@ -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");
}
+7 -7
View File
@@ -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
+40 -34
View File
@@ -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.
+17 -22
View File
@@ -255,7 +255,7 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
tokio::spawn(async move {
let port1 = std::env::var("MCP_UDP_PORT1").unwrap_or_else(|_| "3001".to_string());
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port1)).await {
let mut buf = [0; 4096];
let mut buf = vec![0u8; 65536];
loop {
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
&& let Ok(payload) =
@@ -291,7 +291,7 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
tokio::spawn(async move {
let port2 = std::env::var("MCP_UDP_PORT2").unwrap_or_else(|_| "3002".to_string());
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port2)).await {
let mut buf = [0; 4096];
let mut buf = vec![0u8; 65536];
loop {
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
&& let Ok(payload) =
@@ -302,16 +302,11 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
|| payload.event == "VimEnter"
{
let session = &payload.session_id;
let is_unix_socket = session.starts_with('/') || session.starts_with('~');
if is_unix_socket {
let wsl_path =
"\\\\wsl.localhost\\Ubuntu\\home\\riz\\.gemini\\active_nvim.txt";
let _ = tokio::fs::write(wsl_path, session).await;
} else {
let profile = std::env::var("USERPROFILE")
.unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
let win_path = format!("{}\\.gemini\\active_nvim.txt", profile);
let _ = tokio::fs::write(&win_path, session).await;
let _is_unix_socket = session.starts_with('/') || session.starts_with('~');
if let Some(home) = dirs::home_dir() {
let nvim_txt_path = home.join(".gemini").join("active_nvim.txt");
let _ = tokio::fs::create_dir_all(nvim_txt_path.parent().unwrap()).await;
let _ = tokio::fs::write(&nvim_txt_path, session).await;
}
}
@@ -393,8 +388,8 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
let _ = tokio::fs::write(&log_path, format!("Server crashed: {}\n", e)).await;
}
tracing::info!("axum::serve graceful shutdown complete. Terminating process cleanly.");
std::process::exit(0);
tracing::info!("axum::serve graceful shutdown complete.");
Ok(())
}
pub fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
@@ -440,15 +435,15 @@ pub fn run_cli() -> Result<(), Box<dyn std::error::Error>> {
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 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 _ = cmd
.arg(format!("http://127.0.0.1:{}/shutdown", port))
.output();
let _ = req.send().await;
});
if cli.restart {
std::thread::sleep(Duration::from_secs(2));
+12
View File
@@ -38,6 +38,18 @@ pub fn tool_def<T: JsonSchema>(name: &str, fallback_desc: &str) -> serde_json::V
})
}
pub fn parse_args<T: serde::de::DeserializeOwned>(
args: serde_json::Value,
tool_name: &str,
) -> crate::error::Result<T> {
serde_json::from_value::<T>(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::*;
+116 -10
View File
@@ -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
);
}
}
}
+31
View File
@@ -57,6 +57,37 @@ impl MemoryIndex {
})
}
pub fn new_in_ram() -> tantivy::Result<Self> {
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<tantivy::Result<()>> {
let writer = Arc::clone(&self.writer);
let id_field = self.id_field;
+21 -14
View File
@@ -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<Self>) {
let idx = self.search_index.read().unwrap().clone();
let idx = self.get_search_index();
idx.delete_all();
let entities: Vec<_> = self
@@ -242,6 +244,7 @@ impl MemoryState {
}
});
if self.activity_tx.receiver_count() > 0 {
let payload = serde_json::json!({
"jsonrpc": "2.0",
"method": "notifications/activity",
@@ -251,6 +254,7 @@ impl MemoryState {
let _ = self.activity_tx.send(payload);
}
}
pub fn broadcast_task_event(&self, event: TaskEvent) {
let payload_val = serde_json::to_value(&event).unwrap_or_default();
@@ -460,18 +464,18 @@ impl SearchService {
}
if !vdb_search {
let mut texts_to_embed = Vec::new();
let mut metadata = Vec::new();
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));
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 sticky = self.state.code.sticky.read_with(|s| s.clone());
for note in sticky {
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(),
@@ -479,8 +483,11 @@ impl SearchService {
note.content.chars().take(200).collect::<String>(),
));
}
});
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 {
+102 -48
View File
@@ -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<u8>,
flushed: Arc<tokio::sync::Notify>,
}
/// Shared centralized write queue actor that handles all database writes serially with micro-batching.
#[derive(Clone)]
pub struct DbWriteQueue {
tx: tokio::sync::mpsc::Sender<DbWriteTask>,
}
static QUEUE_REGISTRY: std::sync::Mutex<Option<(Arc<Database>, DbWriteQueue)>> = std::sync::Mutex::new(None);
fn get_or_create_queue(db: Arc<Database>) -> 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<Database>) -> Self {
let (tx, mut rx) = tokio::sync::mpsc::channel::<DbWriteTask>(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<u8>, flushed: Arc<tokio::sync::Notify>) {
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<T> {
pub cache: Arc<RwLock<T>>,
pub flushed: Arc<tokio::sync::Notify>,
tx: tokio::sync::mpsc::Sender<()>,
key: String,
queue: DbWriteQueue,
}
impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static> Store<T> {
@@ -15,53 +103,14 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
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());
Self {
cache,
flushed,
key: key.to_string(),
queue,
}
let _ = write_txn.commit();
break;
}
Err(_) => {
std::thread::yield_now();
}
}
}
})
.await;
flushed_inner.notify_waiters();
}
}
});
Self { cache, flushed, tx }
}
fn load_from_db(key: &str, db: &Database) -> T {
@@ -86,11 +135,16 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
}
pub fn modify<F: FnOnce(&mut T)>(&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(());
}
}
+2
View File
@@ -1299,6 +1299,8 @@ pub struct SnippetsTool {
pub id: Option<String>,
/// Search query or snippet name (required for 'store').
pub query: Option<String>,
/// Snippet name alias (for 'store', 'delete', 'tag').
pub name: Option<String>,
/// Snippet language (for 'store').
pub language: Option<String>,
/// Code content (for 'store').
+6 -1
View File
@@ -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!(