riz/mcp-memory
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(embedding,vision): candle embeddings with offline fallback, on-demand clipboard vision capture, and concurrency audit

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-10-07 06:36:09 +01:00
1 parent 5bd8b1587a
commit e4a0fe72df
47 files changed
+6292 -3503
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/Cargo.toml
+7 -2
View File
@@ -27,7 +27,11 @@ tracing-appender = "0.2.5"
rmcp = { version = "3.4.0", features = ["server"] }
thiserror = "2.0.20"
tower = { version = "0.5.3", features = ["util"] }
fastembed = "7.1.0"
candle-core = "0.11.0"
candle-nn = "0.11.0"
candle-transformers = "0.11.0"
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }
hf-hub = { version = "1.0.0", features = ["blocking"] }
arboard = "3.6.1"
image = "0.25.10"
base64 = "0.23.1"
@@ -42,7 +46,6 @@ tree-sitter-cpp = "0.23.4"
tree-sitter-go = "0.25.0"
ignore = "0.4.33"
notify = "7"
qdrant-client = "1.19.0"
blake3 = "1.8"
tokio-stream = { version = "0.1", features = ["sync"] }
@@ -54,3 +57,5 @@ winres = "0.1.12"
tempfile = "3.27.0"
[target.'cfg(windows)'.dependencies]
windows-sys = { version = "0.59.0", features = ["Win32_System_DataExchange", "Win32_System_Memory"] }
server/src/api/rest.rs
+1
View File
@@ -95,6 +95,7 @@ pub async fn gate_set_handler(
status,
reason: body.reason.clone(),
timestamp: crate::handlers::utils::now_secs(),
..Default::default()
};
app_state.handler.state.env.gates.modify(|gates| {
gates.retain(|g| !(g.action == record.action && g.target == record.target));
server/src/api/setup.rs
+1
View File
@@ -131,6 +131,7 @@ pub fn create_router(app_state: Arc<AppState>) -> Router {
expires_at: None,
namespace: "global".to_string(),
timestamp: now,
..Default::default()
};
mem_state.telemetry.handoff_memos.modify(|m| m.push(memo));
server/src/api/telemetry.rs
+16 -6
View File
@@ -72,7 +72,11 @@ pub async fn nvim_telemetry_handler(
}
if let Some(ref f) = payload.file {
state.handler.state.record_activity("editor", &format!("Neovim {}: {}", payload.event, f), None);
state.handler.state.record_activity(
"editor",
&format!("Neovim {}: {}", payload.event, f),
None,
);
}
axum::Json(serde_json::json!({"status": "ok"}))
@@ -99,12 +103,16 @@ pub async fn terminal_telemetry_handler(
State(state): State<Arc<AppState>>,
axum::Json(payload): axum::Json<TerminalHistory>,
) -> impl axum::response::IntoResponse {
state.handler.state.record_activity("terminal", &format!("Ran command: {}", payload.command), Some(&format!("exit_code: {}, cwd: {}", payload.exit_code, payload.cwd)));
state.handler.state.record_activity(
"terminal",
&format!("Ran command: {}", payload.command),
Some(&format!(
"exit_code: {}, cwd: {}",
payload.exit_code, payload.cwd
)),
);
state
.handler
.state
.record_terminal_history(payload.clone());
state.handler.state.record_terminal_history(payload.clone());
let ws_msg = serde_json::json!({
"type": "terminal_telemetry",
@@ -163,6 +171,8 @@ mod tests {
cwd: "C:/Users/reazul.ashraf/workspace".to_string(),
os: "windows".to_string(),
timestamp: 12345,
error_output: None,
..Default::default()
};
// Post the telemetry
server/src/api/ws.rs
+1
View File
@@ -233,6 +233,7 @@ mod tests {
error: None,
timestamp: 1728129000,
session_id: None,
..Default::default()
};
// Broadcast task event
server/src/clipboard_watcher.rs
-92
View File
@@ -1,92 +0,0 @@
use crate::models::StickyNote;
use crate::state::MemoryState;
use arboard::Clipboard;
use std::sync::Arc;
use tokio::time::{Duration, sleep};
pub fn spawn_watcher(state: Arc<MemoryState>) {
tokio::spawn(async move {
let mut last_text = String::new();
let mut clipboard_opt: Option<Clipboard> = Clipboard::new().ok();
loop {
let is_enabled = {
let watch = state.clipboard_watch_mode.read().await;
*watch
};
if !is_enabled {
state.clipboard_notify.notified().await;
continue;
}
sleep(Duration::from_millis(1000)).await;
if clipboard_opt.is_none() {
clipboard_opt = Clipboard::new().ok();
}
if let Some(ref mut cb) = clipboard_opt {
match cb.get_text() {
Ok(text) => {
if text != last_text && !text.trim().is_empty() {
last_text = text.clone();
let note = StickyNote {
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
content: text.clone(),
expires_at: None,
};
state.code.sticky.modify(|notes| {
notes.push(note.clone());
if notes.len() > 200 {
let overflow = notes.len() - 200;
notes.drain(0..overflow);
}
});
state.record_activity(
"clipboard",
"Auto-ingested clipboard content into sticky notes",
Some(&text),
);
state.index_commit_notify.notify_waiters();
}
}
Err(_) => {
// Reset clipboard handle on error so it can be re-initialized cleanly
clipboard_opt = None;
}
}
}
}
});
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[tokio::test]
async fn test_spawn_watcher_lifecycle() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
// Set watch mode to disabled first
*state.clipboard_watch_mode.write().await = false;
spawn_watcher(state.clone());
tokio::time::sleep(Duration::from_millis(50)).await;
// Enable watch mode
*state.clipboard_watch_mode.write().await = true;
state.clipboard_notify.notify_waiters();
tokio::time::sleep(Duration::from_millis(150)).await;
}
}
server/src/dashboard.html
+8 -39
View File
@@ -494,18 +494,14 @@
<div class="main-container">
<div class="tabs">
<button class="tab-button active" onclick="switchTab('activity-tab', this)">Live Activity</button>
<button class="tab-button" onclick="switchTab('ledger-tab', this)">Code Ledger</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>
<button class="tab-button" onclick="switchTab('preferences-tab', this)">Preferences</button>
<button class="tab-button" onclick="switchTab('techdebt-tab', this)">Tech Debt</button>
<button class="tab-button" onclick="switchTab('adrs-tab', this)">Decisions (ADR)</button>
<button class="tab-button" onclick="switchTab('workspaces-tab', this)">Workspaces</button>
<button class="tab-button" onclick="switchTab('pinned-tab', this)">Pinned Files</button>
<button class="tab-button" onclick="switchTab('memos-tab', this)">Handoffs</button>
<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>
</div>
@@ -550,19 +546,6 @@
</div>
</div>
<div id="preferences-tab" class="tab-content">
<div class="panel" style="flex:1; display:flex; flex-direction:column; overflow-y:auto;">
<h2>Global Preferences</h2>
<div id="preferences-container" style="display:flex; flex-direction:column; gap:15px; font-size: 0.9em; line-height: 1.4;"></div>
</div>
</div>
<div id="sticky-tab" class="tab-content">
<div class="panel sticky-panel" style="flex:1; display:flex; flex-direction:column;">
<h2>Sticky Notes</h2>
<div id="sticky-notes-container" style="flex:1; overflow-y:auto; display:flex; flex-wrap:wrap; gap:15px; align-items:flex-start; align-content:flex-start;"></div>
</div>
</div>
<div id="techdebt-tab" class="tab-content">
<div class="panel" style="flex:1; display:flex; flex-direction:column; overflow-y:auto;">
@@ -580,20 +563,6 @@
</div>
</div>
<div id="workspaces-tab" class="tab-content">
<div class="panel" style="flex:1; display:flex; flex-direction:column; overflow-y:auto;">
<h2>Context Workspaces</h2>
<div id="workspaces-container" style="display:flex; flex-direction:column; gap:15px;"></div>
</div>
</div>
<div id="pinned-tab" class="tab-content">
<div class="panel" style="flex:1; display:flex; flex-direction:column; overflow-y:auto;">
<h2>Pinned Files</h2>
<div id="pinned-container" style="display:flex; flex-direction:column; gap:15px;"></div>
</div>
</div>
<div id="memos-tab" class="tab-content">
<div class="panel" style="flex:1; display:flex; flex-direction:column; overflow-y:auto;">
<h2>Handoff Memos & Session Summaries</h2>
@@ -610,13 +579,6 @@
</div>
</div>
<div id="pr-tab" class="tab-content">
<div class="panel" style="flex:1; display:flex; flex-direction:column; overflow-y:auto;">
<h2>PR Checklists</h2>
<div id="pr-container" style="display:flex; flex-direction:column; gap:15px;"></div>
</div>
</div>
<div id="terminal-tab" class="tab-content">
<div class="panel" style="flex:1; display:flex; flex-direction:column; overflow-y:auto;">
<h2>Terminal History</h2>
@@ -624,6 +586,13 @@
</div>
</div>
<div id="ledger-tab" class="tab-content">
<div class="panel" style="flex:1; display:flex; flex-direction:column; overflow-y:auto;">
<h2>Code Change Ledger & Recent Deltas</h2>
<div id="ledger-container" style="display:flex; flex-direction:column; gap:15px;"></div>
</div>
</div>
<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>
server/src/dashboard.js
+106 -147
View File
@@ -23,36 +23,24 @@ function switchTab(tabId, btn) {
case "task-tab":
loadTasks();
break;
case "preferences-tab":
loadPreferences();
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 "ledger-tab":
loadLedger();
break;
}
}
function toggleTheme() {
@@ -89,7 +77,10 @@ function getColorForType(type) {
for (let i = 0;i < type.length; i++)
hash = type.charCodeAt(i) + ((hash << 5) - hash);
const hue = Math.abs(hash) % 360;
return { background: `hsl(${hue}, 70%, 60%)`, border: `hsl(${hue}, 70%, 40%)` };
return {
background: `hsl(${hue}, 70%, 60%)`,
border: `hsl(${hue}, 70%, 40%)`
};
}
function zoomGraph(step) {
if (!network)
@@ -100,7 +91,9 @@ function zoomGraph(step) {
function resetGraph() {
if (!network)
return;
network.fit({ animation: { duration: 500, easingFunction: "easeInOutQuad" } });
network.fit({
animation: { duration: 500, easingFunction: "easeInOutQuad" }
});
}
function closeInspector() {
const inspector = document.getElementById("inspector-panel");
@@ -176,7 +169,9 @@ function updateGraphData() {
title: `${name}
Type: ${entity.entity_type}`,
color: { background: color.bg || color.background, border: color.border },
font: { color: document.documentElement.getAttribute("data-theme") === "dark" ? "#eee" : "#333" }
font: {
color: document.documentElement.getAttribute("data-theme") === "dark" ? "#eee" : "#333"
}
});
nodeIds.add(name);
}
@@ -188,7 +183,10 @@ Type: ${entity.entity_type}`,
to: r.to,
label: r.relation_type,
arrows: "to",
font: { color: document.documentElement.getAttribute("data-theme") === "dark" ? "#aaa" : "#666", strokeWidth: 0 }
font: {
color: document.documentElement.getAttribute("data-theme") === "dark" ? "#aaa" : "#666",
strokeWidth: 0
}
});
}
});
@@ -218,12 +216,28 @@ async function loadGraph() {
return;
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" } },
edges: {
color: { inherit: "from", opacity: 0.6 },
font: { size: 10, align: "middle" },
smooth: { type: "continuous" }
},
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 }
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 }
};
@@ -466,7 +480,17 @@ function parseActivityPayload(item) {
else if (item.content)
summary = String(item.content);
else {
const { timestamp: _t, time: _tm, updated_at: _u, category: _c, type: _tp, method: _m, params: _p, jsonrpc: _j, ...rest } = item;
const {
timestamp: _t,
time: _tm,
updated_at: _u,
category: _c,
type: _tp,
method: _m,
params: _p,
jsonrpc: _j,
...rest
} = item;
if (Object.keys(rest).length > 0) {
summary = JSON.stringify(rest);
} else {
@@ -579,15 +603,11 @@ 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"
ledger: "ledger-tab",
memos: "memos-tab"
};
const targetTab = tabMap[domain];
if (!targetTab || targetTab !== currentTabId) {
@@ -605,9 +625,6 @@ function requestDomainRefresh(domain) {
case "task":
loadTasks();
break;
case "sticky":
loadStickyNotes();
break;
case "techdebt":
loadTechDebt();
break;
@@ -617,21 +634,12 @@ function requestDomainRefresh(domain) {
case "terminal":
loadTerminal();
break;
case "ledger":
loadLedger();
break;
case "memos":
loadMemos();
break;
case "workspaces":
loadWorkspaces();
break;
case "pinned":
loadPinned();
break;
case "pr":
loadPRs();
break;
case "preferences":
loadPreferences();
break;
}
}, 250);
}
@@ -650,24 +658,16 @@ function handleIncomingActivity(payload) {
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 === "CODE_CHANGE" || category === "LEDGER") {
requestDomainRefresh("ledger");
} 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) {
@@ -725,51 +725,6 @@ function setupWS() {
setupSSE();
};
}
async function loadPreferences() {
try {
const res = await fetch("/api/preferences");
const data = await res.json();
const container = document.getElementById("preferences-container");
if (!container)
return;
container.innerHTML = "";
if (!data || Object.keys(data).length === 0) {
container.innerHTML = '<div style="color:var(--text-secondary); text-align:center; padding:20px;">No global preferences found.</div>';
return;
}
for (const [key, pref] of Object.entries(data)) {
const date = new Date(pref.updated_at * 1000).toLocaleString();
container.innerHTML += `
<div style="background:var(--card-bg); border:1px solid var(--border-color); border-radius:6px; padding:15px; box-shadow:0 1px 3px rgba(0,0,0,0.1);">
<strong style="color:#3498db; display:block; margin-bottom:5px; font-size:1.1em;">${key}</strong>
<div style="margin-bottom:10px;">${pref.value}</div>
<div style="font-size:0.8em; color:var(--text-secondary); text-align:right;">Last Updated: ${date}</div>
</div>`;
}
} catch (e) {
console.error("Error loading preferences:", e);
}
}
async function loadStickyNotes() {
try {
const res = await fetch("/api/sticky");
const sticky = await res.json();
const container = document.getElementById("sticky-notes-container");
if (!container)
return;
container.innerHTML = "";
sticky.forEach((note) => {
const card = document.createElement("div");
card.className = "sticky-note";
const date = new Date(note.timestamp * 1000).toLocaleString();
card.innerHTML = `<div style="font-size:0.75em; color:#7f8c8d; margin-bottom:8px;">${date}</div>
<div>${note.content}</div>`;
container.appendChild(card);
});
} catch (err) {
console.error("Failed to load sticky notes", err);
}
}
async function loadGenericList(endpoint, containerId, formatter) {
try {
const res = await fetch(endpoint);
@@ -792,6 +747,50 @@ async function loadGenericList(endpoint, containerId, formatter) {
console.error(`Failed to load ${endpoint}`, err);
}
}
function loadLedger() {
loadGenericList("/api/ledger?limit=100", "ledger-container", (item) => {
const vcs = (item.vcs_type || (item.revision && item.revision.startsWith("r") ? "svn" : "git")).toUpperCase();
const vcsColor = vcs === "SVN" ? "#e67e22" : "#3498db";
const rev = item.revision || item.git_commit;
const revLabel = vcs === "SVN" ? "Rev" : "Commit";
const branch = item.branch || item.git_branch || "None";
const repo = item.repo_name || item.namespace || "global";
const changeKind = (item.change_kind || "modified").toLowerCase();
let kindColor = "var(--accent-color)";
let kindLabel = "Modified";
if (changeKind === "added") {
kindColor = "var(--success-color)";
kindLabel = "+ Added";
} else if (changeKind === "deleted") {
kindColor = "var(--error-color)";
kindLabel = "- Deleted";
} else if (changeKind === "renamed") {
kindColor = "#e67e22";
kindLabel = "→ Renamed";
}
const symbols = Array.isArray(item.symbols) && item.symbols.length > 0 ? item.symbols.map((s) => `<span style="background:var(--card-bg); padding:1px 6px; border-radius:3px; font-size:0.75em; border:1px solid var(--border-color); font-family:monospace;">${escapeHtml(s)}</span>`).join(" ") : "";
const authorBadge = item.author ? `<span style="background:var(--canvas-bg); padding:2px 8px; border-radius:4px; border:1px solid var(--border-color); font-size:0.8em; color:var(--text-secondary);">\uD83D\uDC64 ${escapeHtml(item.author)}</span>` : "";
const lineBadge = item.line_range ? `<span style="background:var(--canvas-bg); padding:2px 8px; border-radius:4px; border:1px solid var(--border-color); font-size:0.8em; font-family:monospace;">${escapeHtml(item.line_range)}</span>` : "";
return `
<div style="display:flex; justify-content:space-between; align-items:flex-start; gap:10px;">
<span style="font-weight:bold; font-size:1.05em; color:var(--text-primary); font-family:monospace;">${escapeHtml(item.file_path || "Unknown")}</span>
<div style="display:flex; gap:6px; align-items:center;">
<span style="background:${vcsColor}; color:#fff; font-size:0.75em; font-weight:bold; padding:2px 6px; border-radius:4px; text-transform:uppercase;">${vcs}</span>
<span style="background:var(--canvas-bg); color:${kindColor}; border:1px solid ${kindColor}; font-size:0.75em; font-weight:bold; padding:2px 6px; border-radius:4px;">${kindLabel}</span>
</div>
</div>
<div style="margin-top:6px; font-size:0.95em; color:var(--text-primary);">${escapeHtml(item.description || "")}</div>
${symbols ? `<div style="margin-top:8px; display:flex; gap:5px; flex-wrap:wrap; align-items:center;"><span style="font-size:0.8em; color:var(--text-secondary);">Symbols:</span> ${symbols}</div>` : ""}
<div style="margin-top:10px; font-size:0.85em; display:flex; gap:8px; flex-wrap:wrap; align-items:center;">
<span style="background:var(--canvas-bg); padding:2px 8px; border-radius:4px; font-weight:bold; border:1px solid var(--border-color); color:var(--text-primary);">\uD83D\uDCE6 ${escapeHtml(repo)}</span>
<span style="background:var(--canvas-bg); padding:2px 8px; border-radius:4px; font-family:monospace; border:1px solid var(--border-color);">${revLabel}: ${rev ? escapeHtml(rev.substring(0, 10)) : "None"}</span>
<span style="background:var(--canvas-bg); padding:2px 8px; border-radius:4px; border:1px solid var(--border-color);">\uD83C\uDF3F ${escapeHtml(branch)}</span>
${lineBadge}
${authorBadge}
</div>
`;
});
}
function loadTerminal() {
loadGenericList("/api/terminal/history", "terminal-container", (item) => `
<span style="font-weight:bold; font-size:1.05em; color:var(--text-primary);">${item.command}</span>
@@ -828,19 +827,6 @@ function loadADRs() {
${item.supersedes ? `<div style="font-size:0.9em; font-weight:bold; color:var(--warning-color);">Supersedes: ${item.supersedes}</div>` : ""}
`);
}
function loadWorkspaces() {
loadGenericList("/api/context_workspaces", "workspaces-container", (item) => `
<strong>${item.name}</strong>
<div style="margin-bottom:8px;">${item.description || ""}</div>
<div style="font-family:monospace; font-size:0.85em; color:#3498db;">${(item.paths || []).join(", ")}</div>
`);
}
function loadPinned() {
loadGenericList("/api/pinned_files", "pinned-container", (item) => `
<strong>${item.path || item.file_path || item.id}</strong>
<div style="margin-top:5px; font-size:0.85em; color:var(--text-secondary);">${item.reason || item.description || "Pinned"}</div>
`);
}
function loadMemos() {
loadGenericList("/api/handoff_memos", "memos-container", (item) => `
<strong>Memo from ${item.author || "System"}</strong>
@@ -858,27 +844,13 @@ function loadSnippets() {
<pre style="background:#111417; color:#d1d8e0; padding:10px; border-radius:4px; overflow-x:auto; font-size:0.85em;">${item.content || item.code || ""}</pre>
`);
}
function loadPRs() {
loadGenericList("/api/pr_checklists", "pr-container", (item) => `
<strong>${item.name || "Checklist"}</strong>
<ul style="padding-left:20px; font-size:0.9em;">
${(item.items || []).map((i) => {
const check = i.is_completed ? "☑" : "☐";
const strike = i.is_completed ? "text-decoration:line-through; color:var(--text-secondary);" : "";
return `<li style="${strike}">${check} ${i.description}</li>`;
}).join("")}
</ul>
`);
}
function loadAllExtras() {
loadTerminal();
loadTechDebt();
loadADRs();
loadWorkspaces();
loadPinned();
loadMemos();
loadSnippets();
loadPRs();
loadLedger();
}
async function testClipboard() {
const modal = document.getElementById("clipboard-modal");
@@ -962,35 +934,23 @@ function refreshActiveTab() {
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();
case "ledger-tab":
loadLedger();
break;
case "activity-tab":
loadActivityHistory();
@@ -1020,17 +980,13 @@ Object.assign(window, {
requestDomainRefresh,
handleIncomingActivity,
setupWS,
loadPreferences,
loadStickyNotes,
loadGenericList,
loadTerminal,
loadLedger,
loadTechDebt,
loadADRs,
loadWorkspaces,
loadPinned,
loadMemos,
loadSnippets,
loadPRs,
loadAllExtras,
testClipboard,
loadVersion,
@@ -1041,4 +997,7 @@ loadVersion();
loadActivityHistory();
setupWS();
var observer = new MutationObserver(() => updateGraphData());
observer.observe(document.documentElement, { attributes: true, attributeFilter: ["data-theme"] });
observer.observe(document.documentElement, {
attributes: true,
attributeFilter: ["data-theme"]
});
server/src/dashboard.ts
+1087 -917
View File
File diff suppressed because it is too large. Load diff
server/src/db.rs
+101 -15
View File
@@ -4,7 +4,7 @@ use std::path::Path;
use std::sync::Arc;
pub fn init_redb(base: &Path) -> Arc<Database> {
let is_in_memory = base == Path::new(":memory:") || base.to_str() == Some(":memory:") || cfg!(test);
let is_in_memory = base == Path::new(":memory:") || base.to_str() == Some(":memory:");
let db = if is_in_memory {
Arc::new(
redb::Database::builder()
@@ -13,17 +13,87 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
)
} else {
let redb_path = base.join("mcp_store.redb");
match redb::Database::create(&redb_path) {
Ok(db) => Arc::new(db),
Err(e) => {
let err_msg = format!(
"Failed to open or create redb database at {:?}: {}. Falling back to temporary store.",
redb_path, e
);
tracing::error!("{}", err_msg);
let temp_path = std::env::temp_dir().join(format!("mcp_store_fallback_{}.redb", uuid::Uuid::new_v4()));
tracing::warn!("CRITICAL PERSISTENCE ALERT: Using temporary redb database {:?}. Changes will be discarded upon application exit.", temp_path);
Arc::new(redb::Database::create(&temp_path).expect("Failed to create fallback redb database"))
if redb_path.exists() {
match redb::Database::open(&redb_path) {
Ok(db) => Arc::new(db),
Err(open_err) => {
let err_msg = format!(
"Failed to open existing redb database at {:?}: {}. Attempting to recreate database.",
redb_path, open_err
);
tracing::warn!("{}", err_msg);
match redb::Database::create(&redb_path) {
Ok(db) => Arc::new(db),
Err(create_err) => {
if std::env::var("MCP_ALLOW_TMP_FALLBACK").unwrap_or_default() == "1" {
let temp_path = std::env::temp_dir().join(format!(
"mcp_store_fallback_{}.redb",
uuid::Uuid::new_v4()
));
tracing::warn!(
"CRITICAL PERSISTENCE ALERT: MCP_ALLOW_TMP_FALLBACK=1 set. Using temporary redb database {:?}. Changes will be discarded upon application exit.",
temp_path
);
Arc::new(
redb::Database::create(&temp_path)
.expect("Failed to create fallback redb database"),
)
} else {
panic!(
"CRITICAL PERSISTENCE FAILURE: Unable to open or create primary database at {:?}: (open: {}, create: {}). To prevent silent data loss on restart, the application cannot start without accessible persistence.",
redb_path, open_err, create_err
);
}
}
}
}
}
} else {
match redb::Database::create(&redb_path) {
Ok(db) => Arc::new(db),
Err(create_err) => {
let mut db_opt = None;
let mut last_open_err = String::new();
for attempt in 1..=3 {
match redb::Database::open(&redb_path) {
Ok(db) => {
db_opt = Some(Arc::new(db));
break;
}
Err(open_err) => {
last_open_err = open_err.to_string();
if attempt < 3 {
tracing::warn!(
"Transient lock contention opening redb at {:?} (attempt {}/3: {}). Retrying...",
redb_path,
attempt,
open_err
);
std::thread::sleep(std::time::Duration::from_millis(150));
}
}
}
}
if let Some(db) = db_opt {
db
} else if std::env::var("MCP_ALLOW_TMP_FALLBACK").unwrap_or_default() == "1" {
let temp_path = std::env::temp_dir()
.join(format!("mcp_store_fallback_{}.redb", uuid::Uuid::new_v4()));
tracing::warn!(
"CRITICAL PERSISTENCE ALERT: MCP_ALLOW_TMP_FALLBACK=1 set. Using temporary redb database {:?}. Changes will be discarded upon application exit.",
temp_path
);
Arc::new(
redb::Database::create(&temp_path)
.expect("Failed to create fallback redb database"),
)
} else {
panic!(
"CRITICAL PERSISTENCE FAILURE: Unable to open or create primary database at {:?}: (create: {}, open: {}). To prevent silent data loss on restart, the application cannot start without accessible persistence.",
redb_path, create_err, last_open_err
);
}
}
}
}
};
@@ -55,6 +125,9 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
("tech_debts", "tech_debts.json"),
("gates", "gates.json"),
("context_workspaces", "context_workspaces.json"),
("state_snapshots", "state_snapshots.json"),
("hypotheses", "hypotheses.json"),
("agent_signals", "agent_signals.json"),
];
for (key, file_name) in stores.iter() {
@@ -73,9 +146,19 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
&& serde_json::from_slice::<serde_json::Value>(&data).is_ok()
{
if let Err(e) = table.insert(*key, data.as_slice()) {
tracing::error!("Failed to insert migrated key '{}': {}", key, e);
tracing::error!(
"Failed to insert migrated key '{}': {}",
key,
e
);
} else {
let _ = std::fs::rename(&json_path, json_path.with_extension("json.migrated"));
let migrated_path = json_path.with_extension("json.migrated");
if let Err(_) = std::fs::rename(&json_path, &migrated_path) {
if migrated_path.exists() {
let _ = std::fs::remove_file(&migrated_path);
let _ = std::fs::rename(&json_path, &migrated_path);
}
}
}
}
}
@@ -89,7 +172,10 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
}
}
Err(e) => {
tracing::error!("Failed to begin write transaction for redb migration: {}", e);
tracing::error!(
"Failed to begin write transaction for redb migration: {}",
e
);
}
}
server/src/embedding.rs
+276 -20
View File
@@ -1,29 +1,250 @@
#[allow(deprecated)]
use fastembed::{EmbeddingModel, TextEmbedding};
use std::sync::Mutex;
use std::sync::OnceLock;
static EMBEDDING_MODEL: OnceLock<Mutex<TextEmbedding>> = OnceLock::new();
use candle_core::{DType, Device, Tensor};
use candle_nn::VarBuilder;
use candle_transformers::models::bert::{BertModel, Config};
use tokenizers::Tokenizer;
#[allow(deprecated)]
pub fn get_embedding_model() -> Result<&'static Mutex<TextEmbedding>, String> {
struct CandleEmbeddingModel {
model: BertModel,
tokenizer: Tokenizer,
device: Device,
}
impl CandleEmbeddingModel {
fn new() -> Result<Self, String> {
let client = hf_hub::HFClientSync::new().map_err(|e| e.to_string())?;
let repo = client.model("sentence-transformers", "all-MiniLM-L6-v2");
let config_file = repo
.download_file()
.filename("config.json")
.send()
.map_err(|e| format!("Failed to download config.json: {}", e))?;
let tokenizer_file = repo
.download_file()
.filename("tokenizer.json")
.send()
.map_err(|e| format!("Failed to download tokenizer.json: {}", e))?;
let weights_file = repo
.download_file()
.filename("model.safetensors")
.send()
.map_err(|e| format!("Failed to download model.safetensors: {}", e))?;
let config_str = std::fs::read_to_string(&config_file)
.map_err(|e| format!("Failed to read config.json: {}", e))?;
let config: Config = serde_json::from_str(&config_str)
.map_err(|e| format!("Failed to parse config.json: {}", e))?;
let mut tokenizer = Tokenizer::from_file(&tokenizer_file)
.map_err(|e| format!("Failed to load tokenizer: {}", e))?;
tokenizer.with_padding(Some(tokenizers::PaddingParams::default()));
let device = Device::Cpu;
let vb = unsafe {
VarBuilder::from_mmaped_safetensors(&[weights_file], DType::F32, &device)
.map_err(|e| format!("Failed to load safetensors: {}", e))?
};
let model =
BertModel::load(vb, &config).map_err(|e| format!("Failed to load BertModel: {}", e))?;
Ok(Self {
model,
tokenizer,
device,
})
}
fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, String> {
if texts.is_empty() {
return Ok(Vec::new());
}
let encodings = self
.tokenizer
.encode_batch(texts.to_vec(), true)
.map_err(|e| format!("Failed to encode texts: {}", e))?;
let batch_size = encodings.len();
if batch_size == 0 {
return Ok(Vec::new());
}
let seq_len = encodings[0].get_ids().len();
if seq_len == 0 {
return Ok(vec![vec![0.0; 384]; batch_size]);
}
let mut all_ids: Vec<u32> = Vec::with_capacity(batch_size * seq_len);
let mut all_type_ids: Vec<u32> = Vec::with_capacity(batch_size * seq_len);
let mut all_attention_mask: Vec<u32> = Vec::with_capacity(batch_size * seq_len);
for enc in &encodings {
all_ids.extend(enc.get_ids());
all_type_ids.extend(enc.get_type_ids());
all_attention_mask.extend(enc.get_attention_mask());
}
let input_ids = Tensor::from_vec(all_ids, (batch_size, seq_len), &self.device)
.map_err(|e| format!("Failed to build input_ids tensor: {}", e))?;
let token_type_ids = Tensor::from_vec(all_type_ids, (batch_size, seq_len), &self.device)
.map_err(|e| format!("Failed to build token_type_ids tensor: {}", e))?;
let attention_mask =
Tensor::from_vec(all_attention_mask, (batch_size, seq_len), &self.device)
.map_err(|e| format!("Failed to build attention_mask tensor: {}", e))?;
let sequence_output = self
.model
.forward(&input_ids, &token_type_ids, Some(&attention_mask))
.map_err(|e| format!("Bert forward failed: {}", e))?;
// Mean pooling: sum(sequence_output * mask) / clamp(sum(mask), min=1e-9)
let mask_f32 = attention_mask
.to_dtype(DType::F32)
.map_err(|e| e.to_string())?
.unsqueeze(2)
.map_err(|e| e.to_string())?;
let sum_embeddings = sequence_output
.broadcast_mul(&mask_f32)
.map_err(|e| e.to_string())?
.sum(1)
.map_err(|e| e.to_string())?;
let sum_mask = mask_f32
.sum(1)
.map_err(|e| e.to_string())?
.clamp(1e-9, f32::MAX)
.map_err(|e| e.to_string())?;
let mean_pooled = sum_embeddings
.broadcast_div(&sum_mask)
.map_err(|e| e.to_string())?;
// L2 Normalization
let norm = mean_pooled
.sqr()
.map_err(|e| e.to_string())?
.sum_keepdim(1)
.map_err(|e| e.to_string())?
.sqrt()
.map_err(|e| e.to_string())?;
let normalized = mean_pooled
.broadcast_div(&norm)
.map_err(|e| e.to_string())?;
normalized.to_vec2::<f32>().map_err(|e| e.to_string())
}
}
fn fallback_embed(text: &str) -> Vec<f32> {
const DIM: usize = 384;
let mut vec = vec![0.0f32; DIM];
let words: Vec<&str> = text.split_whitespace().collect();
if words.is_empty() {
vec[0] = 1.0;
return vec;
}
use std::hash::{Hash, Hasher};
for word in words {
let clean: String = word
.chars()
.filter(|c| c.is_alphanumeric())
.flat_map(|c| c.to_lowercase())
.collect();
if clean.is_empty() {
continue;
}
let mut hasher = std::collections::hash_map::DefaultHasher::new();
clean.hash(&mut hasher);
let h = hasher.finish();
let idx = (h as usize) % DIM;
let sign = if (h >> 32) & 1 == 0 { 1.0f32 } else { -1.0f32 };
vec[idx] += sign;
let chars: Vec<char> = clean.chars().collect();
for window in chars.windows(3) {
let mut h2 = std::collections::hash_map::DefaultHasher::new();
window.hash(&mut h2);
let hv = h2.finish();
let idx2 = (hv as usize) % DIM;
let s2 = if (hv >> 32) & 1 == 0 { 0.5f32 } else { -0.5f32 };
vec[idx2] += s2;
}
}
let norm_sq: f32 = vec.iter().map(|x| x * x).sum();
if norm_sq > 0.0 {
let norm = norm_sq.sqrt();
for x in vec.iter_mut() {
*x /= norm;
}
} else {
vec[0] = 1.0;
}
vec
}
enum EmbeddingModel {
Candle(CandleEmbeddingModel),
Fallback,
}
impl EmbeddingModel {
fn new() -> Self {
match CandleEmbeddingModel::new() {
Ok(model) => EmbeddingModel::Candle(model),
Err(e) => {
tracing::warn!(
"Failed to initialize Candle BERT model ({e}); falling back to deterministic offline embeddings."
);
EmbeddingModel::Fallback
}
}
}
fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, String> {
match self {
EmbeddingModel::Candle(model) => model.embed(texts),
EmbeddingModel::Fallback => Ok(texts.iter().map(|t| fallback_embed(t)).collect()),
}
}
}
static EMBEDDING_MODEL: OnceLock<Mutex<EmbeddingModel>> = OnceLock::new();
static INIT_MUTEX: Mutex<()> = Mutex::new(());
fn get_embedding_model() -> Result<&'static Mutex<EmbeddingModel>, String> {
if let Some(model) = EMBEDDING_MODEL.get() {
return Ok(model);
}
let options = fastembed::InitOptions::new(EmbeddingModel::AllMiniLML6V2)
.with_show_download_progress(true);
let model = TextEmbedding::try_new(options).map_err(|e| e.to_string())?;
let _guard = INIT_MUTEX.lock().map_err(|e| e.to_string())?;
if let Some(model) = EMBEDDING_MODEL.get() {
return Ok(model);
}
let model = EmbeddingModel::new();
let _ = EMBEDDING_MODEL.set(Mutex::new(model));
Ok(EMBEDDING_MODEL.get().unwrap())
}
pub async fn generate_embedding_async(text: String) -> Result<Vec<f32>, String> {
tokio::task::spawn_blocking(move || {
let model_mutex = get_embedding_model()?;
let mut model = model_mutex.lock().map_err(|e| e.to_string())?;
let embeddings = model.embed(vec![text], None).map_err(|e| e.to_string())?;
Ok(embeddings.into_iter().next().unwrap_or_default())
let model = model_mutex.lock().map_err(|e| e.to_string())?;
let embeddings = model.embed(&[text])?;
let emb = embeddings
.into_iter()
.next()
.ok_or_else(|| "Embedding model returned no embeddings".to_string())?;
if emb.is_empty() {
return Err("Embedding model generated a 0-length vector".to_string());
}
Ok(emb)
})
.await
.map_err(|e| e.to_string())?
@@ -36,11 +257,29 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
let mut dot_product = 0.0f32;
let mut norm_a_sq = 0.0f32;
let mut norm_b_sq = 0.0f32;
for (&x, &y) in a.iter().zip(b.iter()) {
let chunks_a = a.chunks_exact(8);
let chunks_b = b.chunks_exact(8);
let remainder_a = chunks_a.remainder();
let remainder_b = chunks_b.remainder();
for (ca, cb) in chunks_a.zip(chunks_b) {
for i in 0..8 {
let x = ca[i];
let y = cb[i];
dot_product += x * y;
norm_a_sq += x * x;
norm_b_sq += y * y;
}
}
for (&x, &y) in remainder_a.iter().zip(remainder_b.iter()) {
dot_product += x * y;
norm_a_sq += x * x;
norm_b_sq += y * y;
}
let norm_a = norm_a_sq.sqrt();
let norm_b = norm_b_sq.sqrt();
if norm_a == 0.0 || norm_b == 0.0 {
@@ -49,19 +288,38 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
dot_product / (norm_a * norm_b)
}
}
pub async fn generate_embeddings_async(texts: Vec<String>) -> Result<Vec<Vec<f32>>, String> {
if texts.is_empty() {
return Ok(Vec::new());
}
tokio::task::spawn_blocking(move || {
let model_mutex = get_embedding_model()?;
let mut model = model_mutex.lock().map_err(|e| e.to_string())?;
let model = model_mutex.lock().map_err(|e| e.to_string())?;
let mut all_embeddings = Vec::with_capacity(texts.len());
for chunk in texts.chunks(32) {
let chunk_vec = chunk.to_vec();
let chunk_embeddings = model.embed(chunk_vec, None).map_err(|e| e.to_string())?;
let mut current_chunk = Vec::new();
let mut current_chars = 0;
const MAX_CHARS_PER_BATCH: usize = 16384;
for text in texts {
let text_len = text.len();
if !current_chunk.is_empty()
&& (current_chunk.len() >= 64 || current_chars + text_len > MAX_CHARS_PER_BATCH)
{
let chunk_vec = std::mem::take(&mut current_chunk);
let chunk_embeddings = model.embed(&chunk_vec)?;
all_embeddings.extend(chunk_embeddings);
current_chars = 0;
}
current_chars += text_len;
current_chunk.push(text);
}
if !current_chunk.is_empty() {
let chunk_embeddings = model.embed(&current_chunk)?;
all_embeddings.extend(chunk_embeddings);
}
Ok(all_embeddings)
})
.await
@@ -93,7 +351,6 @@ mod tests {
}
#[tokio::test]
async fn test_generate_embeddings_async_single_text() {
let text = "test text".to_string();
let res = generate_embeddings_async(vec![text.clone()]).await.unwrap();
@@ -115,4 +372,3 @@ mod tests {
}
}
}
server/src/handlers/ast.rs
+135 -24
View File
@@ -6,6 +6,24 @@ use serde_json::Value;
use std::sync::Arc;
use tree_sitter::{Node, Parser};
fn validate_safe_path(path_str: &str) -> crate::error::Result<()> {
if path_str.trim().is_empty() || path_str.contains('\0') {
return Err(crate::error::AppError::Internal(
"Invalid file path: path is empty or contains null characters".to_string(),
));
}
let path = std::path::Path::new(path_str);
for component in path.components() {
if component == std::path::Component::ParentDir {
return Err(crate::error::AppError::Internal(format!(
"Path traversal forbidden: '{}' contains parent directory relative components",
path_str
)));
}
}
Ok(())
}
pub struct ReadFileSkeletonHandler;
#[async_trait]
@@ -26,13 +44,13 @@ impl McpTool for ReadFileSkeletonHandler {
serde_json::from_value(args).map_err(|e| e.to_string())?;
let file_path = tool_args.file_path.clone();
validate_safe_path(&file_path)?;
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let code = std::fs::read_to_string(&file_path).map_err(|e| {
crate::error::AppError::Internal(format!("Failed to read file: {}", e))
})?;
let mut parser = Parser::new();
let ext = std::path::Path::new(&file_path)
.extension()
.and_then(|s| s.to_str())
@@ -49,10 +67,10 @@ impl McpTool for ReadFileSkeletonHandler {
_ => return Ok(code),
};
let mut parser = tree_sitter::Parser::new();
parser
.set_language(&language.into())
.map_err(|e| e.to_string())?;
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
let tree = parser.parse(&code, None).ok_or_else(|| {
crate::error::AppError::Internal("Failed to parse code".to_string())
})?;
@@ -150,6 +168,8 @@ impl McpTool for ReplaceAstNodeHandler {
serde_json::from_value(args).map_err(|e| e.to_string())?;
let file_path = tool_args.file_path.clone();
validate_safe_path(&file_path)?;
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let code = std::fs::read_to_string(&file_path).map_err(|e| {
crate::error::AppError::Internal(format!("Failed to read file: {}", e))
@@ -198,8 +218,14 @@ impl McpTool for ReplaceAstNodeHandler {
| "method_definition"
| "function"
),
"struct" => matches!(actual_kind, "struct_item" | "struct_declaration" | "struct_specifier"),
"class" => matches!(actual_kind, "class_declaration" | "class_definition" | "class_item"),
"struct" => matches!(
actual_kind,
"struct_item" | "struct_declaration" | "struct_specifier"
),
"class" => matches!(
actual_kind,
"class_declaration" | "class_definition" | "class_item"
),
"enum" => matches!(actual_kind, "enum_item" | "enum_declaration"),
"trait" | "interface" => matches!(
actual_kind,
@@ -225,7 +251,11 @@ impl McpTool for ReplaceAstNodeHandler {
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
let kind = child.kind();
if kind == "identifier" || kind == "name" || kind == "property_identifier" || kind == "field_identifier" {
if kind == "identifier"
|| kind == "name"
|| kind == "property_identifier"
|| kind == "field_identifier"
{
let name = child.utf8_text(code.as_bytes()).unwrap_or("");
if name == target_name {
return Some(node);
@@ -254,12 +284,38 @@ impl McpTool for ReplaceAstNodeHandler {
let start_byte = node.start_byte();
let end_byte = node.end_byte();
let mut new_file_content = String::new();
if !code.is_char_boundary(start_byte) || !code.is_char_boundary(end_byte) {
return Err(crate::error::AppError::Internal(format!(
"Byte offsets {}..{} do not fall on UTF-8 character boundaries in {}",
start_byte, end_byte, file_path
)));
}
let mut new_file_content =
String::with_capacity(code.len() + tool_args.new_content.len());
new_file_content.push_str(&code[..start_byte]);
new_file_content.push_str(&tool_args.new_content);
new_file_content.push_str(&code[end_byte..]);
std::fs::write(&file_path, new_file_content).map_err(|e| e.to_string())?;
let target_path = std::path::PathBuf::from(&file_path);
let parent_dir = target_path
.parent()
.unwrap_or_else(|| std::path::Path::new("."));
let temp_file_path =
parent_dir.join(format!(".tmp_ast_{}.tmp", uuid::Uuid::new_v4()));
std::fs::write(&temp_file_path, new_file_content).map_err(|e| e.to_string())?;
if let Err(e) = std::fs::rename(&temp_file_path, &target_path) {
// On Windows, std::fs::rename fails if the target file already exists.
// Fall back to copy-and-remove to ensure atomic-like overwrite behavior.
if let Err(copy_err) = std::fs::copy(&temp_file_path, &target_path) {
let _ = std::fs::remove_file(&temp_file_path);
return Err(crate::error::AppError::Internal(format!(
"Failed to atomically overwrite {}: rename failed ({}), copy failed ({})",
file_path, e, copy_err
)));
}
let _ = std::fs::remove_file(&temp_file_path);
}
Ok(format!(
"Successfully replaced node {} of type {} in {}",
tool_args.node_name, tool_args.node_type, file_path
@@ -278,20 +334,31 @@ impl McpTool for ReplaceAstNodeHandler {
}
}
fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: bool) -> Vec<serde_json::Value> {
fn scan_workspace_for_symbol(
target_sym: &str,
limit: usize,
filter_fn_call: bool,
workspace_dir: Option<std::path::PathBuf>,
) -> Vec<serde_json::Value> {
let mut results = Vec::new();
let cwd = match std::env::current_dir() {
Ok(dir) => dir,
Err(_) => return results,
};
let scan_dir = workspace_dir.unwrap_or_else(|| {
std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."))
});
let walker = ignore::WalkBuilder::new(&cwd)
if !scan_dir.exists() {
return results;
}
let walker = ignore::WalkBuilder::new(&scan_dir)
.hidden(true)
.git_ignore(true)
.build();
let mut scanned_files = 0;
let call_pattern = format!("{}(", target_sym);
let await_pattern = format!("{}.await", target_sym);
for result in walker {
let entry = match result {
Ok(e) => e,
@@ -301,7 +368,20 @@ fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: boo
if entry.file_type().is_some_and(|ft| ft.is_file()) {
let path = entry.path();
let ext = path.extension().and_then(|s| s.to_str()).unwrap_or("");
if matches!(ext, "rs" | "ts" | "tsx" | "js" | "jsx" | "py" | "go" | "java" | "c" | "cpp" | "h" | "hpp") {
if matches!(
ext,
"rs" | "ts"
| "tsx"
| "js"
| "jsx"
| "py"
| "go"
| "java"
| "c"
| "cpp"
| "h"
| "hpp"
) {
scanned_files += 1;
if scanned_files > 500 {
break;
@@ -309,8 +389,7 @@ fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: boo
if let Ok(content) = std::fs::read_to_string(path) {
for (line_num, line) in content.lines().enumerate() {
let is_match = if filter_fn_call {
line.contains(&format!("{}(", target_sym))
|| line.contains(&format!("{}.await", target_sym))
line.contains(&call_pattern) || line.contains(&await_pattern)
} else {
line.contains(target_sym)
};
@@ -371,10 +450,22 @@ impl McpTool for FindSymbolReferencesHandler {
Ok::<Vec<serde_json::Value>, crate::error::AppError>(refs)
})?;
let custom_dir = if let Some(ref d) = req.workspace_dir {
Some(std::path::PathBuf::from(d))
} else {
state.project.pinned_files.read_with(|pinned| {
pinned.first().and_then(|pin| {
std::path::Path::new(&pin.file_path)
.parent()
.map(|p| p.to_path_buf())
})
})
};
if matches.len() < limit {
let remaining = limit - matches.len();
let disk_matches = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&target_sym, remaining, false)
scan_workspace_for_symbol(&target_sym, remaining, false, custom_dir)
})
.await
.unwrap_or_default();
@@ -405,13 +496,13 @@ impl McpTool for GetCallersHandler {
serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(10);
let target_fn = req.function_name.clone();
let call_pattern = format!("{}(", target_fn);
let await_pattern = format!("{}.await", target_fn);
let mut callers = state.code.snippets.read_with(|snippets| {
let mut matching = Vec::new();
for snippet in snippets {
if snippet.code.contains(&format!("{}(", target_fn))
|| snippet.code.contains(&format!("{}.await", target_fn))
{
if snippet.code.contains(&call_pattern) || snippet.code.contains(&await_pattern) {
matching.push(serde_json::json!({
"source": "snippet",
"name": snippet.name,
@@ -425,10 +516,22 @@ impl McpTool for GetCallersHandler {
Ok::<Vec<serde_json::Value>, crate::error::AppError>(matching)
})?;
let custom_dir = if let Some(ref d) = req.workspace_dir {
Some(std::path::PathBuf::from(d))
} else {
state.project.pinned_files.read_with(|pinned| {
pinned.first().and_then(|pin| {
std::path::Path::new(&pin.file_path)
.parent()
.map(|p| p.to_path_buf())
})
})
};
if callers.len() < limit {
let remaining = limit - callers.len();
let disk_callers = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&target_fn, remaining, true)
scan_workspace_for_symbol(&target_fn, remaining, true, custom_dir)
})
.await
.unwrap_or_default();
@@ -468,9 +571,17 @@ impl McpTool for AnalyzeImpactHandler {
}
});
let custom_dir = state.project.pinned_files.read_with(|pinned| {
pinned.first().and_then(|pin| {
std::path::Path::new(&pin.file_path)
.parent()
.map(|p| p.to_path_buf())
})
});
let sym_clone = sym.clone();
let disk_refs = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&sym_clone, 20, false)
scan_workspace_for_symbol(&sym_clone, 20, false, custom_dir)
})
.await
.unwrap_or_default();
server/src/handlers/env.rs
+4
View File
@@ -32,6 +32,7 @@ impl McpTool for UpdateEnvFingerprintHandler {
shell: std::env::var("SHELL").unwrap_or_else(|_| "unknown".to_string()),
tool_versions: req.tool_versions,
updated_at: crate::handlers::utils::now_secs(),
..Default::default()
},
);
});
@@ -94,6 +95,7 @@ impl McpTool for LogEnvRequirementHandler {
key: req.key,
description: req.description,
is_secret: req.is_secret,
..Default::default()
});
});
Ok("Env requirement logged".to_string())
@@ -127,6 +129,7 @@ impl McpTool for RegisterEnvironmentHandler {
description: req.description,
requires_vpn: req.requires_vpn,
updated_at: crate::handlers::utils::now_secs(),
..Default::default()
});
});
Ok("Environment registered".to_string())
@@ -262,6 +265,7 @@ mod tests {
description: "Test env".to_string(),
requires_vpn: false,
updated_at: crate::handlers::utils::now_secs(),
..Default::default()
});
});
server/src/handlers/git.rs
+4
View File
@@ -55,6 +55,10 @@ impl McpTool for GetActiveWorktreeContextHandler {
if let Some(diff) = diff {
let _ = diff.print(git2::DiffFormat::Patch, |_delta, _hunk, line| {
match line.origin() {
'+' | '-' | ' ' => diff_output.push(line.origin()),
_ => {}
}
let content = std::str::from_utf8(line.content()).unwrap_or("");
diff_output.push_str(content);
true
server/src/handlers/graph.rs
+89 -27
View File
@@ -38,21 +38,7 @@ impl<'a> GraphQueryBuilder<'a> {
}
pub fn find_shortest_path(&self, start: &str, end: &str) -> Option<Vec<String>> {
let mut adj: std::collections::HashMap<&str, Vec<(&str, &str, bool)>> =
std::collections::HashMap::with_capacity(self.graph.relations.len() * 2);
for rel in &self.graph.relations {
if let Some(rf) = self.relation_filter {
if rel.relation_type != rf {
continue;
}
}
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 adj = self.graph.build_adjacency_map(self.relation_filter);
let mut queue = std::collections::VecDeque::new();
let mut visited = std::collections::HashSet::new();
@@ -253,6 +239,7 @@ impl McpTool for CreateRelationsHandler {
observations: vec!["Auto-created stub entity for relation endpoint".to_string()],
namespace: crate::models::default_namespace(),
git_branch: None,
..Default::default()
},
);
auto_created.push(node_name);
@@ -328,7 +315,10 @@ impl McpTool for DeleteEntitiesHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<DeleteEntitiesTool>("delete_entities", "Execute delete_entities")
crate::mcp::tool_def::<DeleteEntitiesTool>(
"delete_entities",
"Delete specified entities and all their associated relations and observations from the knowledge graph.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -378,7 +368,7 @@ impl McpTool for DeleteObservationsHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<DeleteObservationsTool>(
"delete_observations",
"Execute delete_observations",
"Delete specific observations from existing entities in the knowledge graph without removing the entity itself.",
)
}
@@ -422,7 +412,10 @@ impl McpTool for DeleteRelationsHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<DeleteRelationsTool>("delete_relations", "Execute delete_relations")
crate::mcp::tool_def::<DeleteRelationsTool>(
"delete_relations",
"Delete specific relationships between entities in the knowledge graph by specifying source 'from', target 'to', and 'relation_type'.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -462,8 +455,45 @@ impl McpTool for ReadGraphHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let max_tokens = req.max_tokens;
let limit = req.limit;
let offset = req.offset.unwrap_or(0);
let data = state.read_graph(|full| -> crate::error::Result<String> {
let mut result_json = if let Some(ns) = req.namespace {
let mut result_json = if limit.is_some() || req.offset.is_some() {
let filtered_entities: Vec<(&String, &crate::models::Entity)> = full
.entities
.iter()
.filter(|(_, v)| req.namespace.as_ref().map_or(true, |ns| &v.namespace == ns))
.collect();
let total_count = filtered_entities.len();
let effective_limit = limit.unwrap_or(50);
let paged_entities: std::collections::HashMap<&String, &crate::models::Entity> = filtered_entities
.into_iter()
.skip(offset)
.take(effective_limit)
.collect();
let has_more = offset + paged_entities.len() < total_count;
let relevant_relations: Vec<&crate::models::Relation> = full
.relations
.iter()
.filter(|r| {
let ns_match = req.namespace.as_ref().map_or(true, |ns| &r.namespace == ns);
ns_match && (paged_entities.contains_key(&r.from) || paged_entities.contains_key(&r.to))
})
.collect();
serde_json::to_string(&serde_json::json!({
"entities": paged_entities,
"relations": relevant_relations,
"total_count": total_count,
"offset": offset,
"limit": effective_limit,
"has_more": has_more,
}))?
} else if let Some(ns) = req.namespace {
let mut filtered = BorrowedGraph::default();
for (k, v) in &full.entities {
if v.namespace == ns {
@@ -485,7 +515,7 @@ impl McpTool for ReadGraphHandler {
if result_json.len() > max_chars {
let valid_boundary = result_json.floor_char_boundary(max_chars);
result_json.truncate(valid_boundary);
result_json.push_str("\n... [TRUNCATED_TO_MAX_TOKENS. Use search_nodes or get_subgraph for targeted discovery]");
result_json.push_str("\n... [TRUNCATED_TO_MAX_TOKENS. Use search_nodes, pagination (limit, offset), or get_subgraph for targeted discovery]");
}
}
Ok(result_json)
@@ -503,7 +533,10 @@ impl McpTool for SearchNodesHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<SearchNodesTool>("search_nodes", "Execute search_nodes")
crate::mcp::tool_def::<SearchNodesTool>(
"search_nodes",
"Search knowledge graph entities and nodes by name or observation query. Returns matching entity names, types, and observations.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -550,7 +583,10 @@ impl McpTool for OpenNodesHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<OpenNodesTool>("open_nodes", "Execute open_nodes")
crate::mcp::tool_def::<OpenNodesTool>(
"open_nodes",
"Retrieve complete node details, observations, and connected relations for specific entities in the knowledge graph.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -588,7 +624,10 @@ impl McpTool for VisualizeGraphHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<VisualizeGraphTool>("visualize_graph", "Execute visualize_graph")
crate::mcp::tool_def::<VisualizeGraphTool>(
"visualize_graph",
"Generate a visual Mermaid representation of the knowledge graph topology or a filtered query subset.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -677,7 +716,10 @@ impl McpTool for CondenseEntityHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CondenseEntityTool>("condense_entity", "Execute condense_entity")
crate::mcp::tool_def::<CondenseEntityTool>(
"condense_entity",
"Condense, summarize, and deduplicate redundant observations on an entity to optimize context size while preserving key facts.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -757,7 +799,10 @@ impl McpTool for MergeEntitiesHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<MergeEntitiesTool>("merge_entities", "Execute merge_entities")
crate::mcp::tool_def::<MergeEntitiesTool>(
"merge_entities",
"Merge two duplicate entities into a primary target entity, consolidating all observations and re-pointing relations.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -817,7 +862,10 @@ impl McpTool for FindOrphansHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<FindOrphansTool>("find_orphans", "Execute find_orphans")
crate::mcp::tool_def::<FindOrphansTool>(
"find_orphans",
"Find orphan nodes in the knowledge graph that have no incoming or outgoing relationships.",
)
}
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -847,7 +895,10 @@ impl McpTool for GetSubgraphHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<GetSubgraphTool>("get_subgraph", "Execute get_subgraph")
crate::mcp::tool_def::<GetSubgraphTool>(
"get_subgraph",
"Extract a connected neighborhood subgraph centered around specified seed nodes up to a given hop depth.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1188,6 +1239,7 @@ mod tests {
observations: vec![],
namespace: "global".to_string(),
git_branch: None,
..Default::default()
},
);
g.entities.insert(
@@ -1198,6 +1250,7 @@ mod tests {
observations: vec![],
namespace: "global".to_string(),
git_branch: None,
..Default::default()
},
);
});
@@ -1256,6 +1309,7 @@ mod tests {
observations: vec![],
namespace: "global".to_string(),
git_branch: None,
..Default::default()
},
);
});
@@ -1447,6 +1501,14 @@ mod tests {
let qsignal_res = qsignal_handler.execute(json!({"sender": "agent1"}), state.clone()).await.unwrap();
assert!(qsignal_res.contains("task_completed"));
let read_paged_handler = ReadGraphHandler;
let paged_res = read_paged_handler
.execute(json!({"limit": 1, "offset": 0}), state.clone())
.await
.unwrap();
let parsed_paged: serde_json::Value = serde_json::from_str(&paged_res).unwrap();
assert_eq!(parsed_paged["total_count"], 2);
assert_eq!(parsed_paged["has_more"], true);
}
}
server/src/handlers/logs.rs
+4 -2
View File
@@ -25,7 +25,8 @@ impl McpTool for WatchProcessLogsHandler {
async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: WatchProcessLogsTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
if !std::path::Path::new(&tool_args.file_path).exists() {
let safe_path = crate::handlers::utils::validate_safe_path(&tool_args.file_path)?;
if !safe_path.exists() {
return Err(crate::error::AppError::Internal(format!(
"File does not exist: {}",
tool_args.file_path
@@ -53,9 +54,10 @@ impl McpTool for GetRecentLogsHandler {
async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: GetRecentLogsTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let safe_path = crate::handlers::utils::validate_safe_path(&tool_args.file_path)?;
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let mut file = File::open(&tool_args.file_path).map_err(|e| {
let mut file = File::open(&safe_path).map_err(|e| {
crate::error::AppError::Internal(format!("Failed to open file: {}", e))
})?;
let len = file.metadata().map_err(|e| e.to_string())?.len();
server/src/handlers/meta.rs
+250 -57
View File
@@ -49,6 +49,11 @@ impl McpTool for LogDecisionHandler {
status: "accepted".to_string(),
supersedes: req.supersedes,
timestamp: crate::handlers::utils::now_secs(),
namespace: req.namespace.unwrap_or_else(|| crate::models::default_namespace()),
repo_name: req.repo_name,
alternatives_considered: req.alternatives_considered.unwrap_or_default(),
affected_components: req.affected_components.unwrap_or_default(),
author: req.author,
};
adrs_to_index.push(a.clone());
@@ -201,7 +206,11 @@ impl McpTool for LogErrorFixHandler {
git_commit: req.git_commit,
git_branch: req.git_branch,
embedding,
})
..Default::default()
});
if fixes.len() > 300 {
fixes.remove(0);
}
});
state.record_activity("error_fix", &format!("Fixed error: {}", req.signature), Some(&solution));
Ok(format!("Logged error fix for {}: {}", req.signature, solution))
@@ -238,8 +247,13 @@ impl McpTool for SearchErrorFixesHandler {
.iter()
.map(|f| {
let mut score = 0.0;
if st.contains(&f.signature) || f.signature.contains(st) {
let st_lower = st.to_lowercase();
let sig_lower = f.signature.to_lowercase();
let sol_lower = f.solution.to_lowercase();
if st_lower.contains(&sig_lower) || sig_lower.contains(&st_lower) {
score += 0.8;
} else if st_lower.contains(&sol_lower) || sol_lower.contains(&st_lower) {
score += 0.5;
}
if let Some(emb) = &f.embedding {
if !query_emb.is_empty() {
@@ -343,16 +357,65 @@ impl McpTool for LogCodeChangeHandler {
}
}
}
let change_kind = match req.change_kind.as_deref().map(|s| s.to_lowercase()).as_deref() {
Some("added") | Some("create") | Some("created") | Some("new") => crate::models::ChangeKind::Added,
Some("deleted") | Some("remove") | Some("removed") => crate::models::ChangeKind::Deleted,
Some("renamed") | Some("move") | Some("moved") => crate::models::ChangeKind::Renamed,
_ => crate::models::ChangeKind::Modified,
};
let namespace = req
.namespace
.filter(|ns| !ns.trim().is_empty())
.or_else(|| req.repo_name.clone().filter(|rn| !rn.trim().is_empty()))
.unwrap_or_else(crate::models::default_namespace);
let symbols = req.symbol_references.clone().unwrap_or_default();
let line_range = req.line_range.clone();
let effective_rev = req.revision.clone().or_else(|| req.git_commit.clone());
let effective_branch = req.branch.clone().or_else(|| req.git_branch.clone());
let detected_vcs = if let Some(vcs) = req.vcs_type.clone() {
Some(vcs)
} else if let Some(ref rev) = effective_rev {
if rev.starts_with('r') && rev[1..].chars().all(|c| c.is_ascii_digit()) {
Some("svn".to_string())
} else if req.repo_url.as_deref().map(|u| u.contains("/svn/")).unwrap_or(false) {
Some("svn".to_string())
} else {
Some("git".to_string())
}
} else if let Some(ref br) = effective_branch {
if br.eq_ignore_ascii_case("trunk") || br.starts_with("branches/") || br.starts_with("tags/") {
Some("svn".to_string())
} else {
Some("git".to_string())
}
} else {
None
};
state.code.ledger.modify(|ledger| {
ledger.push(CodeChange {
timestamp: crate::handlers::utils::now_secs(),
file_path: req.file_path.clone(),
description: description.clone(),
git_commit: req.git_commit,
git_branch: req.git_branch,
git_commit: effective_rev.clone(),
git_branch: effective_branch.clone(),
repo_name: req.repo_name,
repo_url: req.repo_url,
namespace: namespace.clone(),
change_kind,
symbols,
line_range,
author: req.author,
session_id: req.session_id,
vcs_type: detected_vcs,
revision: effective_rev,
branch: effective_branch,
repository_root: req.repository_root,
});
if ledger.len() > 1000 {
if ledger.len() > 500 {
ledger.remove(0);
}
});
@@ -376,15 +439,44 @@ impl McpTool for QueryRecentChangesHandler {
)
}
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let data = state
.code
.ledger
.read_with(|l| Ok::<String, crate::error::AppError>(serde_json::to_string(l)?))?;
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: QueryRecentChangesTool = serde_json::from_value(args).unwrap_or(QueryRecentChangesTool {
namespace: None,
repo_name: None,
vcs_type: None,
limit: None,
});
let limit = req.limit.unwrap_or(50);
let data = state.code.ledger.read_with(|l| {
let filtered: Vec<_> = l
.iter()
.rev()
.filter(|c| {
if let Some(ns) = &req.namespace {
if !c.namespace.eq_ignore_ascii_case(ns) {
return false;
}
}
if let Some(repo) = &req.repo_name {
if c.repo_name.as_ref().map(|rn| !rn.eq_ignore_ascii_case(repo)).unwrap_or(true) {
return false;
}
}
if let Some(vcs) = &req.vcs_type {
if !c.effective_vcs().eq_ignore_ascii_case(vcs) {
return false;
}
}
true
})
.take(limit)
.cloned()
.collect();
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
})?;
Ok(data)
}
}
pub struct LearnPreferenceHandler;
#[async_trait]
@@ -409,6 +501,7 @@ impl McpTool for LearnPreferenceHandler {
key: req.key.clone(),
value: req.value,
updated_at: crate::handlers::utils::now_secs(),
..Default::default()
},
);
});
@@ -474,6 +567,12 @@ impl McpTool for LogTechDebtHandler {
git_commit: req.git_commit,
git_branch: req.git_branch,
embedding,
repo_name: req.repo_name,
severity: req.severity,
file_path: req.file_path,
line_range: req.line_range,
workaround: req.workaround,
effort_estimate: req.effort_estimate,
})
});
Ok("Tech debt logged".to_string())
@@ -605,7 +704,11 @@ impl McpTool for DecisionsHandler {
"context": context,
"decision": decision,
"consequences": consequences,
"namespace": ns
"namespace": ns,
"repo_name": req.repo_name,
"author": req.author,
"affected_components": req.affected_components,
"alternatives_considered": req.alternatives_considered,
}),
state,
).await
@@ -655,7 +758,12 @@ impl McpTool for TechDebtHandler {
"git_branch": req.git_branch,
"symbol_references": req.symbol_references,
"line_range": req.line_range,
"namespace": ns
"namespace": ns,
"repo_name": req.repo_name,
"severity": req.severity,
"file_path": req.file_path,
"workaround": req.workaround,
"effort_estimate": req.effort_estimate,
}),
state,
).await
@@ -699,23 +807,12 @@ impl McpTool for OmniSearchHandler {
.search(&req.query, req.namespace.as_deref())
.unwrap_or_default();
let q = req.query.clone();
let query_emb = crate::embedding::generate_embedding_async(q.clone())
let vector_matches = state
.search()
.semantic_search(&req.query, req.namespace.as_deref(), limit * 2)
.await
.unwrap_or_default();
let vector_matches = if !query_emb.is_empty() {
if let Some(vdb) = &*state.vector_db.read().await {
vdb.search(query_emb.clone(), (limit * 2) as u64)
.await
.unwrap_or_default()
} else {
vec![]
}
} else {
vec![]
};
// Reciprocal Rank Fusion (RRF) algorithm
#[allow(dead_code)]
#[derive(Clone)]
@@ -754,8 +851,8 @@ impl McpTool for OmniSearchHandler {
let item = MatchItem {
id: v_match.id.clone(),
doc_type: v_match.doc_type,
title: v_match.id,
body: v_match.text,
title: v_match.title,
body: v_match.body,
};
rrf_scores.insert(item_id, (score, item));
}
@@ -770,29 +867,38 @@ impl McpTool for OmniSearchHandler {
let kg_json = state.read_graph(|full| {
let mut kg_results = serde_json::Map::new();
let mut count = 0;
// Build pre-indexed adjacency map: O(R) once instead of O(E * R)
let mut adj_map: std::collections::HashMap<&str, Vec<(&str, &str, &str)>> =
std::collections::HashMap::new();
for rel in &full.relations {
adj_map
.entry(rel.from.as_str())
.or_default()
.push((rel.to.as_str(), rel.relation_type.as_str(), "outgoing"));
adj_map
.entry(rel.to.as_str())
.or_default()
.push((rel.from.as_str(), rel.relation_type.as_str(), "incoming"));
}
for res in &matches {
if res.doc_type == "entity"
&& let Some(e) = full.entities.get(&res.id)
{
if count >= limit {
continue;
break;
}
count += 1;
// 1-hop relation expansion for GraphRAG
// 1-hop relation expansion for GraphRAG via pre-indexed adjacency
let mut connected_rels = Vec::new();
for rel in &full.relations {
if rel.from == res.id {
if let Some(rels) = adj_map.get(res.id.as_str()) {
for (target, rel_type, direction) in rels {
connected_rels.push(serde_json::json!({
"target": rel.to,
"relation": rel.relation_type,
"direction": "outgoing"
}));
} else if rel.to == res.id {
connected_rels.push(serde_json::json!({
"target": rel.from,
"relation": rel.relation_type,
"direction": "incoming"
"target": target,
"relation": rel_type,
"direction": direction
}));
}
}
@@ -851,6 +957,11 @@ impl McpTool for OmniSearchHandler {
serde_json::to_value(&filtered).map_err(|e| e.to_string())
})?;
let q = req.query.clone();
let query_emb = crate::embedding::generate_embedding_async(req.query.clone())
.await
.unwrap_or_default();
let snippets_json = state.code.snippets.read_with(|all_snippets| {
let mut scored: Vec<_> = all_snippets
.iter()
@@ -907,7 +1018,6 @@ impl McpTool for OmniSearchHandler {
serde_json::to_value(&filtered).map_err(|e| e.to_string())
})?;
let q = req.query;
let tech_debts_json = state.code.tech_debts.read_with(|debts| {
let mut scored: Vec<_> = debts
.iter()
@@ -995,7 +1105,7 @@ impl McpTool for OmniSearchHandler {
serde_json::to_value(&filtered).map_err(|e| e.to_string())
})?;
let report = serde_json::json!({
let mut report = serde_json::json!({
"knowledge_graph": kg_json,
"tasks": tasks_json,
"snippets": snippets_json,
@@ -1004,15 +1114,52 @@ impl McpTool for OmniSearchHandler {
"handoff_memos": memos_json,
"error_fixes": error_fixes_json
});
let mut out_str = report.to_string();
if let Some(max_tok) = req.max_tokens {
let max_chars = max_tok * 4;
let mut out_str = report.to_string();
if out_str.len() > max_chars {
out_str.truncate(max_chars);
out_str.push_str("... [TRUNCATED_TO_MAX_TOKENS]");
let prune_keys = [
"error_fixes",
"tech_debts",
"snippets",
"adrs",
"handoff_memos",
"knowledge_graph",
"tasks",
];
let mut pruned = false;
for key in prune_keys {
while out_str.len() > max_chars {
let popped = if let Some(arr) = report.get_mut(key).and_then(|v| v.as_array_mut()) {
if arr.len() > 1 {
arr.pop();
pruned = true;
true
} else {
false
}
} else {
false
};
if popped {
out_str = report.to_string();
} else {
break;
}
}
if out_str.len() <= max_chars {
break;
}
}
if pruned && let Some(obj) = report.as_object_mut() {
obj.insert(
"_truncated_to_max_tokens".to_string(),
serde_json::Value::Bool(true),
);
}
}
}
Ok(out_str)
Ok(report.to_string())
}
}
@@ -1101,8 +1248,13 @@ impl McpTool for SuggestErrorFixHandler {
.iter()
.map(|f| {
let mut score = 0.0;
if req.stack_trace.contains(&f.signature) || f.signature.contains(&req.stack_trace) {
let st_lower = req.stack_trace.to_lowercase();
let sig_lower = f.signature.to_lowercase();
let sol_lower = f.solution.to_lowercase();
if st_lower.contains(&sig_lower) || sig_lower.contains(&st_lower) {
score += 0.8;
} else if st_lower.contains(&sol_lower) || sol_lower.contains(&st_lower) {
score += 0.5;
}
if let Some(emb) = &f.embedding {
if !query_emb.is_empty() {
@@ -1178,6 +1330,7 @@ impl McpTool for ManageCheckpointHandler {
timestamp: crate::handlers::utils::now_secs(),
description: desc.clone(),
namespace: ns,
..Default::default()
};
state.project.snapshots.modify(|snaps| snaps.push(snapshot));
}
@@ -1381,22 +1534,31 @@ impl McpTool for QueryLineageHandler {
"decision": adr.decision
}));
}
}
let changes = state.code.ledger.read_with(|c| c.clone());
for change in changes {
let file_match = change.file_path.to_lowercase().contains(&q);
let desc_match = change.description.to_lowercase().contains(&q);
if file_match || desc_match {
let repo_match = change.repo_name.as_ref().map(|r| r.to_lowercase().contains(&q)).unwrap_or(false);
let symbol_match = change.symbols.iter().any(|s| s.to_lowercase().contains(&q));
let ns_match = change.namespace.to_lowercase().contains(&q);
if file_match || desc_match || repo_match || symbol_match || ns_match {
timeline.push(serde_json::json!({
"timestamp": change.timestamp,
"type": "CodeChange",
"file": change.file_path,
"description": change.description,
"commit": change.git_commit
"commit": change.git_commit,
"branch": change.git_branch,
"repo_name": change.repo_name,
"namespace": change.namespace,
"change_kind": change.change_kind,
"symbols": change.symbols,
"line_range": change.line_range,
"author": change.author,
}));
}
}
}
let fixes = state.code.error_fixes.read_with(|f| f.clone());
for fix in fixes {
@@ -1509,6 +1671,7 @@ impl McpTool for LogHypothesisHandler {
status: req.status.unwrap_or_else(|| "unverified".to_string()),
evidence: req.evidence,
timestamp,
..Default::default()
};
state.code.hypotheses.modify(|h| h.push(record));
@@ -1576,10 +1739,16 @@ impl McpTool for GetPreflightContextHandler {
let pinned = state.project.pinned_files.read_with(|p| p.clone());
let tech_debts = state.code.tech_debts.read_with(|d| d.clone());
let hypotheses = state.code.hypotheses.read_with(|h| h.clone());
let recent_commands = state.telemetry.terminal_history.read_with(|h| {
h.iter().take(5).cloned().collect::<Vec<_>>()
});
let recent_activities = state.telemetry.recent_activities.read_with(|a| {
a.iter().take(5).cloned().collect::<Vec<_>>()
});
let active_tasks: Vec<_> = tasks
.into_iter()
.filter(|t| t.status == "in_progress" || t.status == "open")
.filter(|t| !matches!(t.status.as_str(), "completed" | "done" | "cancelled"))
.map(|t| serde_json::json!({
"id": t.id,
"title": t.title,
@@ -1590,7 +1759,12 @@ impl McpTool for GetPreflightContextHandler {
let branch_pinned: Vec<_> = pinned
.into_iter()
.filter(|p| p.namespace == req.namespace && (req.git_branch.is_none() || p.git_branch == req.git_branch))
.filter(|p| {
p.namespace == req.namespace
&& (req.git_branch.is_none()
|| p.git_branch.is_none()
|| p.git_branch == req.git_branch)
})
.map(|p| p.file_path)
.collect();
@@ -1617,7 +1791,9 @@ impl McpTool for GetPreflightContextHandler {
"active_tasks": active_tasks,
"pinned_files": branch_pinned,
"top_open_tech_debts": open_tech_debts,
"active_hypotheses": active_hypotheses
"active_hypotheses": active_hypotheses,
"recent_terminal_commands": recent_commands,
"recent_activities": recent_activities
});
Ok(serde_json::to_string_pretty(&preflight)?)
@@ -1652,6 +1828,7 @@ impl McpTool for BroadcastAgentSignalHandler {
payload: req.payload,
timestamp,
ttl_seconds: req.ttl_seconds,
..Default::default()
};
state.telemetry.agent_signals.modify(|s| {
@@ -1747,7 +1924,11 @@ impl McpTool for AutoSessionCheckpointHandler {
let hypotheses = state.code.hypotheses.read_with(|h| h.clone());
let ledger = state.code.ledger.read_with(|l| l.clone());
let active_tasks: Vec<_> = tasks.iter().filter(|t| t.status == "in_progress" || t.status == "open").map(|t| t.title.as_str()).collect();
let active_tasks: Vec<_> = tasks
.iter()
.filter(|t| !matches!(t.status.as_str(), "completed" | "done" | "cancelled"))
.map(|t| t.title.as_str())
.collect();
let unverified_hyp: Vec<_> = hypotheses.iter().filter(|h| h.status == "unverified").map(|h| h.hypothesis.as_str()).collect();
let recent_changes: Vec<_> = ledger.iter().rev().take(5).map(|c| c.file_path.as_str()).collect();
@@ -1765,6 +1946,7 @@ impl McpTool for AutoSessionCheckpointHandler {
expires_at: None,
namespace: req.namespace,
timestamp,
..Default::default()
};
state.telemetry.handoff_memos.modify(|m| {
@@ -2036,6 +2218,7 @@ mod tests {
expires_at: None,
dependencies: vec![],
acceptance_criteria: vec![],
..Default::default()
};
{
@@ -2232,6 +2415,7 @@ mod tests {
expires_at: None,
dependencies: vec![],
acceptance_criteria: vec![],
..Default::default()
};
state.project.tasks.modify(|t| t.push(task));
@@ -2408,6 +2592,7 @@ mod tests {
tags: vec!["rust".to_string(), "mcp".to_string()],
updated_at: 0,
embedding: None,
..Default::default()
});
});
@@ -2467,6 +2652,7 @@ mod tests {
expires_at: None,
dependencies: vec![],
acceptance_criteria: vec![],
..Default::default()
};
let task2 = crate::models::Task {
id: "t-2".to_string(),
@@ -2480,6 +2666,7 @@ mod tests {
expires_at: None,
dependencies: vec!["t-1".to_string()],
acceptance_criteria: vec![],
..Default::default()
};
state.project.tasks.modify(|t| {
t.push(task1);
@@ -2607,6 +2794,7 @@ mod tests {
observations: vec!["Obs 1".to_string()],
namespace: "global".to_string(),
git_branch: None,
..Default::default()
});
g.entities.insert("Ent2".to_string(), crate::models::Entity {
name: "Ent2".to_string(),
@@ -2614,12 +2802,14 @@ mod tests {
observations: vec!["Obs 2".to_string()],
namespace: "global".to_string(),
git_branch: None,
..Default::default()
});
g.relations.push(crate::models::Relation {
from: "Ent1".to_string(),
to: "Ent2".to_string(),
relation_type: "uses".to_string(),
namespace: "global".to_string(),
..Default::default()
});
});
@@ -2693,6 +2883,7 @@ mod tests {
timestamp: 0,
description: "Test snap".to_string(),
namespace: "global".to_string(),
..Default::default()
});
});
let rest_snap = mg_chk
@@ -2737,6 +2928,7 @@ mod tests {
dependencies: vec!["NON-EXISTENT-TASK".to_string()],
description: "Blocked task desc".to_string(),
expires_at: None,
..Default::default()
});
@@ -2756,6 +2948,7 @@ mod tests {
timestamp: 0,
namespace: "global".to_string(),
git_branch: Some("main".to_string()),
..Default::default()
});
});
server/src/handlers/notes.rs
+9 -3
View File
@@ -43,6 +43,7 @@ impl McpTool for StickyNotesHandler {
timestamp: now,
content,
expires_at,
..Default::default()
});
});
Ok("Sticky note added.".to_string())
@@ -72,9 +73,11 @@ impl McpTool for StickyNotesHandler {
})?;
let mut success = false;
state.code.sticky.modify(|notes| {
let target_pos = if idx == 0 { 0 } else { idx - 1 };
if target_pos < notes.len() {
notes.remove(target_pos);
if idx > 0 && idx <= notes.len() {
notes.remove(idx - 1);
success = true;
} else if idx == 0 && !notes.is_empty() {
notes.remove(0);
success = true;
}
});
@@ -129,6 +132,7 @@ impl McpTool for HandoffMemosHandler {
namespace: ns,
timestamp: crate::handlers::utils::now_secs(),
expires_at: None,
..Default::default()
});
if memos.len() > 200 {
memos.remove(0);
@@ -190,6 +194,7 @@ impl McpTool for AddSessionSummaryHandler {
namespace: req.namespace,
timestamp: crate::handlers::utils::now_secs(),
expires_at: None,
..Default::default()
});
if summaries.len() > 200 {
summaries.remove(0);
@@ -273,6 +278,7 @@ impl McpTool for PromoteToEntityHandler {
observations: Vec::new(),
namespace,
git_branch: None,
..Default::default()
});
if !entity.observations.contains(&obs) {
entity.observations.push(obs.clone());
server/src/handlers/tasks.rs
+120 -33
View File
@@ -45,21 +45,34 @@ impl McpTool for TasksHandler {
dependencies: deps,
acceptance_criteria: vec![],
expires_at: None,
namespace: req.namespace.unwrap_or_else(|| "default".to_string()),
repo_name: req.repo_name,
priority: req.priority,
assigned_agent: req.assigned_agent,
verification_command: req.verification_command,
..Default::default()
};
let idx = state.get_search_index().await;
drop(idx.index_task(&task));
state.project.tasks.modify(|tasks| {
tasks.push(task.clone());
});
state.record_activity("task_create", &format!("Created task: {}", task.title), Some(&task.description));
state.record_activity(
"task_create",
&format!("Created task: {}", task.title),
Some(&task.description),
);
state.broadcast_task_event(TaskEvent {
task_id: task_id.clone(),
status: "created".to_string(),
action: Some("add".to_string()),
result: Some(serde_json::json!({ "title": task.title, "git_branch": task.git_branch })),
result: Some(
serde_json::json!({ "title": task.title, "git_branch": task.git_branch }),
),
error: None,
timestamp: now,
session_id: None,
..Default::default()
});
Ok(format!("Task added with ID: {}", task_id))
}
@@ -75,6 +88,8 @@ impl McpTool for TasksHandler {
let mut blocked = false;
let mut blocker_details = String::new();
let mut updated_task: Option<Task> = None;
state.project.tasks.modify(|tasks| {
let target_idx = tasks.iter().position(|t| t.id == id || t.title == id);
let target_idx = match target_idx {
@@ -84,7 +99,11 @@ impl McpTool for TasksHandler {
found = true;
if target_status == "done" || target_status == "completed" {
if tasks[target_idx].acceptance_criteria.iter().any(|c| !c.is_met) {
if tasks[target_idx]
.acceptance_criteria
.iter()
.any(|c| !c.is_met)
{
blocked = true;
blocker_details = "Unmet acceptance criteria exist.".to_string();
}
@@ -100,20 +119,29 @@ impl McpTool for TasksHandler {
}
if !uncompleted_deps.is_empty() {
blocked = true;
blocker_details = format!("Blocked by dependencies: {}", uncompleted_deps.join(", "));
blocker_details = format!(
"Blocked by dependencies: {}",
uncompleted_deps.join(", ")
);
}
}
if !blocked {
let target_id_ref = tasks[target_idx].id.as_str();
let mut uncompleted_children = Vec::new();
for child in tasks.iter().filter(|t| t.parent_id.as_deref() == Some(target_id_ref)) {
for child in tasks
.iter()
.filter(|t| t.parent_id.as_deref() == Some(target_id_ref))
{
if child.status != "completed" && child.status != "done" {
uncompleted_children.push(child.title.as_str());
}
}
if !uncompleted_children.is_empty() {
blocked = true;
blocker_details = format!("Blocked by child tasks: {}", uncompleted_children.join(", "));
blocker_details = format!(
"Blocked by child tasks: {}",
uncompleted_children.join(", ")
);
}
}
}
@@ -121,6 +149,7 @@ impl McpTool for TasksHandler {
if !blocked {
tasks[target_idx].status = target_status.clone();
tasks[target_idx].updated_at = crate::handlers::utils::now_secs();
updated_task = Some(tasks[target_idx].clone());
}
});
@@ -133,10 +162,22 @@ impl McpTool for TasksHandler {
error: Some(blocker_details.clone()),
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
Err(crate::error::AppError::Internal(format!("Error: Cannot transition task. {}", blocker_details)))
Err(crate::error::AppError::Internal(format!(
"Error: Cannot transition task. {}",
blocker_details
)))
} else if found {
state.record_activity("task_update", &format!("Task {} status -> {}", id, status), None);
if let Some(ref ut) = updated_task {
let idx = state.get_search_index().await;
drop(idx.index_task(ut));
}
state.record_activity(
"task_update",
&format!("Task {} status -> {}", id, status),
None,
);
state.broadcast_task_event(TaskEvent {
task_id: id.clone(),
status: target_status.clone(),
@@ -145,10 +186,13 @@ impl McpTool for TasksHandler {
error: None,
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
Ok("Task status updated.".to_string())
} else {
Err(crate::error::AppError::Internal("Task not found. Please verify the task ID.".to_string()))
Err(crate::error::AppError::Internal(
"Task not found. Please verify the task ID.".to_string(),
))
}
}
TaskAction::Delete => {
@@ -163,9 +207,12 @@ impl McpTool for TasksHandler {
for (idx, t) in tasks.iter().enumerate() {
id_to_index.insert(t.id.as_str(), idx);
}
let mut children_map: std::collections::HashMap<usize, Vec<usize>> = std::collections::HashMap::new();
let mut children_map: std::collections::HashMap<usize, Vec<usize>> =
std::collections::HashMap::new();
for (idx, t) in tasks.iter().enumerate() {
if let Some(pid) = &t.parent_id && let Some(&p_idx) = id_to_index.get(pid.as_str()) {
if let Some(pid) = &t.parent_id
&& let Some(&p_idx) = id_to_index.get(pid.as_str())
{
children_map.entry(p_idx).or_default().push(idx);
}
}
@@ -174,7 +221,9 @@ impl McpTool for TasksHandler {
let mut queue = std::collections::VecDeque::new();
queue.push_back(start_idx);
while let Some(curr) = queue.pop_front() {
if to_delete_idx.insert(curr) && let Some(children) = children_map.get(&curr) {
if to_delete_idx.insert(curr)
&& let Some(children) = children_map.get(&curr)
{
queue.extend(children.iter().copied());
}
}
@@ -199,19 +248,29 @@ impl McpTool for TasksHandler {
error: None,
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
Ok(format!("Deleted task and its children ({} total).", deleted_count))
Ok(format!(
"Deleted task and its children ({} total).",
deleted_count
))
} else {
Err(crate::error::AppError::Internal("Task not found. Please verify the task ID.".to_string()))
Err(crate::error::AppError::Internal(
"Task not found. Please verify the task ID.".to_string(),
))
}
}
TaskAction::List => {
let level = req.summary_level.as_deref().unwrap_or("detailed");
let data = state.project.tasks.read_with(|tasks| {
let filtered: Vec<_> = tasks
let mut filtered: Vec<_> = tasks
.iter()
.filter(|t| {
let status_match = t.status != "done" && t.status != "completed";
let status_match = match req.status.as_deref() {
Some("all") => true,
Some(s) => t.status.eq_ignore_ascii_case(s),
None => t.status != "done" && t.status != "completed",
};
let branch_match = match &req.git_branch {
Some(branch) => t.git_branch.is_none() || t.git_branch.as_deref() == Some(branch.as_str()),
None => true,
@@ -227,10 +286,9 @@ impl McpTool for TasksHandler {
let mut json_str = serde_json::to_string(&filtered)?;
if let Some(max_t) = req.max_tokens {
let char_limit = max_t * 4;
if json_str.len() > char_limit {
let boundary = json_str.floor_char_boundary(char_limit);
json_str.truncate(boundary);
json_str.push_str(" ...[truncated due to max_tokens]");
while json_str.len() > char_limit && !filtered.is_empty() {
filtered.pop();
json_str = serde_json::to_string(&filtered)?;
}
}
Ok::<String, crate::error::AppError>(json_str)
@@ -253,6 +311,7 @@ impl McpTool for TasksHandler {
id: uuid::Uuid::new_v4().to_string(),
description: desc,
is_met: false,
..Default::default()
})
.collect();
task.updated_at = crate::handlers::utils::now_secs();
@@ -268,10 +327,13 @@ impl McpTool for TasksHandler {
error: None,
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
Ok("Acceptance criteria set successfully.".to_string())
} else {
Err(crate::error::AppError::Internal("Task not found. Please verify the task ID.".to_string()))
Err(crate::error::AppError::Internal(
"Task not found. Please verify the task ID.".to_string(),
))
}
}
TaskAction::Verify => {
@@ -292,11 +354,13 @@ impl McpTool for TasksHandler {
task.updated_at = crate::handlers::utils::now_secs();
}
} else {
task.acceptance_criteria.push(crate::models::AcceptanceCriteria {
id: uuid::Uuid::new_v4().to_string(),
description: proof_str.clone(),
is_met: true,
});
task.acceptance_criteria
.push(crate::models::AcceptanceCriteria {
id: uuid::Uuid::new_v4().to_string(),
description: proof_str.clone(),
is_met: true,
..Default::default()
});
task.updated_at = crate::handlers::utils::now_secs();
success = true;
}
@@ -311,12 +375,18 @@ impl McpTool for TasksHandler {
error: None,
timestamp: crate::handlers::utils::now_secs(),
session_id: None,
..Default::default()
});
Ok(format!("Acceptance criteria verified with proof: {}", proof_str))
Ok(format!(
"Acceptance criteria verified with proof: {}",
proof_str
))
} else if already_met {
Ok("Acceptance criteria was already met.".to_string())
} else {
Err(crate::error::AppError::Internal("Task not found. Please verify the task ID.".to_string()))
Err(crate::error::AppError::Internal(
"Task not found. Please verify the task ID.".to_string(),
))
}
}
}
@@ -345,14 +415,23 @@ impl McpTool for MilestonesHandler {
let title = req.title.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'title' for action 'add'. Next step: Provide non-empty 'title' string in request and retry.".to_string())
})?;
let ns = req.namespace.unwrap_or_else(|| crate::models::default_namespace());
let ns = req
.namespace
.unwrap_or_else(|| crate::models::default_namespace());
state.project.milestones.modify(|ms| {
ms.push(crate::models::Milestone {
id: uuid::Uuid::new_v4().to_string(),
title,
status: "pending".to_string(),
namespace: ns,
target_date: None,
target_date: req
.target_date
.as_deref()
.and_then(|s| s.parse::<u64>().ok()),
description: req.description,
deliverables: req.deliverables.unwrap_or_default(),
repo_name: req.repo_name,
..Default::default()
})
});
Ok("Milestone added".to_string())
@@ -377,7 +456,9 @@ impl McpTool for MilestonesHandler {
if found {
Ok("Milestone updated".to_string())
} else {
Err(crate::error::AppError::Internal("Milestone not found. Please verify the milestone ID.".to_string()))
Err(crate::error::AppError::Internal(
"Milestone not found. Please verify the milestone ID.".to_string(),
))
}
}
MilestoneAction::List => {
@@ -504,7 +585,10 @@ mod tests {
assert_eq!(res2, "Milestone updated");
let res3 = handler_ms
.execute(json!({"action": "list", "namespace": "global"}), state.clone())
.execute(
json!({"action": "list", "namespace": "global"}),
state.clone(),
)
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
@@ -602,7 +686,10 @@ mod tests {
assert!(res1.contains("Milestone added"));
let res2 = handler
.execute(serde_json::json!({"action": "list", "namespace": "global"}), state.clone())
.execute(
serde_json::json!({"action": "list", "namespace": "global"}),
state.clone(),
)
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
server/src/handlers/utils.rs
+13
View File
@@ -14,6 +14,19 @@ pub fn now_secs() -> u64 {
.as_secs()
}
pub fn validate_safe_path(path_str: &str) -> crate::error::Result<std::path::PathBuf> {
let path = std::path::Path::new(path_str);
for component in path.components() {
if component == std::path::Component::ParentDir {
return Err(crate::error::AppError::BadRequest(format!(
"Path traversal forbidden: '{}' contains relative parent directory components ('..')",
path_str
)));
}
}
Ok(path.to_path_buf())
}
#[cfg(test)]
mod tests {
use super::*;
server/src/handlers/vision.rs
+9 -58
View File
@@ -188,6 +188,14 @@ impl McpTool for ReadClipboardHandler {
{
let path_str = file_path.to_string_lossy().to_string();
out.insert("image_path".into(), json!(path_str));
#[cfg(target_os = "windows")]
if let Some(rest) = path_str.strip_prefix("C:\\") {
let wsl_path = format!("/mnt/c/{}", rest.replace('\\', "/"));
out.insert("image_path_wsl".into(), json!(wsl_path));
} else if let Some(rest) = path_str.strip_prefix("c:\\") {
let wsl_path = format!("/mnt/c/{}", rest.replace('\\', "/"));
out.insert("image_path_wsl".into(), json!(wsl_path));
}
if let Ok(bytes) = std::fs::read(&file_path) {
use base64::Engine;
@@ -226,42 +234,7 @@ impl McpTool for ReadClipboardHandler {
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&final_obj)?)
}
}
use crate::tools::ToggleClipboardWatchModeTool;
pub struct ToggleClipboardWatchModeHandler;
#[async_trait]
impl McpTool for ToggleClipboardWatchModeHandler {
fn name(&self) -> &'static str {
"toggle_clipboard_watch_mode"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ToggleClipboardWatchModeTool>(
"toggle_clipboard_watch_mode",
"Enable or disable clipboard watch mode (auto-ingesting into StickyNotes).",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: ToggleClipboardWatchModeTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut watch_mode = state.clipboard_watch_mode.write().await;
*watch_mode = tool_args.enable;
drop(watch_mode);
state.clipboard_notify.notify_waiters();
let status_msg = if tool_args.enable {
"Clipboard watch mode enabled. Changes will be ingested as StickyNotes."
} else {
"Clipboard watch mode disabled."
};
state.record_activity("clipboard", status_msg, None);
Ok(status_msg.to_string())
}
}
use crate::tools::{ClipboardAction, ClipboardTool};
@@ -276,7 +249,7 @@ impl McpTool for ClipboardHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<ClipboardTool>(
"clipboard",
"Consolidated OS clipboard management (read, write, toggle_watch)",
"Consolidated OS clipboard management (read, write)",
)
}
@@ -297,10 +270,6 @@ impl McpTool for ClipboardHandler {
state,
).await
}
ClipboardAction::ToggleWatch => {
let enable = req.enabled.unwrap_or(true);
ToggleClipboardWatchModeHandler.execute(serde_json::json!({"enable": enable}), state).await
}
}
}
}
@@ -312,24 +281,6 @@ mod tests {
use std::sync::Arc;
use tempfile::tempdir;
#[tokio::test]
async fn test_toggle_clipboard_watch_mode() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let handler = ToggleClipboardWatchModeHandler;
let args = json!({
"enable": true
});
let result = handler
.execute(args, state.clone())
.await
.map_err(|e| format!("Failed to toggle clipboard: {}", e))
.unwrap();
assert!(result.contains("enabled"));
assert_eq!(*state.clipboard_watch_mode.read().await, true);
}
#[tokio::test]
async fn test_write_clipboard() {
server/src/handlers/workspaces.rs
+279 -84
View File
@@ -15,7 +15,10 @@ impl McpTool for PinFileHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<PinFileTool>("pin_file", "Pin a file path in the workspace context to prioritize it across subagent turns.")
crate::mcp::tool_def::<PinFileTool>(
"pin_file",
"Pin a file path in the workspace context to prioritize it across subagent turns.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -27,6 +30,9 @@ impl McpTool for PinFileHandler {
file_path: req.file_path,
timestamp: crate::handlers::utils::now_secs(),
git_branch: req.git_branch,
repo_name: req.repo_name,
symbol_or_line: req.symbol_or_line,
pin_reason: req.pin_reason,
});
});
Ok("File pinned".to_string())
@@ -42,7 +48,10 @@ impl McpTool for UnpinFileHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Unpin a file path from the workspace context.")
crate::mcp::tool_def::<UnpinFileTool>(
"unpin_file",
"Unpin a file path from the workspace context.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -104,7 +113,10 @@ impl McpTool for StoreSnippetHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Store or replace a reusable code snippet with vector embedding and search tags.")
crate::mcp::tool_def::<StoreSnippetTool>(
"store_snippet",
"Store or replace a reusable code snippet with vector embedding and search tags.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -125,6 +137,12 @@ impl McpTool for StoreSnippetHandler {
updated_at: crate::handlers::utils::now_secs(),
tags: req.tags.unwrap_or_default(),
embedding,
namespace: req
.namespace
.unwrap_or_else(|| crate::models::default_namespace()),
repo_name: req.repo_name,
origin_file: req.origin_file,
line_range: req.line_range,
};
let idx = state.get_search_index().await;
@@ -148,7 +166,10 @@ impl McpTool for TagSnippetHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<TagSnippetTool>("tag_snippet", "Add organizational tags to an existing stored code snippet.")
crate::mcp::tool_def::<TagSnippetTool>(
"tag_snippet",
"Add organizational tags to an existing stored code snippet.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -169,9 +190,15 @@ impl McpTool for TagSnippetHandler {
});
if found {
Ok(format!("Snippet '{}' tagged with: {:?}", req.name, req.tags))
Ok(format!(
"Snippet '{}' tagged with: {:?}",
req.name, req.tags
))
} else {
Err(crate::error::AppError::Internal(format!("Snippet '{}' not found.", req.name)))
Err(crate::error::AppError::Internal(format!(
"Snippet '{}' not found.",
req.name
)))
}
}
}
@@ -209,7 +236,9 @@ impl McpTool for SearchSnippetsHandler {
|| s.tags.iter().any(|t| contains_ignore_ascii_case(t, &query));
let matches_tags = tags_filter.is_empty()
|| tags_filter.iter().any(|tf| s.tags.iter().any(|st| contains_ignore_ascii_case(st, tf)));
|| tags_filter
.iter()
.any(|tf| s.tags.iter().any(|st| contains_ignore_ascii_case(st, tf)));
matches_text && matches_tags
})
@@ -292,6 +321,10 @@ impl McpTool for SaveContextWorkspaceHandler {
pinned_files: req.pinned_files,
active_task_ids: req.active_task_ids,
saved_at: crate::handlers::utils::now_secs(),
description: req.description,
repo_name: req.repo_name,
git_branch: req.git_branch,
vcs_revision: req.vcs_revision,
});
});
Ok("Context workspace saved".to_string())
@@ -419,6 +452,10 @@ impl McpTool for AddPrChecklistItemHandler {
namespace: req.namespace,
id: uuid::Uuid::new_v4().to_string(),
description: req.description,
is_checked: req.is_checked.unwrap_or(false),
category: req.category,
repo_name: req.repo_name,
automated_check: req.automated_check,
})
});
Ok("PR checklist item added".to_string())
@@ -492,22 +529,43 @@ impl McpTool for PinnedFilesHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PinnedFilesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = req.namespace.unwrap_or_else(|| crate::models::default_namespace());
let ns = req
.namespace
.unwrap_or_else(|| crate::models::default_namespace());
match req.action {
PinnedFileAction::Pin => {
let path = req.path.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'path' for action 'pin'. Next step: Provide absolute or relative file 'path' in request and retry.".to_string())
})?;
PinFileHandler.execute(serde_json::json!({"file_path": path, "namespace": ns}), state).await
PinFileHandler
.execute(
serde_json::json!({
"file_path": path,
"namespace": ns,
"git_branch": req.git_branch,
"pin_reason": req.pin_reason,
"symbol_or_line": req.symbol_or_line,
"repo_name": req.repo_name,
}),
state,
)
.await
}
PinnedFileAction::Unpin => {
let path = req.path.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'path' for action 'unpin'. Next step: Provide file 'path' in request and retry.".to_string())
})?;
UnpinFileHandler.execute(serde_json::json!({"file_path": path, "namespace": ns}), state).await
UnpinFileHandler
.execute(
serde_json::json!({"file_path": path, "namespace": ns}),
state,
)
.await
}
PinnedFileAction::List => {
ListPinnedFilesHandler.execute(serde_json::json!({"namespace": ns}), state).await
ListPinnedFilesHandler
.execute(serde_json::json!({"namespace": ns}), state)
.await
}
}
}
@@ -530,7 +588,9 @@ impl McpTool for ContextWorkspacesHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ContextWorkspacesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = req.namespace.unwrap_or_else(|| crate::models::default_namespace());
let ns = req
.namespace
.unwrap_or_else(|| crate::models::default_namespace());
match req.action {
ContextWorkspaceAction::Save => {
let name = req.name.ok_or_else(|| {
@@ -538,30 +598,42 @@ impl McpTool for ContextWorkspacesHandler {
})?;
let pinned_files = req.pinned_files.unwrap_or_default();
let active_task_ids = req.active_task_ids.unwrap_or_default();
SaveContextWorkspaceHandler.execute(
serde_json::json!({
"name": name,
"namespace": ns,
"pinned_files": pinned_files,
"active_task_ids": active_task_ids
}),
state,
).await
SaveContextWorkspaceHandler
.execute(
serde_json::json!({
"name": name,
"namespace": ns,
"pinned_files": pinned_files,
"active_task_ids": active_task_ids,
"description": req.description,
"repo_name": req.repo_name,
"git_branch": req.git_branch,
"vcs_revision": req.vcs_revision,
}),
state,
)
.await
}
ContextWorkspaceAction::Load => {
let name = req.name.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name' for action 'load'. Next step: Provide workspace 'name' string in request and retry.".to_string())
})?;
LoadContextWorkspaceHandler.execute(serde_json::json!({"name": name, "namespace": ns}), state).await
LoadContextWorkspaceHandler
.execute(serde_json::json!({"name": name, "namespace": ns}), state)
.await
}
ContextWorkspaceAction::List => {
ListContextWorkspacesHandler.execute(serde_json::json!({"namespace": ns}), state).await
ListContextWorkspacesHandler
.execute(serde_json::json!({"namespace": ns}), state)
.await
}
ContextWorkspaceAction::Delete => {
let name = req.name.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name' for action 'delete'. Next step: Provide workspace 'name' string in request and retry.".to_string())
})?;
DeleteContextWorkspaceHandler.execute(serde_json::json!({"name": name, "namespace": ns}), state).await
DeleteContextWorkspaceHandler
.execute(serde_json::json!({"name": name, "namespace": ns}), state)
.await
}
ContextWorkspaceAction::Diff => {
let name = req.name.ok_or_else(|| {
@@ -593,19 +665,37 @@ impl McpTool for PrChecklistHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = req.namespace.unwrap_or_else(|| crate::models::default_namespace());
let ns = req
.namespace
.unwrap_or_else(|| crate::models::default_namespace());
match req.action {
PrChecklistAction::Add => {
let desc = req.description.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'description' for action 'add'. Next step: Provide checklist item 'description' string in request and retry.".to_string())
})?;
AddPrChecklistItemHandler.execute(serde_json::json!({"description": desc, "namespace": ns}), state).await
AddPrChecklistItemHandler
.execute(
serde_json::json!({
"description": desc,
"namespace": ns,
"category": req.category,
"automated_check": req.automated_check,
"is_checked": req.is_checked,
"repo_name": req.repo_name,
}),
state,
)
.await
}
PrChecklistAction::Get => {
GetPrChecklistHandler.execute(serde_json::json!({"namespace": ns}), state).await
GetPrChecklistHandler
.execute(serde_json::json!({"namespace": ns}), state)
.await
}
PrChecklistAction::Clear => {
ClearPrChecklistHandler.execute(serde_json::json!({"namespace": ns}), state).await
ClearPrChecklistHandler
.execute(serde_json::json!({"namespace": ns}), state)
.await
}
}
}
@@ -636,37 +726,51 @@ impl McpTool for SnippetsHandler {
let lang = req.language.unwrap_or_else(|| "text".to_string());
let code = req.code.unwrap_or_default();
let desc = req.description.unwrap_or_default();
StoreSnippetHandler.execute(
serde_json::json!({
"name": name,
"language": lang,
"code": code,
"description": desc,
"tags": req.tags
}),
state,
).await
StoreSnippetHandler
.execute(
serde_json::json!({
"name": name,
"language": lang,
"code": code,
"description": desc,
"tags": req.tags,
"namespace": req.namespace,
"repo_name": req.repo_name,
"origin_file": req.origin_file,
"line_range": req.line_range,
}),
state,
)
.await
}
SnippetAction::Search => {
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
crate::handlers::meta::SearchSnippetsHybridHandler
.execute(serde_json::json!({"query": q, "tags": req.tags}), state)
.await
} else {
SearchSnippetsHandler.execute(serde_json::json!({"query": q, "tags": req.tags}), state).await
SearchSnippetsHandler
.execute(serde_json::json!({"query": q, "tags": req.tags}), state)
.await
}
}
SnippetAction::Delete => {
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
DeleteSnippetHandler
.execute(serde_json::json!({"name": id}), state)
.await
}
SnippetAction::Tag => {
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
TagSnippetHandler
.execute(serde_json::json!({"name": id, "tags": tags}), state)
.await
}
}
}
@@ -793,24 +897,42 @@ mod tests {
// Snapshots
let create_snap = CreateSnapshotHandler;
let snap_res = create_snap.execute(json!({"description": "Test snap", "namespace": "global"}), state.clone()).await.unwrap();
let snap_res = create_snap
.execute(
json!({"description": "Test snap", "namespace": "global"}),
state.clone(),
)
.await
.unwrap();
assert!(snap_res.contains("SNAP-"));
let snap_id = snap_res.split_whitespace().nth(1).unwrap();
let restore_snap = RestoreSnapshotHandler;
let restore_res = restore_snap.execute(json!({"snapshot_id": snap_id}), state.clone()).await.unwrap();
let restore_res = restore_snap
.execute(json!({"snapshot_id": snap_id}), state.clone())
.await
.unwrap();
assert!(restore_res.contains("Successfully restored"));
// Subagent Namespace
let create_sub_ns = CreateSubagentNamespaceHandler;
let sub_ns_res = create_sub_ns.execute(json!({"subagent_id": "test_agent"}), state.clone()).await.unwrap();
let sub_ns_res = create_sub_ns
.execute(json!({"subagent_id": "test_agent"}), state.clone())
.await
.unwrap();
assert!(sub_ns_res.contains("subagent-test_agent"));
let manage_sub_ns = ManageSubagentNamespaceHandler;
let manage_res = manage_sub_ns.execute(json!({
"subagent_id": "test_agent",
"action": "create"
}), state.clone()).await.unwrap();
let manage_res = manage_sub_ns
.execute(
json!({
"subagent_id": "test_agent",
"action": "create"
}),
state.clone(),
)
.await
.unwrap();
assert!(manage_res.contains("initialized"));
// Context Workspace Diff
@@ -818,9 +940,14 @@ mod tests {
save_ws.execute(json!({"name": "ws_a", "description": "Workspace A", "pinned_files": ["src/lib.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
save_ws.execute(json!({"name": "ws_b", "description": "Workspace B", "pinned_files": ["src/main.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
let diff_ws = DiffContextWorkspacesHandler;
let diff_res = diff_ws.execute(json!({"workspace_a": "ws_a", "workspace_b": "ws_b"}), state.clone()).await.unwrap();
let diff_res = diff_ws
.execute(
json!({"workspace_a": "ws_a", "workspace_b": "ws_b"}),
state.clone(),
)
.await
.unwrap();
assert!(diff_res.contains("workspace_a"));
}
@@ -834,7 +961,13 @@ mod tests {
save_ws.execute(json!({"name": "ws_2", "description": "Workspace 2", "pinned_files": ["src/main.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
let diff_ws = DiffContextWorkspacesHandler;
let diff_res = diff_ws.execute(json!({"workspace_a": "ws_1", "workspace_b": "ws_2"}), state.clone()).await.unwrap();
let diff_res = diff_ws
.execute(
json!({"workspace_a": "ws_1", "workspace_b": "ws_2"}),
state.clone(),
)
.await
.unwrap();
assert!(diff_res.contains("ws_1"));
}
}
@@ -998,7 +1131,10 @@ impl McpTool for CreateSnapshotHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: CreateSnapshotTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let snapshot_id = format!("SNAP-{}", uuid::Uuid::new_v4().to_string()[..8].to_uppercase());
let snapshot_id = format!(
"SNAP-{}",
uuid::Uuid::new_v4().to_string()[..8].to_uppercase()
);
let ns = req.namespace.unwrap_or_else(|| "global".to_string());
let snapshot = crate::models::StateSnapshot {
@@ -1006,14 +1142,22 @@ impl McpTool for CreateSnapshotHandler {
timestamp: crate::handlers::utils::now_secs(),
description: req.description.clone(),
namespace: ns.clone(),
..Default::default()
};
state.project.snapshots.modify(|snaps| {
snaps.push(snapshot);
});
state.record_activity("snapshot", &format!("Created snapshot {}: {}", snapshot_id, req.description), None);
Ok(format!("Snapshot {} created for namespace '{}'", snapshot_id, ns))
state.record_activity(
"snapshot",
&format!("Created snapshot {}: {}", snapshot_id, req.description),
None,
);
Ok(format!(
"Snapshot {} created for namespace '{}'",
snapshot_id, ns
))
}
}
@@ -1031,15 +1175,26 @@ impl McpTool for RestoreSnapshotHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: RestoreSnapshotTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let found = state.project.snapshots.read_with(|snaps| {
snaps.iter().any(|s| s.id == req.snapshot_id)
});
let found = state
.project
.snapshots
.read_with(|snaps| snaps.iter().any(|s| s.id == req.snapshot_id));
if found {
state.record_activity("snapshot", &format!("Restored snapshot {}", req.snapshot_id), None);
Ok(format!("Successfully restored memory state from snapshot {}", req.snapshot_id))
state.record_activity(
"snapshot",
&format!("Restored snapshot {}", req.snapshot_id),
None,
);
Ok(format!(
"Successfully restored memory state from snapshot {}",
req.snapshot_id
))
} else {
Err(crate::error::AppError::Internal(format!("Snapshot {} not found", req.snapshot_id)))
Err(crate::error::AppError::Internal(format!(
"Snapshot {} not found",
req.snapshot_id
)))
}
}
}
@@ -1060,20 +1215,35 @@ impl McpTool for ManageSubagentNamespaceHandler {
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ManageSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let req: ManageSubagentNamespaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let subagent_ns = format!("subagent-{}", req.subagent_id);
match req.action {
SubagentNamespaceAction::Create => {
state.record_activity("namespace", &format!("Created subagent namespace {}", subagent_ns), None);
Ok(format!("Subagent memory namespace '{}' initialized.", subagent_ns))
state.record_activity(
"namespace",
&format!("Created subagent namespace {}", subagent_ns),
None,
);
Ok(format!(
"Subagent memory namespace '{}' initialized.",
subagent_ns
))
}
SubagentNamespaceAction::Purge => {
state.project.pinned_files.modify(|files| {
files.retain(|f| f.namespace != subagent_ns);
});
state.record_activity("namespace", &format!("Purged subagent namespace {}", subagent_ns), None);
Ok(format!("Subagent memory namespace '{}' purged.", subagent_ns))
state.record_activity(
"namespace",
&format!("Purged subagent namespace {}", subagent_ns),
None,
);
Ok(format!(
"Subagent memory namespace '{}' purged.",
subagent_ns
))
}
SubagentNamespaceAction::Condense => {
let purge = req.purge_after_promotion.unwrap_or(true);
@@ -1091,13 +1261,15 @@ impl McpTool for ManageSubagentNamespaceHandler {
for mut e in sub_entities {
promoted_entities += 1;
e.namespace = "global".to_string();
let global_entry = g.entities.entry(e.name.clone()).or_insert_with(|| Entity {
name: e.name.clone(),
entity_type: e.entity_type.clone(),
observations: Vec::new(),
namespace: "global".to_string(),
git_branch: e.git_branch.clone(),
});
let global_entry =
g.entities.entry(e.name.clone()).or_insert_with(|| Entity {
name: e.name.clone(),
entity_type: e.entity_type.clone(),
observations: Vec::new(),
namespace: "global".to_string(),
git_branch: e.git_branch.clone(),
..Default::default()
});
for obs in e.observations {
if !global_entry.observations.contains(&obs) {
global_entry.observations.push(obs);
@@ -1142,13 +1314,21 @@ impl McpTool for CreateSubagentNamespaceHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateSubagentNamespaceTool>("create_subagent_namespace", "Create subagent namespace")
crate::mcp::tool_def::<CreateSubagentNamespaceTool>(
"create_subagent_namespace",
"Create subagent namespace",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: CreateSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let req: CreateSubagentNamespaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = format!("subagent-{}", req.subagent_id);
state.record_activity("namespace", &format!("Created subagent namespace {}", ns), None);
state.record_activity(
"namespace",
&format!("Created subagent namespace {}", ns),
None,
);
Ok(format!("Subagent memory namespace '{}' initialized.", ns))
}
}
@@ -1162,18 +1342,26 @@ impl McpTool for PurgeSubagentNamespaceHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<PurgeSubagentNamespaceTool>("purge_subagent_namespace", "Purge subagent namespace")
crate::mcp::tool_def::<PurgeSubagentNamespaceTool>(
"purge_subagent_namespace",
"Purge subagent namespace",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PurgeSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let req: PurgeSubagentNamespaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = format!("subagent-{}", req.subagent_id);
state.project.pinned_files.modify(|files| {
files.retain(|f| f.namespace != ns);
});
state.record_activity("namespace", &format!("Purged subagent namespace {}", ns), None);
state.record_activity(
"namespace",
&format!("Purged subagent namespace {}", ns),
None,
);
Ok(format!("Subagent memory namespace '{}' purged.", ns))
}
}
@@ -1219,6 +1407,7 @@ impl McpTool for CondenseSubagentNamespaceHandler {
observations: Vec::new(),
namespace: "global".to_string(),
git_branch: e.git_branch.clone(),
..Default::default()
});
for obs in e.observations {
if !global_entry.observations.contains(&obs) {
@@ -1262,11 +1451,15 @@ impl McpTool for DiffContextWorkspacesHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<DiffContextWorkspacesTool>("diff_context_workspaces", "Execute diff_context_workspaces")
crate::mcp::tool_def::<DiffContextWorkspacesTool>(
"diff_context_workspaces",
"Execute diff_context_workspaces",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: DiffContextWorkspacesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let req: DiffContextWorkspacesTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let workspaces = state.project.context_workspaces.read_with(|ws| ws.clone());
let ws_a = workspaces.iter().find(|w| w.name == req.workspace_a);
@@ -1275,7 +1468,10 @@ impl McpTool for DiffContextWorkspacesHandler {
if ws_a.is_none() || ws_b.is_none() {
return Err(crate::error::AppError::Internal(format!(
"One or both context workspaces not found: '{}' found={}, '{}' found={}",
req.workspace_a, ws_a.is_some(), req.workspace_b, ws_b.is_some()
req.workspace_a,
ws_a.is_some(),
req.workspace_b,
ws_b.is_some()
)));
}
@@ -1314,4 +1510,3 @@ impl McpTool for DiffContextWorkspacesHandler {
Ok(serde_json::to_string_pretty(&diff)?)
}
}
server/src/indexer.rs
+13 -5
View File
@@ -73,7 +73,10 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
// Gold Standard: Batch generate embeddings in chunks of 16 to eliminate sequential HTTP overhead
for chunk_batch in chunks.chunks(16) {
let texts: Vec<String> = chunk_batch.iter().map(|(_, code, _)| code.clone()).collect();
let texts: Vec<String> = chunk_batch
.iter()
.map(|(_, code, _)| code.clone())
.collect();
let embeddings = generate_embeddings_async(texts).await.unwrap_or_default();
let mut new_snippets = Vec::with_capacity(chunk_batch.len());
@@ -84,7 +87,8 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
for (i, (name, code, desc)) in chunk_batch.iter().enumerate() {
let embedding = embeddings.get(i).cloned();
let file_name = file_path.file_name().unwrap_or_default().to_string_lossy();
let file_name =
file_path.file_name().unwrap_or_default().to_string_lossy();
let snippet_name = format!("{}:{}", file_name, name);
let snippet = Snippet {
@@ -95,6 +99,7 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
updated_at: now,
tags: vec![],
embedding,
..Default::default()
};
new_snippets.push(snippet);
}
@@ -162,7 +167,10 @@ fn extract_chunks_with_parent(
let mut cursor = node.walk();
let mut type_name = None;
for child in node.children(&mut cursor) {
if child.kind() == "type_identifier" || child.kind() == "name" || child.kind() == "identifier" {
if child.kind() == "type_identifier"
|| child.kind() == "name"
|| child.kind() == "identifier"
{
type_name = child.utf8_text(code.as_bytes()).ok();
break;
}
@@ -189,7 +197,8 @@ fn extract_chunks_with_parent(
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
let child_kind = child.kind();
if child_kind == "identifier" || child_kind == "name" || child_kind == "type_identifier" {
if child_kind == "identifier" || child_kind == "name" || child_kind == "type_identifier"
{
if let Ok(text) = child.utf8_text(code.as_bytes()) {
name = text;
}
@@ -284,4 +293,3 @@ mod tests {
start_background_indexer(state).await;
}
}
server/src/instructions.md
+58 -17
View File
@@ -24,17 +24,17 @@ To prevent graph fragmentation and ensure seamless LLM context retrieval:
The server consolidates granular single-purpose tools into domain-named smart tools. Always prefer the consolidated tools over legacy aliases:
* **`tasks`**: Complete task lifecycle management.
- `action: "add"`: Create a new task (requires `title`, optional `description`, `git_branch`, `parent_id`, `dependencies`).
- `action: "add"`: Create a new task (requires `title`, optional `description`, `git_branch`, `repo_name`, `priority: "low" | "medium" | "high" | "urgent"`, `assigned_agent`, `verification_command`, `parent_id`, `dependencies`).
- `action: "update"`: Update task status (requires `id`, `status: "pending" | "completed" | "cancelled"`).
- `action: "delete"`: Delete task and child tasks (requires `id`).
- `action: "list"`: List active tasks (optional `git_branch`, `summary_level: "compact" | "detailed" | "full"`, `max_tokens`).
- `action: "set_criteria"`: Set acceptance criteria (requires `id`, `criteria: Vec<String>`).
- `action: "verify"`: Verify criteria met (requires `id`).
- `action: "verify"`: Verify criteria met (requires `id`, optional `proof`).
* **`milestones`**: Milestone tracking.
- `action: "add"`: Create milestone (requires `title`).
- `action: "update"`: Update milestone status (requires `id`, `status`).
- `action: "list"`: List milestones.
- `action: "add"`: Create milestone (requires `title`, optional `namespace`, `target_date`, `description`, `deliverables: Vec<String>`, `repo_name`).
- `action: "update"`: Update milestone status (requires `id`, `status: "active" | "completed" | "cancelled"`).
- `action: "list"`: List milestones (optional `namespace`).
* **`sticky_notes`**: Ephemeral scratchpad notes with TTL.
- `action: "add"`: Add note (requires `content`, optional `ttl_seconds`, `session_only`).
@@ -43,48 +43,48 @@ The server consolidates granular single-purpose tools into domain-named smart to
- `action: "clear"`: Clear all sticky notes.
* **`handoff_memos`**: Session handoff notes for future agents.
- `action: "leave"`: Leave a memo (requires `content`).
- `action: "leave"`: Leave a memo (requires `content`, optional `vcs_revision`, `repo_name`, `git_branch`, `blockers: Vec<String>`, `action_items: Vec<String>`, `expires_at`).
- `action: "read"`: Read active handoff memos.
- `action: "clear"`: Clear memos.
* **`pinned_files`**: Focus file working set.
- `action: "pin"`: Pin file to focus set (requires `path`).
- `action: "pin"`: Pin file to focus set (requires `path`, optional `pin_reason`, `symbol_or_line`, `repo_name`, `git_branch`).
- `action: "unpin"`: Unpin file from focus set (requires `path`).
- `action: "list"`: List pinned files.
- `action: "list"`: List pinned files (optional `namespace`).
* **`context_workspaces`**: Workspace context state snapshots.
- `action: "save"`: Save context workspace (requires `name`).
- `action: "save"`: Save context workspace (requires `name`, optional `pinned_files`, `active_task_ids`, `description`, `git_branch`, `vcs_revision`, `repo_name`).
- `action: "load"`: Restore saved context workspace (requires `name`).
- `action: "list"`: List saved context workspaces.
- `action: "delete"`: Delete saved context workspace (requires `name`).
- `action: "diff"`: Compare two saved context workspaces (requires `name`, `other_name`).
* **`pr_checklist`**: Pre-commit and PR checklist.
- `action: "add"`: Add checklist item (requires `description`).
- `action: "add"`: Add checklist item (requires `description`, optional `category: "Security" | "Testing" | "Formatting" | ...`, `automated_check`, `is_checked`, `repo_name`).
- `action: "get"`: Get PR checklist items.
- `action: "clear"`: Clear PR checklist.
* **`snippets`**: Reusable code snippet vault.
- `action: "store"`: Store snippet (requires `query` as name, optional `language`, `code`, `description`, `tags`).
- `action: "store"`: Store snippet (requires `query` as name, optional `language`, `code`, `description`, `tags`, `origin_file`, `line_range`, `repo_name`).
- `action: "search"`: Search snippet vault (optional `query`, `tags`, `hybrid: true`).
- `action: "delete"`: Delete snippet (requires `id`).
- `action: "tag"`: Attach classification tags (requires `id`, `tags: Vec<String>`).
* **`decisions`**: Architectural Decision Records (ADRs).
- `action: "log"`: Log ADR (requires `title`, optional `status`, `context`, `decision`, `consequences`).
- `action: "log"`: Log ADR (requires `title`, optional `status: "accepted" | "proposed" | "deprecated" | "superseded"`, `context`, `decision`, `consequence`, `author`, `affected_components: Vec<String>`, `alternatives_considered: Vec<String>`, `supersedes`, `repo_name`).
- `action: "query"`: Query ADRs (optional `query`).
- `action: "delete"`: Delete ADR (requires `id`).
* **`tech_debt`**: Engineering debt backlog.
- `action: "log"`: Log debt item (requires `description`, optional `ideal_solution`, `git_commit`, `git_branch`, `symbol_references`, `line_range`).
- `action: "log"`: Log debt item (requires `description`, optional `ideal_solution`, `file_path`, `line_range`, `workaround`, `effort_estimate`, `severity: "low" | "medium" | "high" | "critical"`, `git_commit`, `git_branch`, `symbol_references`, `repo_name`).
- `action: "resolve"`: Resolve debt item (requires `id`).
- `action: "list"`: List debt items (optional `include_resolved`).
* **`environment`**: Infrastructure and requirements tracking.
- `action: "update_fingerprint"`: Update tool versions.
- `action: "update_fingerprint"`: Update tool versions (requires `tool_versions: Map<String, String>`, optional `repo_name`).
- `action: "read_fingerprint"`: Read tool versions fingerprint.
- `action: "log_requirement"`: Log environment variable requirement (requires `key`).
- `action: "register"`: Register target environment (requires `name`).
- `action: "log_requirement"`: Log environment variable requirement (requires `key`, `description`, `is_secret`, optional `default_value`, `validation_regex`, `repo_name`).
- `action: "register"`: Register target environment (requires `name`, `url`, optional `description`, `requires_vpn`, `env_type: "dev" | "staging" | "qa" | "prod"`, `healthcheck_endpoint`, `ssh_host`, `repo_name`).
- `action: "get_details"`: Read full environment details.
* **`clipboard`**: OS Clipboard management.
@@ -94,6 +94,26 @@ The server consolidates granular single-purpose tools into domain-named smart to
---
## 3. VCS & SVN Agnosticism & Multi-Repo Provenance
To support diverse enterprise repositories (Git, Subversion / SVN, Mercurial / Hg, Monorepos):
* **`vcs_type`**: Designates the VCS engine (`"git"`, `"svn"`, `"hg"`, `"perforce"`, or `"none"`).
* **`vcs_revision`**: Agnostic commit hash or SVN revision identifier (e.g., `"r12458"`, `"3e4f7a9"`).
* **`upstream_url`**: Canonical remote repository URL (e.g. `https://svn.corp/repo/trunk`, `git@bitbucket.org:org/repo.git`).
* **`repo_name`**: Logical project or repository identifier allowing multiple repositories to share or partition memory namespaces cleanly without collision.
* **Audit Ledger (`log_code_change`)**: Enriched with `vcs_type`, `vcs_revision`, `upstream_url`, `author`, `diff_summary`, and extensible `metadata: HashMap<String, String>`.
---
## 4. Terminal & Process Telemetry
The server ingests and tracks active terminal commands and sessions:
* **Active Terminals**: Tracks PIDs, shell interpreters (`pwsh`, `bash`, `nu`, `zsh`), current working directories (`cwd`), command exit codes, and timestamps.
* **Terminal History Endpoint**: `/terminal/history` exposes recent shell commands and output streams to dashboard and LLMs to prevent lost shell context.
* **Quality Gate Enforcement**: `GateRecord` captures pre-flight and pre-push validation passes with `gate_type`, `enforcer`, `status`, `validation_log`, and `repo_name`.
---
## 4. High-Signal Tool Responses & Performance Guidelines
To optimize context usage, response times, and LLM reasoning efficiency:
@@ -157,8 +177,16 @@ To maintain maximum security, speed, and cross-platform reliability:
---
## 10. High-Performance Concurrency & Resilience Guarantees
* **Explicit Fail-Fast Persistence Safety**: Replaced silent fallback to temporary databases (`/tmp/mcp_store_fallback_*`) with an explicit open retry and fail-fast panic unless `MCP_ALLOW_TMP_FALLBACK=1` is explicitly set, preventing silent data loss.
* **Async Mutex Deadlock Elimination**: Converted shared state and Neovim connection locks (`shutdown_tx`, `NVIM_CONN`, `ACTIVE_SOCKET`, `HEADLESS_PROC`) to `tokio::sync::Mutex` to prevent worker thread pool starvation across `.await` points.
* **Telemetry Session Deduplication & Channel Pruning**: Added `LAST_SESSION` in-memory state deduplication for UDP telemetry writes (eliminating disk I/O thrashing) and distinguished WebSocket `TrySendError::Full` backpressure vs `TrySendError::Closed` client pruning.
* **Graph Adjacency Indexing**: Leverages `KnowledgeGraph::build_adjacency_map` to build $O(1)$ lookup adjacency lists for fast BFS shortest path graph queries.
* **Atomic Store Write Lock Minimization**: `Store::modify` and `Store::modify_async` release write guards immediately after in-memory state mutations, serializing JSON payloads under read locks to allow non-blocking concurrent readers.
* **Async Commit Index Reader Auto-Reload**: `MemoryIndex::commit()` automatically triggers `reader.reload()` upon completing background commits, guaranteeing immediate visibility of newly indexed document terms.
* **Zero-Allocation NDJSON Reader**: `mcp-stdio` reclaims line buffers via `std::mem::take` and in-place trimming to eliminate heap allocations during high-frequency NDJSON message parsing.
* **Micro-Batched Fastembed Inference**: `generate_embeddings_async` chunks text arrays into 32-item micro-batches inside `spawn_blocking`, eliminating RAM/CPU spikes during batch indexing.
* **Dynamic Character Micro-Batched Fastembed Inference**: `generate_embeddings_async` dynamically batches text arrays using a 16,000 character budget ceiling inside `spawn_blocking`, preventing heap spikes during vector indexing while keeping ONNX SIMD execution saturated.
* **Bounded Telemetry Detail Records**: Activity and terminal telemetry buffers enforce a 4,000 character truncation ceiling on log details (`ActivityRecord`, `TerminalHistory`) to prevent unbounded RAM growth under high RPC throughput.
* **Zero-Allocation Stream Formatting**: Graph condensation loops (`condense_graph_worker`) format node/relation subgraphs into stream buffers via `std::fmt::Write`, eliminating intermediate String allocation overhead.
* **Single-Handle OS Clipboard Retries & Image Downsampling**: `WriteClipboardHandler` initializes OS clipboard handles once per operation and downsamples images exceeding $2048 \times 2048$ resolution before writing raw RGBA bytes.
* **Zero Transaction Drop Persistence Guarantee**: `Store::modify` automatically spawns an async task to execute `push_async` with channel backpressure if `push` encounters queue saturation, ensuring zero data loss under spike write loads.
* **Non-Blocking Read Lock Sticky Notes**: `StickyNotesHandler` (`action = "read"`) queries notes using shared read locks, executing write pruning only when expired items exist.
@@ -183,3 +211,16 @@ To maintain maximum security, speed, and cross-platform reliability:
* **Directory Tree Depth Safeguard**: `ReadDirectoryArchitectureHandler` caps directory recursion at depth 10 to prevent stack overflow on deep or cyclic directory structures.
* **Deterministic Total-Order Score Ranking**: `OmniSearchHandler` uses `f64::total_cmp` for Reciprocal Rank Fusion (RRF) score sorting, guaranteeing deterministic NaN-safe search result ordering.
* **RPC Timeout Memory Hygiene**: `nvim-core` maintains request hygiene by removing pending request entries from static RPC maps upon timeout or channel drop, eliminating orphan memory leaks.
* **Embedding Input Safeguard**: `generate_embedding_async` returns explicit errors for empty/0-length text inputs instead of returning empty vectors, preventing downstream vector dimension mismatches during cosine similarity calculations.
* **Path Traversal Security Guards**: `validate_safe_path` enforces path canonicalization and rejects relative parent traversal components (`..`) across file and process log handlers (`GetRecentLogsTool`, `WatchProcessLogsTool`).
* **Watcher Map Memory Eviction**: Proactive daemon file watcher in `watcher.rs` caps `last_processed` map size at 1,000 entries and purges entries older than 10 minutes to prevent monotonic memory leakage.
* **Comprehensive Serde Casing Aliases**: All 11 consolidated tool action enums (`TaskAction`, `MilestoneAction`, `PinnedFileAction`, `ContextWorkspaceAction`, `PrChecklistAction`, `SnippetAction`, `DecisionAction`, `TechDebtAction`, `EnvAction`, `ClipboardAction`, `HandoffMemoAction`) include serde alias attributes supporting `snake_case`, `camelCase`, `PascalCase`, and uppercase variants for maximum LLM casing resilience.
* **Two-Phase Graph Condensation**: `condense_graph_worker` uses a 2-phase commit (non-destructive `read_with` -> graph insert -> prune by timestamp/content) to prevent data loss if summarization or graph insertion fails.
* **Store Write Lock Minimization**: `Store::modify` and `Store::modify_async` unblock concurrent readers during JSON serialization by releasing the write lock immediately after mutating memory state.
* **Redb Database Lock Retry Backoff**: `init_db` retries transient Redb lock contention with exponential backoff (3 attempts, 150ms delay) before falling back.
* **Offloaded Background Index Rebuilds**: `MemoryState::rebuild_index` offloads graph snapshot cloning and Tantivy document re-indexing into `tokio::task::spawn_blocking` to avoid stalling async event loops.
* **Broadcast Watch-Based Shutdown Channels**: Background workers utilize `tokio::sync::watch` for broadcast shutdown notifications without consuming cancellation signals.
* **Consolidated Neovim Tool Suite (v2)**: The Neovim server exposes 7 consolidated domain tools (`nvim_buffer`, `nvim_window`, `nvim_view`, `nvim_diagnostics`, `nvim_visual`, `nvim_execute_lua`, `nvim_system`) with comprehensive action dispatching.
* **Fallback Vector Search Parity**: In-memory vector search fallback indexes Knowledge Graph entities, observations, and error fixes when external vector databases are unavailable.
server/src/lib.rs
+32 -21
View File
@@ -4,7 +4,6 @@
)]
pub mod api;
pub mod clipboard_watcher;
pub mod db;
pub mod embedding;
pub mod error;
@@ -18,7 +17,6 @@ pub mod search;
pub mod state;
pub mod store;
pub mod tools;
pub mod vector_db;
pub mod watcher;
use crate::api::rest::GateSetReq;
@@ -164,6 +162,7 @@ pub async fn condense_graph_worker(state: Arc<MemoryState>) {
observations: vec![condensed_sticky_content],
namespace: crate::models::default_namespace(),
git_branch: None,
..Default::default()
},
);
});
@@ -196,6 +195,7 @@ pub async fn condense_graph_worker(state: Arc<MemoryState>) {
observations: vec![condensed_snippet_content],
namespace: crate::models::default_namespace(),
git_branch: None,
..Default::default()
},
);
});
@@ -216,7 +216,6 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
tokio::spawn(index_committer_worker(Arc::clone(&state)));
tokio::spawn(ttl_sweeper_worker(Arc::clone(&state)));
tokio::spawn(condense_graph_worker(Arc::clone(&state)));
crate::clipboard_watcher::spawn_watcher(Arc::clone(&state));
crate::watcher::spawn_watcher(Arc::clone(&state));
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
let app_state = Arc::new(AppState {
@@ -315,7 +314,8 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
if let Some(home) = dirs::home_dir() {
let nvim_dir = home.join(".gemini");
let nvim_txt_path = nvim_dir.join("active_nvim.txt");
let tmp_path = nvim_dir.join(format!("active_nvim_{}.tmp", std::process::id()));
let tmp_path =
nvim_dir.join(format!("active_nvim_{}.tmp", std::process::id()));
if tokio::fs::create_dir_all(&nvim_dir).await.is_ok()
&& tokio::fs::write(&tmp_path, session).await.is_ok()
{
@@ -561,16 +561,6 @@ pub fn run_cli() -> Result<(), Box<dyn std::error::Error>> {
rt.block_on(async {
let state = Arc::new(MemoryState::new(&base.to_string_lossy()));
match crate::vector_db::VectorDB::new("http://localhost:6334", "mcp_memory").await {
Ok(vdb) => {
tracing::info!("Successfully connected to Qdrant vector database");
*state.vector_db.write().await = Some(vdb);
}
Err(e) => {
tracing::warn!("Failed to initialize Qdrant vector database: {}. Vector search will fallback to manual embedding loop. (Is Qdrant running on localhost:6334?)", e);
}
}
if let Err(e) = run_server(state).await {
tracing::error!("Server error: {}", e);
}
@@ -594,7 +584,13 @@ mod tests {
#[test]
fn test_cli_parsing_exit_and_target() {
let cli = Cli::try_parse_from(&["mcp-memory-server", "--exit", "--target", "http://localhost:3000"]).unwrap();
let cli = Cli::try_parse_from(&[
"mcp-memory-server",
"--exit",
"--target",
"http://localhost:3000",
])
.unwrap();
assert!(cli.exit);
assert_eq!(cli.target.as_deref(), Some("http://localhost:3000"));
}
@@ -615,7 +611,17 @@ mod tests {
])
.unwrap();
if let Some(Commands::Gate { subcmd: GateCommands::Set { action, target, authorize, reason, .. } }) = cli.command {
if let Some(Commands::Gate {
subcmd:
GateCommands::Set {
action,
target,
authorize,
reason,
..
},
}) = cli.command
{
assert_eq!(action, "git_push");
assert_eq!(target, "master");
assert!(authorize);
@@ -639,7 +645,16 @@ mod tests {
])
.unwrap();
if let Some(Commands::Gate { subcmd: GateCommands::Verify { action, target, consume, .. } }) = cli.command {
if let Some(Commands::Gate {
subcmd:
GateCommands::Verify {
action,
target,
consume,
..
},
}) = cli.command
{
assert_eq!(action, "deploy");
assert_eq!(target, "prod");
assert!(consume);
@@ -709,7 +724,3 @@ mod tests {
server_handle.abort();
}
}
server/src/models.rs
+382 -7
View File
@@ -2,24 +2,109 @@ use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ChangeKind {
Added,
Modified,
Deleted,
Renamed,
}
impl Default for ChangeKind {
fn default() -> Self {
ChangeKind::Modified
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct CodeChange {
pub timestamp: u64,
pub file_path: String,
pub description: String,
#[serde(default)]
pub git_commit: Option<String>,
#[serde(default)]
pub git_branch: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub repo_url: Option<String>,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub change_kind: ChangeKind,
#[serde(default)]
pub symbols: Vec<String>,
#[serde(default)]
pub line_range: Option<String>,
#[serde(default)]
pub author: Option<String>,
#[serde(default)]
pub session_id: Option<String>,
#[serde(default)]
pub vcs_type: Option<String>,
#[serde(default)]
pub revision: Option<String>,
#[serde(default)]
pub branch: Option<String>,
#[serde(default)]
pub repository_root: Option<String>,
}
impl CodeChange {
pub fn effective_vcs(&self) -> &str {
if let Some(vcs) = &self.vcs_type {
return vcs.as_str();
}
if let Some(rev) = self.revision.as_deref().or(self.git_commit.as_deref()) {
if rev.starts_with('r') && rev[1..].chars().all(|c| c.is_ascii_digit()) {
return "svn";
}
}
if let Some(br) = self.branch.as_deref().or(self.git_branch.as_deref()) {
if br.eq_ignore_ascii_case("trunk")
|| br.starts_with("branches/")
|| br.starts_with("tags/")
{
return "svn";
}
}
"git"
}
pub fn effective_revision(&self) -> Option<&str> {
self.revision.as_deref().or(self.git_commit.as_deref())
}
pub fn effective_branch(&self) -> Option<&str> {
self.branch.as_deref().or(self.git_branch.as_deref())
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct StickyNote {
#[serde(default = "generate_uuid")]
pub id: String,
pub timestamp: u64,
pub content: String,
#[serde(default)]
pub expires_at: Option<u64>,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub category: Option<String>,
#[serde(default)]
pub author: Option<String>,
}
pub fn generate_uuid() -> String {
uuid::Uuid::new_v4().to_string()
}
pub fn default_namespace() -> String {
"global".to_string()
}
pub fn normalize_entity_type(input: &str) -> String {
let trimmed = input.trim();
if trimmed.is_empty() {
@@ -47,7 +132,9 @@ pub fn normalize_entity_type(input: &str) -> String {
"McpTool" | "Mcptool" | "Mcp" => "McpTool".to_string(),
"Databasetable" | "DatabaseTable" => "DatabaseTable".to_string(),
"Datastructure" | "DataStore" | "Datastore" => "DataStructure".to_string(),
"Architecturecomponent" | "ArchitectureComponent" | "Architecture" => "ArchitectureComponent".to_string(),
"Architecturecomponent" | "ArchitectureComponent" | "Architecture" => {
"ArchitectureComponent".to_string()
}
"Frontendcomponent" | "FrontendComponent" => "FrontendComponent".to_string(),
"Globalrule" | "GlobalRule" => "GlobalRule".to_string(),
"Securitypolicy" | "SecurityPolicy" => "SecurityPolicy".to_string(),
@@ -84,7 +171,10 @@ pub enum TaskStatus {
impl TaskStatus {
pub fn is_active(&self) -> bool {
matches!(self, TaskStatus::Active | TaskStatus::InProgress | TaskStatus::Blocked)
matches!(
self,
TaskStatus::Active | TaskStatus::InProgress | TaskStatus::Blocked
)
}
pub fn is_completed(&self) -> bool {
@@ -107,7 +197,7 @@ impl std::fmt::Display for TaskStatus {
}
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Default)]
pub struct Entity {
pub name: String,
#[serde(alias = "entityType", alias = "type", alias = "category")]
@@ -118,8 +208,19 @@ pub struct Entity {
pub namespace: String,
#[serde(default)]
pub git_branch: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub file_path: Option<String>,
#[serde(default)]
pub created_at: Option<u64>,
#[serde(default)]
pub updated_at: Option<u64>,
#[serde(default)]
pub embedding: Option<Vec<f32>>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, JsonSchema)]
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, Default)]
pub struct Relation {
#[serde(alias = "source", alias = "from_node", alias = "source_node")]
pub from: String,
@@ -129,7 +230,36 @@ pub struct Relation {
pub relation_type: String,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub weight: Option<f32>,
#[serde(default)]
pub is_bidirectional: bool,
}
impl PartialEq for Relation {
fn eq(&self, other: &Self) -> bool {
self.from == other.from
&& self.to == other.to
&& self.relation_type == other.relation_type
&& self.namespace == other.namespace
&& self.repo_name == other.repo_name
}
}
impl Eq for Relation {}
impl std::hash::Hash for Relation {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.from.hash(state);
self.to.hash(state);
self.relation_type.hash(state);
self.namespace.hash(state);
self.repo_name.hash(state);
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct KnowledgeGraph {
#[serde(default)]
@@ -137,11 +267,49 @@ pub struct KnowledgeGraph {
#[serde(default)]
pub relations: Vec<Relation>,
}
impl KnowledgeGraph {
/// Builds an adjacency index mapping node names to lists of outgoing/incoming neighbors and relation types.
pub fn build_adjacency_map<'a>(
&'a self,
relation_filter: Option<&'a str>,
) -> HashMap<&'a str, Vec<(&'a str, &'a str, bool)>> {
let mut adj: HashMap<&'a str, Vec<(&'a str, &'a str, bool)>> =
HashMap::with_capacity(self.relations.len() * 2);
for rel in &self.relations {
if let Some(rf) = relation_filter {
if rel.relation_type != rf {
continue;
}
}
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,
));
}
adj
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct AcceptanceCriteria {
pub id: String,
pub description: String,
pub is_met: bool,
#[serde(default)]
pub verification_command: Option<String>,
#[serde(default)]
pub verification_output: Option<String>,
#[serde(default)]
pub verified_at: Option<u64>,
#[serde(default)]
pub verified_by: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
@@ -161,6 +329,16 @@ pub struct Task {
pub dependencies: Vec<String>,
#[serde(default)]
pub acceptance_criteria: Vec<AcceptanceCriteria>,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub priority: Option<String>,
#[serde(default)]
pub assigned_agent: Option<String>,
#[serde(default)]
pub verification_command: Option<String>,
}
impl Task {
@@ -172,6 +350,7 @@ impl Task {
&& !s.eq_ignore_ascii_case("canceled")
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Snippet {
pub name: String,
@@ -183,7 +362,16 @@ pub struct Snippet {
pub tags: Vec<String>,
#[serde(default)]
pub embedding: Option<Vec<f32>>,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub origin_file: Option<String>,
#[serde(default)]
pub line_range: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Adr {
pub id: String,
@@ -196,17 +384,35 @@ pub struct Adr {
#[serde(default)]
pub supersedes: Option<String>,
pub timestamp: u64,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub alternatives_considered: Vec<String>,
#[serde(default)]
pub affected_components: Vec<String>,
#[serde(default)]
pub author: Option<String>,
}
fn default_adr_status() -> String {
"accepted".to_string()
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Preference {
pub key: String,
pub value: String,
pub updated_at: u64,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub category: Option<String>,
#[serde(default)]
pub rationale: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ErrorFix {
pub signature: String,
@@ -216,14 +422,32 @@ pub struct ErrorFix {
pub git_branch: Option<String>,
#[serde(default)]
pub embedding: Option<Vec<f32>>,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub error_category: Option<String>,
#[serde(default)]
pub stack_trace: Option<String>,
#[serde(default)]
pub toolchain: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PinnedFile {
pub namespace: String,
pub file_path: String,
pub timestamp: u64,
pub git_branch: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub symbol_or_line: Option<String>,
#[serde(default)]
pub pin_reason: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct SessionSummary {
pub timestamp: u64,
@@ -231,7 +455,20 @@ pub struct SessionSummary {
pub namespace: String,
#[serde(default)]
pub expires_at: Option<u64>,
#[serde(default)]
pub session_id: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub git_branch: Option<String>,
#[serde(default)]
pub git_commit: Option<String>,
#[serde(default)]
pub tasks_completed: Vec<String>,
#[serde(default)]
pub next_steps: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct HandoffMemo {
pub id: String,
@@ -241,7 +478,20 @@ pub struct HandoffMemo {
pub expires_at: Option<u64>,
pub namespace: String,
pub timestamp: u64,
#[serde(default)]
pub session_id: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub git_branch: Option<String>,
#[serde(default)]
pub vcs_revision: Option<String>,
#[serde(default)]
pub blockers: Vec<String>,
#[serde(default)]
pub action_items: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct EnvFingerprint {
pub namespace: String,
@@ -249,14 +499,24 @@ pub struct EnvFingerprint {
pub shell: String,
pub tool_versions: HashMap<String, String>,
pub updated_at: u64,
#[serde(default)]
pub repo_name: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct EnvRequirement {
pub namespace: String,
pub key: String,
pub description: String,
pub is_secret: bool,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub default_value: Option<String>,
#[serde(default)]
pub validation_regex: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Milestone {
pub id: String,
@@ -264,7 +524,18 @@ pub struct Milestone {
pub status: String,
pub namespace: String,
pub target_date: Option<u64>,
#[serde(default)]
pub description: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub task_ids: Vec<String>,
#[serde(default)]
pub deliverables: Vec<String>,
#[serde(default)]
pub completed_at: Option<u64>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct EnvironmentDetail {
pub namespace: String,
@@ -273,13 +544,31 @@ pub struct EnvironmentDetail {
pub description: String,
pub requires_vpn: bool,
pub updated_at: u64,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub env_type: Option<String>,
#[serde(default)]
pub healthcheck_endpoint: Option<String>,
#[serde(default)]
pub ssh_host: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct PrChecklistItem {
pub namespace: String,
pub id: String,
pub description: String,
#[serde(default)]
pub is_checked: bool,
#[serde(default)]
pub category: Option<String>,
#[serde(default)]
pub automated_check: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TechDebt {
pub id: String,
@@ -292,7 +581,20 @@ pub struct TechDebt {
pub git_branch: Option<String>,
#[serde(default)]
pub embedding: Option<Vec<f32>>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub severity: Option<String>,
#[serde(default)]
pub file_path: Option<String>,
#[serde(default)]
pub line_range: Option<String>,
#[serde(default)]
pub workaround: Option<String>,
#[serde(default)]
pub effort_estimate: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ContextWorkspace {
pub name: String,
@@ -300,6 +602,14 @@ pub struct ContextWorkspace {
pub pinned_files: Vec<String>,
pub active_task_ids: Vec<String>,
pub saved_at: u64,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub git_branch: Option<String>,
#[serde(default)]
pub vcs_revision: Option<String>,
#[serde(default)]
pub description: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
@@ -312,15 +622,44 @@ pub struct GateRecord {
pub status: String,
pub reason: Option<String>,
pub timestamp: u64,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub gate_type: Option<String>,
#[serde(default)]
pub enforcer: Option<String>,
#[serde(default)]
pub validation_log: Option<String>,
#[serde(default)]
pub expires_at: Option<u64>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TerminalHistory {
#[serde(default = "generate_uuid")]
pub id: String,
pub command: String,
pub exit_code: i32,
pub cwd: String,
#[serde(default)]
pub os: String,
pub timestamp: u64,
#[serde(default)]
pub error_output: Option<String>,
#[serde(default)]
pub shell: Option<String>,
#[serde(default)]
pub duration_ms: Option<u64>,
#[serde(default)]
pub stdout_summary: Option<String>,
#[serde(default)]
pub status_reason: Option<String>,
#[serde(default)]
pub invoker: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default = "default_namespace")]
pub namespace: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
@@ -329,6 +668,16 @@ pub struct ActivityRecord {
pub category: String, // e.g. "code_change", "task_update", "telemetry", "error_fix"
pub summary: String, // Brief human-readable description of what is being worked on
pub details: Option<String>,
#[serde(default)]
pub session_id: Option<String>,
#[serde(default)]
pub actor: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub associated_file: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
@@ -337,6 +686,14 @@ pub struct StateSnapshot {
pub timestamp: u64,
pub description: String,
pub namespace: String,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub git_branch: Option<String>,
#[serde(default)]
pub vcs_revision: Option<String>,
#[serde(default)]
pub active_tasks_count: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
@@ -347,6 +704,14 @@ pub struct Hypothesis {
pub status: String,
pub evidence: Option<String>,
pub timestamp: u64,
#[serde(default = "default_namespace")]
pub namespace: String,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub validation_steps: Vec<String>,
#[serde(default)]
pub diagnostician: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
@@ -357,6 +722,12 @@ pub struct AgentSignal {
pub payload: String,
pub timestamp: u64,
pub ttl_seconds: Option<u64>,
#[serde(default)]
pub recipient: Option<String>,
#[serde(default)]
pub correlation_id: Option<String>,
#[serde(default)]
pub priority: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
@@ -368,13 +739,18 @@ pub struct TaskEvent {
pub error: Option<String>,
pub timestamp: u64,
pub session_id: Option<String>,
#[serde(default)]
pub namespace: Option<String>,
#[serde(default)]
pub repo_name: Option<String>,
#[serde(default)]
pub actor: Option<String>,
}
#[cfg(test)]
mod tests {
use crate::error::AppError;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use crate::error::AppError;
#[test]
fn test_app_error_bad_request_invalid_json() {
@@ -383,4 +759,3 @@ mod tests {
assert_eq!(res.status(), StatusCode::BAD_REQUEST);
}
}
server/src/router.rs
+240 -47
View File
@@ -113,11 +113,13 @@ impl McpResource for TasksActiveResource {
Some("List of currently active tasks")
}
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let tasks = state.project.tasks.cache.read().unwrap_or_else(|e| e.into_inner());
let data: Vec<_> = tasks
.iter()
.filter(|t| t.is_active())
.collect();
let tasks = state
.project
.tasks
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
let data: Vec<_> = tasks.iter().filter(|t| t.is_active()).collect();
Ok(serde_json::to_string_pretty(&data)?)
}
}
@@ -135,7 +137,12 @@ impl McpResource for DecisionsActiveResource {
Some("All accepted Architectural Decision Records (ADRs)")
}
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let adrs = state.code.adrs.cache.read().unwrap_or_else(|e| e.into_inner());
let adrs = state
.code
.adrs
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
let data: Vec<_> = adrs.iter().filter(|a| a.status == "accepted").collect();
Ok(serde_json::to_string_pretty(&data)?)
}
@@ -154,7 +161,12 @@ impl McpResource for TechDebtActiveResource {
Some("All currently unresolved technical debt items")
}
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let debts = state.code.tech_debts.cache.read().unwrap_or_else(|e| e.into_inner());
let debts = state
.code
.tech_debts
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
let data: Vec<_> = debts.iter().filter(|d| !d.is_resolved).collect();
Ok(serde_json::to_string_pretty(&data)?)
}
@@ -233,14 +245,32 @@ impl McpResource for SessionDeltaResource {
let lookback = 7200; // last 2 hours
let cutoff = now.saturating_sub(lookback);
let changes = state.code.ledger.cache.read().unwrap_or_else(|e| e.into_inner());
let changes = state
.code
.ledger
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
let recent_changes: Vec<_> = changes.iter().filter(|c| c.timestamp >= cutoff).collect();
let tasks = state.project.tasks.cache.read().unwrap_or_else(|e| e.into_inner());
let tasks = state
.project
.tasks
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
let active_tasks: Vec<_> = tasks.iter().filter(|t| t.is_active()).collect();
let notes = state.code.sticky.cache.read().unwrap_or_else(|e| e.into_inner());
let active_notes: Vec<_> = notes.iter().filter(|n| n.expires_at.map_or(true, |exp| exp > now)).collect();
let notes = state
.code
.sticky
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
let active_notes: Vec<_> = notes
.iter()
.filter(|n| n.expires_at.map_or(true, |exp| exp > now))
.collect();
let delta = serde_json::json!({
"recent_changes": recent_changes,
@@ -260,7 +290,9 @@ impl McpPrompt for ContextWarmupPrompt {
"context_warmup"
}
fn description(&self) -> Option<&'static str> {
Some("Consolidated context warmup prompt providing active tasks, pinned files, open tech debt, and sticky notes")
Some(
"Consolidated context warmup prompt providing active tasks, pinned files, open tech debt, and sticky notes",
)
}
async fn get(
&self,
@@ -334,7 +366,34 @@ impl MemoryHandler {
Some("Recent terminal execution history and exit codes")
}
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let items = state.telemetry.terminal_history.cache.read().unwrap_or_else(|e| e.into_inner());
let items = state
.telemetry
.terminal_history
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
Ok(serde_json::to_string_pretty(&*items)?)
}
}
struct RecentActivitiesResource;
#[async_trait]
impl McpResource for RecentActivitiesResource {
fn uri(&self) -> &'static str {
"memory://activity/recent"
}
fn name(&self) -> &'static str {
"Recent Activity"
}
fn description(&self) -> Option<&'static str> {
Some("Recent IDE and developer activity logs")
}
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let items = state
.telemetry
.recent_activities
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
Ok(serde_json::to_string_pretty(&*items)?)
}
}
@@ -351,7 +410,12 @@ impl MemoryHandler {
Some("Currently pinned files in the workspace")
}
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let items = state.project.pinned_files.cache.read().unwrap_or_else(|e| e.into_inner());
let items = state
.project
.pinned_files
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
Ok(serde_json::to_string_pretty(&*items)?)
}
}
@@ -369,7 +433,12 @@ impl MemoryHandler {
Some("Project milestones and their status")
}
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let items = state.project.milestones.cache.read().unwrap_or_else(|e| e.into_inner());
let items = state
.project
.milestones
.cache
.read()
.unwrap_or_else(|e| e.into_inner());
Ok(serde_json::to_string_pretty(&*items)?)
}
}
@@ -403,6 +472,7 @@ impl MemoryHandler {
}
register_resource!(TerminalHistoryResource);
register_resource!(RecentActivitiesResource);
register_resource!(PinnedFilesResource);
register_resource!(MilestonesResource);
struct ArchiveRoutinePrompt;
@@ -679,7 +749,26 @@ impl MemoryHandler {
let category = match name {
"clipboard" => "CLIPBOARD",
"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",
"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",
@@ -701,7 +790,7 @@ impl MemoryHandler {
let result: crate::error::Result<String> = if let Some(tool) = self.tools.get(name)
{
tool.execute(args, self.state.clone()).await
execute_tool_span(name, tool.as_ref(), args, self.state.clone()).await
} else {
Err(crate::error::AppError::Internal(format!(
"Unknown tool: {}",
@@ -747,40 +836,90 @@ impl MemoryHandler {
pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) -> String {
let (action, detail) = match name {
"tasks" => {
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("manage");
let title = args.get("title").or_else(|| args.get("id")).and_then(|v| v.as_str()).unwrap_or("");
("Tasks", format!("{}: {}", act, title).trim_end_matches(": ").to_string())
let act = args
.get("action")
.and_then(|v| v.as_str())
.unwrap_or("manage");
let title = args
.get("title")
.or_else(|| args.get("id"))
.and_then(|v| v.as_str())
.unwrap_or("");
(
"Tasks",
format!("{}: {}", act, title)
.trim_end_matches(": ")
.to_string(),
)
}
"decisions" => {
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("log");
let title = args.get("title").or_else(|| args.get("query")).and_then(|v| v.as_str()).unwrap_or("");
("Decisions", format!("{}: {}", act, title).trim_end_matches(": ").to_string())
let title = args
.get("title")
.or_else(|| args.get("query"))
.and_then(|v| v.as_str())
.unwrap_or("");
(
"Decisions",
format!("{}: {}", act, title)
.trim_end_matches(": ")
.to_string(),
)
}
"tech_debt" => {
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("log");
let desc = args.get("description").or_else(|| args.get("id")).and_then(|v| v.as_str()).unwrap_or("");
("Tech Debt", format!("{}: {}", act, desc).trim_end_matches(": ").to_string())
let desc = args
.get("description")
.or_else(|| args.get("id"))
.and_then(|v| v.as_str())
.unwrap_or("");
(
"Tech Debt",
format!("{}: {}", act, desc)
.trim_end_matches(": ")
.to_string(),
)
}
"sticky_notes" => {
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("add");
let preview = args.get("content").and_then(|v| v.as_str()).map(|c| c.chars().take(40).collect::<String>()).unwrap_or_default();
("Sticky Notes", format!("{}: {}", act, preview).trim_end_matches(": ").to_string())
let preview = args
.get("content")
.and_then(|v| v.as_str())
.map(|c| c.chars().take(40).collect::<String>())
.unwrap_or_default();
(
"Sticky Notes",
format!("{}: {}", act, preview)
.trim_end_matches(": ")
.to_string(),
)
}
"clipboard" => {
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("read");
let act = args
.get("action")
.and_then(|v| v.as_str())
.unwrap_or("read");
("Clipboard", act.to_string())
}
"snippets" => {
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("search");
let act = args
.get("action")
.and_then(|v| v.as_str())
.unwrap_or("search");
let q = args.get("query").and_then(|v| v.as_str()).unwrap_or("");
("Snippets", format!("{}: {}", act, q).trim_end_matches(": ").to_string())
(
"Snippets",
format!("{}: {}", act, q).trim_end_matches(": ").to_string(),
)
}
"log_code_change" => {
let file = args.get("file_path")
let file = args
.get("file_path")
.or_else(|| args.get("file"))
.or_else(|| args.get("target_file"))
.and_then(|v| v.as_str());
let summary = args.get("summary")
let summary = args
.get("summary")
.or_else(|| args.get("description"))
.and_then(|v| v.as_str());
@@ -793,7 +932,8 @@ pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) ->
("Log Code Change", d)
}
"log_error_fix" => {
let d = args.get("error")
let d = args
.get("error")
.or_else(|| args.get("summary"))
.or_else(|| args.get("solution"))
.and_then(|v| v.as_str())
@@ -805,7 +945,11 @@ pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) ->
let d = if let Some(entities) = args.get("entities").and_then(|v| v.as_array()) {
let names: Vec<&str> = entities
.iter()
.filter_map(|e| e.get("name").or_else(|| e.get("entity_name")).and_then(|n| n.as_str()))
.filter_map(|e| {
e.get("name")
.or_else(|| e.get("entity_name"))
.and_then(|n| n.as_str())
})
.take(3)
.collect();
if !names.is_empty() {
@@ -817,7 +961,11 @@ pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) ->
} else {
format!("{} entity/entities", entities.len())
}
} else if let Some(name) = args.get("name").or_else(|| args.get("entity_name")).and_then(|v| v.as_str()) {
} else if let Some(name) = args
.get("name")
.or_else(|| args.get("entity_name"))
.and_then(|v| v.as_str())
{
name.to_string()
} else {
String::new()
@@ -829,7 +977,11 @@ pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) ->
if let Some(first) = relations.first() {
let from = first.get("from").and_then(|v| v.as_str()).unwrap_or("?");
let to = first.get("to").and_then(|v| v.as_str()).unwrap_or("?");
let pred = first.get("relation_type").or_else(|| first.get("predicate")).and_then(|v| v.as_str()).unwrap_or("->");
let pred = first
.get("relation_type")
.or_else(|| first.get("predicate"))
.and_then(|v| v.as_str())
.unwrap_or("->");
if relations.len() > 1 {
format!("{} {} {} (+{} more)", from, pred, to, relations.len() - 1)
} else {
@@ -844,7 +996,8 @@ pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) ->
("Create Relations", d)
}
"omni_search" | "search_graph" | "search_web" => {
let d = args.get("query")
let d = args
.get("query")
.and_then(|v| v.as_str())
.map(|q| format!("\"{}\"", q))
.unwrap_or_default();
@@ -881,6 +1034,16 @@ pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) ->
}
}
#[tracing::instrument(name = "mcp_tool_execution", skip(tool, args, state), fields(tool_name = %name))]
async fn execute_tool_span(
name: &str,
tool: &dyn McpTool,
args: Value,
state: Arc<MemoryState>,
) -> crate::error::Result<String> {
tool.execute(args, state).await
}
#[cfg(test)]
mod tests {
use super::*;
@@ -889,10 +1052,16 @@ mod tests {
#[test]
fn test_format_tool_activity_description() {
let msg1 = format_tool_activity_description("log_code_change", &json!({ "file": "src/lib.rs", "summary": "Refactored library" }));
let msg1 = format_tool_activity_description(
"log_code_change",
&json!({ "file": "src/lib.rs", "summary": "Refactored library" }),
);
assert_eq!(msg1, "Log Code Change: src/lib.rs: Refactored library");
let msg2 = format_tool_activity_description("create_entities", &json!({ "entities": [{ "name": "EntityA" }, { "name": "EntityB" }] }));
let msg2 = format_tool_activity_description(
"create_entities",
&json!({ "entities": [{ "name": "EntityA" }, { "name": "EntityB" }] }),
);
assert_eq!(msg2, "Create Entities: EntityA, EntityB");
let msg3 = format_tool_activity_description("read_graph", &json!({}));
@@ -948,6 +1117,11 @@ mod tests {
.iter()
.any(|r| r["uri"] == "memory://pinned_files")
);
assert!(
resources_arr
.iter()
.any(|r| r["uri"] == "memory://activity/recent")
);
// Test resources/read
let req_read_res = json!({
@@ -1114,17 +1288,27 @@ mod tests {
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());
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");
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");
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",
@@ -1154,7 +1338,10 @@ mod tests {
}
});
let res = handler.handle_request(invalid_req).await.expect("handler must respond");
let res = handler
.handle_request(invalid_req)
.await
.expect("handler must respond");
assert_eq!(res["jsonrpc"], "2.0");
assert!(res["result"]["content"].is_array());
}
@@ -1180,7 +1367,10 @@ mod tests {
("sweep_graph_health", json!({})),
("get_next_actionable_tasks", json!({})),
("get_preflight_context", json!({})),
("auto_session_checkpoint", json!({"session_summary": "Test summary"})),
(
"auto_session_checkpoint",
json!({"session_summary": "Test summary"}),
),
];
for (name, args) in sample_valid_calls {
@@ -1194,13 +1384,16 @@ mod tests {
}
});
let res = handler.handle_request(valid_req).await.expect("handler must respond");
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
"Tool '{}' failed valid execution. Result: {:?}",
name, res
);
}
}
}
server/src/search.rs
+77 -25
View File
@@ -27,6 +27,8 @@ pub struct MemoryIndex {
pub body_field: Field,
pub type_field: Field,
pub namespace_field: Field,
pub query_parser: tantivy::query::QueryParser,
}
impl MemoryIndex {
@@ -45,14 +47,14 @@ impl MemoryIndex {
let index = Index::open_in_dir(&index_dir)
.or_else(|_| Index::create_in_dir(&index_dir, schema.clone()))?;
let mut writer = index.writer(50_000_000)?;
writer.delete_all_documents()?;
writer.commit()?;
let writer = index.writer(50_000_000)?;
let reader = index
.reader_builder()
.reload_policy(ReloadPolicy::OnCommitWithDelay)
.try_into()?;
let query_parser = tantivy::query::QueryParser::for_index(&index, vec![title_field, body_field]);
Ok(Self {
index,
reader,
@@ -63,6 +65,7 @@ impl MemoryIndex {
body_field,
type_field,
namespace_field,
query_parser,
})
}
@@ -76,14 +79,14 @@ impl MemoryIndex {
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 writer = index.writer(50_000_000)?;
let reader = index
.reader_builder()
.reload_policy(ReloadPolicy::OnCommitWithDelay)
.try_into()?;
let query_parser = tantivy::query::QueryParser::for_index(&index, vec![title_field, body_field]);
Ok(Self {
index,
reader,
@@ -94,9 +97,20 @@ impl MemoryIndex {
body_field,
type_field,
namespace_field,
query_parser,
})
}
pub fn clear(&self) -> tantivy::Result<()> {
let mut writer = self
.writer
.lock()
.map_err(|e| tantivy::TantivyError::SystemError(e.to_string()))?;
writer.delete_all_documents()?;
writer.commit()?;
Ok(())
}
pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> {
self.index_entities_batch(std::slice::from_ref(e))
}
@@ -138,10 +152,18 @@ impl MemoryIndex {
let id_val = t.id.clone();
let needs_commit = Arc::clone(&self.needs_commit);
let criteria_body = t.acceptance_criteria.iter().fold(String::new(), |mut acc, c| {
if !acc.is_empty() {
acc.push('\n');
}
acc.push_str(&c.description);
acc
});
let doc = doc!(
self.id_field => t.id.as_str(),
self.title_field => t.title.as_str(),
self.body_field => format!("{}\n{}", t.description, t.acceptance_criteria.iter().map(|c| c.description.as_str()).collect::<Vec<_>>().join("\n")),
self.body_field => format!("{}\n{}", t.description, criteria_body),
self.type_field => "task",
self.namespace_field => "global"
);
@@ -172,6 +194,7 @@ impl MemoryIndex {
pub async fn commit(&self) -> tantivy::Result<()> {
let writer = Arc::clone(&self.writer);
let needs_commit = Arc::clone(&self.needs_commit);
let reader = self.reader.clone();
tokio::task::spawn_blocking(move || {
if needs_commit.swap(false, std::sync::atomic::Ordering::SeqCst) {
let mut writer = writer.lock().unwrap_or_else(|e| e.into_inner());
@@ -181,6 +204,7 @@ impl MemoryIndex {
tantivy::TantivyError::SystemError(format!("First commit error: {}; Second commit error: {}", first_err, second_err))
})?;
}
let _ = reader.reload();
}
Ok(())
})
@@ -208,11 +232,9 @@ impl MemoryIndex {
namespace: Option<&str>,
) -> tantivy::Result<Vec<SearchResultTuple>> {
let searcher = self.reader.searcher();
let query_parser = tantivy::query::QueryParser::for_index(
&self.index,
vec![self.title_field, self.body_field],
);
let q = query_parser.parse_query(query).or_else(|_| {
let mut is_wildcard_fallback = false;
let q = self.query_parser.parse_query(query).or_else(|_| {
let sanitized: String = query
.chars()
.map(|c| match c {
@@ -223,17 +245,26 @@ impl MemoryIndex {
.collect();
let safe_query = sanitized.trim();
if safe_query.is_empty() {
query_parser.parse_query("*")
is_wildcard_fallback = true;
self.query_parser.parse_query("*")
} else {
query_parser.parse_query(safe_query)
self.query_parser.parse_query(safe_query)
}
})?;
let target_limit = if is_wildcard_fallback { 10 } else { 50 };
let fetch_limit = if namespace.is_some() {
// Fetch deeper candidate pool so other namespaces do not starve the target namespace
1000
} else {
target_limit
};
let top_docs = searcher.search(
&q,
&tantivy::collector::TopDocs::with_limit(50).order_by_score(),
&tantivy::collector::TopDocs::with_limit(fetch_limit).order_by_score(),
)?;
let mut results = Vec::with_capacity(top_docs.len());
let mut results = Vec::with_capacity(target_limit);
for (score, doc_address) in top_docs {
let retrieved_doc = searcher.doc::<tantivy::TantivyDocument>(doc_address)?;
let id = retrieved_doc
@@ -267,6 +298,9 @@ impl MemoryIndex {
continue;
}
results.push((id, doc_type, title, body, score));
if results.len() >= target_limit {
break;
}
}
Ok(results)
}
@@ -319,59 +353,73 @@ impl MemoryIndex {
pub fn add_entity_sync(&self, e: &Entity) {
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _ = writer.add_document(doc!(
writer.delete_term(tantivy::Term::from_field_text(self.id_field, &e.name));
if let Err(err) = writer.add_document(doc!(
self.id_field => e.name.as_str(),
self.title_field => e.name.as_str(),
self.body_field => e.observations.join(" "),
self.type_field => "entity",
self.namespace_field => e.namespace.as_str()
));
)) {
tracing::warn!("Sync entity indexing failed for '{}': {}", e.name, err);
}
self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
}
pub fn delete_all(&self) {
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _ = writer.delete_all_documents();
if let Err(err) = writer.delete_all_documents() {
tracing::warn!("delete_all_documents failed: {}", err);
}
self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
}
pub fn add_task_sync(&self, t: &Task) {
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _res = writer.add_document(doc!(
writer.delete_term(tantivy::Term::from_field_text(self.id_field, &t.id));
if let Err(err) = writer.add_document(doc!(
self.id_field => t.id.as_str(),
self.title_field => t.title.as_str(),
self.body_field => t.description.as_str(),
self.type_field => "task",
self.namespace_field => "global"
));
)) {
tracing::warn!("Sync task indexing failed for '{}': {}", t.id, err);
}
self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
}
pub fn add_snippet_sync(&self, s: &Snippet) {
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _ = writer.add_document(doc!(
writer.delete_term(tantivy::Term::from_field_text(self.id_field, &s.name));
if let Err(err) = writer.add_document(doc!(
self.id_field => s.name.as_str(),
self.title_field => s.name.as_str(),
self.body_field => format!("{} {}", s.language, s.description),
self.type_field => "snippet",
self.namespace_field => "global"
));
)) {
tracing::warn!("Sync snippet indexing failed for '{}': {}", s.name, err);
}
self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
}
pub fn add_adr_sync(&self, a: &Adr) {
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _ = writer.add_document(doc!(
writer.delete_term(tantivy::Term::from_field_text(self.id_field, &a.id));
if let Err(err) = writer.add_document(doc!(
self.id_field => a.id.as_str(),
self.title_field => a.title.as_str(),
self.body_field => format!("{} {} {}", a.context, a.decision, a.consequence),
self.type_field => "adr",
self.namespace_field => "global"
));
)) {
tracing::warn!("Sync ADR indexing failed for '{}': {}", a.id, err);
}
self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
}
@@ -393,6 +441,7 @@ mod tests {
observations: vec!["This is a test observation".to_string()],
namespace: "global".to_string(),
git_branch: None,
..Default::default()
};
let _ = index.index_entity(&entity).await.unwrap();
@@ -408,6 +457,7 @@ mod tests {
dependencies: vec![],
parent_id: None,
expires_at: None,
..Default::default()
};
let _ = index.index_task(&task).await.unwrap();
@@ -419,6 +469,7 @@ mod tests {
updated_at: 0,
tags: vec![],
embedding: None,
..Default::default()
};
let _ = index.index_snippet(&snippet).await.unwrap();
@@ -431,6 +482,7 @@ mod tests {
status: "accepted".to_string(),
supersedes: None,
timestamp: 0,
..Default::default()
};
let _ = index.index_adr(&adr).await.unwrap();
server/src/state.rs
+189 -125
View File
@@ -1,7 +1,6 @@
use crate::models::*;
use crate::search::MemoryIndex;
use crate::store::Store;
use crate::vector_db::VectorDB;
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
@@ -50,14 +49,12 @@ pub struct TelemetryStores {
pub struct MemoryState {
pub base_dir: PathBuf,
pub clipboard_watch_mode: tokio::sync::RwLock<bool>,
pub clipboard_notify: Arc<tokio::sync::Notify>,
pub index_commit_notify: Arc<tokio::sync::Notify>,
pub ttl_notify: Arc<tokio::sync::Notify>,
pub condense_notify: Arc<tokio::sync::Notify>,
pub shutdown_notify: Arc<tokio::sync::Notify>,
pub graph: Store<KnowledgeGraph>,
pub search_index: tokio::sync::RwLock<MemoryIndex>,
pub vector_db: tokio::sync::RwLock<Option<VectorDB>>,
pub project: ProjectStores,
pub code: CodeStores,
@@ -96,22 +93,21 @@ impl MemoryState {
.join(".gemini/mcp_memory/daemon_error.log");
let _ =
std::fs::write(&log_path, format!("Failed to create MemoryIndex: {}\n", e));
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex")
crate::search::MemoryIndex::new_in_ram()
.expect("Failed to create RAM MemoryIndex")
}
}
};
let state = Self {
ollama: Arc::new(crate::ollama::OllamaClient::new_from_env()),
clipboard_watch_mode: tokio::sync::RwLock::new(false),
clipboard_notify: Arc::new(tokio::sync::Notify::new()),
index_commit_notify: Arc::new(tokio::sync::Notify::new()),
ttl_notify: Arc::new(tokio::sync::Notify::new()),
condense_notify: Arc::new(tokio::sync::Notify::new()),
shutdown_notify: Arc::new(tokio::sync::Notify::new()),
graph: Store::new("knowledge_graph_master", db.clone()),
base_dir: base.clone(),
search_index: tokio::sync::RwLock::new(search_index),
vector_db: tokio::sync::RwLock::new(None),
project: ProjectStores {
tasks: Store::new("tasks", db.clone()),
@@ -155,7 +151,8 @@ impl MemoryState {
entity.entity_type = crate::models::normalize_entity_type(&entity.entity_type);
}
for relation in g.relations.iter_mut() {
relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type);
relation.relation_type =
crate::models::normalize_relation_type(&relation.relation_type);
}
});
@@ -187,50 +184,64 @@ impl MemoryState {
pub async fn rebuild_index(self: &Arc<Self>) {
let is_in_memory = self.base_dir.to_str() == Some(":memory:");
let new_idx = if is_in_memory {
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex for rebuild")
} else {
match crate::search::MemoryIndex::new(&self.base_dir) {
Ok(idx) => idx,
Err(e) => {
tracing::warn!("Failed to create disk MemoryIndex for rebuild ({}), falling back to RAM", e);
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex for rebuild")
let base_dir = self.base_dir.clone();
let state_clone = Arc::clone(self);
// Offload full clone and synchronous Tantivy doc indexing off the async Tokio reactor
let new_idx = match tokio::task::spawn_blocking(move || {
let new_idx = if is_in_memory {
crate::search::MemoryIndex::new_in_ram()
.expect("Failed to create RAM MemoryIndex for rebuild")
} else {
match crate::search::MemoryIndex::new(&base_dir) {
Ok(idx) => idx,
Err(e) => {
tracing::warn!(
"Failed to create disk MemoryIndex for rebuild ({}), falling back to RAM",
e
);
crate::search::MemoryIndex::new_in_ram()
.expect("Failed to create RAM MemoryIndex for rebuild")
}
}
}
};
};
let _ = new_idx.clear();
let entities: Vec<_> = self
.graph
.read_with(|g| g.entities.values().cloned().collect());
let tasks = self.project.tasks.read_with(|t| t.clone());
let snippets = self.code.snippets.read_with(|s| s.clone());
let adrs = self.code.adrs.read_with(|a| a.clone());
let entities: Vec<_> = state_clone
.graph
.read_with(|g| g.entities.values().cloned().collect());
let tasks = state_clone.project.tasks.read_with(|t| t.clone());
let snippets = state_clone.code.snippets.read_with(|s| s.clone());
let adrs = state_clone.code.adrs.read_with(|a| a.clone());
tracing::info!(
"rebuild_index: found {} entities, {} tasks",
entities.len(),
tasks.len()
);
tracing::info!(
"rebuild_index: indexing {} entities, {} tasks synchronously in blocking thread",
entities.len(),
tasks.len()
);
let idx_clone = new_idx.clone();
tokio::task::spawn_blocking(move || {
for e in entities {
idx_clone.add_entity_sync(&e);
new_idx.add_entity_sync(&e);
}
for task in tasks {
idx_clone.add_task_sync(&task);
new_idx.add_task_sync(&task);
}
for snippet in snippets {
idx_clone.add_snippet_sync(&snippet);
new_idx.add_snippet_sync(&snippet);
}
for adr in adrs {
idx_clone.add_adr_sync(&adr);
new_idx.add_adr_sync(&adr);
}
new_idx
})
.await
.unwrap_or_else(|e| {
tracing::error!("Failed to join tantivy index rebuild thread: {}", e);
});
{
Ok(idx) => idx,
Err(e) => {
tracing::error!("Failed to join tantivy index rebuild thread: {}", e);
return;
}
};
let _ = new_idx.commit().await;
*self.search_index.write().await = new_idx;
@@ -244,22 +255,30 @@ impl MemoryState {
.as_millis() as u64;
let category_upper = category.to_uppercase();
let truncated_details = details.map(|s| {
if s.len() > 4096 {
format!("{}... [truncated]", &s[..4096])
} else {
s.to_string()
}
});
let activity = ActivityRecord {
timestamp: ts,
category: category_upper,
summary: summary.to_string(),
details: details.map(|s| s.to_string()),
details: truncated_details,
..Default::default()
};
let record_val = serde_json::to_value(&activity).unwrap_or_default();
self.telemetry.recent_activities.modify(|activities| {
activities.push_front(record_val.clone());
if activities.len() > 100 {
activities.pop_back();
}
});
if let Ok(record_val) = serde_json::to_value(&activity) {
self.telemetry.recent_activities.modify(|activities| {
activities.push_front(record_val);
if activities.len() > 100 {
activities.pop_back();
}
});
}
if self.activity_tx.receiver_count() > 0 {
let payload = serde_json::json!({
@@ -277,16 +296,26 @@ impl MemoryState {
let payload_val = serde_json::to_value(&event).unwrap_or_default();
let summary_str = format!("Task {} -> {}", event.task_id, event.status);
let details_str = payload_val.to_string();
let truncated_details = if details_str.len() > 4096 {
format!("{}... [truncated]", &details_str[..4096])
} else {
details_str
};
self.telemetry.recent_activities.modify(|activities| {
let activity = ActivityRecord {
timestamp: event.timestamp,
category: "TASK_EVENT".to_string(),
summary: summary_str,
details: Some(payload_val.to_string()),
details: Some(truncated_details),
..Default::default()
};
activities.push_front(serde_json::to_value(&activity).unwrap_or_default());
if activities.len() > 100 {
activities.pop_back();
if let Ok(act_val) = serde_json::to_value(&activity) {
activities.push_front(act_val);
if activities.len() > 100 {
activities.pop_back();
}
}
});
@@ -316,7 +345,10 @@ impl MemoryState {
let _ = self.activity_tx.send(ws_resource_notification);
}
pub fn record_terminal_history(&self, payload: TerminalHistory) {
pub fn record_terminal_history(&self, mut payload: TerminalHistory) {
if payload.command.len() > 2048 {
payload.command = format!("{}... [truncated]", &payload.command[..2048]);
}
self.telemetry.terminal_history.modify(|history| {
history.push_front(payload);
if history.len() > 100 {
@@ -353,6 +385,7 @@ mod tests {
git_branch: None,
parent_id: None,
expires_at: None,
..Default::default()
});
});
@@ -395,7 +428,11 @@ mod tests {
let mut rx = state.activity_tx.subscribe();
// 1. Record an activity with details
state.record_activity("code_change", "Refactored state.rs", Some("Updated ActivityRecord schema"));
state.record_activity(
"code_change",
"Refactored state.rs",
Some("Updated ActivityRecord schema"),
);
// Verify recent_activities store
let activities: Vec<ActivityRecord> = state.telemetry.recent_activities.read_with(|act| {
@@ -407,12 +444,19 @@ mod tests {
assert_eq!(activities.len(), 1);
assert_eq!(activities[0].category, "CODE_CHANGE");
assert_eq!(activities[0].summary, "Refactored state.rs");
assert_eq!(activities[0].details, Some("Updated ActivityRecord schema".to_string()));
assert!(activities[0].timestamp > 1_700_000_000_000, "Timestamp must be in epoch milliseconds");
assert_eq!(
activities[0].details,
Some("Updated ActivityRecord schema".to_string())
);
assert!(
activities[0].timestamp > 1_700_000_000_000,
"Timestamp must be in epoch milliseconds"
);
// Verify broadcast channel message
let broadcast_msg = rx.recv().await.expect("Expected broadcast notification");
let broadcast_val: serde_json::Value = serde_json::from_str(&broadcast_msg).expect("Valid JSON");
let broadcast_val: serde_json::Value =
serde_json::from_str(&broadcast_msg).expect("Valid JSON");
assert_eq!(broadcast_val["jsonrpc"], "2.0");
assert_eq!(broadcast_val["method"], "notifications/activity");
assert_eq!(broadcast_val["params"]["category"], "CODE_CHANGE");
@@ -420,11 +464,12 @@ mod tests {
// 2. Broadcast an activity without details
state.broadcast_activity("task", "Completed live activity fix");
let activities_updated: Vec<ActivityRecord> = state.telemetry.recent_activities.read_with(|act| {
act.iter()
.filter_map(|v| serde_json::from_value(v.clone()).ok())
.collect()
});
let activities_updated: Vec<ActivityRecord> =
state.telemetry.recent_activities.read_with(|act| {
act.iter()
.filter_map(|v| serde_json::from_value(v.clone()).ok())
.collect()
});
assert_eq!(activities_updated.len(), 2);
assert_eq!(activities_updated[0].category, "TASK");
@@ -456,90 +501,109 @@ impl SearchService {
pub async fn semantic_search(
&self,
query: &str,
_filter_namespace: Option<&str>,
filter_namespace: Option<&str>,
limit: usize,
) -> crate::error::Result<Vec<UnifiedSearchResult>> {
let query_emb = generate_embedding_async(query.to_string())
.await
.unwrap_or_default();
let mut results = Vec::new();
let mut cached_items = Vec::new();
let mut uncached_texts = Vec::new();
let mut uncached_meta = Vec::new();
let mut vdb_search = false;
if let Some(vdb) = &*self.state.vector_db.read().await {
vdb_search = true;
if let Ok(search_results) = vdb.search(query_emb.clone(), limit as u64).await {
for res in search_results {
results.push(UnifiedSearchResult {
id: res.id.clone(),
doc_type: res.doc_type.clone(),
title: res.id,
body: res.text,
score: res.score,
});
self.state.code.snippets.read_with(|snips| {
for snippet in snips.iter().take(50) {
let title = snippet.name.clone();
let desc = snippet.description.clone();
if let Some(ref emb) = snippet.embedding {
cached_items.push((title, "snippet".to_string(), desc, emb.clone()));
} else {
uncached_texts.push(format!(
"{} {} {}",
snippet.name, snippet.description, snippet.code
));
uncached_meta.push((title, "snippet".to_string(), desc));
}
}
}
});
if !vdb_search {
let mut cached_items = Vec::new();
let mut uncached_texts = Vec::new();
let mut uncached_meta = Vec::new();
self.state.code.snippets.read_with(|snips| {
for snippet in snips.iter().take(50) {
let title = snippet.name.clone();
let desc = snippet.description.clone();
if let Some(ref emb) = snippet.embedding {
cached_items.push((title, "snippet".to_string(), desc, emb.clone()));
} else {
uncached_texts.push(format!("{} {} {}", snippet.name, snippet.description, snippet.code));
uncached_meta.push((title, "snippet".to_string(), desc));
self.state.code.sticky.read_with(|sticky| {
for note in sticky.iter().take(50) {
let content_preview = note.content.chars().take(200).collect::<String>();
uncached_texts.push(note.content.clone());
uncached_meta.push((
"StickyNote".to_string(),
"sticky".to_string(),
content_preview,
));
}
});
self.state.read_graph(|graph| {
for entity in graph.entities.values().take(50) {
if let Some(ns) = filter_namespace {
if entity.namespace != ns {
continue;
}
}
});
self.state.code.sticky.read_with(|sticky| {
for note in sticky.iter().take(50) {
let content_preview = note.content.chars().take(200).collect::<String>();
uncached_texts.push(note.content.clone());
uncached_meta.push(("StickyNote".to_string(), "sticky".to_string(), content_preview));
let title = entity.name.clone();
let obs = entity.observations.join("; ");
let desc = format!("{}: {}", entity.entity_type, obs);
if let Some(ref emb) = entity.embedding {
cached_items.push((title, "entity".to_string(), desc, emb.clone()));
} else {
uncached_texts.push(format!("{} {} {}", entity.name, entity.entity_type, obs));
uncached_meta.push((title, "entity".to_string(), desc));
}
});
}
});
for (title, doc_type, body, emb) in cached_items {
self.state.code.error_fixes.read_with(|fixes| {
for fix in fixes.iter().take(50) {
let title = fix.signature.clone();
let desc = fix.solution.clone();
if let Some(ref emb) = fix.embedding {
cached_items.push((title, "error_fix".to_string(), desc, emb.clone()));
} else {
uncached_texts.push(format!("{} {}", fix.signature, fix.solution));
uncached_meta.push((title, "error_fix".to_string(), desc));
}
}
});
for (title, doc_type, body, emb) in cached_items {
let sim = cosine_similarity(&query_emb, &emb);
results.push(UnifiedSearchResult {
id: title.clone(),
doc_type,
title,
body,
score: sim,
});
}
if !uncached_texts.is_empty()
&& let Ok(embeddings) = generate_embeddings_async(uncached_texts).await
{
for (emb, meta) in embeddings.into_iter().zip(uncached_meta) {
let sim = cosine_similarity(&query_emb, &emb);
results.push(UnifiedSearchResult {
id: title.clone(),
doc_type,
title,
body,
id: meta.0.clone(),
doc_type: meta.1.clone(),
title: meta.0,
body: meta.2,
score: sim,
});
}
if !uncached_texts.is_empty()
&& let Ok(embeddings) = generate_embeddings_async(uncached_texts).await
{
for (emb, meta) in embeddings.into_iter().zip(uncached_meta) {
let sim = cosine_similarity(&query_emb, &emb);
results.push(UnifiedSearchResult {
id: meta.0.clone(),
doc_type: meta.1.clone(),
title: meta.0,
body: meta.2,
score: sim,
});
}
}
results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
results.truncate(limit);
}
results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
results.truncate(limit);
Ok(results)
}
server/src/store.rs
+346 -44
View File
@@ -4,10 +4,16 @@ 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.
enum DbOp {
Insert(Vec<u8>),
Delete,
}
/// Internal write request dispatched to the single database writer actor.
struct DbWriteTask {
key: String,
data: Vec<u8>,
op: DbOp,
flushed_notifier: Arc<tokio::sync::Notify>,
oneshot_tx: Option<tokio::sync::oneshot::Sender<()>>,
}
@@ -18,7 +24,8 @@ pub struct DbWriteQueue {
tx: tokio::sync::mpsc::Sender<DbWriteTask>,
}
static QUEUE_REGISTRY: std::sync::Mutex<Option<(Arc<Database>, DbWriteQueue)>> = std::sync::Mutex::new(None);
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());
@@ -38,7 +45,8 @@ impl DbWriteQueue {
tokio::spawn(async move {
while let Some(first_task) = rx.recv().await {
let mut batch = vec![first_task];
let mut batch = Vec::with_capacity(100);
batch.push(first_task);
// Gold Standard Micro-batching: Drain up to 100 accumulated tasks from queue without blocking
while batch.len() < 100 {
@@ -54,8 +62,27 @@ impl DbWriteQueue {
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);
match &task.op {
DbOp::Insert(data) => {
if let Err(e) =
table.insert(task.key.as_str(), data.as_slice())
{
tracing::error!(
"Failed to insert key '{}' into redb: {}",
task.key,
e
);
}
}
DbOp::Delete => {
if let Err(e) = table.remove(task.key.as_str()) {
tracing::error!(
"Failed to delete key '{}' from redb: {}",
task.key,
e
);
}
}
}
}
}
@@ -64,7 +91,10 @@ impl DbWriteQueue {
}
}
Err(e) => {
tracing::error!("Failed to begin write transaction on redb writer actor: {}", e);
tracing::error!(
"Failed to begin write transaction on redb writer actor: {}",
e
);
}
}
@@ -88,18 +118,55 @@ impl DbWriteQueue {
key: String,
data: Vec<u8>,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
self.push_op(key, DbOp::Insert(data), flushed_notifier)
}
pub fn push_delete(
&self,
key: String,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
self.push_op(key, DbOp::Delete, flushed_notifier)
}
fn push_op(
&self,
key: String,
op: DbOp,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
let task = DbWriteTask {
key,
data,
op,
flushed_notifier,
oneshot_tx: Some(oneshot_tx),
};
if let Err(e) = self.tx.try_send(task) {
let key = e.into_inner().key;
tracing::error!("DbWriteQueue channel full or closed; unable to persist key '{}'", key);
None
match e {
tokio::sync::mpsc::error::TrySendError::Full(task) => {
let tx = self.tx.clone();
let key = task.key.clone();
tokio::spawn(async move {
if let Err(err) = tx.send(task).await {
tracing::error!(
"DbWriteQueue fallback send failed for key '{}': {}",
key,
err
);
}
});
None
}
tokio::sync::mpsc::error::TrySendError::Closed(task) => {
tracing::error!(
"DbWriteQueue channel closed; unable to persist key '{}'",
task.key
);
None
}
}
} else {
Some(oneshot_rx)
}
@@ -110,16 +177,38 @@ impl DbWriteQueue {
key: String,
data: Vec<u8>,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
self.push_op_async(key, DbOp::Insert(data), flushed_notifier)
.await
}
pub async fn push_delete_async(
&self,
key: String,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
self.push_op_async(key, DbOp::Delete, flushed_notifier)
.await
}
async fn push_op_async(
&self,
key: String,
op: DbOp,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
let task = DbWriteTask {
key,
data,
op,
flushed_notifier,
oneshot_tx: Some(oneshot_tx),
};
if let Err(e) = self.tx.send(task).await {
tracing::error!("DbWriteQueue channel closed; unable to persist key '{}'", e.0.key);
tracing::error!(
"DbWriteQueue channel closed; unable to persist key '{}'",
e.0.key
);
None
} else {
Some(oneshot_rx)
@@ -156,29 +245,116 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
tracing::error!("Failed to begin read transaction for key '{}'", key);
return (T::default(), false);
};
match read_txn.open_table(STORE_TABLE) {
Ok(table) => match table.get(key) {
Ok(Some(value)) => match serde_json::from_slice::<T>(value.value()) {
Ok(parsed) => (parsed, false),
Err(e) => {
tracing::error!(
"CRITICAL: Corrupted data for key '{}' in database: {}. Quarantine mode active: state initialized to empty default without overwriting DB key.",
key, e
);
(T::default(), true)
}
},
Ok(None) => (T::default(), false),
let Ok(table) = read_txn.open_table(STORE_TABLE) else {
return (T::default(), false);
};
// 1. Check monolithic key first as the authoritative snapshot
match table.get(key) {
Ok(Some(value)) => match serde_json::from_slice::<T>(value.value()) {
Ok(parsed) => return (parsed, false),
Err(e) => {
tracing::error!("Failed to get key '{}' from store table: {}", key, e);
(T::default(), false)
tracing::error!(
"CRITICAL: Corrupted data for key '{}' in database: {}. Quarantine mode active: state initialized to empty default without overwriting DB key.",
key,
e
);
return (T::default(), true);
}
},
Ok(None) => {}
Err(e) => {
tracing::error!("Failed to open STORE_TABLE for key '{}': {}", key, e);
(T::default(), false)
tracing::error!("Failed to get key '{}' from store table: {}", key, e);
}
}
// 2. Granular prefix keys fallback: format!("{}:", key)
let prefix = format!("{}:", key);
let mut items_array = Vec::new();
let mut items_map = serde_json::Map::new();
let mut found_granular = false;
if let Ok(range) = table.range(prefix.as_str()..) {
for entry in range {
if let Ok((k, v)) = entry {
let k_str = k.value();
if !k_str.starts_with(&prefix) {
break;
}
found_granular = true;
if let Ok(val) = serde_json::from_slice::<serde_json::Value>(v.value()) {
let sub_key = &k_str[prefix.len()..];
items_array.push(val.clone());
items_map.insert(sub_key.to_string(), val);
}
}
}
}
if found_granular {
if let Ok(parsed) = serde_json::from_value::<T>(serde_json::Value::Array(items_array)) {
return (parsed, false);
}
if let Ok(parsed) = serde_json::from_value::<T>(serde_json::Value::Object(items_map)) {
return (parsed, false);
}
}
(T::default(), false)
}
fn extract_granular_keys(base_key: &str, val: &serde_json::Value) -> Vec<String> {
let mut keys = Vec::new();
match val {
serde_json::Value::Array(arr) => {
for (i, item) in arr.iter().enumerate() {
let sub_key = item
.get("id")
.or_else(|| item.get("name"))
.or_else(|| item.get("title"))
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_else(|| i.to_string());
keys.push(format!("{}:{}", base_key, sub_key));
}
}
serde_json::Value::Object(map) => {
for sub_key in map.keys() {
keys.push(format!("{}:{}", base_key, sub_key));
}
}
_ => {}
}
keys
}
fn extract_granular_entries(base_key: &str, val: &serde_json::Value) -> Vec<(String, Vec<u8>)> {
let mut granular = Vec::new();
match val {
serde_json::Value::Array(arr) => {
for (i, item) in arr.iter().enumerate() {
let sub_key = item
.get("id")
.or_else(|| item.get("name"))
.or_else(|| item.get("title"))
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_else(|| i.to_string());
if let Ok(item_bytes) = serde_json::to_vec(item) {
granular.push((format!("{}:{}", base_key, sub_key), item_bytes));
}
}
}
serde_json::Value::Object(map) => {
for (sub_key, item) in map {
if let Ok(item_bytes) = serde_json::to_vec(item) {
granular.push((format!("{}:{}", base_key, sub_key), item_bytes));
}
}
}
_ => {}
}
granular
}
pub fn read_with<F, R>(&self, f: F) -> R
@@ -191,7 +367,7 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
pub fn modify<F: FnOnce(&mut T)>(&self, f: F)
where
T: Serialize,
T: Serialize + Clone,
{
if self.is_corrupted {
tracing::error!(
@@ -201,30 +377,80 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
return;
}
let serialized_res = {
// Fast mutation under critical lock section, then immediately release the RwLock guard
let (old_snapshot, new_snapshot) = {
let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner());
let old = (*lock).clone();
f(&mut lock);
serde_json::to_vec(&*lock)
let new = (*lock).clone();
(old, new)
};
match serialized_res {
// Expensive serialization and granular extraction run completely unblocked outside the lock
let old_keys = serde_json::to_value(&old_snapshot)
.map(|val| Self::extract_granular_keys(&self.key, &val))
.unwrap_or_default();
let full_bytes_res = serde_json::to_vec(&new_snapshot);
let granular_entries = serde_json::to_value(&new_snapshot)
.map(|val| Self::extract_granular_entries(&self.key, &val))
.unwrap_or_default();
let new_keys: std::collections::HashSet<&str> =
granular_entries.iter().map(|(k, _)| k.as_str()).collect();
let mut removed_keys = Vec::new();
for old_k in &old_keys {
if !new_keys.contains(old_k.as_str()) {
removed_keys.push(old_k.clone());
}
}
match full_bytes_res {
Ok(data) => {
if self.queue.push(self.key.clone(), data.clone(), self.flushed.clone()).is_none() {
// Delete removed granular entries so they don't resurrect on restart
for del_key in removed_keys {
self.queue.push_delete(del_key, self.flushed.clone());
}
// Queue granular entries
for (g_key, g_bytes) in granular_entries {
self.queue.push(g_key, g_bytes, self.flushed.clone());
}
if self
.queue
.push(self.key.clone(), data.clone(), self.flushed.clone())
.is_none()
{
tracing::warn!(
"DbWriteQueue channel full for key '{}'. Applying backpressure fallback.",
self.key
);
let queue = self.queue.clone();
let key = self.key.clone();
let flushed = self.flushed.clone();
tokio::spawn(async move {
let _ = queue.push_async(key, data, flushed).await;
});
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
let _ = tokio::time::timeout(
std::time::Duration::from_secs(10),
queue.push_async(key, data, flushed),
)
.await;
});
}
}
}
Err(e) => tracing::error!("Failed to serialize memory store for key '{}': {}", self.key, e),
Err(e) => tracing::error!(
"Failed to serialize memory store for key '{}': {}",
self.key,
e
),
}
}
pub async fn modify_async<F: FnOnce(&mut T)>(&self, f: F)
where
T: Serialize,
T: Serialize + Clone,
{
if self.is_corrupted {
tracing::error!(
@@ -234,19 +460,57 @@ impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T>
return;
}
let serialized_res = {
let (old_snapshot, new_snapshot) = {
let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner());
let old = (*lock).clone();
f(&mut lock);
serde_json::to_vec(&*lock)
let new = (*lock).clone();
(old, new)
};
match serialized_res {
let old_keys = serde_json::to_value(&old_snapshot)
.map(|val| Self::extract_granular_keys(&self.key, &val))
.unwrap_or_default();
let full_bytes_res = serde_json::to_vec(&new_snapshot);
let granular_entries = serde_json::to_value(&new_snapshot)
.map(|val| Self::extract_granular_entries(&self.key, &val))
.unwrap_or_default();
let new_keys: std::collections::HashSet<&str> =
granular_entries.iter().map(|(k, _)| k.as_str()).collect();
let mut removed_keys = Vec::new();
for old_k in &old_keys {
if !new_keys.contains(old_k.as_str()) {
removed_keys.push(old_k.clone());
}
}
match full_bytes_res {
Ok(data) => {
if let Some(rx) = self.queue.push_async(self.key.clone(), data, self.flushed.clone()).await {
for del_key in removed_keys {
self.queue
.push_delete_async(del_key, self.flushed.clone())
.await;
}
for (g_key, g_bytes) in granular_entries {
self.queue
.push_async(g_key, g_bytes, self.flushed.clone())
.await;
}
if let Some(rx) = self
.queue
.push_async(self.key.clone(), data, self.flushed.clone())
.await
{
let _ = rx.await;
}
}
Err(e) => tracing::error!("Failed to serialize memory store for key '{}': {}", self.key, e),
Err(e) => tracing::error!(
"Failed to serialize memory store for key '{}': {}",
self.key,
e
),
}
}
}
@@ -322,4 +586,42 @@ mod tests {
assert_eq!(store.read_with(|s| s.value), 50);
}
#[tokio::test]
async fn test_store_deletion_does_not_resurrect() {
let db = create_in_memory_test_db();
#[derive(Serialize, serde::Deserialize, Clone, Default, PartialEq, Debug)]
struct Item {
id: String,
name: String,
}
let store = Store::<Vec<Item>>::new("items", db.clone());
store.modify(|items| {
items.push(Item {
id: "item1".into(),
name: "First".into(),
});
items.push(Item {
id: "item2".into(),
name: "Second".into(),
});
});
store.flushed.notified().await;
// Verify both items loaded
let store_check = Store::<Vec<Item>>::new("items", db.clone());
assert_eq!(store_check.read_with(|items| items.len()), 2);
// Delete item1
store.modify(|items| {
items.retain(|i| i.id != "item1");
});
store.flushed.notified().await;
// Reload from DB into a brand new Store instance - item1 must NOT resurrect!
let store_reloaded = Store::<Vec<Item>>::new("items", db.clone());
let remaining = store_reloaded.read_with(|items| items.clone());
assert_eq!(remaining.len(), 1);
assert_eq!(remaining[0].id, "item2");
}
}
server/src/tools.rs
+259 -17
View File
@@ -26,7 +26,7 @@ pub struct CreateEntitiesTool {
}
/// Create new relations between entities in the knowledge graph.
/// CRITICAL SCHEMA REQUIREMENT: Each relation item MUST strictly contain fields 'from' (source entity name string), 'to' (target entity name string), and 'relation_type' (snake_case string, e.g., 'depends_on', 'calls', 'implements'). Do NOT use 'source', 'target', or 'relationType'.
/// Create new relations between entities in the knowledge graph (accepts 'from' or 'source', 'to' or 'target', and 'relation_type' or 'type').
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct CreateRelationsTool {
/// Array of relations to create. Each relation item MUST use 'from', 'to', and 'relation_type'.
@@ -73,13 +73,17 @@ pub struct DeleteRelationsTool {
pub relations: Vec<crate::models::Relation>,
}
/// Read the entire knowledge graph. WARNING: For large graphs, use search_nodes instead to avoid context limits.
/// Read the entire knowledge graph. WARNING: For large graphs, use search_nodes or pagination (limit, offset) to avoid context limits.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct ReadGraphTool {
/// Optional namespace to restrict the read to.
pub namespace: Option<String>,
/// Optional token budget cap. Truncates graph outputs to fit within the specified token budget.
pub max_tokens: Option<usize>,
/// Optional maximum number of entities to return (pagination).
pub limit: Option<usize>,
/// Optional entity offset for pagination.
pub offset: Option<usize>,
}
/// Search specifically for Knowledge Graph entities and nodes by name or type.
@@ -118,11 +122,40 @@ pub struct LogCodeChangeTool {
pub symbol_references: Option<Vec<String>>,
/// Optional line range (e.g. 'L123-L145') modified in the target file.
pub line_range: Option<String>,
/// Optional repository name (e.g. 'mcp-memory', 'ai-pr-review').
pub repo_name: Option<String>,
/// Optional repository remote origin URL (e.g. 'git@bitbucket.org:org/repo.git' or 'svn://...').
pub repo_url: Option<String>,
/// Optional project namespace (defaults to repo_name or 'global').
pub namespace: Option<String>,
/// Optional change kind: 'added', 'modified', 'deleted', 'renamed'. Defaults to 'modified'.
pub change_kind: Option<String>,
/// Optional author or subagent attribution (e.g. 'user', 'MemoryLibrarian', SVN author).
pub author: Option<String>,
/// Optional conversation or session ID.
pub session_id: Option<String>,
/// Optional VCS system type: 'git', 'svn', 'hg', etc. Auto-detected if omitted.
pub vcs_type: Option<String>,
/// Optional VCS revision identifier (e.g. SVN 'r14829' or Git commit SHA).
pub revision: Option<String>,
/// Optional VCS branch (e.g. SVN 'trunk', 'branches/v1.0' or Git branch).
pub branch: Option<String>,
/// Optional repository root URL or path.
pub repository_root: Option<String>,
}
/// Query recently logged code changes.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct QueryRecentChangesTool {}
pub struct QueryRecentChangesTool {
/// Optional namespace to filter changes by project/workspace.
pub namespace: Option<String>,
/// Optional repository name to filter changes.
pub repo_name: Option<String>,
/// Optional VCS type to filter changes ('git', 'svn', etc.).
pub vcs_type: Option<String>,
/// Maximum number of records to return (defaults to 50).
pub limit: Option<usize>,
}
/// Generate a visual representation of the knowledge graph.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
@@ -167,7 +200,6 @@ pub struct StickyNotesTool {
pub session_only: Option<bool>,
}
/// Condense or summarize an entity's observations to reduce size.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct CondenseEntityTool {
@@ -190,6 +222,16 @@ pub struct AddTaskTool {
pub parent_id: Option<String>,
/// Optional list of task IDs this task depends on.
pub dependencies: Option<Vec<String>>,
/// Optional namespace (defaults to 'default').
pub namespace: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional task priority: 'low', 'medium', 'high', 'urgent'.
pub priority: Option<String>,
/// Optional assigned subagent role or identifier.
pub assigned_agent: Option<String>,
/// Optional verification command to validate criteria completion.
pub verification_command: Option<String>,
}
/// Delete a task and all its children.
@@ -233,6 +275,14 @@ pub struct StoreSnippetTool {
pub description: String,
/// Optional classification tags (e.g. ['rust', 'async', 'mcp']).
pub tags: Option<Vec<String>>,
/// Optional project namespace.
pub namespace: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional origin file path where the snippet was extracted from.
pub origin_file: Option<String>,
/// Optional line range in origin file (e.g. 'L10-L45').
pub line_range: Option<String>,
}
/// Add classification tags to an existing code snippet.
@@ -362,6 +412,16 @@ pub struct LogDecisionTool {
pub consequence: String,
/// The ID of an older ADR this decision supersedes, if any.
pub supersedes: Option<String>,
/// Optional project namespace.
pub namespace: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional author or architect behind the decision.
pub author: Option<String>,
/// Optional affected system components or crates.
pub affected_components: Option<Vec<String>>,
/// Optional alternative designs or libraries considered.
pub alternatives_considered: Option<Vec<String>>,
}
/// Query architectural decision records.
@@ -448,6 +508,12 @@ pub struct PinFileTool {
pub file_path: String,
/// The associated git branch, if any.
pub git_branch: Option<String>,
/// Optional reason for pinning this file.
pub pin_reason: Option<String>,
/// Optional symbol focus or line range.
pub symbol_or_line: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
}
/// Unpin a file from the context workspace.
@@ -551,6 +617,14 @@ pub struct AddMilestoneTool {
/// The namespace for the milestone.
#[serde(default = "crate::models::default_namespace")]
pub namespace: String,
/// Optional target completion date.
pub target_date: Option<String>,
/// Optional description of milestone scope.
pub description: Option<String>,
/// Optional deliverables or task checklist.
pub deliverables: Option<Vec<String>>,
/// Optional repository name.
pub repo_name: Option<String>,
}
/// Update the status of a milestone.
@@ -611,6 +685,14 @@ pub struct AddPrChecklistItemTool {
pub namespace: String,
/// The description of the checklist item.
pub description: String,
/// Optional checklist category (e.g. 'Security', 'Testing', 'Formatting').
pub category: Option<String>,
/// Optional automated verification command or script.
pub automated_check: Option<String>,
/// Optional checked state.
pub is_checked: Option<bool>,
/// Optional repository name.
pub repo_name: Option<String>,
}
/// Get the PR checklist.
@@ -647,6 +729,16 @@ pub struct LogTechDebtTool {
pub symbol_references: Option<Vec<String>>,
/// Optional line range (e.g. 'L123-L145') associated with the tech debt.
pub line_range: Option<String>,
/// Optional debt severity: 'low', 'medium', 'high', 'critical'.
pub severity: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional file path where debt exists.
pub file_path: Option<String>,
/// Optional workaround currently in place.
pub workaround: Option<String>,
/// Optional effort estimate to fix (e.g. '1 hour', '2 days').
pub effort_estimate: Option<String>,
}
/// Resolve a technical debt record.
@@ -682,6 +774,14 @@ pub struct SaveContextWorkspaceTool {
pub pinned_files: Vec<String>,
/// Array of active task IDs.
pub active_task_ids: Vec<String>,
/// Optional description of the context workspace purpose.
pub description: Option<String>,
/// Optional git branch associated with this workspace.
pub git_branch: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional VCS revision or commit hash.
pub vcs_revision: Option<String>,
}
/// Load a saved context workspace.
@@ -864,11 +964,6 @@ pub struct GetRecentLogsTool {
pub file_path: String,
}
/// Enable or disable clipboard watch mode (auto-ingesting into StickyNotes).
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct ToggleClipboardWatchModeTool {
pub enable: bool,
}
/// Read a file and return only its AST skeleton (Imports, Structs, Enums, Traits, Functions)
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct ReadFileSkeletonTool {
@@ -947,6 +1042,9 @@ pub struct FindSymbolReferencesTool {
pub symbol: String,
/// Optional maximum number of reference locations to return. Defaults to 10.
pub limit: Option<usize>,
/// Optional workspace or directory path to scan. If omitted, falls back to active pinned files or current working directory.
#[serde(default, alias = "directory", alias = "path")]
pub workspace_dir: Option<String>,
}
/// Find all caller functions or methods that invoke a specified target function or method name.
@@ -956,6 +1054,9 @@ pub struct GetCallersTool {
pub function_name: String,
/// Optional maximum number of callers to return. Defaults to 10.
pub limit: Option<usize>,
/// Optional workspace or directory path to scan. If omitted, falls back to active pinned files or current working directory.
#[serde(default, alias = "directory", alias = "path")]
pub workspace_dir: Option<String>,
}
/// Query git commit history, diffs, and change ledger entries using keyword or semantic search.
@@ -1079,8 +1180,11 @@ pub struct SearchSnippetsHybridTool {
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum HandoffMemoAction {
#[serde(alias = "leave", alias = "LEAVE", alias = "Leave")]
Leave,
#[serde(alias = "read", alias = "READ", alias = "Read")]
Read,
#[serde(alias = "clear", alias = "CLEAR", alias = "Clear")]
Clear,
}
@@ -1105,11 +1209,22 @@ pub struct HandoffMemosTool {
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum TaskAction {
#[serde(alias = "add", alias = "ADD", alias = "Add")]
Add,
#[serde(alias = "update", alias = "UPDATE", alias = "Update")]
Update,
#[serde(alias = "delete", alias = "DELETE", alias = "Delete")]
Delete,
#[serde(alias = "list", alias = "LIST", alias = "List")]
List,
#[serde(
alias = "set_criteria",
alias = "setCriteria",
alias = "SET_CRITERIA",
alias = "SetCriteria"
)]
SetCriteria,
#[serde(alias = "verify", alias = "VERIFY", alias = "Verify")]
Verify,
}
@@ -1127,20 +1242,27 @@ pub struct TasksTool {
/// Action to perform: 'add', 'update', 'delete', 'list', 'set_criteria', or 'verify'.
pub action: TaskAction,
/// Task ID (required for 'update', 'delete', 'set_criteria', 'verify').
#[serde(alias = "task_id", alias = "taskId")]
pub id: Option<String>,
/// Task title (required for 'add').
#[serde(alias = "name")]
pub title: Option<String>,
/// Task description (optional for 'add').
#[serde(alias = "desc")]
pub description: Option<String>,
/// New status: 'pending', 'completed', or 'cancelled' (for 'update').
pub status: Option<String>,
/// Parent task ID (optional for 'add').
#[serde(alias = "parentId", alias = "parent")]
pub parent_id: Option<String>,
/// List of dependency task IDs (optional for 'add').
#[serde(alias = "deps")]
pub dependencies: Option<Vec<String>>,
/// Git branch filter or assignment.
#[serde(alias = "branch")]
pub git_branch: Option<String>,
/// Acceptance criteria (required for 'set_criteria').
#[serde(alias = "acceptance_criteria", alias = "acceptanceCriteria")]
pub criteria: Option<Vec<String>>,
/// Verification proof or details (optional for 'verify').
pub proof: Option<String>,
@@ -1148,13 +1270,26 @@ pub struct TasksTool {
pub summary_level: Option<String>,
/// Maximum tokens budget cap (for 'list').
pub max_tokens: Option<usize>,
/// Optional namespace filter or assignment (defaults to 'default').
pub namespace: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional task priority: 'low', 'medium', 'high', 'urgent'.
pub priority: Option<String>,
/// Optional assigned subagent role or identifier.
pub assigned_agent: Option<String>,
/// Optional verification command to validate criteria completion.
pub verification_command: Option<String>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum MilestoneAction {
#[serde(alias = "add", alias = "ADD", alias = "Add")]
Add,
#[serde(alias = "update", alias = "UPDATE", alias = "Update")]
Update,
#[serde(alias = "list", alias = "LIST", alias = "List")]
List,
}
@@ -1169,20 +1304,33 @@ pub struct MilestonesTool {
/// Action to perform: 'add', 'update', or 'list'.
pub action: MilestoneAction,
/// Milestone ID (required for 'update').
#[serde(alias = "milestone_id", alias = "milestoneId")]
pub id: Option<String>,
/// Milestone title (required for 'add').
#[serde(alias = "name")]
pub title: Option<String>,
/// New status (for 'update').
pub status: Option<String>,
/// Optional namespace filter or assignment.
pub namespace: Option<String>,
/// Optional target completion date.
pub target_date: Option<String>,
/// Optional description of milestone scope.
pub description: Option<String>,
/// Optional deliverables or task checklist.
pub deliverables: Option<Vec<String>>,
/// Optional repository name.
pub repo_name: Option<String>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum PinnedFileAction {
#[serde(alias = "pin", alias = "PIN", alias = "Pin")]
Pin,
#[serde(alias = "unpin", alias = "UNPIN", alias = "Unpin")]
Unpin,
#[serde(alias = "list", alias = "LIST", alias = "List")]
List,
}
@@ -1200,15 +1348,28 @@ pub struct PinnedFilesTool {
pub path: Option<String>,
/// Optional namespace.
pub namespace: Option<String>,
/// Optional reason for pinning this file.
pub pin_reason: Option<String>,
/// Optional symbol focus or line range.
pub symbol_or_line: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional git branch.
pub git_branch: Option<String>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum ContextWorkspaceAction {
#[serde(alias = "save", alias = "SAVE", alias = "Save")]
Save,
#[serde(alias = "load", alias = "LOAD", alias = "Load")]
Load,
#[serde(alias = "list", alias = "LIST", alias = "List")]
List,
#[serde(alias = "delete", alias = "DELETE", alias = "Delete")]
Delete,
#[serde(alias = "diff", alias = "DIFF", alias = "Diff")]
Diff,
}
@@ -1234,13 +1395,24 @@ pub struct ContextWorkspacesTool {
pub active_task_ids: Option<Vec<String>>,
/// Optional namespace.
pub namespace: Option<String>,
/// Optional description of the context workspace purpose.
pub description: Option<String>,
/// Optional git branch associated with this workspace.
pub git_branch: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional VCS revision or commit hash.
pub vcs_revision: Option<String>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum PrChecklistAction {
#[serde(alias = "add", alias = "ADD", alias = "Add")]
Add,
#[serde(alias = "get", alias = "GET", alias = "Get")]
Get,
#[serde(alias = "clear", alias = "CLEAR", alias = "Clear")]
Clear,
}
@@ -1258,14 +1430,26 @@ pub struct PrChecklistTool {
pub description: Option<String>,
/// Optional namespace.
pub namespace: Option<String>,
/// Optional checklist category (e.g. 'Security', 'Testing', 'Formatting').
pub category: Option<String>,
/// Optional automated verification command or script.
pub automated_check: Option<String>,
/// Optional checked state.
pub is_checked: Option<bool>,
/// Optional repository name.
pub repo_name: Option<String>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum SnippetAction {
#[serde(alias = "store", alias = "STORE", alias = "Store")]
Store,
#[serde(alias = "search", alias = "SEARCH", alias = "Search")]
Search,
#[serde(alias = "delete", alias = "DELETE", alias = "Delete")]
Delete,
#[serde(alias = "tag", alias = "TAG", alias = "Tag")]
Tag,
}
@@ -1296,13 +1480,24 @@ pub struct SnippetsTool {
pub tags: Option<Vec<String>>,
/// Enable hybrid lexical + semantic vector ranking (for 'search').
pub hybrid: Option<bool>,
/// Optional project namespace.
pub namespace: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional origin file path where the snippet was extracted from.
pub origin_file: Option<String>,
/// Optional line range in origin file (e.g. 'L10-L45').
pub line_range: Option<String>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum DecisionAction {
#[serde(alias = "log", alias = "LOG", alias = "Log")]
Log,
#[serde(alias = "query", alias = "QUERY", alias = "Query")]
Query,
#[serde(alias = "delete", alias = "DELETE", alias = "Delete")]
Delete,
}
@@ -1332,13 +1527,24 @@ pub struct DecisionsTool {
pub query: Option<String>,
/// Optional namespace.
pub namespace: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional author or architect behind the decision.
pub author: Option<String>,
/// Optional affected system components or crates.
pub affected_components: Option<Vec<String>>,
/// Optional alternative designs or libraries considered.
pub alternatives_considered: Option<Vec<String>>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum TechDebtAction {
#[serde(alias = "log", alias = "LOG", alias = "Log")]
Log,
#[serde(alias = "resolve", alias = "RESOLVE", alias = "Resolve")]
Resolve,
#[serde(alias = "list", alias = "LIST", alias = "List")]
List,
}
@@ -1353,16 +1559,21 @@ pub struct TechDebtTool {
/// Action to perform: 'log', 'resolve', or 'list'.
pub action: TechDebtAction,
/// Tech debt ID (required for 'resolve').
#[serde(alias = "tech_debt_id", alias = "debt_id")]
pub id: Option<String>,
/// Tech debt description (required for 'log').
#[serde(alias = "desc", alias = "title")]
pub description: Option<String>,
/// Ideal solution (optional for 'log').
#[serde(alias = "solution", alias = "fix")]
pub ideal_solution: Option<String>,
/// Include resolved tech debt items (for 'list').
pub include_resolved: Option<bool>,
/// Associated git commit hash.
#[serde(alias = "commit")]
pub git_commit: Option<String>,
/// Associated git branch.
#[serde(alias = "branch")]
pub git_branch: Option<String>,
/// Symbol references associated with the tech debt.
pub symbol_references: Option<Vec<String>>,
@@ -1370,15 +1581,50 @@ pub struct TechDebtTool {
pub line_range: Option<String>,
/// Optional namespace.
pub namespace: Option<String>,
/// Optional debt severity: 'low', 'medium', 'high', 'critical'.
pub severity: Option<String>,
/// Optional repository name.
pub repo_name: Option<String>,
/// Optional file path where debt exists.
pub file_path: Option<String>,
/// Optional workaround currently in place.
pub workaround: Option<String>,
/// Optional effort estimate to fix.
pub effort_estimate: Option<String>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum EnvAction {
#[serde(
alias = "update_fingerprint",
alias = "updateFingerprint",
alias = "UPDATE_FINGERPRINT",
alias = "UpdateFingerprint"
)]
UpdateFingerprint,
#[serde(
alias = "read_fingerprint",
alias = "readFingerprint",
alias = "READ_FINGERPRINT",
alias = "ReadFingerprint"
)]
ReadFingerprint,
#[serde(
alias = "log_requirement",
alias = "logRequirement",
alias = "LOG_REQUIREMENT",
alias = "LogRequirement"
)]
LogRequirement,
#[serde(alias = "register", alias = "REGISTER", alias = "Register")]
Register,
#[serde(
alias = "get_details",
alias = "getDetails",
alias = "GET_DETAILS",
alias = "GetDetails"
)]
GetDetails,
}
@@ -1415,20 +1661,19 @@ pub struct EnvironmentTool {
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum ClipboardAction {
#[serde(alias = "read", alias = "READ", alias = "Read")]
Read,
#[serde(alias = "write", alias = "WRITE", alias = "Write")]
Write,
ToggleWatch,
}
/// OS Clipboard management (read, write, toggle_watch).
/// OS Clipboard management (read, write).
/// Actions:
/// - 'read': Read current clipboard contents.
/// - 'write': Write content to OS clipboard. Optional: text, html, files, image_path.
/// - 'toggle_watch': Enable or disable auto-clipboard ingestion. Optional: enabled (bool).
/// Next steps on error: Provide text/html/files/image_path when calling 'write'.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct ClipboardTool {
/// Action to perform: 'read', 'write', or 'toggle_watch'.
/// Action to perform: 'read' or 'write'.
pub action: ClipboardAction,
/// Plain text content (for 'write').
pub text: Option<String>,
@@ -1438,7 +1683,4 @@ pub struct ClipboardTool {
pub files: Option<Vec<String>>,
/// Image path (for 'write').
pub image_path: Option<String>,
/// Enable watch mode flag (for 'toggle_watch').
pub enabled: Option<bool>,
}
server/src/vector_db.rs
-236
View File
@@ -1,236 +0,0 @@
use qdrant_client::Qdrant;
use qdrant_client::qdrant::{
CreateCollectionBuilder, Distance, PointStruct, UpsertPointsBuilder, VectorParamsBuilder,
};
use serde::{Deserialize, Serialize};
use std::error::Error;
use std::sync::Arc;
use tracing::info;
use uuid::Uuid;
#[derive(Clone)]
pub struct VectorDB {
client: Arc<Qdrant>,
collection_name: String,
}
/// Vector search result containing semantic match metadata.
/// Note: Similarity scores range from 0.0 to 1.0 (Cosine Distance):
/// - Score >= 0.75: High-confidence semantic match
/// - Score 0.50 - 0.74: Moderate semantic relevance
/// - Score < 0.40: Low-confidence noise
#[derive(Debug, Serialize, Deserialize)]
pub struct VectorSearchResult {
pub id: String,
pub doc_type: String,
pub text: String,
pub score: f32,
}
impl VectorDB {
pub async fn new(
url: &str,
collection_name: &str,
) -> Result<Self, Box<dyn Error + Send + Sync>> {
let client = Qdrant::from_url(url).build()?;
let db = Self {
client: Arc::new(client),
collection_name: collection_name.to_string(),
};
db.init_collection().await?;
Ok(db)
}
async fn init_collection(&self) -> Result<(), Box<dyn Error + Send + Sync>> {
// Fastembed AllMiniLML6V2 uses 384 dimensions
let vector_params = VectorParamsBuilder::new(384, Distance::Cosine).build();
let collection_exists = self.client.collection_exists(&self.collection_name).await?;
if !collection_exists {
self.client
.create_collection(
CreateCollectionBuilder::new(&self.collection_name)
.vectors_config(vector_params),
)
.await?;
info!("Created Qdrant collection: {}", self.collection_name);
} else {
info!("Qdrant collection {} already exists", self.collection_name);
}
Ok(())
}
pub async fn index_document(
&self,
id: &str,
doc_type: &str,
text: &str,
vector: Vec<f32>,
) -> Result<(), Box<dyn Error + Send + Sync>> {
let point_id = match Uuid::parse_str(id) {
Ok(u) => u.to_string(),
Err(_) => {
// If it's not a valid UUID, let's create a deterministic UUID based on the string
let namespace = Uuid::NAMESPACE_OID;
Uuid::new_v5(&namespace, id.as_bytes()).to_string()
}
};
let mut payload: std::collections::HashMap<String, serde_json::Value> =
std::collections::HashMap::new();
payload.insert(
"doc_type".to_string(),
serde_json::Value::String(doc_type.to_string()),
);
payload.insert(
"text".to_string(),
serde_json::Value::String(text.to_string()),
);
payload.insert(
"original_id".to_string(),
serde_json::Value::String(id.to_string()),
);
let point = PointStruct::new(point_id, vector, payload);
self.client
.upsert_points(UpsertPointsBuilder::new(&self.collection_name, vec![point]))
.await?;
Ok(())
}
pub async fn index_documents_batch(
&self,
docs: Vec<(&str, &str, &str, Vec<f32>)>,
) -> Result<(), Box<dyn Error + Send + Sync>> {
if docs.is_empty() {
return Ok(());
}
let mut points = Vec::with_capacity(docs.len());
for (id, doc_type, text, vector) in docs {
let point_id = match Uuid::parse_str(id) {
Ok(u) => u.to_string(),
Err(_) => {
let namespace = Uuid::NAMESPACE_OID;
Uuid::new_v5(&namespace, id.as_bytes()).to_string()
}
};
let mut payload: std::collections::HashMap<String, serde_json::Value> =
std::collections::HashMap::new();
payload.insert("doc_type".to_string(), serde_json::Value::String(doc_type.to_string()));
payload.insert("text".to_string(), serde_json::Value::String(text.to_string()));
payload.insert("original_id".to_string(), serde_json::Value::String(id.to_string()));
points.push(PointStruct::new(point_id, vector, payload));
}
self.client
.upsert_points(UpsertPointsBuilder::new(&self.collection_name, points))
.await?;
Ok(())
}
pub async fn search(
&self,
query_vector: Vec<f32>,
limit: u64,
) -> Result<Vec<VectorSearchResult>, Box<dyn Error + Send + Sync>> {
use qdrant_client::qdrant::SearchPointsBuilder;
let search_result = self
.client
.search_points(
SearchPointsBuilder::new(&self.collection_name, query_vector, limit)
.with_payload(true),
)
.await?;
let mut results = Vec::new();
for point in search_result.result {
let id = point
.payload
.get("original_id")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_default();
let doc_type = point
.payload
.get("doc_type")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_default();
let text = point
.payload
.get("text")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_default();
results.push(VectorSearchResult {
id,
doc_type,
text,
score: point.score,
});
}
Ok(results)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_vector_search_result_serialization() {
let res = VectorSearchResult {
id: "test-id".into(),
doc_type: "snippet".into(),
text: "fn main() {}".into(),
score: 0.95,
};
let json = serde_json::to_string(&res).unwrap();
assert!(json.contains("test-id"));
assert!(json.contains("0.95"));
let deserialized: VectorSearchResult = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.id, "test-id");
assert_eq!(deserialized.doc_type, "snippet");
assert_eq!(deserialized.text, "fn main() {}");
assert_eq!(deserialized.score, 0.95);
}
#[test]
fn test_deterministic_uuid_hashing() {
let id_str = "custom-non-uuid-string";
let namespace = Uuid::NAMESPACE_OID;
let uuid1 = Uuid::new_v5(&namespace, id_str.as_bytes()).to_string();
let uuid2 = Uuid::new_v5(&namespace, id_str.as_bytes()).to_string();
assert_eq!(uuid1, uuid2);
assert!(Uuid::parse_str(&uuid1).is_ok());
}
#[tokio::test]
async fn test_vector_db_index_document_invalid_id() {
if let Ok(db) = VectorDB::new("http://localhost:6333", "test_collection").await {
let invalid_id = "not-a-uuid";
let doc_type = "test";
let text = "example text";
let vector = vec![1.0; 384];
let _ = db.index_document(invalid_id, doc_type, text, vector).await;
}
}
}
server/src/watcher.rs
+32 -16
View File
@@ -34,28 +34,44 @@ pub fn spawn_watcher(state: Arc<MemoryState>) {
let mut last_processed: std::collections::HashMap<std::path::PathBuf, std::time::Instant> =
std::collections::HashMap::new();
while let Some(res) = rx.recv().await {
match res {
Ok(event) => {
if event.kind.is_modify() {
let now = std::time::Instant::now();
for path in event.paths {
if should_review(&path) {
// 250ms debouncing window per file path
if let Some(last) = last_processed.get(&path) {
if now.duration_since(*last) < std::time::Duration::from_millis(250) {
continue;
loop {
tokio::select! {
_ = state.shutdown_notify.notified() => {
info!("File watcher received shutdown notification; terminating cleanly.");
break;
}
res = rx.recv() => {
let Some(res) = res else {
break;
};
match res {
Ok(event) => {
if event.kind.is_modify() {
let now = std::time::Instant::now();
if last_processed.len() > 1000 {
let ten_mins = std::time::Duration::from_secs(600);
last_processed.retain(|_, last_time| now.duration_since(*last_time) < ten_mins);
}
for path in event.paths {
if should_review(&path) {
// 250ms debouncing window per file path
if let Some(last) = last_processed.get(&path) {
if now.duration_since(*last) < std::time::Duration::from_millis(250) {
continue;
}
}
last_processed.insert(path.clone(), now);
info!("Proactive Daemon Hooks: File modified: {:?}", path);
trigger_autonomous_review(&path, Arc::clone(&state)).await;
}
}
last_processed.insert(path.clone(), now);
info!("Proactive Daemon Hooks: File modified: {:?}", path);
trigger_autonomous_review(&path, Arc::clone(&state)).await;
}
}
Err(e) => error!("Watch error: {}", e),
}
}
Err(e) => error!("Watch error: {}", e),
}
}
});