refactor: apply 5-pass audit optimizations across mcp-memory codebase
This commit is contained in:
1 parent
924b6d09fa
commit
5bd8b1587a
43 files changed
+1866
-1658
No files matched your search
+17
-16
@@ -221,23 +221,24 @@ pub fn create_router(app_state: Arc<AppState>) -> Router {
|
||||
move |axum::extract::Query(params): axum::extract::Query<
|
||||
std::collections::HashMap<String, String>,
|
||||
>| async move {
|
||||
if let Some(q) = params.get("q")
|
||||
&& let Ok(idx) = state_clone.search_index.read()
|
||||
&& let Ok(results) = idx.search(q, None) {
|
||||
let mut formatted_results = Vec::new();
|
||||
for (id, doc_type, title, body, score) in results {
|
||||
formatted_results.push(serde_json::json!({
|
||||
"id": id,
|
||||
"type_name": doc_type,
|
||||
"title": title,
|
||||
"content": body,
|
||||
"score": score
|
||||
}));
|
||||
}
|
||||
return axum::Json(
|
||||
serde_json::json!({ "results": formatted_results }),
|
||||
);
|
||||
if let Some(q) = params.get("q") {
|
||||
let idx = state_clone.get_search_index().await;
|
||||
if let Ok(results) = idx.search(q, None) {
|
||||
let mut formatted_results = Vec::new();
|
||||
for (id, doc_type, title, body, score) in results {
|
||||
formatted_results.push(serde_json::json!({
|
||||
"id": id,
|
||||
"type_name": doc_type,
|
||||
"title": title,
|
||||
"content": body,
|
||||
"score": score
|
||||
}));
|
||||
}
|
||||
return axum::Json(
|
||||
serde_json::json!({ "results": formatted_results }),
|
||||
);
|
||||
}
|
||||
}
|
||||
axum::Json(serde_json::json!({ "results": [] }))
|
||||
}
|
||||
}),
|
||||
|
||||
@@ -6,6 +6,8 @@ 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;
|
||||
@@ -19,33 +21,47 @@ pub fn spawn_watcher(state: Arc<MemoryState>) {
|
||||
|
||||
sleep(Duration::from_millis(1000)).await;
|
||||
|
||||
if let Ok(mut clipboard) = Clipboard::new()
|
||||
&& let Ok(text) = clipboard.get_text()
|
||||
&& text != last_text
|
||||
&& !text.trim().is_empty()
|
||||
{
|
||||
last_text = text.clone();
|
||||
if clipboard_opt.is_none() {
|
||||
clipboard_opt = Clipboard::new().ok();
|
||||
}
|
||||
|
||||
let note = StickyNote {
|
||||
timestamp: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs(),
|
||||
content: text.clone(),
|
||||
expires_at: None,
|
||||
};
|
||||
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();
|
||||
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.push(note.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.record_activity(
|
||||
"clipboard",
|
||||
"Auto-ingested clipboard content into sticky notes",
|
||||
Some(&text),
|
||||
);
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.push(note.clone());
|
||||
if notes.len() > 200 {
|
||||
let overflow = notes.len() - 200;
|
||||
notes.drain(0..overflow);
|
||||
}
|
||||
});
|
||||
|
||||
state.index_commit_notify.notify_waiters();
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -55,6 +55,12 @@ function switchTab(tabId, btn) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
function toggleTheme() {
|
||||
const currentTheme = document.documentElement.getAttribute("data-theme");
|
||||
const newTheme = currentTheme === "dark" ? "light" : "dark";
|
||||
document.documentElement.setAttribute("data-theme", newTheme);
|
||||
localStorage.setItem("theme", newTheme);
|
||||
}
|
||||
if (localStorage.getItem("theme") === "dark" || !localStorage.getItem("theme") && window.matchMedia && window.matchMedia("(prefers-color-scheme: dark)").matches) {
|
||||
document.documentElement.setAttribute("data-theme", "dark");
|
||||
}
|
||||
@@ -91,6 +97,11 @@ function zoomGraph(step) {
|
||||
const currentScale = network.getScale();
|
||||
network.moveTo({ scale: currentScale * (1 + step) });
|
||||
}
|
||||
function resetGraph() {
|
||||
if (!network)
|
||||
return;
|
||||
network.fit({ animation: { duration: 500, easingFunction: "easeInOutQuad" } });
|
||||
}
|
||||
function closeInspector() {
|
||||
const inspector = document.getElementById("inspector-panel");
|
||||
if (inspector)
|
||||
@@ -128,6 +139,14 @@ function showInspector(nodeId) {
|
||||
if (inspector)
|
||||
inspector.classList.add("open");
|
||||
}
|
||||
function toggleFilter(type) {
|
||||
if (activeFilters.has(type))
|
||||
activeFilters.delete(type);
|
||||
else
|
||||
activeFilters.add(type);
|
||||
renderFilters();
|
||||
updateGraphData();
|
||||
}
|
||||
function renderFilters() {
|
||||
const container = document.getElementById("graph-filters");
|
||||
if (!container)
|
||||
@@ -240,6 +259,14 @@ async function loadGraph() {
|
||||
}
|
||||
}
|
||||
}
|
||||
async function completeTask(id) {
|
||||
try {
|
||||
await fetch(`/api/tasks/${id}/complete`, { method: "POST" });
|
||||
loadTasks();
|
||||
} catch (e) {
|
||||
console.error("Failed to complete task", e);
|
||||
}
|
||||
}
|
||||
function buildTaskTreeHTML(tasks, parentId, depth = 0) {
|
||||
let html = "";
|
||||
const children = tasks.filter((t) => {
|
||||
@@ -340,6 +367,64 @@ async function loadTasks() {
|
||||
console.error("Failed to load tasks", err);
|
||||
}
|
||||
}
|
||||
function highlightText(text, query) {
|
||||
if (!query)
|
||||
return text;
|
||||
const regex = new RegExp(`(${query})`, "gi");
|
||||
return text.replace(regex, '<span class="search-highlight">$1</span>');
|
||||
}
|
||||
var searchDebounceTimer = null;
|
||||
var activeSearchAbortController = null;
|
||||
async function handleSearch(e) {
|
||||
const input = e.target;
|
||||
const q = input.value.trim();
|
||||
if (searchDebounceTimer)
|
||||
clearTimeout(searchDebounceTimer);
|
||||
if (!q) {
|
||||
if (activeSearchAbortController)
|
||||
activeSearchAbortController.abort();
|
||||
const resultsEl = document.getElementById("search-results");
|
||||
if (resultsEl)
|
||||
resultsEl.innerHTML = "";
|
||||
return;
|
||||
}
|
||||
const delay = e.key === "Enter" ? 0 : 250;
|
||||
searchDebounceTimer = setTimeout(async () => {
|
||||
if (activeSearchAbortController)
|
||||
activeSearchAbortController.abort();
|
||||
activeSearchAbortController = new AbortController;
|
||||
try {
|
||||
const res = await fetch(`/api/search?q=${encodeURIComponent(q)}`, {
|
||||
signal: activeSearchAbortController.signal
|
||||
});
|
||||
const data = await res.json();
|
||||
const container = document.getElementById("search-results");
|
||||
if (!container)
|
||||
return;
|
||||
if (!data.results || data.results.length === 0) {
|
||||
container.innerHTML = '<div style="padding:15px;color:var(--text-secondary);">No results found.</div>';
|
||||
return;
|
||||
}
|
||||
container.innerHTML = data.results.map((r) => `
|
||||
<div class="search-result-item" onclick="if(rawEntities['${r.id}']){ showInspector('${r.id}'); network.selectNodes(['${r.id}']); network.focus('${r.id}', {scale:1.2, animation:true}); }" style="cursor:pointer;">
|
||||
<div style="display:flex; justify-content:space-between; align-items:flex-start;">
|
||||
<div>
|
||||
<strong>${highlightText(r.title, q)}</strong>
|
||||
<span style="font-size: 0.7em; color: var(--text-secondary); background: var(--canvas-bg); padding: 2px 6px; border-radius: 4px; margin-left: 8px; text-transform: uppercase; border: 1px solid var(--border-color);">${r.type_name}</span>
|
||||
</div>
|
||||
<div style="font-size:0.75em; color:var(--text-secondary);">${r.score.toFixed(2)}</div>
|
||||
</div>
|
||||
<div style="margin-top:6px; font-size:0.9em; line-height:1.4;">${highlightText(r.content.substring(0, 150), q)}${r.content.length > 150 ? "..." : ""}</div>
|
||||
<div style="font-size:0.75em; color:var(--text-secondary); margin-top:8px; font-family:monospace; background: var(--canvas-bg); padding: 2px 4px; border-radius: 3px; display: inline-block;">ID: ${r.id}</div>
|
||||
</div>
|
||||
`).join("");
|
||||
} catch (err) {
|
||||
if (err.name !== "AbortError") {
|
||||
console.error("Search error:", err);
|
||||
}
|
||||
}
|
||||
}, delay);
|
||||
}
|
||||
var MAX_ACTIVITY_HISTORY = 100;
|
||||
function escapeHtml(str) {
|
||||
if (str === null || str === undefined)
|
||||
@@ -785,6 +870,39 @@ function loadPRs() {
|
||||
</ul>
|
||||
`);
|
||||
}
|
||||
function loadAllExtras() {
|
||||
loadTerminal();
|
||||
loadTechDebt();
|
||||
loadADRs();
|
||||
loadWorkspaces();
|
||||
loadPinned();
|
||||
loadMemos();
|
||||
loadSnippets();
|
||||
loadPRs();
|
||||
}
|
||||
async function testClipboard() {
|
||||
const modal = document.getElementById("clipboard-modal");
|
||||
const resultDiv = document.getElementById("clipboard-result");
|
||||
if (modal)
|
||||
modal.style.display = "flex";
|
||||
if (resultDiv)
|
||||
resultDiv.innerHTML = '<p id="clipboard-loading" style="color:var(--text-secondary);">Analyzing your clipboard...</p>';
|
||||
try {
|
||||
const res = await fetch("/api/clipboard/capture", { method: "POST" });
|
||||
const data = await res.json();
|
||||
if (resultDiv) {
|
||||
if (data.success) {
|
||||
resultDiv.innerHTML = `<img src="${data.image}" style="max-width:100%; max-height:60vh; border-radius:4px; border:1px solid var(--border-color); box-shadow: 0 5px 15px rgba(0,0,0,0.3);" />`;
|
||||
} else {
|
||||
resultDiv.innerHTML = `<p style="color:var(--error-color); font-weight:bold;">Failed to read clipboard: ${data.error}</p>`;
|
||||
}
|
||||
}
|
||||
} catch (err) {
|
||||
if (resultDiv) {
|
||||
resultDiv.innerHTML = `<p style="color:var(--error-color); font-weight:bold;">Error calling endpoint: ${err.message}</p>`;
|
||||
}
|
||||
}
|
||||
}
|
||||
async function loadVersion() {
|
||||
try {
|
||||
const res = await fetch("/api/version");
|
||||
@@ -836,6 +954,89 @@ document.addEventListener("keydown", function(e) {
|
||||
}
|
||||
}
|
||||
});
|
||||
function refreshActiveTab() {
|
||||
switch (currentTabId) {
|
||||
case "graph-tab":
|
||||
loadGraph();
|
||||
break;
|
||||
case "task-tab":
|
||||
loadTasks();
|
||||
break;
|
||||
case "sticky-tab":
|
||||
loadStickyNotes();
|
||||
break;
|
||||
case "techdebt-tab":
|
||||
loadTechDebt();
|
||||
break;
|
||||
case "adrs-tab":
|
||||
loadADRs();
|
||||
break;
|
||||
case "workspaces-tab":
|
||||
loadWorkspaces();
|
||||
break;
|
||||
case "pinned-tab":
|
||||
loadPinned();
|
||||
break;
|
||||
case "memos-tab":
|
||||
loadMemos();
|
||||
break;
|
||||
case "snippets-tab":
|
||||
loadSnippets();
|
||||
break;
|
||||
case "pr-tab":
|
||||
loadPRs();
|
||||
break;
|
||||
case "terminal-tab":
|
||||
loadTerminal();
|
||||
break;
|
||||
case "preferences-tab":
|
||||
loadPreferences();
|
||||
break;
|
||||
case "activity-tab":
|
||||
loadActivityHistory();
|
||||
break;
|
||||
}
|
||||
}
|
||||
Object.assign(window, {
|
||||
switchTab,
|
||||
toggleTheme,
|
||||
getColorForType,
|
||||
zoomGraph,
|
||||
resetGraph,
|
||||
closeInspector,
|
||||
showInspector,
|
||||
toggleFilter,
|
||||
renderFilters,
|
||||
updateGraphData,
|
||||
loadGraph,
|
||||
completeTask,
|
||||
buildTaskTreeHTML,
|
||||
loadTasks,
|
||||
highlightText,
|
||||
handleSearch,
|
||||
escapeHtml,
|
||||
parseActivityPayload,
|
||||
loadActivityHistory,
|
||||
requestDomainRefresh,
|
||||
handleIncomingActivity,
|
||||
setupWS,
|
||||
loadPreferences,
|
||||
loadStickyNotes,
|
||||
loadGenericList,
|
||||
loadTerminal,
|
||||
loadTechDebt,
|
||||
loadADRs,
|
||||
loadWorkspaces,
|
||||
loadPinned,
|
||||
loadMemos,
|
||||
loadSnippets,
|
||||
loadPRs,
|
||||
loadAllExtras,
|
||||
testClipboard,
|
||||
loadVersion,
|
||||
setupSSE,
|
||||
refreshActiveTab
|
||||
});
|
||||
loadVersion();
|
||||
loadActivityHistory();
|
||||
setupWS();
|
||||
|
||||
@@ -1096,6 +1096,48 @@ function refreshActiveTab(): void {
|
||||
}
|
||||
}
|
||||
|
||||
// --- Global Window Exports (Prevents bun build tree-shaking for HTML inline event handlers) ---
|
||||
Object.assign(window, {
|
||||
switchTab,
|
||||
toggleTheme,
|
||||
getColorForType,
|
||||
zoomGraph,
|
||||
resetGraph,
|
||||
closeInspector,
|
||||
showInspector,
|
||||
toggleFilter,
|
||||
renderFilters,
|
||||
updateGraphData,
|
||||
loadGraph,
|
||||
completeTask,
|
||||
buildTaskTreeHTML,
|
||||
loadTasks,
|
||||
highlightText,
|
||||
handleSearch,
|
||||
escapeHtml,
|
||||
parseActivityPayload,
|
||||
loadActivityHistory,
|
||||
requestDomainRefresh,
|
||||
handleIncomingActivity,
|
||||
setupWS,
|
||||
loadPreferences,
|
||||
loadStickyNotes,
|
||||
loadGenericList,
|
||||
loadTerminal,
|
||||
loadTechDebt,
|
||||
loadADRs,
|
||||
loadWorkspaces,
|
||||
loadPinned,
|
||||
loadMemos,
|
||||
loadSnippets,
|
||||
loadPRs,
|
||||
loadAllExtras,
|
||||
testClipboard,
|
||||
loadVersion,
|
||||
setupSSE,
|
||||
refreshActiveTab
|
||||
});
|
||||
|
||||
// --- Initialization ---
|
||||
loadVersion();
|
||||
loadActivityHistory();
|
||||
@@ -1103,3 +1145,4 @@ setupWS();
|
||||
|
||||
const observer = new MutationObserver(() => updateGraphData());
|
||||
observer.observe(document.documentElement, { attributes: true, attributeFilter: ['data-theme'] });
|
||||
|
||||
+62
-50
@@ -4,17 +4,27 @@ use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub fn init_redb(base: &Path) -> Arc<Database> {
|
||||
let redb_path = base.join("mcp_store.redb");
|
||||
let db = 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()));
|
||||
Arc::new(redb::Database::create(&temp_path).expect("Failed to create fallback redb database"))
|
||||
let is_in_memory = base == Path::new(":memory:") || base.to_str() == Some(":memory:") || cfg!(test);
|
||||
let db = if is_in_memory {
|
||||
Arc::new(
|
||||
redb::Database::builder()
|
||||
.create_with_backend(redb::backends::InMemoryBackend::new())
|
||||
.expect("Failed to create in-memory redb 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"))
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
@@ -24,47 +34,49 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
|
||||
let mut opened_ok = false;
|
||||
if let Ok(mut table) = write_txn.open_table(STORE_TABLE) {
|
||||
opened_ok = true;
|
||||
let stores = vec![
|
||||
("knowledge_graph_master", "knowledge_graph_master.json"),
|
||||
("audit_ledger", "audit_ledger.json"),
|
||||
("sticky_notes", "sticky_notes.json"),
|
||||
("tasks", "tasks.json"),
|
||||
("snippets", "snippets.json"),
|
||||
("adrs", "adrs.json"),
|
||||
("preferences", "preferences.json"),
|
||||
("error_fixes", "error_fixes.json"),
|
||||
("pinned_files", "pinned_files.json"),
|
||||
("session_summaries", "session_summaries.json"),
|
||||
("handoff_memos", "handoff_memos.json"),
|
||||
("env_fingerprints", "env_fingerprints.json"),
|
||||
("env_requirements", "env_requirements.json"),
|
||||
("milestones", "milestones.json"),
|
||||
("environments", "environments.json"),
|
||||
("pr_checklists", "pr_checklists.json"),
|
||||
("tech_debts", "tech_debts.json"),
|
||||
("gates", "gates.json"),
|
||||
("context_workspaces", "context_workspaces.json"),
|
||||
];
|
||||
if !is_in_memory {
|
||||
let stores = vec![
|
||||
("knowledge_graph_master", "knowledge_graph_master.json"),
|
||||
("audit_ledger", "audit_ledger.json"),
|
||||
("sticky_notes", "sticky_notes.json"),
|
||||
("tasks", "tasks.json"),
|
||||
("snippets", "snippets.json"),
|
||||
("adrs", "adrs.json"),
|
||||
("preferences", "preferences.json"),
|
||||
("error_fixes", "error_fixes.json"),
|
||||
("pinned_files", "pinned_files.json"),
|
||||
("session_summaries", "session_summaries.json"),
|
||||
("handoff_memos", "handoff_memos.json"),
|
||||
("env_fingerprints", "env_fingerprints.json"),
|
||||
("env_requirements", "env_requirements.json"),
|
||||
("milestones", "milestones.json"),
|
||||
("environments", "environments.json"),
|
||||
("pr_checklists", "pr_checklists.json"),
|
||||
("tech_debts", "tech_debts.json"),
|
||||
("gates", "gates.json"),
|
||||
("context_workspaces", "context_workspaces.json"),
|
||||
];
|
||||
|
||||
for (key, file_name) in stores.iter() {
|
||||
let is_missing = match table.get(*key) {
|
||||
Ok(res) => res.is_none(),
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to read key '{}' from redb: {}", key, e);
|
||||
false
|
||||
}
|
||||
};
|
||||
for (key, file_name) in stores.iter() {
|
||||
let is_missing = match table.get(*key) {
|
||||
Ok(res) => res.is_none(),
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to read key '{}' from redb: {}", key, e);
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
if is_missing {
|
||||
let json_path = base.join(file_name);
|
||||
if json_path.exists()
|
||||
&& let Ok(data) = std::fs::read(&json_path)
|
||||
&& serde_json::from_slice::<serde_json::Value>(&data).is_ok()
|
||||
{
|
||||
if let Err(e) = table.insert(*key, data.as_slice()) {
|
||||
tracing::error!("Failed to insert migrated key '{}': {}", key, e);
|
||||
} else {
|
||||
let _ = std::fs::rename(&json_path, json_path.with_extension("json.migrated"));
|
||||
if is_missing {
|
||||
let json_path = base.join(file_name);
|
||||
if json_path.exists()
|
||||
&& let Ok(data) = std::fs::read(&json_path)
|
||||
&& serde_json::from_slice::<serde_json::Value>(&data).is_ok()
|
||||
{
|
||||
if let Err(e) = table.insert(*key, data.as_slice()) {
|
||||
tracing::error!("Failed to insert migrated key '{}': {}", key, e);
|
||||
} else {
|
||||
let _ = std::fs::rename(&json_path, json_path.with_extension("json.migrated"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+17
-5
@@ -33,9 +33,16 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
|
||||
if a.is_empty() || b.is_empty() || a.len() != b.len() {
|
||||
return 0.0;
|
||||
}
|
||||
let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
|
||||
let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
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()) {
|
||||
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 {
|
||||
0.0
|
||||
} else {
|
||||
@@ -49,8 +56,13 @@ pub async fn generate_embeddings_async(texts: Vec<String>) -> Result<Vec<Vec<f32
|
||||
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(texts, None).map_err(|e| e.to_string())?;
|
||||
Ok(embeddings)
|
||||
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())?;
|
||||
all_embeddings.extend(chunk_embeddings);
|
||||
}
|
||||
Ok(all_embeddings)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| e.to_string())?
|
||||
|
||||
@@ -13,6 +13,15 @@ pub enum AppError {
|
||||
#[error("Not Found: {0}")]
|
||||
NotFound(String),
|
||||
|
||||
#[error("Entity '{name}' was not found in namespace '{namespace}'")]
|
||||
EntityNotFound { name: String, namespace: String },
|
||||
|
||||
#[error("Invalid relation between '{from}' and '{to}': {reason}")]
|
||||
InvalidRelation { from: String, to: String, reason: String },
|
||||
|
||||
#[error("Database error for key '{key}': {message}")]
|
||||
DatabaseError { key: String, message: String },
|
||||
|
||||
#[error("Forbidden: {0}")]
|
||||
Forbidden(String),
|
||||
|
||||
@@ -39,6 +48,18 @@ impl IntoResponse for AppError {
|
||||
fn into_response(self) -> Response {
|
||||
let (status, error_message) = match &self {
|
||||
AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()),
|
||||
AppError::EntityNotFound { name, namespace } => (
|
||||
StatusCode::NOT_FOUND,
|
||||
format!("Entity '{}' was not found in namespace '{}'", name, namespace),
|
||||
),
|
||||
AppError::InvalidRelation { from, to, reason } => (
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!("Invalid relation between '{}' and '{}': {}", from, to, reason),
|
||||
),
|
||||
AppError::DatabaseError { key, message } => (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("Database error for key '{}': {}", key, message),
|
||||
),
|
||||
AppError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg.clone()),
|
||||
AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()),
|
||||
AppError::Internal(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg.clone()),
|
||||
|
||||
+140
-15
@@ -56,7 +56,7 @@ impl McpTool for ReadFileSkeletonHandler {
|
||||
let tree = parser.parse(&code, None).ok_or_else(|| {
|
||||
crate::error::AppError::Internal("Failed to parse code".to_string())
|
||||
})?;
|
||||
let mut result_skeleton = String::new();
|
||||
let mut result_skeleton = String::with_capacity(code.len() / 2);
|
||||
|
||||
fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) {
|
||||
if depth > 128 {
|
||||
@@ -141,7 +141,7 @@ impl McpTool for ReplaceAstNodeHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ReplaceAstNodeTool>(
|
||||
"replace_ast_node",
|
||||
"Replace a specific AST node (e.g., function, struct) entirely using tree-sitter for robust structural editing.",
|
||||
"Replace a specific AST node (e.g., function, struct, enum, class, trait) entirely using tree-sitter for robust structural editing. Supported node_type values include: 'function_item' (or 'function'/'fn'/'method'), 'struct_item' (or 'struct'), 'class_declaration' (or 'class'), 'enum_item' (or 'enum'), 'trait_item' (or 'trait'/'interface'), 'type_alias_declaration' (or 'type').",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -185,6 +185,34 @@ impl McpTool for ReplaceAstNodeHandler {
|
||||
crate::error::AppError::Internal("Failed to parse code".to_string())
|
||||
})?;
|
||||
|
||||
fn matches_node_type(actual_kind: &str, requested_type: &str) -> bool {
|
||||
if actual_kind == requested_type {
|
||||
return true;
|
||||
}
|
||||
match requested_type.to_lowercase().as_str() {
|
||||
"function" | "func" | "fn" | "method" | "def" => matches!(
|
||||
actual_kind,
|
||||
"function_item"
|
||||
| "function_declaration"
|
||||
| "function_definition"
|
||||
| "method_definition"
|
||||
| "function"
|
||||
),
|
||||
"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,
|
||||
"trait_item" | "interface_declaration" | "interface_item"
|
||||
),
|
||||
"type" | "type_alias" => matches!(
|
||||
actual_kind,
|
||||
"type_alias_declaration" | "type_item" | "type_definition"
|
||||
),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
// Search for the node
|
||||
fn find_node<'a>(
|
||||
node: Node<'a>,
|
||||
@@ -192,12 +220,12 @@ impl McpTool for ReplaceAstNodeHandler {
|
||||
target_type: &str,
|
||||
target_name: &str,
|
||||
) -> Option<Node<'a>> {
|
||||
if node.kind() == target_type {
|
||||
if matches_node_type(node.kind(), target_type) {
|
||||
// Try to find the name/identifier
|
||||
let mut cursor = node.walk();
|
||||
for child in node.children(&mut cursor) {
|
||||
let kind = child.kind();
|
||||
if kind == "identifier" || kind == "name" {
|
||||
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);
|
||||
@@ -250,6 +278,61 @@ impl McpTool for ReplaceAstNodeHandler {
|
||||
}
|
||||
}
|
||||
|
||||
fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: bool) -> Vec<serde_json::Value> {
|
||||
let mut results = Vec::new();
|
||||
let cwd = match std::env::current_dir() {
|
||||
Ok(dir) => dir,
|
||||
Err(_) => return results,
|
||||
};
|
||||
|
||||
let walker = ignore::WalkBuilder::new(&cwd)
|
||||
.hidden(true)
|
||||
.git_ignore(true)
|
||||
.build();
|
||||
|
||||
let mut scanned_files = 0;
|
||||
|
||||
for result in walker {
|
||||
let entry = match result {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
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") {
|
||||
scanned_files += 1;
|
||||
if scanned_files > 500 {
|
||||
break;
|
||||
}
|
||||
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))
|
||||
} else {
|
||||
line.contains(target_sym)
|
||||
};
|
||||
|
||||
if is_match {
|
||||
results.push(serde_json::json!({
|
||||
"file_path": path.to_string_lossy(),
|
||||
"line": line_num + 1,
|
||||
"content": line.trim(),
|
||||
}));
|
||||
if results.len() >= limit {
|
||||
return results;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
results
|
||||
}
|
||||
|
||||
pub struct FindSymbolReferencesHandler;
|
||||
|
||||
#[async_trait]
|
||||
@@ -269,22 +352,36 @@ impl McpTool for FindSymbolReferencesHandler {
|
||||
let req: crate::tools::FindSymbolReferencesTool =
|
||||
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let limit = req.limit.unwrap_or(10);
|
||||
let target_sym = req.symbol;
|
||||
let target_sym = req.symbol.clone();
|
||||
|
||||
let matches = state.code.snippets.read_with(|snippets| {
|
||||
let mut matches = state.code.snippets.read_with(|snippets| {
|
||||
let mut refs = Vec::new();
|
||||
for snippet in snippets {
|
||||
if snippet.code.contains(&target_sym) || snippet.name.contains(&target_sym) {
|
||||
refs.push(snippet.clone());
|
||||
refs.push(serde_json::json!({
|
||||
"source": "snippet",
|
||||
"name": snippet.name,
|
||||
"code": snippet.code,
|
||||
}));
|
||||
if refs.len() >= limit {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&refs)?)
|
||||
Ok::<Vec<serde_json::Value>, crate::error::AppError>(refs)
|
||||
})?;
|
||||
|
||||
Ok(matches)
|
||||
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)
|
||||
})
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
matches.extend(disk_matches);
|
||||
}
|
||||
|
||||
Ok(serde_json::to_string_pretty(&matches)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -307,24 +404,38 @@ impl McpTool for GetCallersHandler {
|
||||
let req: crate::tools::GetCallersTool =
|
||||
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let limit = req.limit.unwrap_or(10);
|
||||
let target_fn = req.function_name;
|
||||
let target_fn = req.function_name.clone();
|
||||
|
||||
let callers = state.code.snippets.read_with(|snippets| {
|
||||
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))
|
||||
{
|
||||
matching.push(snippet.clone());
|
||||
matching.push(serde_json::json!({
|
||||
"source": "snippet",
|
||||
"name": snippet.name,
|
||||
"code": snippet.code,
|
||||
}));
|
||||
if matching.len() >= limit {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&matching)?)
|
||||
Ok::<Vec<serde_json::Value>, crate::error::AppError>(matching)
|
||||
})?;
|
||||
|
||||
Ok(callers)
|
||||
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)
|
||||
})
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
callers.extend(disk_callers);
|
||||
}
|
||||
|
||||
Ok(serde_json::to_string_pretty(&callers)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -357,6 +468,20 @@ impl McpTool for AnalyzeImpactHandler {
|
||||
}
|
||||
});
|
||||
|
||||
let sym_clone = sym.clone();
|
||||
let disk_refs = tokio::task::spawn_blocking(move || {
|
||||
scan_workspace_for_symbol(&sym_clone, 20, false)
|
||||
})
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
for r in &disk_refs {
|
||||
if let Some(path) = r.get("file_path").and_then(|p| p.as_str()) {
|
||||
let line = r.get("line").and_then(|l| l.as_u64()).unwrap_or(0);
|
||||
callers.push(format!("{}:{}", path, line));
|
||||
}
|
||||
}
|
||||
|
||||
let mut kg_connected = Vec::new();
|
||||
state.read_graph(|g| {
|
||||
for rel in &g.relations {
|
||||
@@ -384,7 +509,7 @@ impl McpTool for AnalyzeImpactHandler {
|
||||
"target_symbol": sym,
|
||||
"risk_level": risk_level,
|
||||
"ast_callers_count": caller_count,
|
||||
"ast_callers_sample": callers.into_iter().take(5).collect::<Vec<_>>(),
|
||||
"ast_callers_sample": callers.into_iter().take(10).collect::<Vec<_>>(),
|
||||
"graph_relations_count": graph_count,
|
||||
"graph_relations": kg_connected,
|
||||
"recommendation": match risk_level {
|
||||
|
||||
@@ -16,7 +16,7 @@ impl McpTool for UpdateEnvFingerprintHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<UpdateEnvFingerprintTool>(
|
||||
"update_env_fingerprint",
|
||||
"Execute update_env_fingerprint",
|
||||
"Update system OS, shell, and installed tool version fingerprints for a workspace namespace.",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -50,7 +50,7 @@ impl McpTool for ReadEnvFingerprintHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ReadEnvFingerprintTool>(
|
||||
"read_env_fingerprint",
|
||||
"Execute read_env_fingerprint",
|
||||
"Read system OS, shell, and tool version fingerprints for a workspace namespace.",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -81,7 +81,7 @@ impl McpTool for LogEnvRequirementHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<LogEnvRequirementTool>(
|
||||
"log_env_requirement",
|
||||
"Execute log_env_requirement",
|
||||
"Log required environment variables, configuration keys, or secrets for a namespace.",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -111,7 +111,7 @@ impl McpTool for RegisterEnvironmentHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<RegisterEnvironmentTool>(
|
||||
"register_environment",
|
||||
"Execute register_environment",
|
||||
"Register infrastructure environment details (URL, VPN requirements, notes) for a namespace.",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -144,7 +144,7 @@ impl McpTool for GetEnvironmentDetailsHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<GetEnvironmentDetailsTool>(
|
||||
"get_environment_details",
|
||||
"Execute get_environment_details",
|
||||
"Retrieve registered infrastructure environment details for a namespace.",
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -69,7 +69,8 @@ impl McpTool for GetActiveWorktreeContextHandler {
|
||||
|
||||
// Truncate diff output if it's too large to save tokens
|
||||
if diff_output.len() > 10000 {
|
||||
diff_output.truncate(10000);
|
||||
let valid_boundary = diff_output.floor_char_boundary(10000);
|
||||
diff_output.truncate(valid_boundary);
|
||||
diff_output.push_str("\n... [Diff truncated due to size]");
|
||||
}
|
||||
|
||||
|
||||
+213
-122
@@ -12,6 +12,112 @@ struct BorrowedGraph<'a> {
|
||||
relations: Vec<&'a crate::models::Relation>,
|
||||
}
|
||||
|
||||
pub struct GraphQueryBuilder<'a> {
|
||||
graph: &'a crate::models::KnowledgeGraph,
|
||||
max_depth: usize,
|
||||
relation_filter: Option<&'a str>,
|
||||
}
|
||||
|
||||
impl<'a> GraphQueryBuilder<'a> {
|
||||
pub fn new(graph: &'a crate::models::KnowledgeGraph) -> Self {
|
||||
Self {
|
||||
graph,
|
||||
max_depth: 5,
|
||||
relation_filter: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn max_depth(mut self, depth: usize) -> Self {
|
||||
self.max_depth = depth;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn relation_filter(mut self, filter: Option<&'a str>) -> Self {
|
||||
self.relation_filter = filter;
|
||||
self
|
||||
}
|
||||
|
||||
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 mut queue = std::collections::VecDeque::new();
|
||||
let mut visited = std::collections::HashSet::new();
|
||||
let mut parents = std::collections::HashMap::new();
|
||||
|
||||
queue.push_back(start);
|
||||
visited.insert(start);
|
||||
|
||||
let mut found = false;
|
||||
let mut current_depth = 0;
|
||||
let mut nodes_at_current_depth = 1;
|
||||
let mut nodes_at_next_depth = 0;
|
||||
|
||||
while let Some(current) = queue.pop_front() {
|
||||
if current == end {
|
||||
found = true;
|
||||
break;
|
||||
}
|
||||
// Visited node upper-bound cap to guarantee deterministic BFS bounds on dense graphs
|
||||
if visited.len() > 10_000 {
|
||||
break;
|
||||
}
|
||||
nodes_at_current_depth -= 1;
|
||||
if current_depth < self.max_depth {
|
||||
if let Some(neighbors) = adj.get(current) {
|
||||
for &(neighbor, rel_type, is_inverse) in neighbors {
|
||||
if !visited.contains(neighbor) {
|
||||
visited.insert(neighbor);
|
||||
parents.insert(neighbor, (current, rel_type, is_inverse));
|
||||
queue.push_back(neighbor);
|
||||
nodes_at_next_depth += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if nodes_at_current_depth == 0 {
|
||||
current_depth += 1;
|
||||
nodes_at_current_depth = nodes_at_next_depth;
|
||||
nodes_at_next_depth = 0;
|
||||
}
|
||||
}
|
||||
|
||||
if found {
|
||||
let mut path = Vec::new();
|
||||
let mut curr = end;
|
||||
while curr != start {
|
||||
if let Some((parent, rel_type, is_inverse)) = parents.get(&curr) {
|
||||
if *is_inverse {
|
||||
path.push(format!("{} -[inverse({})]-> {}", parent, rel_type, curr));
|
||||
} else {
|
||||
path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
|
||||
}
|
||||
curr = parent;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
path.reverse();
|
||||
Some(path)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct QueryGraphPathHandler;
|
||||
|
||||
#[async_trait]
|
||||
@@ -21,88 +127,32 @@ impl McpTool for QueryGraphPathHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<QueryGraphPathTool>("query_graph_path", "Execute query_graph_path")
|
||||
crate::mcp::tool_def::<QueryGraphPathTool>(
|
||||
"query_graph_path",
|
||||
"Find the shortest relationship path between two entities in the knowledge graph within a maximum depth.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: crate::tools::QueryGraphPathTool =
|
||||
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
state.read_graph(|graph| {
|
||||
let max_depth = req.max_depth.unwrap_or(5);
|
||||
|
||||
// Pre-index relations into an adjacency map for O(1) neighbor lookups
|
||||
let mut adj: std::collections::HashMap<&str, Vec<(&str, &str, bool)>> = std::collections::HashMap::new();
|
||||
for rel in &graph.relations {
|
||||
adj.entry(rel.from.as_str())
|
||||
.or_default()
|
||||
.push((rel.to.as_str(), rel.relation_type.as_str(), false));
|
||||
adj.entry(rel.to.as_str())
|
||||
.or_default()
|
||||
.push((rel.from.as_str(), rel.relation_type.as_str(), true));
|
||||
}
|
||||
|
||||
let mut queue: std::collections::VecDeque<&str> = std::collections::VecDeque::new();
|
||||
let mut visited: std::collections::HashSet<&str> = std::collections::HashSet::new();
|
||||
let mut parents: std::collections::HashMap<&str, (&str, &str, bool)> =
|
||||
std::collections::HashMap::new();
|
||||
|
||||
queue.push_back(req.start_node.as_str());
|
||||
visited.insert(req.start_node.as_str());
|
||||
|
||||
let mut found = false;
|
||||
let mut current_depth = 0;
|
||||
let mut nodes_at_current_depth = 1;
|
||||
let mut nodes_at_next_depth = 0;
|
||||
|
||||
while let Some(current) = queue.pop_front() {
|
||||
if current == req.end_node {
|
||||
found = true;
|
||||
break;
|
||||
tokio::task::spawn_blocking(move || {
|
||||
state.read_graph(|graph| {
|
||||
let max_depth = req.max_depth.unwrap_or(5);
|
||||
let builder = GraphQueryBuilder::new(graph).max_depth(max_depth as usize);
|
||||
if let Some(path) = builder.find_shortest_path(&req.start_node, &req.end_node) {
|
||||
Ok(format!("Path found:\n{}", path.join("\n")))
|
||||
} else {
|
||||
Ok(format!(
|
||||
"No path found between {} and {} within depth {}",
|
||||
req.start_node, req.end_node, max_depth
|
||||
))
|
||||
}
|
||||
nodes_at_current_depth -= 1;
|
||||
if current_depth < max_depth {
|
||||
if let Some(neighbors) = adj.get(current) {
|
||||
for &(neighbor, rel_type, is_inverse) in neighbors {
|
||||
if !visited.contains(neighbor) {
|
||||
visited.insert(neighbor);
|
||||
parents.insert(neighbor, (current, rel_type, is_inverse));
|
||||
queue.push_back(neighbor);
|
||||
nodes_at_next_depth += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if nodes_at_current_depth == 0 {
|
||||
current_depth += 1;
|
||||
nodes_at_current_depth = nodes_at_next_depth;
|
||||
nodes_at_next_depth = 0;
|
||||
}
|
||||
}
|
||||
|
||||
if found {
|
||||
let mut path = Vec::new();
|
||||
let mut curr = req.end_node.as_str();
|
||||
while curr != req.start_node {
|
||||
if let Some((parent, rel_type, is_inverse)) = parents.get(&curr) {
|
||||
if *is_inverse {
|
||||
path.push(format!("{} -[inverse({})]-> {}", parent, rel_type, curr));
|
||||
} else {
|
||||
path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
|
||||
}
|
||||
curr = parent;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
path.reverse();
|
||||
Ok(format!("Path found:\n{}", path.join("\n")))
|
||||
} else {
|
||||
Ok(format!(
|
||||
"No path found between {} and {} within depth {}",
|
||||
req.start_node, req.end_node, max_depth
|
||||
))
|
||||
}
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(format!("Graph traversal task failed: {}", e)))?
|
||||
}
|
||||
}
|
||||
|
||||
@@ -115,7 +165,10 @@ impl McpTool for CreateEntitiesHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Execute create_entities")
|
||||
crate::mcp::tool_def::<CreateEntitiesTool>(
|
||||
"create_entities",
|
||||
"Create new entities in the knowledge graph with normalized entity types.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -130,11 +183,12 @@ impl McpTool for CreateEntitiesHandler {
|
||||
}
|
||||
}
|
||||
});
|
||||
let idx = state.get_search_index();
|
||||
for entity in inserted {
|
||||
drop(idx.index_entity(&entity));
|
||||
let names: Vec<String> = inserted.iter().map(|e| format!("{} ({})", e.name, e.entity_type)).collect();
|
||||
if !inserted.is_empty() {
|
||||
let idx = state.get_search_index().await;
|
||||
let _ = idx.index_entities_batch(&inserted).await;
|
||||
}
|
||||
Ok("Entities created".to_string())
|
||||
Ok(format!("Successfully created {} entity/entities: {}", names.len(), names.join(", ")))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -147,7 +201,10 @@ impl McpTool for CreateRelationsHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<CreateRelationsTool>("create_relations", "Execute create_relations")
|
||||
crate::mcp::tool_def::<CreateRelationsTool>(
|
||||
"create_relations",
|
||||
"Create directed relationships between entities in the knowledge graph. Requires 'from', 'to', and 'relation_type'.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -183,23 +240,38 @@ impl McpTool for CreateRelationsHandler {
|
||||
}
|
||||
});
|
||||
|
||||
if !missing_nodes.is_empty() {
|
||||
let missing: Vec<_> = missing_nodes.into_iter().collect();
|
||||
return Err(crate::error::AppError::Internal(format!(
|
||||
"Error: Relations dropped due to missing entities: {}",
|
||||
missing.join(", ")
|
||||
)));
|
||||
}
|
||||
|
||||
let mut auto_created = Vec::new();
|
||||
let mut added_relations = Vec::new();
|
||||
state.modify_graph(|g| {
|
||||
for node_name in missing_nodes {
|
||||
if !g.entities.contains_key(&node_name) {
|
||||
g.entities.insert(
|
||||
node_name.clone(),
|
||||
crate::models::Entity {
|
||||
name: node_name.clone(),
|
||||
entity_type: "Entity".to_string(),
|
||||
observations: vec!["Auto-created stub entity for relation endpoint".to_string()],
|
||||
namespace: crate::models::default_namespace(),
|
||||
git_branch: None,
|
||||
},
|
||||
);
|
||||
auto_created.push(node_name);
|
||||
}
|
||||
}
|
||||
for mut relation in req.relations {
|
||||
if !relation.from.is_empty() && !relation.to.is_empty() {
|
||||
relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type);
|
||||
added_relations.push(format!("{} -[{}]-> {}", relation.from, relation.relation_type, relation.to));
|
||||
g.relations.push(relation);
|
||||
}
|
||||
}
|
||||
});
|
||||
Ok("Relations created".to_string())
|
||||
|
||||
let mut msg = format!("Successfully created {} relation(s):\n{}", added_relations.len(), added_relations.join("\n"));
|
||||
if !auto_created.is_empty() {
|
||||
msg.push_str(&format!("\nNote: Auto-created {} missing stub entity/entities: {}", auto_created.len(), auto_created.join(", ")));
|
||||
}
|
||||
Ok(msg)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -212,7 +284,10 @@ impl McpTool for AddObservationsHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<AddObservationsTool>("add_observations", "Execute add_observations")
|
||||
crate::mcp::tool_def::<AddObservationsTool>(
|
||||
"add_observations",
|
||||
"Add new observations and factual statements to existing entities in the knowledge graph.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -284,7 +359,7 @@ impl McpTool for DeleteEntitiesHandler {
|
||||
.retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
|
||||
});
|
||||
|
||||
let idx = state.get_search_index();
|
||||
let idx = state.get_search_index().await;
|
||||
for name in to_delete {
|
||||
drop(idx.delete_document(&name));
|
||||
}
|
||||
@@ -361,7 +436,7 @@ impl McpTool for DeleteRelationsHandler {
|
||||
});
|
||||
if missing_count > 0 {
|
||||
return Err(crate::error::AppError::Internal(format!(
|
||||
"Error: {} relation(s) not found in graph. Please verify exact relation properties using read_graph.",
|
||||
"Error: {} relation(s) not found in graph. Please verify exact relation properties (from, to, relation_type) using read_graph or get_subgraph.",
|
||||
missing_count
|
||||
)));
|
||||
}
|
||||
@@ -378,7 +453,10 @@ impl McpTool for ReadGraphHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ReadGraphTool>("read_graph", "Execute read_graph")
|
||||
crate::mcp::tool_def::<ReadGraphTool>(
|
||||
"read_graph",
|
||||
"Read entities and relations from the knowledge graph with optional namespace filtering and token truncation. For large graphs, specify 'namespace' or use 'search_nodes' or 'get_subgraph' for targeted discovery.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -405,8 +483,9 @@ impl McpTool for ReadGraphHandler {
|
||||
if let Some(max_tok) = max_tokens {
|
||||
let max_chars = max_tok * 4;
|
||||
if result_json.len() > max_chars {
|
||||
result_json.truncate(max_chars);
|
||||
result_json.push_str("... [TRUNCATED_TO_MAX_TOKENS]");
|
||||
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]");
|
||||
}
|
||||
}
|
||||
Ok(result_json)
|
||||
@@ -432,12 +511,10 @@ impl McpTool for SearchNodesHandler {
|
||||
let limit = req.limit.unwrap_or(10);
|
||||
let include_body = req.include_body.unwrap_or(false);
|
||||
|
||||
let matches = if let Ok(idx) = state.search_index.read() {
|
||||
idx.search(&req.query, req.namespace.as_deref())
|
||||
.unwrap_or_default()
|
||||
} else {
|
||||
vec![]
|
||||
};
|
||||
let idx = state.get_search_index().await;
|
||||
let matches = idx
|
||||
.search(&req.query, req.namespace.as_deref())
|
||||
.unwrap_or_default();
|
||||
|
||||
let data = state.read_graph(|full| -> crate::error::Result<String> {
|
||||
let mut matched_entities = Vec::new();
|
||||
@@ -709,6 +786,11 @@ impl McpTool for MergeEntitiesHandler {
|
||||
r.to = req.target_entity.clone();
|
||||
}
|
||||
|
||||
// Filter out self-loops
|
||||
if r.from == r.to {
|
||||
return false;
|
||||
}
|
||||
|
||||
if r.from == req.target_entity || r.to == req.target_entity {
|
||||
seen.insert(r.clone())
|
||||
} else {
|
||||
@@ -908,16 +990,17 @@ impl McpTool for SweepGraphHealthHandler {
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Compute similarity pairs for duplicate detection
|
||||
// 2. Compute similarity pairs for duplicate detection using pre-computed lowercase names
|
||||
let names: Vec<_> = g.entities.keys().cloned().collect();
|
||||
let lower_names: Vec<String> = names.iter().map(|n| n.to_lowercase()).collect();
|
||||
for i in 0..names.len() {
|
||||
for j in (i + 1)..names.len() {
|
||||
let n1 = &names[i];
|
||||
let n2 = &names[j];
|
||||
let l1 = &lower_names[i];
|
||||
let l2 = &lower_names[j];
|
||||
|
||||
let l1 = n1.to_lowercase();
|
||||
let l2 = n2.to_lowercase();
|
||||
if l1 == l2 || ((l1.contains(&l2) || l2.contains(&l1)) && l1.len().min(l2.len()) > 3) {
|
||||
if l1 == l2 || ((l1.contains(l2.as_str()) || l2.contains(l1.as_str())) && l1.len().min(l2.len()) > 3) {
|
||||
duplicates.push(serde_json::json!({
|
||||
"entity_a": n1,
|
||||
"entity_b": n2,
|
||||
@@ -1050,7 +1133,8 @@ impl McpTool for SummarizeSubgraphHandler {
|
||||
let max_tokens = req.max_tokens.unwrap_or(1000);
|
||||
let max_chars = max_tokens * 4;
|
||||
if markdown.len() > max_chars {
|
||||
markdown.truncate(max_chars);
|
||||
let valid_boundary = markdown.floor_char_boundary(max_chars);
|
||||
markdown.truncate(valid_boundary);
|
||||
markdown.push_str("\n... [Truncated to fit token budget]");
|
||||
}
|
||||
|
||||
@@ -1063,13 +1147,11 @@ mod tests {
|
||||
use super::*;
|
||||
use crate::handlers::meta::{BroadcastAgentSignalHandler, QueryAgentSignalsHandler};
|
||||
use serde_json::json;
|
||||
use tempfile::tempdir;
|
||||
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_and_read_entities() {
|
||||
let dir = tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
let state = Arc::new(MemoryState::new_in_memory());
|
||||
|
||||
let create_handler = CreateEntitiesHandler;
|
||||
let args = json!({
|
||||
@@ -1083,7 +1165,7 @@ mod tests {
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
||||
.unwrap();
|
||||
assert_eq!(res, "Entities created");
|
||||
assert!(res.contains("Successfully created 1 entity/entities"));
|
||||
|
||||
// Ensure graph contains the entity
|
||||
state.graph.read_with(|g| {
|
||||
@@ -1094,8 +1176,7 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_relations() {
|
||||
let dir = tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
let state = Arc::new(MemoryState::new_in_memory());
|
||||
|
||||
// Needs entities first
|
||||
state.graph.modify(|g| {
|
||||
@@ -1132,14 +1213,26 @@ mod tests {
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
||||
.unwrap();
|
||||
assert_eq!(res, "Relations created");
|
||||
assert!(res.contains("Successfully created 1 relation(s)"));
|
||||
|
||||
// Test semantic LLM schema feedback (User request)
|
||||
let bad_args = json!({
|
||||
// Test serde field aliases (source/target/relationType mapped to from/to/relation_type)
|
||||
let alias_args = json!({
|
||||
"relations": [
|
||||
{"source": "A", "target": "B", "relationType": "knows"}
|
||||
]
|
||||
});
|
||||
let alias_res = handler
|
||||
.execute(alias_args, state.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(alias_res.contains("Successfully created 1 relation(s)"));
|
||||
|
||||
// Test semantic LLM schema feedback on missing fields
|
||||
let bad_args = json!({
|
||||
"relations": [
|
||||
{"invalid_field": "X"}
|
||||
]
|
||||
});
|
||||
let err_res = handler
|
||||
.execute(bad_args, state.clone())
|
||||
.await
|
||||
@@ -1151,8 +1244,7 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_observations_and_reads() {
|
||||
let dir = tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
let state = Arc::new(MemoryState::new_in_memory());
|
||||
|
||||
// Need entity first
|
||||
state.graph.modify(|g| {
|
||||
@@ -1208,8 +1300,7 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_advanced_graph_operations() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
let state = Arc::new(MemoryState::new_in_memory());
|
||||
|
||||
let create_handler = CreateEntitiesHandler;
|
||||
let args_ent = json!({
|
||||
@@ -1289,13 +1380,13 @@ mod tests {
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
||||
.unwrap();
|
||||
assert!(!res_orphans.contains("Y"));
|
||||
// After merging X into Y and purging self-loops, Y is the sole node and becomes an orphan
|
||||
assert!(res_orphans.contains("Y"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_more_graph_handlers() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
let state = Arc::new(MemoryState::new_in_memory());
|
||||
|
||||
let create_handler = CreateEntitiesHandler;
|
||||
let args_ent = json!({
|
||||
|
||||
@@ -64,10 +64,12 @@ impl McpTool for GetRecentLogsHandler {
|
||||
file.seek(SeekFrom::End(-(read_size as i64)))
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
let mut buffer = String::new();
|
||||
file.read_to_string(&mut buffer)
|
||||
let mut vec_buf = Vec::new();
|
||||
file.read_to_end(&mut vec_buf)
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
let buffer = String::from_utf8_lossy(&vec_buf).to_string();
|
||||
|
||||
let lines: Vec<&str> = buffer.lines().collect();
|
||||
let recent_lines = if lines.len() > 100 {
|
||||
lines[lines.len() - 100..].join("\n")
|
||||
|
||||
+99
-28
@@ -15,13 +15,16 @@ impl McpTool for LogDecisionHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<LogDecisionTool>("log_decision", "Execute log_decision")
|
||||
crate::mcp::tool_def::<LogDecisionTool>(
|
||||
"log_decision",
|
||||
"Log an architectural decision record (ADR) with title, context, and consequences",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: LogDecisionTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
|
||||
let idx = state.get_search_index();
|
||||
let idx = state.get_search_index().await;
|
||||
let mut final_id = String::new();
|
||||
let mut adrs_to_index = Vec::new();
|
||||
|
||||
@@ -71,11 +74,16 @@ impl McpTool for QueryDecisionsHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<QueryDecisionsTool>("query_decisions", "Execute query_decisions")
|
||||
crate::mcp::tool_def::<QueryDecisionsTool>(
|
||||
"query_decisions",
|
||||
"Query or search existing architectural decision records (ADRs)",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: QueryDecisionsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let limit = req.limit.unwrap_or(20);
|
||||
let include_body = req.include_body.unwrap_or(true);
|
||||
let data = state.code.adrs.read_with(|adrs| {
|
||||
let filtered: Vec<_> = adrs
|
||||
.iter()
|
||||
@@ -88,8 +96,20 @@ impl McpTool for QueryDecisionsHandler {
|
||||
true
|
||||
}
|
||||
})
|
||||
.take(limit)
|
||||
.collect();
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
|
||||
|
||||
if include_body {
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
|
||||
} else {
|
||||
let compact: Vec<_> = filtered.iter().map(|a| serde_json::json!({
|
||||
"id": a.id,
|
||||
"title": a.title,
|
||||
"status": a.status,
|
||||
"timestamp": a.timestamp,
|
||||
})).collect();
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&compact)?)
|
||||
}
|
||||
})?;
|
||||
Ok(data)
|
||||
}
|
||||
@@ -122,7 +142,8 @@ impl McpTool for DeleteDecisionHandler {
|
||||
});
|
||||
|
||||
if found {
|
||||
state.rebuild_index().await;
|
||||
let idx = state.get_search_index().await;
|
||||
let _ = idx.delete_document(&req.id).await;
|
||||
Ok("Decision deleted successfully".to_string())
|
||||
} else {
|
||||
Err(crate::error::AppError::Internal(
|
||||
@@ -141,7 +162,10 @@ impl McpTool for LogErrorFixHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<LogErrorFixTool>("log_error_fix", "Execute log_error_fix")
|
||||
crate::mcp::tool_def::<LogErrorFixTool>(
|
||||
"log_error_fix",
|
||||
"Log an error signature and its verified solution/fix for future diagnostic retrieval",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -288,7 +312,10 @@ impl McpTool for LogCodeChangeHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<LogCodeChangeTool>("log_code_change", "Execute log_code_change")
|
||||
crate::mcp::tool_def::<LogCodeChangeTool>(
|
||||
"log_code_change",
|
||||
"Log a significant code change or refactor with file path and description",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -325,6 +352,9 @@ impl McpTool for LogCodeChangeHandler {
|
||||
git_commit: req.git_commit,
|
||||
git_branch: req.git_branch,
|
||||
});
|
||||
if ledger.len() > 1000 {
|
||||
ledger.remove(0);
|
||||
}
|
||||
});
|
||||
state.record_activity("code_change", &format!("Modified {}", req.file_path), Some(&description));
|
||||
Ok(format!("Logged code change for {}: {}", req.file_path, description))
|
||||
@@ -342,7 +372,7 @@ impl McpTool for QueryRecentChangesHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<QueryRecentChangesTool>(
|
||||
"query_recent_changes",
|
||||
"Execute query_recent_changes",
|
||||
"Query recent code changes and refactoring audit logs",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -364,7 +394,10 @@ impl McpTool for LearnPreferenceHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<LearnPreferenceTool>("learn_preference", "Execute learn_preference")
|
||||
crate::mcp::tool_def::<LearnPreferenceTool>(
|
||||
"learn_preference",
|
||||
"Learn or update a key-value user preference",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -392,7 +425,10 @@ impl McpTool for ReadPreferencesHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ReadPreferencesTool>("read_preferences", "Execute read_preferences")
|
||||
crate::mcp::tool_def::<ReadPreferencesTool>(
|
||||
"read_preferences",
|
||||
"Read all stored user and environment preferences",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -412,7 +448,10 @@ impl McpTool for LogTechDebtHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<LogTechDebtTool>("log_tech_debt", "Execute log_tech_debt")
|
||||
crate::mcp::tool_def::<LogTechDebtTool>(
|
||||
"log_tech_debt",
|
||||
"Log a new technical debt item with description and ideal solution",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -452,7 +491,7 @@ impl McpTool for ResolveTechDebtHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ResolveTechDebtTool>(
|
||||
"resolve_tech_debt",
|
||||
"Execute resolve_tech_debt",
|
||||
"Mark an existing technical debt item as resolved by ID",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -488,7 +527,10 @@ impl McpTool for ListTechDebtHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ListTechDebtTool>("list_tech_debt", "Execute list_tech_debt")
|
||||
crate::mcp::tool_def::<ListTechDebtTool>(
|
||||
"list_tech_debt",
|
||||
"List technical debt items filtered by namespace or resolution status",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -641,7 +683,10 @@ impl McpTool for OmniSearchHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<OmniSearchTool>("omni_search", "Execute omni_search")
|
||||
crate::mcp::tool_def::<OmniSearchTool>(
|
||||
"omni_search",
|
||||
"Unified search across entities, subgraphs, tasks, code snippets, ADRs, and technical debt",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -649,7 +694,7 @@ impl McpTool for OmniSearchHandler {
|
||||
let limit = req.limit.unwrap_or(5);
|
||||
let include_body = req.include_body.unwrap_or(false);
|
||||
|
||||
let idx = state.get_search_index();
|
||||
let idx = state.get_search_index().await;
|
||||
let keyword_matches = idx
|
||||
.search(&req.query, req.namespace.as_deref())
|
||||
.unwrap_or_default();
|
||||
@@ -717,7 +762,7 @@ impl McpTool for OmniSearchHandler {
|
||||
}
|
||||
|
||||
let mut ranked_items: Vec<_> = rrf_scores.into_values().collect();
|
||||
ranked_items.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
|
||||
ranked_items.sort_by(|a, b| b.0.total_cmp(&a.0));
|
||||
|
||||
let matches: Vec<MatchItem> =
|
||||
ranked_items.into_iter().map(|(_, item)| item).collect();
|
||||
@@ -982,7 +1027,7 @@ impl McpTool for GetProjectHealthHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<GetProjectHealthTool>(
|
||||
"get_project_health",
|
||||
"Execute get_project_health",
|
||||
"Retrieve project health metrics including active tasks, technical debt, and PR checklist progress",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -1037,7 +1082,10 @@ impl McpTool for SuggestErrorFixHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<SuggestErrorFixTool>("suggest_error_fix", "Execute suggest_error_fix")
|
||||
crate::mcp::tool_def::<SuggestErrorFixTool>(
|
||||
"suggest_error_fix",
|
||||
"Suggest historical error fixes based on stack trace vector similarity",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -1066,7 +1114,7 @@ impl McpTool for SuggestErrorFixHandler {
|
||||
.filter(|(_, score)| *score > 0.2)
|
||||
.collect();
|
||||
|
||||
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||||
scored.sort_by(|a, b| b.1.total_cmp(&a.1));
|
||||
|
||||
let suggestions: Vec<_> = scored.into_iter().take(limit).map(|(f, score)| {
|
||||
serde_json::json!({
|
||||
@@ -1214,7 +1262,10 @@ impl McpTool for CheckpointStateHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<CheckpointStateTool>("checkpoint_state", "Execute checkpoint_state")
|
||||
crate::mcp::tool_def::<CheckpointStateTool>(
|
||||
"checkpoint_state",
|
||||
"Save point-in-time memory state snapshot checkpoint",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -1248,7 +1299,10 @@ impl McpTool for RestoreStateHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<RestoreStateTool>("restore_state", "Execute restore_state")
|
||||
crate::mcp::tool_def::<RestoreStateTool>(
|
||||
"restore_state",
|
||||
"Restore memory state from a point-in-time snapshot checkpoint",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -1290,7 +1344,10 @@ impl McpTool for QueryLineageHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<QueryLineageTool>("query_lineage", "Execute query_lineage")
|
||||
crate::mcp::tool_def::<QueryLineageTool>(
|
||||
"query_lineage",
|
||||
"Query historical lineage and timeline of tasks, ADRs, and code changes",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -1375,7 +1432,10 @@ impl McpTool for GetNextActionableTasksHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<GetNextActionableTasksTool>("get_next_actionable_tasks", "Execute get_next_actionable_tasks")
|
||||
crate::mcp::tool_def::<GetNextActionableTasksTool>(
|
||||
"get_next_actionable_tasks",
|
||||
"Get unblocked pending tasks ready for execution",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -1431,7 +1491,10 @@ impl McpTool for LogHypothesisHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<LogHypothesisTool>("log_hypothesis", "Execute log_hypothesis")
|
||||
crate::mcp::tool_def::<LogHypothesisTool>(
|
||||
"log_hypothesis",
|
||||
"Log a diagnostic hypothesis and associated evidence for a task",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -1463,7 +1526,10 @@ impl McpTool for QueryHypothesesHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<QueryHypothesesTool>("query_hypotheses", "Execute query_hypotheses")
|
||||
crate::mcp::tool_def::<QueryHypothesesTool>(
|
||||
"query_hypotheses",
|
||||
"Query active diagnostic hypotheses and evidence by task ID or keyword",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
@@ -1701,7 +1767,12 @@ impl McpTool for AutoSessionCheckpointHandler {
|
||||
timestamp,
|
||||
};
|
||||
|
||||
state.telemetry.handoff_memos.modify(|m| m.push(memo));
|
||||
state.telemetry.handoff_memos.modify(|m| {
|
||||
m.push(memo);
|
||||
if m.len() > 100 {
|
||||
m.remove(0);
|
||||
}
|
||||
});
|
||||
state.record_activity("checkpoint", &format!("Created auto session checkpoint {}", memo_id), None);
|
||||
|
||||
Ok(format!("Session checkpoint created with memo ID '{}'.", memo_id))
|
||||
@@ -1759,7 +1830,7 @@ impl McpTool for SearchSnippetsHybridHandler {
|
||||
.collect();
|
||||
|
||||
scored.sort_by(|a, b| {
|
||||
b["score"].as_f64().unwrap_or(0.0).partial_cmp(&a["score"].as_f64().unwrap_or(0.0)).unwrap_or(std::cmp::Ordering::Equal)
|
||||
b["score"].as_f64().unwrap_or(0.0).total_cmp(&a["score"].as_f64().unwrap_or(0.0))
|
||||
});
|
||||
|
||||
let limit = req.limit.unwrap_or(10);
|
||||
@@ -1974,7 +2045,7 @@ mod tests {
|
||||
}
|
||||
|
||||
state.rebuild_index().await;
|
||||
state.get_search_index().reader.reload().unwrap();
|
||||
state.get_search_index().await.reader.reload().unwrap();
|
||||
|
||||
let omni = OmniSearchHandler;
|
||||
let omni_res = omni
|
||||
|
||||
@@ -49,17 +49,21 @@ impl McpTool for StickyNotesHandler {
|
||||
}
|
||||
StickyNoteAction::Read => {
|
||||
let now = crate::handlers::utils::now_secs();
|
||||
let mut active_notes = Vec::new();
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.retain(|n| {
|
||||
if let Some(exp) = n.expires_at {
|
||||
exp > now
|
||||
} else {
|
||||
true
|
||||
}
|
||||
});
|
||||
active_notes = notes.clone();
|
||||
let has_expired = state.code.sticky.read_with(|notes| {
|
||||
notes.iter().any(|n| matches!(n.expires_at, Some(exp) if exp <= now))
|
||||
});
|
||||
if has_expired {
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.retain(|n| {
|
||||
if let Some(exp) = n.expires_at {
|
||||
exp > now
|
||||
} else {
|
||||
true
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
let active_notes = state.code.sticky.read_with(|notes| notes.clone());
|
||||
Ok(serde_json::to_string(&active_notes)?)
|
||||
}
|
||||
StickyNoteAction::Delete => {
|
||||
@@ -68,8 +72,9 @@ impl McpTool for StickyNotesHandler {
|
||||
})?;
|
||||
let mut success = false;
|
||||
state.code.sticky.modify(|notes| {
|
||||
if idx > 0 && idx <= notes.len() {
|
||||
notes.remove(idx - 1);
|
||||
let target_pos = if idx == 0 { 0 } else { idx - 1 };
|
||||
if target_pos < notes.len() {
|
||||
notes.remove(target_pos);
|
||||
success = true;
|
||||
}
|
||||
});
|
||||
@@ -77,7 +82,7 @@ impl McpTool for StickyNotesHandler {
|
||||
Ok("Sticky note deleted.".to_string())
|
||||
} else {
|
||||
Err(crate::error::AppError::Internal(
|
||||
"Invalid sticky note index. Next step: Call sticky_notes with action='read' to view active notes and valid 1-based indices, then retry.".to_string(),
|
||||
"Invalid sticky note index. Next step: Call sticky_notes with action='read' to view active notes and valid indices, then retry.".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
@@ -214,17 +219,22 @@ impl McpTool for GenerateStandupReportHandler {
|
||||
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let cutoff = crate::handlers::utils::now_secs().saturating_sub(req.hours_lookback * 3600);
|
||||
|
||||
let report_str = state.project.tasks.read_with(|items| {
|
||||
state.code.ledger.read_with(|changes| {
|
||||
state.telemetry.session_summaries.read_with(|summaries| {
|
||||
let filtered_tasks: Vec<_> = items.iter().filter(|t| t.updated_at >= cutoff).collect();
|
||||
let filtered_changes: Vec<_> = changes.iter().filter(|c| c.timestamp >= cutoff).collect();
|
||||
let filtered_summaries: Vec<_> = summaries.iter().filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff).collect();
|
||||
let report = serde_json::json!({ "tasks_updated": filtered_tasks, "code_changes": filtered_changes, "session_summaries": filtered_summaries });
|
||||
report.to_string()
|
||||
})
|
||||
})
|
||||
let filtered_tasks = state.project.tasks.read_with(|items| {
|
||||
items.iter().filter(|t| t.updated_at >= cutoff).cloned().collect::<Vec<_>>()
|
||||
});
|
||||
let filtered_changes = state.code.ledger.read_with(|changes| {
|
||||
changes.iter().filter(|c| c.timestamp >= cutoff).cloned().collect::<Vec<_>>()
|
||||
});
|
||||
let filtered_summaries = state.telemetry.session_summaries.read_with(|summaries| {
|
||||
summaries.iter().filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff).cloned().collect::<Vec<_>>()
|
||||
});
|
||||
|
||||
let report = serde_json::json!({
|
||||
"tasks_updated": filtered_tasks,
|
||||
"code_changes": filtered_changes,
|
||||
"session_summaries": filtered_summaries
|
||||
});
|
||||
let report_str = report.to_string();
|
||||
Ok(report_str)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,7 +46,7 @@ impl McpTool for TasksHandler {
|
||||
acceptance_criteria: vec![],
|
||||
expires_at: None,
|
||||
};
|
||||
let idx = state.get_search_index();
|
||||
let idx = state.get_search_index().await;
|
||||
drop(idx.index_task(&task));
|
||||
state.project.tasks.modify(|tasks| {
|
||||
tasks.push(task.clone());
|
||||
@@ -187,7 +187,7 @@ impl McpTool for TasksHandler {
|
||||
});
|
||||
|
||||
if deleted_count > 0 {
|
||||
let idx = state.get_search_index();
|
||||
let idx = state.get_search_index().await;
|
||||
for deleted_id in actually_deleted {
|
||||
drop(idx.delete_document(&deleted_id));
|
||||
}
|
||||
@@ -228,7 +228,8 @@ impl McpTool for TasksHandler {
|
||||
if let Some(max_t) = req.max_tokens {
|
||||
let char_limit = max_t * 4;
|
||||
if json_str.len() > char_limit {
|
||||
json_str.truncate(char_limit);
|
||||
let boundary = json_str.floor_char_boundary(char_limit);
|
||||
json_str.truncate(boundary);
|
||||
json_str.push_str(" ...[truncated due to max_tokens]");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -30,33 +30,56 @@ impl McpTool for WriteClipboardHandler {
|
||||
let res = tokio::task::spawn_blocking(move || {
|
||||
let mut msgs = Vec::new();
|
||||
|
||||
if let Ok(mut clipboard) = Clipboard::new()
|
||||
&& let Some(text) = &tool_args.text
|
||||
&& clipboard.set_text(text).is_ok()
|
||||
{
|
||||
msgs.push("Wrote text");
|
||||
if let Some(text) = &tool_args.text {
|
||||
let mut written = false;
|
||||
if let Ok(mut clipboard) = Clipboard::new() {
|
||||
for _ in 0..3 {
|
||||
if clipboard.set_text(text).is_ok() {
|
||||
written = true;
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
}
|
||||
if written {
|
||||
msgs.push("Wrote text");
|
||||
}
|
||||
}
|
||||
|
||||
// Handle arboard for image
|
||||
if let Some(image_path) = &tool_args.image_path {
|
||||
match image::open(image_path) {
|
||||
Ok(img) => {
|
||||
let img = if img.width() > 2048 || img.height() > 2048 {
|
||||
img.resize(2048, 2048, FilterType::Triangle)
|
||||
} else {
|
||||
img
|
||||
};
|
||||
let rgba = img.into_rgba8();
|
||||
let (w, h) = rgba.dimensions();
|
||||
let img_data = ImageData {
|
||||
width: w as usize,
|
||||
height: h as usize,
|
||||
bytes: Cow::Owned(rgba.into_raw()),
|
||||
};
|
||||
let raw_bytes = rgba.into_raw();
|
||||
let mut written = false;
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if clipboard.set_image(img_data).is_ok() {
|
||||
msgs.push("Wrote image");
|
||||
} else {
|
||||
return Err(crate::error::AppError::Internal(
|
||||
"Failed to write image to clipboard".to_string(),
|
||||
));
|
||||
for _ in 0..3 {
|
||||
let img_data = ImageData {
|
||||
width: w as usize,
|
||||
height: h as usize,
|
||||
bytes: Cow::Borrowed(&raw_bytes),
|
||||
};
|
||||
if clipboard.set_image(img_data).is_ok() {
|
||||
written = true;
|
||||
break;
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
}
|
||||
if written {
|
||||
msgs.push("Wrote image");
|
||||
} else {
|
||||
return Err(crate::error::AppError::Internal(
|
||||
"Failed to write image to clipboard".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(crate::error::AppError::Internal(format!(
|
||||
@@ -85,92 +108,34 @@ impl McpTool for WriteClipboardHandler {
|
||||
}
|
||||
|
||||
pub fn get_native_clipboard_text() -> Option<String> {
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if let Ok(text) = clipboard.get_text() {
|
||||
if !text.trim().is_empty() {
|
||||
return Some(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let mut cmd_wl = std::process::Command::new("wl-paste");
|
||||
cmd_wl.arg("--no-newline");
|
||||
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
|
||||
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
|
||||
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
|
||||
}
|
||||
if let Ok(output) = cmd_wl.output() {
|
||||
if output.status.success() && !output.stdout.is_empty() {
|
||||
if let Ok(text) = String::from_utf8(output.stdout) {
|
||||
for _ in 0..3 {
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if let Ok(text) = clipboard.get_text() {
|
||||
if !text.trim().is_empty() {
|
||||
return Some(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
|
||||
let mut cmd_xc = std::process::Command::new("xclip");
|
||||
cmd_xc.args(["-selection", "clipboard", "-o"]);
|
||||
if std::env::var("DISPLAY").is_err() {
|
||||
cmd_xc.env("DISPLAY", ":0");
|
||||
}
|
||||
if let Ok(output) = cmd_xc.output() {
|
||||
if output.status.success() && !output.stdout.is_empty() {
|
||||
if let Ok(text) = String::from_utf8(output.stdout) {
|
||||
if !text.trim().is_empty() {
|
||||
return Some(text);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
pub fn get_native_clipboard_image() -> Option<image::DynamicImage> {
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if let Ok(image_data) = clipboard.get_image() {
|
||||
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
|
||||
image_data.width as u32,
|
||||
image_data.height as u32,
|
||||
image_data.bytes.into_owned(),
|
||||
) {
|
||||
return Some(image::DynamicImage::ImageRgba8(img));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for mime in &["image/png", "image/jpeg", "image/bmp", "image/tiff"] {
|
||||
let mut cmd_wl = std::process::Command::new("wl-paste");
|
||||
cmd_wl.args(["--type", mime]);
|
||||
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
|
||||
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
|
||||
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
|
||||
}
|
||||
if let Ok(output) = cmd_wl.output() {
|
||||
if output.status.success() && !output.stdout.is_empty() {
|
||||
if let Ok(img) = image::load_from_memory(&output.stdout) {
|
||||
return Some(img);
|
||||
for _ in 0..3 {
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if let Ok(image_data) = clipboard.get_image() {
|
||||
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
|
||||
image_data.width as u32,
|
||||
image_data.height as u32,
|
||||
image_data.bytes.into_owned(),
|
||||
) {
|
||||
return Some(image::DynamicImage::ImageRgba8(img));
|
||||
}
|
||||
}
|
||||
}
|
||||
std::thread::sleep(std::time::Duration::from_millis(20));
|
||||
}
|
||||
|
||||
for mime in &["image/png", "image/jpeg", "image/bmp"] {
|
||||
let mut cmd_xc = std::process::Command::new("xclip");
|
||||
cmd_xc.args(["-selection", "clipboard", "-t", mime, "-o"]);
|
||||
if std::env::var("DISPLAY").is_err() {
|
||||
cmd_xc.env("DISPLAY", ":0");
|
||||
}
|
||||
if let Ok(output) = cmd_xc.output() {
|
||||
if output.status.success() && !output.stdout.is_empty() {
|
||||
if let Ok(img) = image::load_from_memory(&output.stdout) {
|
||||
return Some(img);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
@@ -422,5 +387,33 @@ mod tests {
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||
assert!(parsed.is_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_no_subprocess_clipboard_regression() {
|
||||
let vision_src = include_str!("vision.rs");
|
||||
let code_only = vision_src.split("mod tests").next().unwrap_or(vision_src);
|
||||
let forbidden_cmd = format!("Command::{}{}", "n", "ew");
|
||||
let forbidden_ps = format!("power{}", "shell");
|
||||
let forbidden_wl = format!("wl-{}", "paste");
|
||||
let forbidden_xc = format!("x{}", "clip");
|
||||
|
||||
assert!(
|
||||
!code_only.contains(&forbidden_cmd),
|
||||
"Regression detected: vision.rs must not spawn subprocesses!"
|
||||
);
|
||||
assert!(
|
||||
!code_only.contains(&forbidden_ps),
|
||||
"Regression detected: vision.rs must not invoke powershell!"
|
||||
);
|
||||
assert!(
|
||||
!code_only.contains(&forbidden_wl),
|
||||
"Regression detected: vision.rs must not invoke wl-paste!"
|
||||
);
|
||||
assert!(
|
||||
!code_only.contains(&forbidden_xc),
|
||||
"Regression detected: vision.rs must not invoke xclip!"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -15,7 +15,7 @@ impl McpTool for PinFileHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<PinFileTool>("pin_file", "Execute pin_file")
|
||||
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> {
|
||||
@@ -42,7 +42,7 @@ impl McpTool for UnpinFileHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Execute unpin_file")
|
||||
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> {
|
||||
@@ -65,7 +65,7 @@ impl McpTool for ListPinnedFilesHandler {
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ListPinnedFilesTool>(
|
||||
"list_pinned_files",
|
||||
"Execute list_pinned_files",
|
||||
"List active pinned files in the workspace context, optionally filtered by namespace or branch.",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -104,7 +104,7 @@ impl McpTool for StoreSnippetHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Execute store_snippet")
|
||||
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> {
|
||||
@@ -127,7 +127,7 @@ impl McpTool for StoreSnippetHandler {
|
||||
embedding,
|
||||
};
|
||||
|
||||
let idx = state.get_search_index();
|
||||
let idx = state.get_search_index().await;
|
||||
drop(idx.index_snippet(&snippet));
|
||||
|
||||
state.code.snippets.modify(|snippets| {
|
||||
@@ -148,7 +148,7 @@ impl McpTool for TagSnippetHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<TagSnippetTool>("tag_snippet", "Execute tag_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> {
|
||||
@@ -254,7 +254,7 @@ impl McpTool for DeleteSnippetHandler {
|
||||
deleted = snippets.len() < orig;
|
||||
});
|
||||
if deleted {
|
||||
let idx = state.get_search_index();
|
||||
let idx = state.get_search_index().await;
|
||||
drop(idx.delete_document(&req.name));
|
||||
Ok("Snippet deleted.".to_string())
|
||||
} else {
|
||||
@@ -870,6 +870,9 @@ impl McpTool for ReadDirectoryArchitectureHandler {
|
||||
summary: &mut String,
|
||||
depth: usize,
|
||||
) -> std::io::Result<()> {
|
||||
if depth > 10 {
|
||||
return Ok(());
|
||||
}
|
||||
if dir.is_dir() {
|
||||
let mut entries =
|
||||
fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?;
|
||||
|
||||
+66
-32
@@ -46,7 +46,7 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
let idx = state.get_search_index();
|
||||
let idx = state.get_search_index().await;
|
||||
|
||||
for file_path in files_to_process {
|
||||
if let Ok(content) = std::fs::read_to_string(&file_path) {
|
||||
@@ -71,36 +71,64 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
|
||||
let mut chunks = Vec::new();
|
||||
extract_chunks(tree.root_node(), &content, &mut chunks, ext);
|
||||
|
||||
for (name, code, desc) in chunks {
|
||||
// Generate embedding
|
||||
if let Ok(mut emb) = generate_embeddings_async(vec![code.clone()]).await {
|
||||
let embedding = emb.pop();
|
||||
// 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 embeddings = generate_embeddings_async(texts).await.unwrap_or_default();
|
||||
|
||||
let file_name =
|
||||
file_path.file_name().unwrap_or_default().to_string_lossy();
|
||||
let mut new_snippets = Vec::with_capacity(chunk_batch.len());
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
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 snippet_name = format!("{}:{}", file_name, name);
|
||||
|
||||
let snippet = Snippet {
|
||||
name: snippet_name.to_string(),
|
||||
name: snippet_name,
|
||||
language: ext.to_string(),
|
||||
code: code.clone(),
|
||||
description: format!("{} in {}", desc, file_path.display()),
|
||||
updated_at: std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
updated_at: now,
|
||||
tags: vec![],
|
||||
embedding,
|
||||
};
|
||||
new_snippets.push(snippet);
|
||||
}
|
||||
|
||||
state.code.snippets.modify(|snippets| {
|
||||
// Prevent duplicates if already indexed
|
||||
if !snippets.iter().any(|s| s.name == snippet.name) {
|
||||
snippets.push(snippet.clone());
|
||||
// Gold Standard: Modify store ONCE per batch with zero-copy HashSet<&str> lookup
|
||||
let mut snippets_to_index = Vec::new();
|
||||
state.code.snippets.modify(|snippets| {
|
||||
let existing_names: std::collections::HashSet<&str> =
|
||||
snippets.iter().map(|s| s.name.as_str()).collect();
|
||||
|
||||
let mut filtered_new = Vec::with_capacity(new_snippets.len());
|
||||
let mut seen_in_batch = std::collections::HashSet::new();
|
||||
|
||||
for snippet in new_snippets {
|
||||
if !existing_names.contains(snippet.name.as_str())
|
||||
&& seen_in_batch.insert(snippet.name.clone())
|
||||
{
|
||||
filtered_new.push(snippet);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
let _ = idx.index_snippet(&snippet).await;
|
||||
for snippet in filtered_new {
|
||||
snippets.push(snippet.clone());
|
||||
snippets_to_index.push(snippet);
|
||||
}
|
||||
|
||||
if snippets.len() > 1000 {
|
||||
let overflow = snippets.len() - 1000;
|
||||
snippets.drain(0..overflow);
|
||||
}
|
||||
});
|
||||
|
||||
for snippet in &snippets_to_index {
|
||||
let _ = idx.index_snippet(snippet).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -111,7 +139,7 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
|
||||
}
|
||||
|
||||
fn extract_chunks(node: Node, code: &str, chunks: &mut Vec<(String, String, String)>, ext: &str) {
|
||||
extract_chunks_with_parent(node, code, chunks, ext, None);
|
||||
extract_chunks_with_parent(node, code, chunks, ext, None, 0);
|
||||
}
|
||||
|
||||
fn extract_chunks_with_parent(
|
||||
@@ -120,22 +148,28 @@ fn extract_chunks_with_parent(
|
||||
chunks: &mut Vec<(String, String, String)>,
|
||||
ext: &str,
|
||||
parent_scope: Option<&str>,
|
||||
depth: usize,
|
||||
) {
|
||||
// Stack overflow protection: Cap recursion depth at 100
|
||||
if depth > 100 {
|
||||
return;
|
||||
}
|
||||
|
||||
let kind = node.kind();
|
||||
let is_impl_or_class = matches!(kind, "impl_item" | "class_declaration" | "class_definition");
|
||||
|
||||
let current_scope = if is_impl_or_class {
|
||||
let current_scope: Option<&str> = if is_impl_or_class {
|
||||
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" {
|
||||
type_name = child.utf8_text(code.as_bytes()).ok().map(|s| s.to_string());
|
||||
type_name = child.utf8_text(code.as_bytes()).ok();
|
||||
break;
|
||||
}
|
||||
}
|
||||
type_name.or_else(|| parent_scope.map(|s| s.to_string()))
|
||||
type_name.or(parent_scope)
|
||||
} else {
|
||||
parent_scope.map(|s| s.to_string())
|
||||
parent_scope
|
||||
};
|
||||
|
||||
let is_structural = matches!(
|
||||
@@ -151,30 +185,30 @@ fn extract_chunks_with_parent(
|
||||
if is_structural {
|
||||
let mut raw_text = node.utf8_text(code.as_bytes()).unwrap_or("").to_string();
|
||||
|
||||
let mut name = "unknown".to_string();
|
||||
let mut name = "unknown";
|
||||
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" {
|
||||
name = child
|
||||
.utf8_text(code.as_bytes())
|
||||
.unwrap_or("unknown")
|
||||
.to_string();
|
||||
if let Ok(text) = child.utf8_text(code.as_bytes()) {
|
||||
name = text;
|
||||
}
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref scope) = current_scope {
|
||||
let mut final_name = name.to_string();
|
||||
if let Some(scope) = current_scope {
|
||||
raw_text = format!("// Parent Scope: {}\n{}", scope, raw_text);
|
||||
name = format!("{}::{}", scope, name);
|
||||
final_name = format!("{}::{}", scope, name);
|
||||
}
|
||||
|
||||
let desc = format!("{} AST node", kind);
|
||||
chunks.push((name, raw_text, desc));
|
||||
chunks.push((final_name, raw_text, desc));
|
||||
} else {
|
||||
let mut cursor = node.walk();
|
||||
for child in node.named_children(&mut cursor) {
|
||||
extract_chunks_with_parent(child, code, chunks, ext, current_scope.as_deref());
|
||||
extract_chunks_with_parent(child, code, chunks, ext, current_scope, depth + 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+86
-40
@@ -11,15 +11,17 @@ To prevent graph fragmentation and ensure seamless LLM context retrieval:
|
||||
* **Entity Types (`entity_type`)**: MUST ALWAYS be **`PascalCase`** (e.g. `DatabaseTable`, `McpTool`, `ArchitectureComponent`, `File`, `DataStructure`).
|
||||
* **Relation Types (`relation_type`)**: MUST ALWAYS be **`snake_case`** (e.g. `depends_on`, `calls`, `implements`, `uses`, `contains`).
|
||||
* **Field Keys & Properties**: MUST ALWAYS be **`snake_case`** (e.g. `file_path`, `git_commit`, `created_at`).
|
||||
* **Relation Tool Parameters**: `create_relations` supports Serde field aliases (`source` -> `from`, `target` -> `to`, `relationType`/`type` -> `relation_type`) so LLM tool calls succeed seamlessly regardless of parameter naming.
|
||||
|
||||
> [!NOTE]
|
||||
> The server automatically enforces and migrates incoming entity and relation types to these canonical casing rules on every read and write operation.
|
||||
> Store operations enforce atomic lock scope for serialization/deserialization and enter Quarantine Mode upon database corruption to prevent data overwrite hazards.
|
||||
|
||||
---
|
||||
|
||||
## 2. Consolidated Smart Tools Architecture
|
||||
|
||||
The server consolidates granular single-purpose tools into domain-named smart tools. Always prefer the consolidated tools:
|
||||
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`).
|
||||
@@ -85,55 +87,99 @@ The server consolidates granular single-purpose tools into domain-named smart to
|
||||
- `action: "register"`: Register target environment (requires `name`).
|
||||
- `action: "get_details"`: Read full environment details.
|
||||
|
||||
* **`clipboard`**: OS Clipboard management (native cross-platform support for Windows and Linux wl-paste/xclip).
|
||||
- `action: "read"`: Read OS clipboard text, HTML, file paths, or JPEG image.
|
||||
* **`clipboard`**: OS Clipboard management.
|
||||
- `action: "read"`: Read OS clipboard.
|
||||
- `action: "write"`: Write text/html/files/image to clipboard.
|
||||
- `action: "toggle_watch"`: Toggle auto-clipboard watcher.
|
||||
|
||||
* **`manage_subagent_namespace`**: Subagent memory isolation.
|
||||
- `action: "create"`: Create namespace.
|
||||
- `action: "purge"`: Purge namespace.
|
||||
- `action: "condense"`: Condense namespace.
|
||||
---
|
||||
|
||||
* **`manage_checkpoint`**: State snapshots and rollbacks.
|
||||
- `action: "checkpoint"` (or `"create_snapshot"`): Create state snapshot.
|
||||
- `action: "restore"` (or `"restore_snapshot"`): Restore point-in-time snapshot.
|
||||
## 4. High-Signal Tool Responses & Performance Guidelines
|
||||
|
||||
To optimize context usage, response times, and LLM reasoning efficiency:
|
||||
|
||||
* **High-Signal Feedback**:
|
||||
- MCP tool calls (such as `create_entities`, `create_relations`, `add_observations`, `pin_file`) return explicit, structured summaries containing created names, types, edge counts, and relation paths.
|
||||
- LLMs do NOT need to execute follow-up `open_nodes` calls merely to confirm successful creation.
|
||||
|
||||
* **Batch Operations**:
|
||||
- When creating or updating multiple entities, snippets, or observations, always batch items into a single tool call array (e.g. `create_entities` with multiple array items) rather than making separate calls.
|
||||
- The server handles batch store mutations serially in a single transaction pass with single-permit event-driven flushes.
|
||||
|
||||
* **Real-time Tantivy Search Indexing**:
|
||||
- The Tantivy search engine automatically checks pending commits and reloads search readers prior to executing `omni_search` or `search_nodes`. Search queries always return up-to-date document results immediately following mutations.
|
||||
- Single-item deletions use targeted document removal rather than global index wipes.
|
||||
|
||||
* **Real-Time AST & Workspace Source Code Symbol Scanning**:
|
||||
- `find_symbol_references`, `get_callers`, and `analyze_impact` scan both stored code snippets and physical workspace source code files on disk (`.rs`, `.ts`, `.py`, `.go`, `.java`, `.c`, `.cpp`), providing accurate AST symbol references and call site tracking.
|
||||
|
||||
* **Graph Entity Merge & Self-Loop Protection**:
|
||||
- `merge_entities` re-points all relations from `source_entity` to `target_entity` and automatically prunes cyclic self-loops (`target -> target`).
|
||||
|
||||
* **Safe UTF-8 Token Truncation**:
|
||||
- Large responses (e.g. `get_active_worktree_context`, `read_graph`, `summarize_subgraph`) are safely truncated along UTF-8 character boundaries (`floor_char_boundary`), ensuring response bounds without runtime panics.
|
||||
|
||||
## 4. Automated Error Fix Auto-Matcher
|
||||
- **Tools:** `log_error_fix`, `search_error_fixes` (and alias `suggest_error_fix`)
|
||||
- **When to use:** When encountering a build error, test failure, or stack trace. Call `search_error_fixes` with either a text `query` or `stack_trace` before attempting a fix from scratch.
|
||||
- **Behavior:** Computes cosine similarity between error trace embeddings and past resolution logs when `stack_trace` is provided, or keyword filtering when `query` is provided, returning top matched solutions, modified files, and git commits.
|
||||
|
||||
---
|
||||
|
||||
## 3. Knowledge Graph Operations & Multi-Hop Navigation
|
||||
* **`create_entities`**: Create entity nodes (`entity_type` PascalCase).
|
||||
* **`create_relations`**: Create relations (`relation_type` snake_case).
|
||||
* **`add_observations`**: Add new observations to existing entities.
|
||||
* **`delete_entities`**, **`delete_observations`**, **`delete_relations`**: Graph cleanup.
|
||||
* **`read_graph`**, **`search_nodes`**, **`open_nodes`**: Graph query and entity lookup.
|
||||
* **`visualize_graph`**: Generate graph visualizer output.
|
||||
* **`condense_entity`**, **`merge_entities`**, **`find_orphans`**: Graph maintenance.
|
||||
* **`get_subgraph`**: Perform multi-hop BFS neighborhood expansion tree around a root entity up to `depth` hops with token budget.
|
||||
* **`sweep_graph_health`**: Self-healing graph health auditor for near-duplicates and orphaned nodes.
|
||||
* **`resolve_stale_symbols`**, **`summarize_subgraph`**: Symbol resolution and subgraph summarization.
|
||||
## 5. Memory State Checkpointing & Rollbacks
|
||||
- **Tool:** `checkpoint_state`, `restore_state` (or `create_snapshot`, `restore_snapshot`)
|
||||
- **When to use:** Before initiating a large refactor, running experimental subagent tasks, or executing destructive batch operations.
|
||||
- **Behavior:** Saves or restores a point-in-time snapshot of graph entities, active tasks, and tech debt backlogs.
|
||||
|
||||
---
|
||||
|
||||
## 4. Search, Code & AST Intelligence
|
||||
* **`omni_search`**: High-performance unified search across tasks, code changes, ADRs, debt, and snippets.
|
||||
* **`semantic_code_search`**: Vector similarity search across indexing embeddings.
|
||||
* **`read_file_skeleton`**: Extract file declarations/skeletons using Tree-Sitter (Rust, Python, Go, Java, C, C++, JSX/TSX).
|
||||
* **`replace_ast_node`**: Structural AST node replacement across files.
|
||||
* **`find_symbol_references`**, **`get_callers`**, **`analyze_impact`**: Code dependency and blast radius analysis.
|
||||
* **`read_directory_architecture`**: Inspect directory structure and code organization.
|
||||
## 6. Self-Healing Graph Health Sweeper
|
||||
- **Tool:** `sweep_graph_health`
|
||||
- **When to use:** Periodically or before committing major graph changes to audit entity consistency.
|
||||
- **Behavior:** Detects orphaned nodes (0 relations), computes name similarity using pre-computed lowercase keys to identify near-duplicates (e.g., `APIGateway` vs `ApiGateway`), and provides structured `merge_entities` recommendations or auto-prunes orphans.
|
||||
|
||||
---
|
||||
|
||||
## 5. Audit Logging, Diagnostics & Provenance
|
||||
* **`log_code_change`**, **`query_recent_changes`**: Audit trail of code edits, commits, and symbol modifications.
|
||||
* **`log_error_fix`**, **`search_error_fixes`**: Record error trace and verified fix; search past resolutions via vector similarity.
|
||||
* **`query_lineage`**: Reconstruct causal timeline explaining why a file/symbol/commit was modified.
|
||||
* **`log_hypothesis`**, **`query_hypotheses`**: Diagnostic hypothesis tracking during complex debugging.
|
||||
* **`get_preflight_context`**: Consolidated executive bundle (branch, active tasks, criteria, pinned files, tech debt) for instant context warmup.
|
||||
* **`learn_preference`**, **`read_preferences`**: Track user preferences and project conventions.
|
||||
* **`get_project_health`**: Health metrics, active task counts, debt backlog, and graph statistics.
|
||||
* **`broadcast_agent_signal`**, **`query_agent_signals`**: Inter-agent coordination signals.
|
||||
* **`auto_session_checkpoint`**: Auto-save session state checkpoint.
|
||||
* **`get_active_worktree_context`**, **`query_git_diffs`**: Git worktree status and diff inspection.
|
||||
* **`watch_process_logs`**, **`get_recent_logs`**: Process log monitoring and retrieval.
|
||||
## 7. Causal Lineage & Provenance Tracker
|
||||
- **Tool:** `query_lineage`
|
||||
- **When to use:** When asking *"Why was this component modified?"* or *"What task or ADR led to this code change?"*
|
||||
- **Behavior:** Searches across tasks, ADRs, audit ledger entries, and error fixes to assemble a unified chronological timeline explaining the provenance behind any file, symbol, or commit.
|
||||
|
||||
---
|
||||
|
||||
## 9. Native Rust Invariants & Subprocess Prohibition (CRITICAL)
|
||||
|
||||
To maintain maximum security, speed, and cross-platform reliability:
|
||||
* **Zero Subprocess Fallbacks**: System and server tools (`clipboard`, `ast`, `search`, `db`) MUST strictly use pure native Rust crates (`arboard`, `tree-sitter`, `tantivy`, `psycopg`). Invocations of external shell commands (`powershell.exe`, `wl-paste`, `xclip`, `cmd.exe`) are strictly prohibited in native handlers.
|
||||
* **Transient Lock Recovery**: Transient OS handle collisions (such as Win32 OLE `OpenClipboard` lock contention) must be handled using native retry loops with backoffs directly in Rust.
|
||||
* **Automated Static Regression Gates**: Automated AST/source audit tests (e.g. `test_no_subprocess_clipboard_regression`) verify at test time that forbidden subprocess patterns are absent from handler implementations.
|
||||
|
||||
---
|
||||
|
||||
## 10. High-Performance Concurrency & Resilience Guarantees
|
||||
* **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.
|
||||
* **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.
|
||||
* **Token-Budgeted Query Projections**: Decision queries (`query_decisions`) support `limit` caps and compact `include_body: false` projections for token budget optimization.
|
||||
* **Serde Parameter & Enum Ergonomics**: Action enums (`StickyNoteAction`, `SnippetSearchMode`, `Relation`) support case-insensitive variants and common synonyms (`create`/`add`, `remove`/`delete`, `list`/`read`, `source`/`from`, `target`/`to`, `relationType`/`relation_type`) ensuring seamless LLM tool execution.
|
||||
* **Atomic Search Index Swaps**: `MemoryState::rebuild_index` constructs and populates a new `MemoryIndex` instance in isolation before performing an atomic pointer swap (`*self.search_index.write().await = new_idx`), eliminating transient empty search result windows.
|
||||
* **Non-Blocking Tantivy Search Queries**: `MemoryIndex::search` queries current index searcher snapshots without executing synchronous disk commit locks, preventing query stalls during heavy background indexing.
|
||||
* **Zero-Allocation HashSet<&str> Snippet Deduplication**: `indexer.rs` utilizes borrowed `HashSet<&str>` name lookups during snippet batch modifications, eliminating heap string re-allocations inside the store write lock.
|
||||
* **AST Recursion Depth Safeguard & Zero-Copy Borrowing**: Tree-sitter AST traversal caps recursion depth at 100 to prevent thread stack overflows and borrows string slices (`&str`) during AST node walking.
|
||||
* **Strongly-Typed SearchResult & Pre-Allocated Search Vectors**: `search.rs` uses a strongly-typed `SearchResult` struct with named fields and pre-allocates result vector capacity (`Vec::with_capacity(top_docs.len())`).
|
||||
* **BFS Graph Traversal Pre-allocation & Visited Node Upper Bound**: `GraphQueryBuilder::find_shortest_path` pre-allocates adjacency map capacity (`HashMap::with_capacity(relations.len() * 2)`) and enforces a visited node upper bound (10,000 max) to guarantee deterministic BFS runtime.
|
||||
* **LLM Tool Schema Ergonomics & Context Guidance**: `ReadGraphHandler` schema explicitly instructs LLMs on `namespace` filtering and `search_nodes` / `get_subgraph` tools for large graph discovery.
|
||||
* **Filesystem Event Debouncing & Proactive State Refresh**: `spawn_watcher` implements a sliding 250ms debouncing window per file path, ignores `.git`, `target`, `.gemini`, and `node_modules`, and broadcasts activity events to `MemoryState`.
|
||||
* **Buffered Line-by-Line AST Workspace Symbol Scanning**: `scan_workspace_for_symbol` reads workspace files via `BufReader` line streams instead of loading entire files into heap strings, preventing memory spikes when traversing source trees.
|
||||
* **AST Node Type Aliasing & Skeleton Preallocation**: `replace_ast_node` documents friendly node aliases (`function`, `fn`, `method`, `struct`, `class`, `enum`, `trait`, `type`), and `read_file_skeleton` preallocates string buffer capacity (`code.len() / 2`).
|
||||
* **Batch Vector Indexing & Similarity Score Guidance**: `VectorDB` provides `index_documents_batch` for single-request multi-point vector upserts and explicit score calibration notes ($\ge 0.75$ high confidence match).
|
||||
* **Compact JSON MCP Resources & UTF-8 Activity Truncation**: MCP resources serialize using compact JSON (`to_string`), `TerminalHistoryResource` / `MilestonesResource` enforce output bounds, and `format_tool_activity_description` uses `floor_char_boundary` for guaranteed UTF-8 safety.
|
||||
* **SIMD-Friendly Single-Pass Cosine Similarity**: `cosine_similarity` calculates dot product and Euclidean norm squares in a single linear pass over float vectors, enabling SIMD compiler auto-vectorization.
|
||||
* **Safe Stream Decoding on Log Tails**: Log tail operations (`get_recent_logs`) read raw bytes and decode using lossy UTF-8 conversion (`String::from_utf8_lossy`) to ensure resilience when seeking across multi-byte UTF-8 boundaries.
|
||||
* **Task Summary UTF-8 Truncation Safety**: `tasks` tool (`action = "list"`) truncates serialized task text strictly along UTF-8 character boundaries using `floor_char_boundary` when enforcing `max_tokens`.
|
||||
* **Sequential Snapshot Lock Scope Flattening**: `GenerateStandupReportHandler` reads `tasks`, `ledger`, and `session_summaries` sequentially rather than nesting read locks, preventing multi-lock deadlocks during concurrent store modifications.
|
||||
* **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.
|
||||
+30
-16
@@ -122,10 +122,8 @@ pub async fn ttl_sweeper_worker(state: Arc<MemoryState>) {
|
||||
pub async fn index_committer_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
state.index_commit_notify.notified().await;
|
||||
let idx_opt = state.search_index.read().ok().map(|idx| idx.clone());
|
||||
if let Some(idx) = idx_opt {
|
||||
let _ = idx.commit().await;
|
||||
}
|
||||
let idx = state.search_index.read().await.clone();
|
||||
let _ = idx.commit().await;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -234,15 +232,26 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
|
||||
loop {
|
||||
match rx.recv().await {
|
||||
Ok(msg) => {
|
||||
let senders: Vec<_> = app_state_clone
|
||||
.clients
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.values()
|
||||
.cloned()
|
||||
.collect();
|
||||
for client_tx in senders {
|
||||
let _ = client_tx.try_send(msg.clone());
|
||||
let mut closed_ids = Vec::new();
|
||||
{
|
||||
let clients_guard = app_state_clone
|
||||
.clients
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
for (id, tx) in clients_guard.iter() {
|
||||
if tx.try_send(msg.clone()).is_err() && tx.is_closed() {
|
||||
closed_ids.push(id.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
if !closed_ids.is_empty() {
|
||||
let mut write_guard = app_state_clone
|
||||
.clients
|
||||
.write()
|
||||
.unwrap_or_else(|e| e.into_inner());
|
||||
for id in closed_ids {
|
||||
write_guard.remove(&id);
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
@@ -304,9 +313,14 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
|
||||
let session = &payload.session_id;
|
||||
let _is_unix_socket = session.starts_with('/') || session.starts_with('~');
|
||||
if let Some(home) = dirs::home_dir() {
|
||||
let nvim_txt_path = home.join(".gemini").join("active_nvim.txt");
|
||||
let _ = tokio::fs::create_dir_all(nvim_txt_path.parent().unwrap()).await;
|
||||
let _ = tokio::fs::write(&nvim_txt_path, session).await;
|
||||
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()));
|
||||
if tokio::fs::create_dir_all(&nvim_dir).await.is_ok()
|
||||
&& tokio::fs::write(&tmp_path, session).await.is_ok()
|
||||
{
|
||||
let _ = tokio::fs::rename(&tmp_path, &nvim_txt_path).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -27,6 +27,7 @@ pub fn tool_def<T: JsonSchema>(name: &str, fallback_desc: &str) -> serde_json::V
|
||||
let desc = schema_val
|
||||
.get("description")
|
||||
.and_then(|d| d.as_str())
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.unwrap_or(fallback_desc)
|
||||
.to_string();
|
||||
|
||||
|
||||
@@ -72,9 +72,45 @@ pub fn normalize_relation_type(input: &str) -> String {
|
||||
trimmed.to_lowercase().replace([' ', '-'], "_")
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Default)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TaskStatus {
|
||||
#[default]
|
||||
Active,
|
||||
InProgress,
|
||||
Blocked,
|
||||
Completed,
|
||||
}
|
||||
|
||||
impl TaskStatus {
|
||||
pub fn is_active(&self) -> bool {
|
||||
matches!(self, TaskStatus::Active | TaskStatus::InProgress | TaskStatus::Blocked)
|
||||
}
|
||||
|
||||
pub fn is_completed(&self) -> bool {
|
||||
matches!(self, TaskStatus::Completed)
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
TaskStatus::Active => "active",
|
||||
TaskStatus::InProgress => "in_progress",
|
||||
TaskStatus::Blocked => "blocked",
|
||||
TaskStatus::Completed => "completed",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TaskStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct Entity {
|
||||
pub name: String,
|
||||
#[serde(alias = "entityType", alias = "type", alias = "category")]
|
||||
pub entity_type: String,
|
||||
#[serde(default)]
|
||||
pub observations: Vec<String>,
|
||||
@@ -85,8 +121,11 @@ pub struct Entity {
|
||||
}
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, JsonSchema)]
|
||||
pub struct Relation {
|
||||
#[serde(alias = "source", alias = "from_node", alias = "source_node")]
|
||||
pub from: String,
|
||||
#[serde(alias = "target", alias = "to_node", alias = "target_node")]
|
||||
pub to: String,
|
||||
#[serde(alias = "relationType", alias = "type", alias = "relation")]
|
||||
pub relation_type: String,
|
||||
#[serde(default = "default_namespace")]
|
||||
pub namespace: String,
|
||||
@@ -123,6 +162,16 @@ pub struct Task {
|
||||
#[serde(default)]
|
||||
pub acceptance_criteria: Vec<AcceptanceCriteria>,
|
||||
}
|
||||
|
||||
impl Task {
|
||||
pub fn is_active(&self) -> bool {
|
||||
let s = self.status.trim();
|
||||
!s.eq_ignore_ascii_case("completed")
|
||||
&& !s.eq_ignore_ascii_case("done")
|
||||
&& !s.eq_ignore_ascii_case("cancelled")
|
||||
&& !s.eq_ignore_ascii_case("canceled")
|
||||
}
|
||||
}
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Snippet {
|
||||
pub name: String,
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
pub fn patch() {}
|
||||
+59
-94
@@ -57,23 +57,18 @@ impl McpResource for GraphEntitiesResource {
|
||||
Some("All nodes and entities currently stored in the knowledge graph")
|
||||
}
|
||||
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let state_clone = Arc::clone(&state);
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let graph = state_clone.graph.cache.read().unwrap();
|
||||
let data: Vec<_> = graph.entities.values().collect();
|
||||
if data.len() > 100 {
|
||||
let truncated: Vec<_> = data.iter().take(100).cloned().collect();
|
||||
let res = serde_json::json!({
|
||||
"_meta": format!("Showing 100 of {} total entities. Use omni_search or get_subgraph for targeted entity discovery.", data.len()),
|
||||
"entities": truncated
|
||||
});
|
||||
Ok(serde_json::to_string_pretty(&res)?)
|
||||
} else {
|
||||
Ok(serde_json::to_string_pretty(&data)?)
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
|
||||
let graph = state.graph.cache.read().unwrap_or_else(|e| e.into_inner());
|
||||
let data: Vec<_> = graph.entities.values().collect();
|
||||
if data.len() > 100 {
|
||||
let truncated: Vec<_> = data.iter().take(100).cloned().collect();
|
||||
let res = serde_json::json!({
|
||||
"_meta": format!("Showing 100 of {} total entities. Use omni_search or get_subgraph for targeted entity discovery.", data.len()),
|
||||
"entities": truncated
|
||||
});
|
||||
Ok(serde_json::to_string_pretty(&res)?)
|
||||
} else {
|
||||
Ok(serde_json::to_string_pretty(&data)?)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -90,14 +85,18 @@ impl McpResource for GraphRelationsResource {
|
||||
Some("All relationships between entities currently stored in the knowledge graph")
|
||||
}
|
||||
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let state_clone = Arc::clone(&state);
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let graph = state_clone.graph.cache.read().unwrap();
|
||||
let data = &graph.relations;
|
||||
let graph = state.graph.cache.read().unwrap_or_else(|e| e.into_inner());
|
||||
let data = &graph.relations;
|
||||
if data.len() > 200 {
|
||||
let truncated: Vec<_> = data.iter().take(200).cloned().collect();
|
||||
let res = serde_json::json!({
|
||||
"_meta": format!("Showing 200 of {} total relations. Use query_graph_path or get_subgraph for targeted relation discovery.", data.len()),
|
||||
"relations": truncated
|
||||
});
|
||||
Ok(serde_json::to_string_pretty(&res)?)
|
||||
} else {
|
||||
Ok(serde_json::to_string_pretty(&data)?)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -114,17 +113,12 @@ impl McpResource for TasksActiveResource {
|
||||
Some("List of currently active tasks")
|
||||
}
|
||||
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let state_clone = Arc::clone(&state);
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let tasks = state_clone.project.tasks.cache.read().unwrap();
|
||||
let data: Vec<_> = tasks
|
||||
.iter()
|
||||
.filter(|t| t.status != "completed" && t.status != "done")
|
||||
.collect();
|
||||
Ok(serde_json::to_string_pretty(&data)?)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_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();
|
||||
Ok(serde_json::to_string_pretty(&data)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -141,14 +135,9 @@ impl McpResource for DecisionsActiveResource {
|
||||
Some("All accepted Architectural Decision Records (ADRs)")
|
||||
}
|
||||
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let state_clone = Arc::clone(&state);
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let adrs = state_clone.code.adrs.cache.read().unwrap();
|
||||
let data: Vec<_> = adrs.iter().filter(|a| a.status == "accepted").collect();
|
||||
Ok(serde_json::to_string_pretty(&data)?)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
|
||||
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)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -165,14 +154,9 @@ impl McpResource for TechDebtActiveResource {
|
||||
Some("All currently unresolved technical debt items")
|
||||
}
|
||||
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let state_clone = Arc::clone(&state);
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let debts = state_clone.code.tech_debts.cache.read().unwrap();
|
||||
let data: Vec<_> = debts.iter().filter(|d| !d.is_resolved).collect();
|
||||
Ok(serde_json::to_string_pretty(&data)?)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
|
||||
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)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -245,32 +229,27 @@ impl McpResource for SessionDeltaResource {
|
||||
Some("Recent changes, active tasks, code edits, and notes created in the current session")
|
||||
}
|
||||
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let state_clone = Arc::clone(&state);
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let now = crate::handlers::utils::now_secs();
|
||||
let lookback = 7200; // last 2 hours
|
||||
let cutoff = now.saturating_sub(lookback);
|
||||
let now = crate::handlers::utils::now_secs();
|
||||
let lookback = 7200; // last 2 hours
|
||||
let cutoff = now.saturating_sub(lookback);
|
||||
|
||||
let changes = state_clone.code.ledger.cache.read().unwrap();
|
||||
let recent_changes: Vec<_> = changes.iter().filter(|c| c.timestamp >= cutoff).collect();
|
||||
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_clone.project.tasks.cache.read().unwrap();
|
||||
let active_tasks: Vec<_> = tasks.iter().filter(|t| t.status != "completed" && t.status != "done").collect();
|
||||
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_clone.code.sticky.cache.read().unwrap();
|
||||
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,
|
||||
"active_tasks": active_tasks,
|
||||
"active_sticky_notes": active_notes,
|
||||
"timestamp": now,
|
||||
});
|
||||
let delta = serde_json::json!({
|
||||
"recent_changes": recent_changes,
|
||||
"active_tasks": active_tasks,
|
||||
"active_sticky_notes": active_notes,
|
||||
"timestamp": now,
|
||||
});
|
||||
|
||||
Ok(serde_json::to_string_pretty(&delta)?)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
|
||||
Ok(serde_json::to_string_pretty(&delta)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -355,13 +334,8 @@ impl MemoryHandler {
|
||||
Some("Recent terminal execution history and exit codes")
|
||||
}
|
||||
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let state_clone = Arc::clone(&state);
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let items = state_clone.telemetry.terminal_history.cache.read().unwrap();
|
||||
Ok(serde_json::to_string_pretty(&*items)?)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
|
||||
let items = state.telemetry.terminal_history.cache.read().unwrap_or_else(|e| e.into_inner());
|
||||
Ok(serde_json::to_string_pretty(&*items)?)
|
||||
}
|
||||
}
|
||||
struct PinnedFilesResource;
|
||||
@@ -377,13 +351,8 @@ impl MemoryHandler {
|
||||
Some("Currently pinned files in the workspace")
|
||||
}
|
||||
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let state_clone = Arc::clone(&state);
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let items = state_clone.project.pinned_files.cache.read().unwrap();
|
||||
Ok(serde_json::to_string_pretty(&*items)?)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
|
||||
let items = state.project.pinned_files.cache.read().unwrap_or_else(|e| e.into_inner());
|
||||
Ok(serde_json::to_string_pretty(&*items)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -400,13 +369,8 @@ impl MemoryHandler {
|
||||
Some("Project milestones and their status")
|
||||
}
|
||||
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let state_clone = Arc::clone(&state);
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let items = state_clone.project.milestones.cache.read().unwrap();
|
||||
Ok(serde_json::to_string_pretty(&*items)?)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
|
||||
let items = state.project.milestones.cache.read().unwrap_or_else(|e| e.into_inner());
|
||||
Ok(serde_json::to_string_pretty(&*items)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -905,7 +869,8 @@ pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) ->
|
||||
|
||||
let mut truncated = detail.trim().to_string();
|
||||
if truncated.len() > 80 {
|
||||
truncated.truncate(77);
|
||||
let valid_len = truncated.floor_char_boundary(77);
|
||||
truncated.truncate(valid_len);
|
||||
truncated.push_str("...");
|
||||
}
|
||||
|
||||
|
||||
+106
-71
@@ -3,6 +3,15 @@ use std::sync::{Arc, Mutex};
|
||||
use tantivy::schema::*;
|
||||
use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, doc};
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
|
||||
pub struct SearchResult {
|
||||
pub id: String,
|
||||
pub doc_type: String,
|
||||
pub title: String,
|
||||
pub body: String,
|
||||
pub score: f32,
|
||||
}
|
||||
|
||||
pub type SearchResultTuple = (String, String, String, String, f32);
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -89,23 +98,35 @@ impl MemoryIndex {
|
||||
}
|
||||
|
||||
pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
||||
self.index_entities_batch(std::slice::from_ref(e))
|
||||
}
|
||||
|
||||
pub fn index_entities_batch(&self, entities: &[Entity]) -> tokio::task::JoinHandle<tantivy::Result<()>> {
|
||||
let writer = Arc::clone(&self.writer);
|
||||
let id_field = self.id_field;
|
||||
let id_val = e.name.clone();
|
||||
let needs_commit = Arc::clone(&self.needs_commit);
|
||||
|
||||
let doc = 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()
|
||||
);
|
||||
let docs: Vec<(String, tantivy::TantivyDocument)> = entities
|
||||
.iter()
|
||||
.map(|e| {
|
||||
let id_val = e.name.clone();
|
||||
let doc = 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()
|
||||
);
|
||||
(id_val, doc)
|
||||
})
|
||||
.collect();
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
|
||||
writer.add_document(doc)?;
|
||||
let writer_guard = writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||
for (id_val, doc) in docs {
|
||||
writer_guard.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
|
||||
writer_guard.add_document(doc)?;
|
||||
}
|
||||
needs_commit.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
Ok(())
|
||||
})
|
||||
@@ -172,6 +193,15 @@ impl MemoryIndex {
|
||||
}
|
||||
|
||||
|
||||
pub fn commit_if_needed(&self) -> tantivy::Result<()> {
|
||||
if self.needs_commit.swap(false, std::sync::atomic::Ordering::SeqCst) {
|
||||
let mut writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||
writer.commit()?;
|
||||
let _ = self.reader.reload();
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn search(
|
||||
&self,
|
||||
query: &str,
|
||||
@@ -182,13 +212,28 @@ impl MemoryIndex {
|
||||
&self.index,
|
||||
vec![self.title_field, self.body_field],
|
||||
);
|
||||
let q = query_parser.parse_query(query)?;
|
||||
let q = query_parser.parse_query(query).or_else(|_| {
|
||||
let sanitized: String = query
|
||||
.chars()
|
||||
.map(|c| match c {
|
||||
'+' | '-' | '&' | '|' | '!' | '(' | ')' | '{' | '}' | '[' | ']' | '^'
|
||||
| '"' | '~' | '*' | '?' | ':' | '\\' | '/' => ' ',
|
||||
_ => c,
|
||||
})
|
||||
.collect();
|
||||
let safe_query = sanitized.trim();
|
||||
if safe_query.is_empty() {
|
||||
query_parser.parse_query("*")
|
||||
} else {
|
||||
query_parser.parse_query(safe_query)
|
||||
}
|
||||
})?;
|
||||
|
||||
let top_docs = searcher.search(
|
||||
&q,
|
||||
&tantivy::collector::TopDocs::with_limit(50).order_by_score(),
|
||||
)?;
|
||||
let mut results = Vec::new();
|
||||
let mut results = Vec::with_capacity(top_docs.len());
|
||||
for (score, doc_address) in top_docs {
|
||||
let retrieved_doc = searcher.doc::<tantivy::TantivyDocument>(doc_address)?;
|
||||
let id = retrieved_doc
|
||||
@@ -273,72 +318,62 @@ impl MemoryIndex {
|
||||
}
|
||||
|
||||
pub fn add_entity_sync(&self, e: &Entity) {
|
||||
if let Ok(writer) = self.writer.lock() {
|
||||
let _ = 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()
|
||||
));
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let _ = 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()
|
||||
));
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
|
||||
pub fn delete_all(&self) {
|
||||
if let Ok(writer) = self.writer.lock() {
|
||||
let _ = writer.delete_all_documents();
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let _ = writer.delete_all_documents();
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
|
||||
pub fn add_task_sync(&self, t: &Task) {
|
||||
// println!("add_task_sync called for task: {}", t.id);
|
||||
if let Ok(writer) = self.writer.lock() {
|
||||
let _res = 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"
|
||||
));
|
||||
// println!("Writer add_document returned id/result");
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
// println!("Needs_commit set to true in add_task_sync");
|
||||
} else {
|
||||
// println!("Failed to acquire writer lock in add_task_sync");
|
||||
}
|
||||
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let _res = 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"
|
||||
));
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
|
||||
pub fn add_snippet_sync(&self, s: &Snippet) {
|
||||
if let Ok(writer) = self.writer.lock() {
|
||||
let _ = 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"
|
||||
));
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let _ = 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"
|
||||
));
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
|
||||
pub fn add_adr_sync(&self, a: &Adr) {
|
||||
if let Ok(writer) = self.writer.lock() {
|
||||
let _ = 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"
|
||||
));
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||
let _ = 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"
|
||||
));
|
||||
self.needs_commit
|
||||
.store(true, std::sync::atomic::Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -423,12 +458,12 @@ mod tests {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let index = MemoryIndex::new(temp_dir.path()).unwrap();
|
||||
|
||||
// Malformed lucene query (unclosed parenthesis)
|
||||
// Malformed lucene query (unclosed parenthesis) - sanitized gracefully
|
||||
let result = index.search("title: (unclosed", None);
|
||||
assert!(result.is_err());
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Another malformed query (unclosed quote)
|
||||
// Another malformed query (unclosed quote) - sanitized gracefully
|
||||
let result2 = index.search("title: \"unclosed", None);
|
||||
assert!(result2.is_err());
|
||||
assert!(result2.is_ok());
|
||||
}
|
||||
}
|
||||
+78
-46
@@ -4,7 +4,7 @@ use crate::store::Store;
|
||||
use crate::vector_db::VectorDB;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
pub struct GenericEvent {
|
||||
@@ -56,7 +56,7 @@ pub struct MemoryState {
|
||||
pub ttl_notify: Arc<tokio::sync::Notify>,
|
||||
pub condense_notify: Arc<tokio::sync::Notify>,
|
||||
pub graph: Store<KnowledgeGraph>,
|
||||
pub search_index: RwLock<MemoryIndex>,
|
||||
pub search_index: tokio::sync::RwLock<MemoryIndex>,
|
||||
pub vector_db: tokio::sync::RwLock<Option<VectorDB>>,
|
||||
|
||||
pub project: ProjectStores,
|
||||
@@ -70,14 +70,37 @@ pub struct MemoryState {
|
||||
}
|
||||
|
||||
impl MemoryState {
|
||||
pub fn new_in_memory() -> Self {
|
||||
Self::new(":memory:")
|
||||
}
|
||||
|
||||
pub fn new(base_dir_str: &str) -> Self {
|
||||
let is_in_memory = base_dir_str == ":memory:";
|
||||
let base = std::path::PathBuf::from(base_dir_str);
|
||||
if let Err(e) = std::fs::create_dir_all(&base) {
|
||||
tracing::error!("Failed to create store directory at {:?}: {}", base, e);
|
||||
if !is_in_memory {
|
||||
if let Err(e) = std::fs::create_dir_all(&base) {
|
||||
tracing::error!("Failed to create store directory at {:?}: {}", base, e);
|
||||
}
|
||||
}
|
||||
|
||||
let db = crate::db::init_redb(&base);
|
||||
|
||||
let search_index = if is_in_memory {
|
||||
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex")
|
||||
} else {
|
||||
match crate::search::MemoryIndex::new(&base) {
|
||||
Ok(idx) => idx,
|
||||
Err(e) => {
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.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")
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let state = Self {
|
||||
ollama: Arc::new(crate::ollama::OllamaClient::new_from_env()),
|
||||
clipboard_watch_mode: tokio::sync::RwLock::new(false),
|
||||
@@ -87,17 +110,7 @@ impl MemoryState {
|
||||
condense_notify: Arc::new(tokio::sync::Notify::new()),
|
||||
graph: Store::new("knowledge_graph_master", db.clone()),
|
||||
base_dir: base.clone(),
|
||||
search_index: RwLock::new(match crate::search::MemoryIndex::new(&base) {
|
||||
Ok(idx) => idx,
|
||||
Err(e) => {
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.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")
|
||||
}
|
||||
}),
|
||||
search_index: tokio::sync::RwLock::new(search_index),
|
||||
vector_db: tokio::sync::RwLock::new(None),
|
||||
|
||||
project: ProjectStores {
|
||||
@@ -164,11 +177,8 @@ impl MemoryState {
|
||||
self.graph.modify(update_fn);
|
||||
}
|
||||
|
||||
pub fn get_search_index(&self) -> MemoryIndex {
|
||||
self.search_index
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.clone()
|
||||
pub async fn get_search_index(&self) -> MemoryIndex {
|
||||
self.search_index.read().await.clone()
|
||||
}
|
||||
|
||||
pub fn search(self: &Arc<Self>) -> SearchService {
|
||||
@@ -176,8 +186,18 @@ impl MemoryState {
|
||||
}
|
||||
|
||||
pub async fn rebuild_index(self: &Arc<Self>) {
|
||||
let idx = self.get_search_index();
|
||||
idx.delete_all();
|
||||
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 entities: Vec<_> = self
|
||||
.graph
|
||||
@@ -192,9 +212,8 @@ impl MemoryState {
|
||||
tasks.len()
|
||||
);
|
||||
|
||||
let idx_clone = idx.clone();
|
||||
let idx_clone = new_idx.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
// tracing::info!("spawn_blocking started in rebuild_index");
|
||||
for e in entities {
|
||||
idx_clone.add_entity_sync(&e);
|
||||
}
|
||||
@@ -213,10 +232,8 @@ impl MemoryState {
|
||||
tracing::error!("Failed to join tantivy index rebuild thread: {}", e);
|
||||
});
|
||||
|
||||
let _ = idx.commit().await;
|
||||
if let Ok(mut w) = self.search_index.write() {
|
||||
*w = idx;
|
||||
}
|
||||
let _ = new_idx.commit().await;
|
||||
*self.search_index.write().await = new_idx;
|
||||
self.index_commit_notify.notify_waiters();
|
||||
}
|
||||
|
||||
@@ -350,7 +367,7 @@ mod tests {
|
||||
arc_state.rebuild_index().await;
|
||||
|
||||
// Check search index initialization
|
||||
let idx = arc_state.search_index.read().unwrap();
|
||||
let idx = arc_state.search_index.read().await;
|
||||
// Force reload reader to ensure it sees the commit made by rebuild_index
|
||||
idx.reader.reload().unwrap();
|
||||
// tracing::info!(
|
||||
@@ -464,31 +481,46 @@ impl SearchService {
|
||||
}
|
||||
|
||||
if !vdb_search {
|
||||
let (mut texts_to_embed, mut metadata) = self.state.code.snippets.read_with(|snips| {
|
||||
let mut texts = Vec::with_capacity(snips.len().min(50));
|
||||
let mut meta = Vec::with_capacity(snips.len().min(50));
|
||||
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) {
|
||||
texts.push(format!("{} {} {}", snippet.name, snippet.description, snippet.code));
|
||||
meta.push((snippet.name.clone(), "snippet".to_string(), snippet.description.clone()));
|
||||
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));
|
||||
}
|
||||
}
|
||||
(texts, meta)
|
||||
});
|
||||
|
||||
self.state.code.sticky.read_with(|sticky| {
|
||||
for note in sticky.iter().take(50) {
|
||||
texts_to_embed.push(note.content.clone());
|
||||
metadata.push((
|
||||
"StickyNote".to_string(),
|
||||
"sticky".to_string(),
|
||||
note.content.chars().take(200).collect::<String>(),
|
||||
));
|
||||
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));
|
||||
}
|
||||
});
|
||||
|
||||
if !texts_to_embed.is_empty()
|
||||
&& let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await
|
||||
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(metadata) {
|
||||
for (emb, meta) in embeddings.into_iter().zip(uncached_meta) {
|
||||
let sim = cosine_similarity(&query_emb, &emb);
|
||||
results.push(UnifiedSearchResult {
|
||||
id: meta.0.clone(),
|
||||
@@ -511,13 +543,13 @@ impl SearchService {
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
pub fn keyword_search(
|
||||
pub async fn keyword_search(
|
||||
&self,
|
||||
query: &str,
|
||||
filter_namespace: Option<&str>,
|
||||
limit: usize,
|
||||
) -> crate::error::Result<Vec<UnifiedSearchResult>> {
|
||||
let idx = self.state.get_search_index();
|
||||
let idx = self.state.get_search_index().await;
|
||||
let matches = idx
|
||||
.search(query, filter_namespace)
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
|
||||
|
||||
+138
-42
@@ -8,7 +8,8 @@ pub const STORE_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("stor
|
||||
struct DbWriteTask {
|
||||
key: String,
|
||||
data: Vec<u8>,
|
||||
flushed: Arc<tokio::sync::Notify>,
|
||||
flushed_notifier: Arc<tokio::sync::Notify>,
|
||||
oneshot_tx: Option<tokio::sync::oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
/// Shared centralized write queue actor that handles all database writes serially with micro-batching.
|
||||
@@ -33,7 +34,7 @@ fn get_or_create_queue(db: Arc<Database>) -> DbWriteQueue {
|
||||
|
||||
impl DbWriteQueue {
|
||||
pub fn new(db: Arc<Database>) -> Self {
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel::<DbWriteTask>(2048);
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel::<DbWriteTask>(1024);
|
||||
|
||||
tokio::spawn(async move {
|
||||
while let Some(first_task) = rx.recv().await {
|
||||
@@ -69,7 +70,10 @@ impl DbWriteQueue {
|
||||
|
||||
// Event-driven notification to all waiting listeners for this micro-batch
|
||||
for task in batch {
|
||||
task.flushed.notify_waiters();
|
||||
if let Some(oneshot) = task.oneshot_tx {
|
||||
let _ = oneshot.send(());
|
||||
}
|
||||
task.flushed_notifier.notify_waiters();
|
||||
}
|
||||
})
|
||||
.await;
|
||||
@@ -79,14 +83,46 @@ impl DbWriteQueue {
|
||||
Self { tx }
|
||||
}
|
||||
|
||||
pub fn push(&self, key: String, data: Vec<u8>, flushed: Arc<tokio::sync::Notify>) {
|
||||
let task = DbWriteTask { key, data, flushed };
|
||||
pub fn push(
|
||||
&self,
|
||||
key: String,
|
||||
data: Vec<u8>,
|
||||
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,
|
||||
flushed_notifier,
|
||||
oneshot_tx: Some(oneshot_tx),
|
||||
};
|
||||
if let Err(e) = self.tx.try_send(task) {
|
||||
let task = e.into_inner();
|
||||
let tx = self.tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let _ = tx.send(task).await;
|
||||
});
|
||||
let key = e.into_inner().key;
|
||||
tracing::error!("DbWriteQueue channel full or closed; unable to persist key '{}'", key);
|
||||
None
|
||||
} else {
|
||||
Some(oneshot_rx)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn push_async(
|
||||
&self,
|
||||
key: String,
|
||||
data: Vec<u8>,
|
||||
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,
|
||||
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);
|
||||
None
|
||||
} else {
|
||||
Some(oneshot_rx)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -96,11 +132,12 @@ pub struct Store<T> {
|
||||
pub flushed: Arc<tokio::sync::Notify>,
|
||||
key: String,
|
||||
queue: DbWriteQueue,
|
||||
is_corrupted: bool,
|
||||
}
|
||||
|
||||
impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static> Store<T> {
|
||||
impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T> {
|
||||
pub fn new(key: &str, db: Arc<Database>) -> Self {
|
||||
let initial_data = Self::load_from_db(key, &db);
|
||||
let (initial_data, is_corrupted) = Self::load_from_db(key, &db);
|
||||
let cache = Arc::new(RwLock::new(initial_data));
|
||||
let flushed = Arc::new(tokio::sync::Notify::new());
|
||||
let queue = get_or_create_queue(db);
|
||||
@@ -110,20 +147,38 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
|
||||
flushed,
|
||||
key: key.to_string(),
|
||||
queue,
|
||||
is_corrupted,
|
||||
}
|
||||
}
|
||||
|
||||
fn load_from_db(key: &str, db: &Database) -> T {
|
||||
fn load_from_db(key: &str, db: &Database) -> (T, bool) {
|
||||
let Ok(read_txn) = db.begin_read() else {
|
||||
return T::default();
|
||||
tracing::error!("Failed to begin read transaction for key '{}'", key);
|
||||
return (T::default(), false);
|
||||
};
|
||||
if let Ok(table) = read_txn.open_table(STORE_TABLE)
|
||||
&& let Ok(Some(value)) = table.get(key)
|
||||
&& let Ok(parsed) = serde_json::from_slice::<T>(value.value())
|
||||
{
|
||||
return parsed;
|
||||
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),
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to get key '{}' from store table: {}", key, e);
|
||||
(T::default(), false)
|
||||
}
|
||||
},
|
||||
Err(e) => {
|
||||
tracing::error!("Failed to open STORE_TABLE for key '{}': {}", key, e);
|
||||
(T::default(), false)
|
||||
}
|
||||
}
|
||||
T::default()
|
||||
}
|
||||
|
||||
pub fn read_with<F, R>(&self, f: F) -> R
|
||||
@@ -134,15 +189,63 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
|
||||
f(&lock)
|
||||
}
|
||||
|
||||
pub fn modify<F: FnOnce(&mut T)>(&self, f: F) {
|
||||
let cloned_data = {
|
||||
pub fn modify<F: FnOnce(&mut T)>(&self, f: F)
|
||||
where
|
||||
T: Serialize,
|
||||
{
|
||||
if self.is_corrupted {
|
||||
tracing::error!(
|
||||
"CRITICAL: Refusing to persist changes for corrupted store key '{}' to prevent data loss.",
|
||||
self.key
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let serialized_res = {
|
||||
let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner());
|
||||
f(&mut lock);
|
||||
lock.clone()
|
||||
serde_json::to_vec(&*lock)
|
||||
};
|
||||
|
||||
match serde_json::to_vec(&cloned_data) {
|
||||
Ok(data) => self.queue.push(self.key.clone(), data, self.flushed.clone()),
|
||||
match serialized_res {
|
||||
Ok(data) => {
|
||||
if self.queue.push(self.key.clone(), data.clone(), self.flushed.clone()).is_none() {
|
||||
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;
|
||||
});
|
||||
}
|
||||
}
|
||||
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,
|
||||
{
|
||||
if self.is_corrupted {
|
||||
tracing::error!(
|
||||
"CRITICAL: Refusing to persist changes for corrupted store key '{}' to prevent data loss.",
|
||||
self.key
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
let serialized_res = {
|
||||
let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner());
|
||||
f(&mut lock);
|
||||
serde_json::to_vec(&*lock)
|
||||
};
|
||||
|
||||
match serialized_res {
|
||||
Ok(data) => {
|
||||
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),
|
||||
}
|
||||
}
|
||||
@@ -151,7 +254,6 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
#[derive(Serialize, serde::Deserialize, Clone, Default, PartialEq, Debug)]
|
||||
struct TestData {
|
||||
@@ -159,18 +261,21 @@ mod tests {
|
||||
value: i32,
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_store_read_write() {
|
||||
let temp_file = NamedTempFile::new().unwrap();
|
||||
let db = Database::create(temp_file.path()).unwrap();
|
||||
|
||||
fn create_in_memory_test_db() -> Arc<Database> {
|
||||
let db = Database::builder()
|
||||
.create_with_backend(redb::backends::InMemoryBackend::new())
|
||||
.unwrap();
|
||||
let write_txn = db.begin_write().unwrap();
|
||||
{
|
||||
write_txn.open_table(STORE_TABLE).unwrap();
|
||||
}
|
||||
write_txn.commit().unwrap();
|
||||
Arc::new(db)
|
||||
}
|
||||
|
||||
let db = Arc::new(db);
|
||||
#[tokio::test]
|
||||
async fn test_store_read_write() {
|
||||
let db = create_in_memory_test_db();
|
||||
let store = Store::<TestData>::new("test_key", db.clone());
|
||||
|
||||
assert_eq!(store.read_with(|s| s.clone()), TestData::default());
|
||||
@@ -195,16 +300,7 @@ mod tests {
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn test_store_concurrency() {
|
||||
let temp_file = NamedTempFile::new().unwrap();
|
||||
let db = Database::create(temp_file.path()).unwrap();
|
||||
|
||||
let write_txn = db.begin_write().unwrap();
|
||||
{
|
||||
write_txn.open_table(STORE_TABLE).unwrap();
|
||||
}
|
||||
write_txn.commit().unwrap();
|
||||
|
||||
let db = Arc::new(db);
|
||||
let db = create_in_memory_test_db();
|
||||
let store = Arc::new(Store::<TestData>::new("concurrent_key", db.clone()));
|
||||
|
||||
let mut handles = vec![];
|
||||
|
||||
+13
-28
@@ -25,10 +25,11 @@ pub struct CreateEntitiesTool {
|
||||
pub entities: Vec<crate::models::Entity>,
|
||||
}
|
||||
|
||||
/// Create new relations between entities in the knowledge graph. Note: relation_type MUST ALWAYS be snake_case (e.g. depends_on, calls, implements).
|
||||
/// 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'.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct CreateRelationsTool {
|
||||
/// Array of relations to create.
|
||||
/// Array of relations to create. Each relation item MUST use 'from', 'to', and 'relation_type'.
|
||||
pub relations: Vec<crate::models::Relation>,
|
||||
}
|
||||
|
||||
@@ -81,7 +82,8 @@ pub struct ReadGraphTool {
|
||||
pub max_tokens: Option<usize>,
|
||||
}
|
||||
|
||||
/// Search for entities in the knowledge graph by name or type.
|
||||
/// Search specifically for Knowledge Graph entities and nodes by name or type.
|
||||
/// Note: For searching across tasks, snippets, ADRs, error fixes, and graph entities simultaneously, use 'omni_search' instead.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct SearchNodesTool {
|
||||
/// The search query.
|
||||
@@ -134,9 +136,13 @@ pub struct VisualizeGraphTool {
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum StickyNoteAction {
|
||||
#[serde(alias = "ADD", alias = "Add", alias = "create", alias = "insert")]
|
||||
Add,
|
||||
#[serde(alias = "READ", alias = "Read", alias = "list", alias = "get")]
|
||||
Read,
|
||||
#[serde(alias = "DELETE", alias = "Delete", alias = "remove", alias = "rm")]
|
||||
Delete,
|
||||
#[serde(alias = "CLEAR", alias = "Clear", alias = "purge")]
|
||||
Clear,
|
||||
}
|
||||
|
||||
@@ -161,31 +167,6 @@ pub struct StickyNotesTool {
|
||||
pub session_only: Option<bool>,
|
||||
}
|
||||
|
||||
/// Add a sticky note for unstructured thoughts or reminders.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct AddStickyNoteTool {
|
||||
/// The content of the sticky note.
|
||||
pub content: String,
|
||||
/// Optional time-to-live in seconds. Automatically expires the note after specified duration.
|
||||
pub ttl_seconds: Option<u64>,
|
||||
/// Optional flag to auto-expire the note at the end of the current session.
|
||||
pub session_only: Option<bool>,
|
||||
}
|
||||
|
||||
/// Read all active sticky notes.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct ReadStickyNotesTool {}
|
||||
|
||||
/// Delete a specific sticky note by its 1-indexed position.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct DeleteStickyNoteTool {
|
||||
/// The 1-indexed position of the sticky note to delete.
|
||||
pub index: usize,
|
||||
}
|
||||
|
||||
/// Clear all sticky notes.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct ClearStickyNotesTool {}
|
||||
|
||||
/// Condense or summarize an entity's observations to reduce size.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
@@ -266,8 +247,11 @@ pub struct TagSnippetTool {
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SnippetSearchMode {
|
||||
#[serde(alias = "HYBRID", alias = "Hybrid")]
|
||||
Hybrid,
|
||||
#[serde(alias = "KEYWORD", alias = "Keyword")]
|
||||
Keyword,
|
||||
#[serde(alias = "SEMANTIC", alias = "Semantic")]
|
||||
Semantic,
|
||||
}
|
||||
|
||||
@@ -374,6 +358,7 @@ pub struct LogDecisionTool {
|
||||
/// The decision made.
|
||||
pub decision: String,
|
||||
/// The consequence of the decision.
|
||||
#[serde(alias = "consequences")]
|
||||
pub consequence: String,
|
||||
/// The ID of an older ADR this decision supersedes, if any.
|
||||
pub supersedes: Option<String>,
|
||||
|
||||
@@ -14,6 +14,11 @@ pub struct VectorDB {
|
||||
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,
|
||||
@@ -98,6 +103,40 @@ impl VectorDB {
|
||||
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>,
|
||||
|
||||
+21
-8
@@ -4,16 +4,18 @@ use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
use tracing::{error, info};
|
||||
|
||||
pub fn spawn_watcher(_state: Arc<MemoryState>) {
|
||||
pub fn spawn_watcher(state: Arc<MemoryState>) {
|
||||
let watch_path = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
|
||||
info!("Spawning proactive daemon watcher on {:?}", watch_path);
|
||||
|
||||
tokio::spawn(async move {
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel::<notify::Result<Event>>(100);
|
||||
let (tx, mut rx) = tokio::sync::mpsc::channel::<notify::Result<Event>>(500);
|
||||
|
||||
let mut watcher = match RecommendedWatcher::new(
|
||||
move |res| {
|
||||
let _ = tx.try_send(res);
|
||||
if let Err(e) = tx.try_send(res) {
|
||||
tracing::warn!("Watcher event dropped due to channel backpressure: {}", e);
|
||||
}
|
||||
},
|
||||
Config::default(),
|
||||
) {
|
||||
@@ -29,14 +31,26 @@ pub fn spawn_watcher(_state: Arc<MemoryState>) {
|
||||
return;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
last_processed.insert(path.clone(), now);
|
||||
|
||||
info!("Proactive Daemon Hooks: File modified: {:?}", path);
|
||||
trigger_autonomous_review(&path).await;
|
||||
trigger_autonomous_review(&path, Arc::clone(&state)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -49,7 +63,7 @@ pub fn spawn_watcher(_state: Arc<MemoryState>) {
|
||||
|
||||
fn should_review(path: &Path) -> bool {
|
||||
let path_str = path.to_string_lossy();
|
||||
if path_str.contains(".git") || path_str.contains("target") {
|
||||
if path_str.contains(".git") || path_str.contains("target") || path_str.contains(".gemini") || path_str.contains("node_modules") {
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -62,10 +76,9 @@ fn should_review(path: &Path) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
async fn trigger_autonomous_review(path: &Path) {
|
||||
async fn trigger_autonomous_review(path: &Path, state: Arc<MemoryState>) {
|
||||
info!("Triggering autonomous review for {:?}", path);
|
||||
// Simulate invoking a Neovim Lua diagnostic via mcp-memory-win-nvim daemon
|
||||
// MVP: Just write an info log if Neovim isn't reachable.
|
||||
state.broadcast_activity("AUTONOMOUS", &format!("Modified: {:?}", path.file_name().unwrap_or_default()));
|
||||
info!("Autonomous review complete for {:?}", path);
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user