refactor: consolidate nvim crates, extract server library, and update workspace dependencies
This commit is contained in:
1 parent
a083719cf1
commit
533adfd41b
53 files changed
+5967
-1230
No files matched your search
+252
-8
@@ -120,8 +120,20 @@ pub fn create_router(app_state: Arc<AppState>) -> Router {
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
"Received shutdown request via /shutdown endpoint. Initiating graceful shutdown."
|
||||
"Received shutdown request via /shutdown endpoint. Creating session checkpoint and initiating graceful shutdown."
|
||||
);
|
||||
let mem_state = state.handler.state.clone();
|
||||
let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs();
|
||||
let memo = crate::models::HandoffMemo {
|
||||
id: format!("chk_shutdown_{}", now),
|
||||
author: "ShutdownHook".to_string(),
|
||||
content: format!("Automatic session checkpoint triggered on graceful server shutdown at epoch {}", now),
|
||||
expires_at: None,
|
||||
namespace: "global".to_string(),
|
||||
timestamp: now,
|
||||
};
|
||||
mem_state.telemetry.handoff_memos.modify(|m| m.push(memo));
|
||||
|
||||
if let Some(tx) = state.shutdown_tx.lock().unwrap().take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
@@ -389,15 +401,14 @@ mod tests {
|
||||
use crate::router::MemoryHandler;
|
||||
use crate::state::MemoryState;
|
||||
use axum::body::Body;
|
||||
use axum::http::Request;
|
||||
use axum::http::{Request, StatusCode};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::RwLock;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use tempfile::tempdir;
|
||||
use tower::ServiceExt;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_router_health() {
|
||||
async fn setup_app() -> (axum::Router, Arc<AppState>, tempfile::TempDir) {
|
||||
let dir = tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
let (shutdown_tx, _) = tokio::sync::oneshot::channel();
|
||||
@@ -408,16 +419,249 @@ mod tests {
|
||||
shutdown_tx: std::sync::Mutex::new(Some(shutdown_tx)),
|
||||
});
|
||||
|
||||
let app = create_router(app_state);
|
||||
let app = create_router(app_state.clone());
|
||||
(app, app_state, dir)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_router_health() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
// Test health endpoint
|
||||
let request = Request::builder()
|
||||
.uri("/health")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), 200);
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_ping_endpoint() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
let request = Request::builder()
|
||||
.uri("/ping")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_version_endpoint() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
let request = Request::builder()
|
||||
.uri("/api/version")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_metrics_endpoint() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
let request = Request::builder()
|
||||
.uri("/metrics")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_activity_endpoint() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
let request = Request::builder()
|
||||
.uri("/api/activity")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_nvim_telemetry_endpoint() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
let body_json = serde_json::json!({
|
||||
"session_id": "test_session",
|
||||
"event": "BufEnter",
|
||||
"file": "test.rs"
|
||||
});
|
||||
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/nvim/telemetry")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body_json.to_string()))
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_terminal_history_endpoint() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
let request = Request::builder()
|
||||
.uri("/api/terminal/history")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_terminal_telemetry_endpoint() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
let body_json = serde_json::json!({
|
||||
"command": "cargo build",
|
||||
"exit_code": 0,
|
||||
"cwd": "/tmp",
|
||||
"os": "windows",
|
||||
"timestamp": 1234567890
|
||||
});
|
||||
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/terminal/telemetry")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body_json.to_string()))
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_post_event_endpoint() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
let body_json = serde_json::json!({
|
||||
"topic": "test_topic",
|
||||
"session_id": "s123",
|
||||
"payload": { "key": "value" }
|
||||
});
|
||||
|
||||
let request = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/events")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(body_json.to_string()))
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_gate_verify_endpoint() {
|
||||
let (app, _, _dir) = setup_app().await;
|
||||
|
||||
let request = Request::builder()
|
||||
.uri("/gate/verify?action=deploy&target=prod")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(request).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_gate_set_and_verify_success() {
|
||||
let (app, app_state, _dir) = setup_app().await;
|
||||
|
||||
let set_req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/gate/set")
|
||||
.header("content-type", "application/json")
|
||||
.body(Body::from(serde_json::json!({
|
||||
"action": "deploy",
|
||||
"target": "prod",
|
||||
"authorize": true,
|
||||
"reason": "Tests passed"
|
||||
}).to_string()))
|
||||
.unwrap();
|
||||
|
||||
let response = app.oneshot(set_req).await.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
|
||||
let app2 = create_router(app_state);
|
||||
let verify_req = Request::builder()
|
||||
.uri("/gate/verify?action=deploy&target=prod")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
|
||||
let verify_resp = app2.oneshot(verify_req).await.unwrap();
|
||||
assert_eq!(verify_resp.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_all_rest_endpoints_comprehensive() {
|
||||
let (_app, app_state, _dir) = setup_app().await;
|
||||
|
||||
let endpoints = vec![
|
||||
"/",
|
||||
"/api/graph",
|
||||
"/api/tasks",
|
||||
"/api/sticky",
|
||||
"/api/search?q=test",
|
||||
"/api/tech_debts",
|
||||
"/api/adrs",
|
||||
"/api/context_workspaces",
|
||||
"/api/handoff_memos",
|
||||
"/api/milestones",
|
||||
"/api/snippets",
|
||||
"/api/pr_checklists",
|
||||
"/api/error_fixes",
|
||||
"/api/stats",
|
||||
];
|
||||
|
||||
for ep in endpoints {
|
||||
let app_inst = create_router(app_state.clone());
|
||||
let req = Request::builder()
|
||||
.uri(ep)
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let resp = app_inst.oneshot(req).await.unwrap();
|
||||
assert_eq!(resp.status(), StatusCode::OK, "Failed endpoint: {}", ep);
|
||||
}
|
||||
|
||||
// Test task completion endpoint
|
||||
let app_task = create_router(app_state.clone());
|
||||
let complete_req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/tasks/task_123/complete")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let complete_resp = app_task.oneshot(complete_req).await.unwrap();
|
||||
assert_eq!(complete_resp.status(), StatusCode::OK);
|
||||
|
||||
// Test /shutdown unauthorized without token
|
||||
let app_shut = create_router(app_state.clone());
|
||||
let shut_req = Request::builder()
|
||||
.method("POST")
|
||||
.uri("/shutdown")
|
||||
.body(Body::empty())
|
||||
.unwrap();
|
||||
let shut_resp = app_shut.oneshot(shut_req).await.unwrap();
|
||||
assert_eq!(shut_resp.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
}
|
||||
// Force recompile
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -104,14 +104,7 @@ pub async fn terminal_telemetry_handler(
|
||||
state
|
||||
.handler
|
||||
.state
|
||||
.telemetry
|
||||
.terminal_history
|
||||
.modify(|history| {
|
||||
history.push_front(payload.clone());
|
||||
if history.len() > 100 {
|
||||
history.pop_back();
|
||||
}
|
||||
});
|
||||
.record_terminal_history(payload.clone());
|
||||
|
||||
let ws_msg = serde_json::json!({
|
||||
"type": "terminal_telemetry",
|
||||
|
||||
@@ -1,10 +0,0 @@
|
||||
use arboard::Clipboard;
|
||||
|
||||
fn main() {
|
||||
if let Ok(mut clipboard) = Clipboard::new() {
|
||||
let text = clipboard.get_text();
|
||||
println!("Text: {:?}", text.ok());
|
||||
} else {
|
||||
println!("Failed to open clipboard");
|
||||
}
|
||||
}
|
||||
@@ -30,7 +30,7 @@ pub fn spawn_watcher(state: Arc<MemoryState>) {
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs(),
|
||||
content: text,
|
||||
content: text.clone(),
|
||||
expires_at: None,
|
||||
};
|
||||
|
||||
@@ -38,9 +38,38 @@ pub fn spawn_watcher(state: Arc<MemoryState>) {
|
||||
notes.push(note.clone());
|
||||
});
|
||||
|
||||
state.record_activity(
|
||||
"clipboard",
|
||||
"Auto-ingested clipboard content into sticky notes",
|
||||
Some(&text),
|
||||
);
|
||||
|
||||
// We use rebuild_index to index the new sticky note
|
||||
state.rebuild_index().await;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_spawn_watcher_lifecycle() {
|
||||
let dir = tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
|
||||
// Set watch mode to disabled first
|
||||
*state.clipboard_watch_mode.write().await = false;
|
||||
spawn_watcher(state.clone());
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
|
||||
// Enable watch mode
|
||||
*state.clipboard_watch_mode.write().await = true;
|
||||
tokio::time::sleep(Duration::from_millis(150)).await;
|
||||
}
|
||||
}
|
||||
|
||||
+226
-54
@@ -633,29 +633,61 @@
|
||||
|
||||
<script>
|
||||
// --- Tabs ---
|
||||
let currentTabId = 'graph-tab';
|
||||
|
||||
function switchTab(tabId, btn) {
|
||||
currentTabId = tabId;
|
||||
|
||||
document.querySelectorAll('.tab-content').forEach(el => el.classList.remove('active'));
|
||||
document.querySelectorAll('.tab-button').forEach(el => el.classList.remove('active'));
|
||||
|
||||
document.getElementById(tabId).classList.add('active');
|
||||
btn.classList.add('active');
|
||||
if (btn) btn.classList.add('active');
|
||||
|
||||
if (tabId === 'graph-tab' && network) {
|
||||
// Resize graph when becoming visible again
|
||||
network.redraw();
|
||||
network.fit();
|
||||
loadGraph();
|
||||
} else if (tabId === 'activity-tab') {
|
||||
const feed = document.getElementById('activity-feed');
|
||||
feed.scrollTop = 0;
|
||||
} else if (tabId === 'task-tab') {
|
||||
loadTasks();
|
||||
} else if (tabId === 'preferences-tab') {
|
||||
loadPreferences();
|
||||
} else if (tabId === 'sticky-tab') {
|
||||
loadStickyNotes();
|
||||
} else {
|
||||
loadAllExtras();
|
||||
switch (tabId) {
|
||||
case 'graph-tab':
|
||||
if (network) {
|
||||
network.redraw();
|
||||
} else {
|
||||
loadGraph();
|
||||
}
|
||||
break;
|
||||
case 'activity-tab':
|
||||
loadActivityHistory();
|
||||
break;
|
||||
case 'task-tab':
|
||||
loadTasks();
|
||||
break;
|
||||
case 'preferences-tab':
|
||||
loadPreferences();
|
||||
break;
|
||||
case 'sticky-tab':
|
||||
loadStickyNotes();
|
||||
break;
|
||||
case 'techdebt-tab':
|
||||
loadTechDebt();
|
||||
break;
|
||||
case 'adrs-tab':
|
||||
loadADRs();
|
||||
break;
|
||||
case 'workspaces-tab':
|
||||
loadWorkspaces();
|
||||
break;
|
||||
case 'pinned-tab':
|
||||
loadPinned();
|
||||
break;
|
||||
case 'memos-tab':
|
||||
loadMemos();
|
||||
break;
|
||||
case 'snippets-tab':
|
||||
loadSnippets();
|
||||
break;
|
||||
case 'pr-tab':
|
||||
loadPRs();
|
||||
break;
|
||||
case 'terminal-tab':
|
||||
loadTerminal();
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1028,23 +1060,99 @@
|
||||
// --- WebSocket Activity Feed ---
|
||||
const MAX_ACTIVITY_HISTORY = 100;
|
||||
|
||||
function parseActivityPayload(item) {
|
||||
let category = 'activity';
|
||||
let message = '';
|
||||
let timestamp = Date.now();
|
||||
|
||||
if (typeof item === 'string') {
|
||||
try {
|
||||
const parsed = JSON.parse(item);
|
||||
return parseActivityPayload(parsed);
|
||||
} catch (e) {
|
||||
message = item;
|
||||
}
|
||||
} else if (typeof item === 'object' && item !== null) {
|
||||
if (item.method === 'notifications/activity' && item.params) {
|
||||
message = item.params.message || JSON.stringify(item.params);
|
||||
timestamp = item.params.time || item.params.timestamp || timestamp;
|
||||
category = item.params.category || item.params.type || 'tool';
|
||||
} else {
|
||||
category = item.category || item.type || category;
|
||||
message = item.summary || item.message || item.description || item.data || JSON.stringify(item);
|
||||
timestamp = item.timestamp || item.time || item.updated_at || timestamp;
|
||||
if (typeof timestamp === 'number' && timestamp < 10000000000) {
|
||||
timestamp = timestamp * 1000;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback inference from tool name if category is generic
|
||||
const lowerCategory = category.toLowerCase();
|
||||
if (lowerCategory === 'activity' || lowerCategory === 'mcp_tool' || lowerCategory === 'tool') {
|
||||
const match = message.match(/Agent executed tool:\s*([a-z0-9_]+)/i);
|
||||
if (match && match[1]) {
|
||||
const tool = match[1].toLowerCase();
|
||||
if (tool.includes('clipboard')) category = 'CLIPBOARD';
|
||||
else if (tool.includes('graph') || tool.includes('entities') || tool.includes('relations') || tool.includes('schema')) category = 'GRAPH';
|
||||
else if (tool.includes('decision')) category = 'DECISION';
|
||||
else if (tool.includes('code_change') || tool.includes('ast')) category = 'CODE';
|
||||
else if (tool.includes('task') || tool.includes('milestone')) category = 'TASK';
|
||||
else if (tool.includes('note')) category = 'STICKY_NOTE';
|
||||
else if (tool.includes('checkpoint') || tool.includes('snapshot')) category = 'CHECKPOINT';
|
||||
else if (tool.includes('subagent') || tool.includes('namespace')) category = 'SUBAGENT';
|
||||
else if (tool.includes('snippet')) category = 'SNIPPET';
|
||||
else if (tool.includes('error')) category = 'ERROR_FIX';
|
||||
else category = tool.toUpperCase();
|
||||
}
|
||||
}
|
||||
|
||||
const catUpper = category.toUpperCase();
|
||||
const colorMap = {
|
||||
'GRAPH': '#2ecc71',
|
||||
'DECISION': '#f39c12',
|
||||
'CODE': '#3498db',
|
||||
'TASK': '#1abc9c',
|
||||
'CLIPBOARD': '#9b59b6',
|
||||
'STICKY_NOTE': '#e67e22',
|
||||
'CHECKPOINT': '#e74c3c',
|
||||
'ERROR_FIX': '#e74c3c',
|
||||
'TECH_DEBT': '#d35400',
|
||||
'SUBAGENT': '#8e44ad',
|
||||
'SNIPPET': '#16a085',
|
||||
'TOOL': '#3498db',
|
||||
'SYSTEM': '#95a5a6'
|
||||
};
|
||||
const badgeColor = colorMap[catUpper] || '#3498db';
|
||||
|
||||
const dateObj = new Date(timestamp);
|
||||
const validDate = isNaN(dateObj.getTime()) ? new Date() : dateObj;
|
||||
const timeStr = validDate.toLocaleString([], {
|
||||
month: '2-digit', day: '2-digit', hour: '2-digit', minute:'2-digit', second:'2-digit'
|
||||
});
|
||||
|
||||
return `<span class="time">[${timeStr}]</span> <span style="background:var(--canvas-bg); color:${badgeColor}; padding:2px 6px; border-radius:3px; font-weight:bold; margin-right:6px; border:1px solid ${badgeColor}60; text-transform:uppercase; font-size:0.75em;">${catUpper}</span> ${message}`;
|
||||
}
|
||||
|
||||
async function loadActivityHistory() {
|
||||
try {
|
||||
const response = await fetch('/api/activity');
|
||||
const history = await response.json();
|
||||
const feed = document.getElementById('activity-feed');
|
||||
if (!feed) return;
|
||||
feed.innerHTML = '';
|
||||
[...history].reverse().forEach(item => {
|
||||
const div = document.createElement('div');
|
||||
div.className = 'feed-entry';
|
||||
const timeStr = new Date(item.time).toLocaleString([], {month: '2-digit', day: '2-digit', hour: '2-digit', minute:'2-digit', second:'2-digit'});
|
||||
div.innerHTML = `<span class="time">[${timeStr}]</span> ${item.message || item.data}`;
|
||||
div.innerHTML = parseActivityPayload(item);
|
||||
feed.appendChild(div);
|
||||
});
|
||||
if (history.length > 0) {
|
||||
feed.scrollTop = 0;
|
||||
}
|
||||
} catch(e) {}
|
||||
} catch(e) {
|
||||
console.error("Failed to load activity history", e);
|
||||
}
|
||||
}
|
||||
|
||||
function setupWS() {
|
||||
@@ -1056,27 +1164,21 @@
|
||||
try {
|
||||
const data = JSON.parse(event.data);
|
||||
if (data.type === 'activity') {
|
||||
// Check scroll state before appending
|
||||
const isScrolledToTop = feed.scrollTop <= 20;
|
||||
const isScrolledToTop = feed ? feed.scrollTop <= 20 : true;
|
||||
|
||||
const div = document.createElement('div');
|
||||
div.className = 'feed-entry';
|
||||
const timeStr = new Date(data.data.time).toLocaleString([], {month: '2-digit', day: '2-digit', hour: '2-digit', minute:'2-digit', second:'2-digit'});
|
||||
div.innerHTML = `<span class="time">[${timeStr}]</span> ${data.data.message || data.data.data || data.data}`;
|
||||
feed.prepend(div);
|
||||
|
||||
// Enforce max history
|
||||
while (feed.children.length > MAX_ACTIVITY_HISTORY) {
|
||||
feed.removeChild(feed.lastChild);
|
||||
}
|
||||
|
||||
// Auto-scroll logic
|
||||
if (isScrolledToTop) {
|
||||
feed.scrollTop = 0;
|
||||
div.innerHTML = parseActivityPayload(data.data);
|
||||
if (feed) {
|
||||
feed.prepend(div);
|
||||
while (feed.children.length > MAX_ACTIVITY_HISTORY) {
|
||||
feed.removeChild(feed.lastChild);
|
||||
}
|
||||
if (isScrolledToTop) feed.scrollTop = 0;
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
// Ignore non-JSON or other messages for now
|
||||
console.error("WebSocket message parse error", e);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -1161,7 +1263,7 @@
|
||||
}
|
||||
}
|
||||
|
||||
function loadAllExtras() {
|
||||
function loadTerminal() {
|
||||
loadGenericList('/api/terminal/history', 'terminal-container', item => `
|
||||
<span style="font-weight:bold; font-size:1.05em; color:var(--text-primary);">${item.command}</span>
|
||||
<span style="margin-left:15px; font-size:0.85em;"><span style="font-weight:bold; color:var(--text-secondary);">Exit Code:</span> <span style="color:${item.exit_code === 0 ? 'var(--success-color)' : 'var(--error-color)'}; font-weight:bold;">${item.exit_code}</span></span>
|
||||
@@ -1170,7 +1272,9 @@
|
||||
<span style="background:var(--canvas-bg); padding:2px 6px; border-radius:4px; font-family:monospace; border:1px solid var(--border-color);">CWD: ${item.cwd || 'Unknown'}</span>
|
||||
</div>
|
||||
`);
|
||||
}
|
||||
|
||||
function loadTechDebt() {
|
||||
loadGenericList('/api/tech_debts', 'techdebt-container', item => `
|
||||
<strong>${item.id}</strong> ${item.is_resolved ? '<span style="color:var(--success-color); font-weight:bold; background: rgba(46, 204, 113, 0.1); padding: 2px 6px; border-radius: 4px; border: 1px solid var(--success-color); margin-left: 8px;">(Resolved)</span>' : '<span style="color:var(--open-color); font-weight:bold; background: rgba(230, 126, 34, 0.1); padding: 2px 6px; border-radius: 4px; border: 1px solid var(--open-color); margin-left: 8px;">(Open)</span>'}
|
||||
<div style="margin-top:5px;"><strong>Description:</strong> ${item.description}</div>
|
||||
@@ -1186,7 +1290,9 @@
|
||||
<div style="font-size:0.9em; margin-bottom:8px; margin-top:5px;"><em>Solution:</em> ${item.solution || ''}</div>
|
||||
<div style="font-family:monospace; font-size:0.85em; color:#e74c3c;">Commit: ${item.git_commit || 'None'}</div>
|
||||
`);
|
||||
|
||||
}
|
||||
|
||||
function loadADRs() {
|
||||
loadGenericList('/api/adrs', 'adrs-container', item => `
|
||||
<strong>${item.id} | ${item.title}</strong>
|
||||
<span style="float:right; background:var(--canvas-bg); padding:2px 6px; border-radius:4px; font-size:0.8em; border:1px solid var(--border-color); ${item.status === 'superseded' ? 'text-decoration: line-through; color:var(--error-color);' : 'color:var(--success-color);'}">${item.status}</span>
|
||||
@@ -1195,24 +1301,44 @@
|
||||
<div style="font-size:0.9em; margin-bottom:8px;"><em>Consequence:</em> ${item.consequence || ''}</div>
|
||||
${item.supersedes ? `<div style="font-size:0.9em; font-weight:bold; color:var(--warning-color);">Supersedes: ${item.supersedes}</div>` : ''}
|
||||
`);
|
||||
}
|
||||
|
||||
function loadWorkspaces() {
|
||||
loadGenericList('/api/context_workspaces', 'workspaces-container', item => `
|
||||
<strong>${item.name}</strong>
|
||||
<div style="margin-bottom:8px;">${item.description || ''}</div>
|
||||
<div style="font-family:monospace; font-size:0.85em; color:#3498db;">${(item.paths || []).join(', ')}</div>
|
||||
`);
|
||||
}
|
||||
|
||||
function loadPinned() {
|
||||
loadGenericList('/api/pinned_files', 'pinned-container', item => `
|
||||
<strong>${item.path || item.file_path || item.id}</strong>
|
||||
<div style="margin-top:5px; font-size:0.85em; color:var(--text-secondary);">${item.reason || item.description || 'Pinned'}</div>
|
||||
`);
|
||||
}
|
||||
|
||||
function loadMemos() {
|
||||
loadGenericList('/api/handoff_memos', 'memos-container', item => `
|
||||
<strong>Memo from ${item.author || 'System'}</strong>
|
||||
<div style="white-space:pre-wrap; font-size:0.9em; line-height:1.4;">${item.content || item.summary || ''}</div>
|
||||
`);
|
||||
|
||||
loadGenericList('/api/milestones', 'milestones-container', item => `
|
||||
<strong>${item.title || item.name}</strong>
|
||||
<div style="margin-top:5px; font-size:0.85em;">${item.description || ''}</div>
|
||||
`);
|
||||
}
|
||||
|
||||
function loadSnippets() {
|
||||
loadGenericList('/api/snippets', 'snippets-container', item => `
|
||||
<strong>${item.description || 'Snippet'}</strong>
|
||||
<div style="margin-bottom:8px; font-size:0.8em;">Language: ${item.language || 'txt'} | Tags: ${(item.tags || []).join(', ')}</div>
|
||||
<pre style="background:#111417; color:#d1d8e0; padding:10px; border-radius:4px; overflow-x:auto; font-size:0.85em;">${item.content || item.code || ''}</pre>
|
||||
`);
|
||||
}
|
||||
|
||||
function loadPRs() {
|
||||
loadGenericList('/api/pr_checklists', 'pr-container', item => `
|
||||
<strong>${item.name || 'Checklist'}</strong>
|
||||
<ul style="padding-left:20px; font-size:0.9em;">
|
||||
@@ -1225,6 +1351,17 @@
|
||||
`);
|
||||
}
|
||||
|
||||
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');
|
||||
@@ -1264,16 +1401,17 @@
|
||||
sse.onmessage = function(event) {
|
||||
if (event.data) {
|
||||
try {
|
||||
const rec = JSON.parse(event.data);
|
||||
const feed = document.getElementById('activity-feed');
|
||||
if (feed) {
|
||||
const item = document.createElement('div');
|
||||
item.style.padding = '8px 12px';
|
||||
item.style.borderBottom = '1px solid var(--border-color)';
|
||||
item.style.fontSize = '0.85em';
|
||||
item.style.background = 'rgba(39, 174, 96, 0.08)';
|
||||
item.innerHTML = `<strong>${rec.category || 'activity'}</strong>: ${rec.summary || event.data}`;
|
||||
feed.insertBefore(item, feed.firstChild);
|
||||
const isScrolledToTop = feed.scrollTop <= 20;
|
||||
const div = document.createElement('div');
|
||||
div.className = 'feed-entry';
|
||||
div.innerHTML = parseActivityPayload(event.data);
|
||||
feed.prepend(div);
|
||||
while (feed.children.length > MAX_ACTIVITY_HISTORY) {
|
||||
feed.removeChild(feed.lastChild);
|
||||
}
|
||||
if (isScrolledToTop) feed.scrollTop = 0;
|
||||
}
|
||||
} catch(e) {}
|
||||
}
|
||||
@@ -1300,10 +1438,6 @@
|
||||
// --- Start ---
|
||||
loadVersion();
|
||||
loadGraph();
|
||||
loadTasks();
|
||||
loadStickyNotes();
|
||||
loadActivityHistory();
|
||||
loadAllExtras();
|
||||
setupWS();
|
||||
setupSSE();
|
||||
|
||||
@@ -1311,12 +1445,50 @@
|
||||
const observer = new MutationObserver(() => updateGraphData());
|
||||
observer.observe(document.documentElement, { attributes: true, attributeFilter: ['data-theme'] });
|
||||
|
||||
// Smart background polling: refresh only the currently active tab
|
||||
setInterval(() => {
|
||||
loadGraph();
|
||||
loadTasks();
|
||||
loadStickyNotes();
|
||||
loadAllExtras();
|
||||
}, 5000);
|
||||
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;
|
||||
}
|
||||
}, 10000);
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -79,4 +79,28 @@ mod tests {
|
||||
let res = generate_embeddings_async(vec![]).await.unwrap();
|
||||
assert!(res.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
async fn test_generate_embeddings_async_single_text() {
|
||||
let text = "test text".to_string();
|
||||
let res = generate_embeddings_async(vec![text.clone()]).await.unwrap();
|
||||
assert_eq!(res.len(), 1);
|
||||
assert_eq!(res[0].len(), 384);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_generate_embeddings_async_multiple_texts() {
|
||||
let texts = vec![
|
||||
"test text 1".to_string(),
|
||||
"test text 2".to_string(),
|
||||
"test text 3".to_string(),
|
||||
];
|
||||
let res = generate_embeddings_async(texts.clone()).await.unwrap();
|
||||
assert_eq!(res.len(), 3);
|
||||
for embedding in &res {
|
||||
assert_eq!(embedding.len(), 384);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -247,6 +247,154 @@ impl McpTool for ReplaceAstNodeHandler {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct FindSymbolReferencesHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for FindSymbolReferencesHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"find_symbol_references"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<crate::tools::FindSymbolReferencesTool>(
|
||||
"find_symbol_references",
|
||||
"Find all source locations and AST chunks where a specific symbol is referenced or called.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
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 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());
|
||||
if refs.len() >= limit {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&refs)?)
|
||||
})?;
|
||||
|
||||
Ok(matches)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct GetCallersHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for GetCallersHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"get_callers"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<crate::tools::GetCallersTool>(
|
||||
"get_callers",
|
||||
"Find all caller functions or methods that invoke a specified target function name.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
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 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());
|
||||
if matching.len() >= limit {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&matching)?)
|
||||
})?;
|
||||
|
||||
Ok(callers)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct AnalyzeImpactHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for AnalyzeImpactHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"analyze_impact"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<crate::tools::AnalyzeImpactTool>(
|
||||
"analyze_impact",
|
||||
"Analyze the potential downstream breaking impact of modifying a function, struct, or file.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: crate::tools::AnalyzeImpactTool =
|
||||
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let sym = req.target_symbol.clone();
|
||||
|
||||
let mut callers = Vec::new();
|
||||
state.code.snippets.read_with(|snippets| {
|
||||
for snippet in snippets {
|
||||
if snippet.code.contains(&sym) {
|
||||
callers.push(snippet.name.clone());
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let mut kg_connected = Vec::new();
|
||||
state.read_graph(|g| {
|
||||
for rel in &g.relations {
|
||||
if rel.from == sym {
|
||||
kg_connected.push(format!("Outgoing: {} -> {}", rel.relation_type, rel.to));
|
||||
} else if rel.to == sym {
|
||||
kg_connected.push(format!("Incoming: {} <- {}", rel.relation_type, rel.from));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let caller_count = callers.len();
|
||||
let graph_count = kg_connected.len();
|
||||
let risk_level = if caller_count > 10 || graph_count > 5 {
|
||||
"CRITICAL"
|
||||
} else if caller_count > 3 || graph_count > 2 {
|
||||
"HIGH"
|
||||
} else if caller_count > 0 || graph_count > 0 {
|
||||
"MEDIUM"
|
||||
} else {
|
||||
"LOW"
|
||||
};
|
||||
|
||||
let result = serde_json::json!({
|
||||
"target_symbol": sym,
|
||||
"risk_level": risk_level,
|
||||
"ast_callers_count": caller_count,
|
||||
"ast_callers_sample": callers.into_iter().take(5).collect::<Vec<_>>(),
|
||||
"graph_relations_count": graph_count,
|
||||
"graph_relations": kg_connected,
|
||||
"recommendation": match risk_level {
|
||||
"CRITICAL" | "HIGH" => "Requires comprehensive unit test verification and backwards compatibility checks before modifying.",
|
||||
"MEDIUM" => "Verify direct call sites and run affected module tests.",
|
||||
_ => "Safe to modify with standard unit test verification.",
|
||||
}
|
||||
});
|
||||
|
||||
Ok(serde_json::to_string_pretty(&result)?)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
@@ -270,4 +270,26 @@ mod tests {
|
||||
assert!(res3.contains("prod.local"));
|
||||
assert!(!res3.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_empty_namespace_in_update_env_fingerprint() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
|
||||
let update_handler = UpdateEnvFingerprintHandler;
|
||||
let args = serde_json::json!({
|
||||
"namespace": "",
|
||||
"tool_versions": {
|
||||
"rustc": "1.70.0"
|
||||
}
|
||||
});
|
||||
|
||||
let res = update_handler
|
||||
.execute(args, state.clone())
|
||||
.await;
|
||||
|
||||
assert!(res.is_ok());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -86,6 +86,88 @@ impl McpTool for GetActiveWorktreeContextHandler {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct QueryGitDiffsHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for QueryGitDiffsHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"query_git_diffs"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<crate::tools::QueryGitDiffsTool>(
|
||||
"query_git_diffs",
|
||||
"Query recent git commit history, diffs, and change ledger entries.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: crate::tools::QueryGitDiffsTool =
|
||||
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let limit = req.limit.unwrap_or(5);
|
||||
let q = req.query.to_lowercase();
|
||||
|
||||
let mut matches = Vec::new();
|
||||
|
||||
state.code.ledger.read_with(|ledger| {
|
||||
for change in ledger {
|
||||
if change.file_path.to_lowercase().contains(&q)
|
||||
|| change.description.to_lowercase().contains(&q)
|
||||
|| change.git_commit.as_ref().is_some_and(|c| c.contains(&q))
|
||||
{
|
||||
matches.push(json!({
|
||||
"file_path": change.file_path,
|
||||
"description": change.description,
|
||||
"commit": change.git_commit,
|
||||
"branch": change.git_branch,
|
||||
"timestamp": change.timestamp,
|
||||
}));
|
||||
if matches.len() >= limit {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if matches.len() < limit {
|
||||
let remaining = limit - matches.len();
|
||||
let git_matches = tokio::task::spawn_blocking(move || {
|
||||
let mut results = Vec::new();
|
||||
let cwd = env::current_dir().unwrap_or_default();
|
||||
if let Ok(repo) = git2::Repository::discover(&cwd) {
|
||||
if let Ok(mut revwalk) = repo.revwalk() {
|
||||
let _ = revwalk.push_head();
|
||||
let mut count = 0;
|
||||
for oid in revwalk.flatten() {
|
||||
if count >= remaining {
|
||||
break;
|
||||
}
|
||||
if let Ok(commit) = repo.find_commit(oid) {
|
||||
let summary = commit.summary().unwrap_or("");
|
||||
if summary.to_lowercase().contains(&q) {
|
||||
count += 1;
|
||||
results.push(json!({
|
||||
"commit_id": oid.to_string(),
|
||||
"author": commit.author().name().unwrap_or("unknown"),
|
||||
"message": summary,
|
||||
"timestamp": commit.time().seconds(),
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
results
|
||||
})
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
matches.extend(git_matches);
|
||||
}
|
||||
|
||||
Ok(serde_json::to_string_pretty(&matches)?)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -110,4 +192,22 @@ mod tests {
|
||||
assert!(parsed.get("modified_files").is_some());
|
||||
assert!(parsed.get("diff").is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_active_worktree_context_empty_git_repo() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
|
||||
let handler = GetActiveWorktreeContextHandler;
|
||||
|
||||
let result = handler
|
||||
.execute(serde_json::json!({}), state)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to get worktree context: {}", e))
|
||||
.unwrap();
|
||||
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||
assert!(parsed.get("branch").is_some() || parsed.is_object());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -119,8 +119,9 @@ impl McpTool for CreateEntitiesHandler {
|
||||
let req: CreateEntitiesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let mut inserted = Vec::new();
|
||||
state.modify_graph(|g| {
|
||||
for entity in req.entities {
|
||||
for mut entity in req.entities {
|
||||
if !entity.name.is_empty() {
|
||||
entity.entity_type = crate::models::normalize_entity_type(&entity.entity_type);
|
||||
inserted.push(entity.clone());
|
||||
g.entities.insert(entity.name.clone(), entity);
|
||||
}
|
||||
@@ -165,8 +166,9 @@ impl McpTool for CreateRelationsHandler {
|
||||
};
|
||||
let mut missing_nodes = std::collections::HashSet::new();
|
||||
state.modify_graph(|g| {
|
||||
for relation in req.relations {
|
||||
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);
|
||||
let from_exists = g.entities.contains_key(&relation.from);
|
||||
let to_exists = g.entities.contains_key(&relation.to);
|
||||
if from_exists && to_exists {
|
||||
@@ -353,8 +355,9 @@ impl McpTool for ReadGraphHandler {
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: ReadGraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let max_tokens = req.max_tokens;
|
||||
let data = state.read_graph(|full| -> crate::error::Result<String> {
|
||||
if let Some(ns) = req.namespace {
|
||||
let mut result_json = if let Some(ns) = req.namespace {
|
||||
let mut filtered = BorrowedGraph::default();
|
||||
for (k, v) in &full.entities {
|
||||
if v.namespace == ns {
|
||||
@@ -366,10 +369,19 @@ impl McpTool for ReadGraphHandler {
|
||||
filtered.relations.push(r);
|
||||
}
|
||||
}
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
|
||||
serde_json::to_string(&filtered)?
|
||||
} else {
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(full)?)
|
||||
serde_json::to_string(full)?
|
||||
};
|
||||
|
||||
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]");
|
||||
}
|
||||
}
|
||||
Ok(result_json)
|
||||
})?;
|
||||
Ok(data)
|
||||
}
|
||||
@@ -553,6 +565,54 @@ impl McpTool for CondenseEntityHandler {
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: CondenseEntityTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let mut missing = false;
|
||||
|
||||
if req.summarized_observations.is_empty() {
|
||||
let mut obs_to_condense = Vec::new();
|
||||
state.read_graph(|g| {
|
||||
if let Some(e) = g.entities.get(&req.entity_name) {
|
||||
obs_to_condense = e.observations.clone();
|
||||
} else {
|
||||
missing = true;
|
||||
}
|
||||
});
|
||||
if missing {
|
||||
return Err(crate::error::AppError::Internal(format!(
|
||||
"Error: Entity '{}' not found. Please verify the exact entity name using search_nodes.",
|
||||
req.entity_name
|
||||
)));
|
||||
}
|
||||
let mut unique_obs: Vec<String> = Vec::new();
|
||||
for obs in obs_to_condense {
|
||||
let clean = obs.trim().to_string();
|
||||
if !clean.is_empty() && !unique_obs.contains(&clean) {
|
||||
unique_obs.push(clean);
|
||||
}
|
||||
}
|
||||
if unique_obs.len() > 3 && state.ollama.is_available().await {
|
||||
let prompt = format!(
|
||||
"Condense the following observations for entity '{}' into 2-3 concise architectural facts:\n- {}",
|
||||
req.entity_name,
|
||||
unique_obs.join("\n- ")
|
||||
);
|
||||
if let Ok(summary) = state.ollama.generate(&prompt, None, None).await {
|
||||
let lines: Vec<String> = summary
|
||||
.lines()
|
||||
.map(|l| l.trim().trim_start_matches('-').trim().to_string())
|
||||
.filter(|l| !l.is_empty())
|
||||
.collect();
|
||||
if !lines.is_empty() {
|
||||
unique_obs = lines;
|
||||
}
|
||||
}
|
||||
}
|
||||
state.modify_graph(|master| {
|
||||
if let Some(e) = master.entities.get_mut(&req.entity_name) {
|
||||
e.observations = unique_obs;
|
||||
}
|
||||
});
|
||||
return Ok(format!("Entity '{}' observations condensed with deduplication/AI summary.", req.entity_name));
|
||||
}
|
||||
|
||||
state.modify_graph(|master| {
|
||||
if let Some(e) = master.entities.get_mut(&req.entity_name) {
|
||||
e.observations = req.summarized_observations;
|
||||
@@ -655,14 +715,316 @@ impl McpTool for FindOrphansHandler {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct GetSubgraphHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for GetSubgraphHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"get_subgraph"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<GetSubgraphTool>("get_subgraph", "Execute get_subgraph")
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: GetSubgraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let root = req
|
||||
.root_entity
|
||||
.or(req.root_node)
|
||||
.ok_or_else(|| crate::error::AppError::Internal("root_entity or root_node is required".to_string()))?;
|
||||
let depth = req.depth.unwrap_or(2);
|
||||
let format = req.format.unwrap_or(SubgraphFormat::Json);
|
||||
|
||||
if format == SubgraphFormat::MarkdownTree {
|
||||
let (sub_entities, sub_relations) = state.read_graph(|g| {
|
||||
let mut visited = std::collections::HashSet::new();
|
||||
let mut queue = std::collections::VecDeque::new();
|
||||
queue.push_back((root.clone(), 0));
|
||||
visited.insert(root.clone());
|
||||
|
||||
while let Some((curr, d)) = queue.pop_front() {
|
||||
if d >= depth {
|
||||
continue;
|
||||
}
|
||||
for r in &g.relations {
|
||||
if r.from == curr && !visited.contains(&r.to) {
|
||||
visited.insert(r.to.clone());
|
||||
queue.push_back((r.to.clone(), d + 1));
|
||||
} else if r.to == curr && !visited.contains(&r.from) {
|
||||
visited.insert(r.from.clone());
|
||||
queue.push_back((r.from.clone(), d + 1));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ent: Vec<_> = g.entities.values().filter(|e| visited.contains(&e.name)).cloned().collect();
|
||||
let rel: Vec<_> = g.relations.iter().filter(|r| visited.contains(&r.from) && visited.contains(&r.to)).cloned().collect();
|
||||
(ent, rel)
|
||||
});
|
||||
|
||||
let mut markdown = format!("# Subgraph Topology for `{}`\n\n", root);
|
||||
markdown.push_str("## Entities\n");
|
||||
for e in &sub_entities {
|
||||
markdown.push_str(&format!("- **{}** (`{}`)\n", e.name, e.entity_type));
|
||||
for obs in &e.observations {
|
||||
markdown.push_str(&format!(" - {}\n", obs));
|
||||
}
|
||||
}
|
||||
|
||||
markdown.push_str("\n## Relations\n");
|
||||
for r in &sub_relations {
|
||||
markdown.push_str(&format!("- `{}` --[{}]--> `{}`\n", r.from, r.relation_type, r.to));
|
||||
}
|
||||
return Ok(markdown);
|
||||
}
|
||||
|
||||
let data = state.read_graph(|full| {
|
||||
let mut visited_nodes: std::collections::HashSet<String> = std::collections::HashSet::new();
|
||||
let mut matched_entities = std::collections::HashMap::new();
|
||||
let mut matched_relations = Vec::new();
|
||||
|
||||
let mut queue = std::collections::VecDeque::new();
|
||||
queue.push_back((root.clone(), 0));
|
||||
visited_nodes.insert(root.clone());
|
||||
|
||||
while let Some((curr, current_depth)) = queue.pop_front() {
|
||||
if let Some(entity) = full.entities.get(&curr) {
|
||||
matched_entities.insert(curr.clone(), entity.clone());
|
||||
}
|
||||
|
||||
if current_depth < depth {
|
||||
for rel in &full.relations {
|
||||
if rel.from == curr {
|
||||
matched_relations.push(rel.clone());
|
||||
if !visited_nodes.contains(&rel.to) {
|
||||
visited_nodes.insert(rel.to.clone());
|
||||
queue.push_back((rel.to.clone(), current_depth + 1));
|
||||
}
|
||||
} else if rel.to == curr {
|
||||
matched_relations.push(rel.clone());
|
||||
if !visited_nodes.contains(&rel.from) {
|
||||
visited_nodes.insert(rel.from.clone());
|
||||
queue.push_back((rel.from.clone(), current_depth + 1));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let result = serde_json::json!({
|
||||
"root_node": root,
|
||||
"depth": depth,
|
||||
"entities": matched_entities,
|
||||
"relations": matched_relations,
|
||||
});
|
||||
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&result)?)
|
||||
})?;
|
||||
|
||||
Ok(data)
|
||||
}
|
||||
}
|
||||
|
||||
use crate::handlers::utils::*;
|
||||
|
||||
pub struct SweepGraphHealthHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for SweepGraphHealthHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"sweep_graph_health"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<SweepGraphHealthTool>("sweep_graph_health", "Audit and sweep graph health")
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: SweepGraphHealthTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let auto_prune = req.auto_prune_orphans.unwrap_or(false);
|
||||
|
||||
let mut orphans = Vec::new();
|
||||
let mut duplicates = Vec::new();
|
||||
|
||||
state.modify_graph(|g| {
|
||||
// 1. Identify Orphans
|
||||
let mut connected = std::collections::HashSet::new();
|
||||
for r in &g.relations {
|
||||
connected.insert(r.from.clone());
|
||||
connected.insert(r.to.clone());
|
||||
}
|
||||
|
||||
for name in g.entities.keys() {
|
||||
if !connected.contains(name) {
|
||||
orphans.push(name.clone());
|
||||
}
|
||||
}
|
||||
|
||||
if auto_prune {
|
||||
for orphan in &orphans {
|
||||
g.entities.remove(orphan);
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Compute similarity pairs for duplicate detection
|
||||
let names: Vec<_> = g.entities.keys().cloned().collect();
|
||||
for i in 0..names.len() {
|
||||
for j in (i + 1)..names.len() {
|
||||
let n1 = &names[i];
|
||||
let n2 = &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) {
|
||||
duplicates.push(serde_json::json!({
|
||||
"entity_a": n1,
|
||||
"entity_b": n2,
|
||||
"suggested_action": format!("merge_entities(source: '{}', target: '{}')", n1, n2)
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let report = serde_json::json!({
|
||||
"orphaned_entities": orphans,
|
||||
"orphans_pruned": auto_prune,
|
||||
"potential_duplicates": duplicates,
|
||||
"health_score": if orphans.is_empty() && duplicates.is_empty() { "100%" } else { "Needs Maintenance" }
|
||||
});
|
||||
|
||||
Ok(serde_json::to_string_pretty(&report)?)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ResolveStaleSymbolsHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for ResolveStaleSymbolsHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"resolve_stale_symbols"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ResolveStaleSymbolsTool>("resolve_stale_symbols", "Inspect Knowledge Graph entities and tech debt symbol/line references against files on disk, healing stale pointers.")
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: ResolveStaleSymbolsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let auto_heal = req.auto_heal.unwrap_or(true);
|
||||
|
||||
let mut checked = 0;
|
||||
let mut stale_entities = Vec::new();
|
||||
|
||||
state.modify_graph(|g| {
|
||||
for (name, _entity) in g.entities.iter_mut() {
|
||||
if let Some(target) = &req.target {
|
||||
if name != target && !name.contains(target) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
checked += 1;
|
||||
let path = std::path::Path::new(name);
|
||||
if (name.contains('/') || name.contains('\\') || name.contains('.')) && !path.exists() {
|
||||
stale_entities.push(name.clone());
|
||||
}
|
||||
}
|
||||
|
||||
if auto_heal {
|
||||
for stale in &stale_entities {
|
||||
g.entities.remove(stale);
|
||||
g.relations.retain(|r| &r.from != stale && &r.to != stale);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let report = serde_json::json!({
|
||||
"checked_count": checked,
|
||||
"stale_entities_found": stale_entities,
|
||||
"auto_healed": auto_heal,
|
||||
"status": "success"
|
||||
});
|
||||
|
||||
Ok(serde_json::to_string_pretty(&report)?)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SummarizeSubgraphHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for SummarizeSubgraphHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"summarize_subgraph"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<SummarizeSubgraphTool>("summarize_subgraph", "Generate compact, LLM-optimized Markdown topology of a graph component capped within a token budget.")
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: SummarizeSubgraphTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let depth = req.depth.unwrap_or(2);
|
||||
|
||||
let (sub_entities, sub_relations) = state.read_graph(|g| {
|
||||
let mut visited = std::collections::HashSet::new();
|
||||
let mut queue = std::collections::VecDeque::new();
|
||||
queue.push_back((req.root_entity.clone(), 0));
|
||||
visited.insert(req.root_entity.clone());
|
||||
|
||||
while let Some((curr, d)) = queue.pop_front() {
|
||||
if d >= depth {
|
||||
continue;
|
||||
}
|
||||
for r in &g.relations {
|
||||
if r.from == curr && !visited.contains(&r.to) {
|
||||
visited.insert(r.to.clone());
|
||||
queue.push_back((r.to.clone(), d + 1));
|
||||
} else if r.to == curr && !visited.contains(&r.from) {
|
||||
visited.insert(r.from.clone());
|
||||
queue.push_back((r.from.clone(), d + 1));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let ent: Vec<_> = g.entities.values().filter(|e| visited.contains(&e.name)).cloned().collect();
|
||||
let rel: Vec<_> = g.relations.iter().filter(|r| visited.contains(&r.from) && visited.contains(&r.to)).cloned().collect();
|
||||
(ent, rel)
|
||||
});
|
||||
|
||||
let mut markdown = format!("# Subgraph Topology for `{}`\n\n", req.root_entity);
|
||||
markdown.push_str("## Entities\n");
|
||||
for e in &sub_entities {
|
||||
markdown.push_str(&format!("- **{}** (`{}`)\n", e.name, e.entity_type));
|
||||
for obs in &e.observations {
|
||||
markdown.push_str(&format!(" - {}\n", obs));
|
||||
}
|
||||
}
|
||||
|
||||
markdown.push_str("\n## Relations\n");
|
||||
for r in &sub_relations {
|
||||
markdown.push_str(&format!("- `{}` --[{}]--> `{}`\n", r.from, r.relation_type, r.to));
|
||||
}
|
||||
|
||||
let max_tokens = req.max_tokens.unwrap_or(1000);
|
||||
let max_chars = max_tokens * 4;
|
||||
if markdown.len() > max_chars {
|
||||
markdown.truncate(max_chars);
|
||||
markdown.push_str("\n... [Truncated to fit token budget]");
|
||||
}
|
||||
|
||||
Ok(markdown)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
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();
|
||||
@@ -888,4 +1250,71 @@ mod tests {
|
||||
.unwrap();
|
||||
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 create_handler = CreateEntitiesHandler;
|
||||
let args_ent = json!({
|
||||
"entities": [
|
||||
{"name": "Alpha", "entity_type": "Node", "observations": ["Obs1", "Obs2"], "namespace": "global"},
|
||||
{"name": "Beta", "entity_type": "Node", "observations": ["Obs3"], "namespace": "global"}
|
||||
]
|
||||
});
|
||||
create_handler.execute(args_ent, state.clone()).await.unwrap();
|
||||
|
||||
let rel_handler = CreateRelationsHandler;
|
||||
let args_rel = json!({
|
||||
"relations": [
|
||||
{"from": "Alpha", "to": "Beta", "relation_type": "links_to", "namespace": "global"}
|
||||
]
|
||||
});
|
||||
rel_handler.execute(args_rel, state.clone()).await.unwrap();
|
||||
|
||||
let path_handler = QueryGraphPathHandler;
|
||||
let path_res = path_handler.execute(json!({"start_node": "Alpha", "end_node": "Beta"}), state.clone()).await.unwrap();
|
||||
assert!(path_res.contains("Alpha"));
|
||||
|
||||
let sub_handler = GetSubgraphHandler;
|
||||
let sub_res = sub_handler.execute(json!({"root_entity": "Alpha"}), state.clone()).await.unwrap();
|
||||
assert!(sub_res.contains("Alpha"));
|
||||
|
||||
let sum_handler = SummarizeSubgraphHandler;
|
||||
let sum_res = sum_handler.execute(json!({"root_entity": "Alpha"}), state.clone()).await.unwrap();
|
||||
assert!(!sum_res.is_empty());
|
||||
|
||||
|
||||
let sweep_handler = SweepGraphHealthHandler;
|
||||
let sweep_res = sweep_handler.execute(json!({}), state.clone()).await.unwrap();
|
||||
assert!(!sweep_res.is_empty());
|
||||
|
||||
let del_obs_handler = DeleteObservationsHandler;
|
||||
let del_obs_res = del_obs_handler.execute(json!({
|
||||
"deletions": [{"entity_name": "Alpha", "observations": ["Obs1"]}]
|
||||
}), state.clone()).await.unwrap();
|
||||
assert_eq!(del_obs_res, "Observations deleted");
|
||||
|
||||
let del_rel_handler = DeleteRelationsHandler;
|
||||
let del_rel_res = del_rel_handler.execute(json!({
|
||||
"relations": [{"from": "Alpha", "to": "Beta", "relation_type": "links_to"}]
|
||||
}), state.clone()).await.unwrap();
|
||||
assert_eq!(del_rel_res, "Relations deleted");
|
||||
|
||||
let bcast_handler = BroadcastAgentSignalHandler;
|
||||
let bcast_res = bcast_handler.execute(json!({
|
||||
"sender": "agent1",
|
||||
"signal_type": "task_completed",
|
||||
"payload": "fix_bug"
|
||||
}), state.clone()).await.unwrap();
|
||||
assert!(bcast_res.contains("Broadcasted signal"));
|
||||
|
||||
|
||||
let qsignal_handler = QueryAgentSignalsHandler;
|
||||
let qsignal_res = qsignal_handler.execute(json!({"sender": "agent1"}), state.clone()).await.unwrap();
|
||||
assert!(qsignal_res.contains("task_completed"));
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
@@ -133,4 +133,30 @@ mod tests {
|
||||
assert!(result.contains("line1"));
|
||||
assert!(result.contains("line3"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_recent_logs_with_large_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
let handler = GetRecentLogsHandler;
|
||||
|
||||
let log_file = dir.path().join("large_test.log");
|
||||
let mut buffer = String::new();
|
||||
for _ in 0..1000 {
|
||||
buffer.push_str("line\n");
|
||||
}
|
||||
std::fs::write(&log_file, buffer).unwrap();
|
||||
|
||||
let args = serde_json::json!({
|
||||
"file_path": log_file.to_str().unwrap()
|
||||
});
|
||||
|
||||
let result = handler
|
||||
.execute(args, state)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to get recent logs: {}", e))
|
||||
.unwrap();
|
||||
assert!(result.contains("line"));
|
||||
}
|
||||
}
|
||||
|
||||
+1650
-15
File diff suppressed because it is too large.
Load diff
@@ -7,6 +7,90 @@ use serde_json::Value;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::Arc;
|
||||
|
||||
pub struct ManageStickyNotesHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for ManageStickyNotesHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"manage_sticky_notes"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ManageStickyNotesTool>(
|
||||
"manage_sticky_notes",
|
||||
"Manage ephemeral sticky notes with TTL (add, read, delete, clear)",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: ManageStickyNotesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
match req.action {
|
||||
StickyNoteAction::Add => {
|
||||
let content = req.content.ok_or_else(|| {
|
||||
crate::error::AppError::Internal("content is required for 'add' action".to_string())
|
||||
})?;
|
||||
let now = crate::handlers::utils::now_secs();
|
||||
let expires_at = if let Some(ttl) = req.ttl_seconds {
|
||||
Some(now + ttl)
|
||||
} else if req.session_only.unwrap_or(false) {
|
||||
Some(now + 14400)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.push(StickyNote {
|
||||
timestamp: now,
|
||||
content,
|
||||
expires_at,
|
||||
});
|
||||
});
|
||||
Ok("Sticky note added.".to_string())
|
||||
}
|
||||
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();
|
||||
});
|
||||
Ok(serde_json::to_string(&active_notes)?)
|
||||
}
|
||||
StickyNoteAction::Delete => {
|
||||
let idx = req.index.ok_or_else(|| {
|
||||
crate::error::AppError::Internal("index is required for 'delete' action".to_string())
|
||||
})?;
|
||||
let mut success = false;
|
||||
state.code.sticky.modify(|notes| {
|
||||
if idx > 0 && idx <= notes.len() {
|
||||
notes.remove(idx - 1);
|
||||
success = true;
|
||||
}
|
||||
});
|
||||
if success {
|
||||
Ok("Sticky note deleted.".to_string())
|
||||
} else {
|
||||
Err(crate::error::AppError::Internal(
|
||||
"Invalid sticky note index.".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
StickyNoteAction::Clear => {
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.clear();
|
||||
});
|
||||
Ok("All sticky notes cleared.".to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct AddStickyNoteHandler;
|
||||
|
||||
#[async_trait]
|
||||
@@ -21,11 +105,20 @@ impl McpTool for AddStickyNoteHandler {
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: AddStickyNoteTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let now = crate::handlers::utils::now_secs();
|
||||
let expires_at = if let Some(ttl) = req.ttl_seconds {
|
||||
Some(now + ttl)
|
||||
} else if req.session_only.unwrap_or(false) {
|
||||
Some(now + 14400) // Default 4-hour session TTL
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.push(StickyNote {
|
||||
timestamp: crate::handlers::utils::now_secs(),
|
||||
timestamp: now,
|
||||
content: req.content,
|
||||
expires_at: None,
|
||||
expires_at,
|
||||
});
|
||||
});
|
||||
Ok("Sticky note added.".to_string())
|
||||
@@ -48,11 +141,21 @@ impl McpTool for ReadStickyNotesHandler {
|
||||
}
|
||||
|
||||
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let data = state
|
||||
.code
|
||||
.sticky
|
||||
.read_with(|s| Ok::<String, crate::error::AppError>(serde_json::to_string(s)?))?;
|
||||
Ok(data)
|
||||
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();
|
||||
});
|
||||
|
||||
Ok(serde_json::to_string(&active_notes)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -138,7 +241,10 @@ impl McpTool for LeaveHandoffMemoHandler {
|
||||
namespace: req.namespace,
|
||||
timestamp: crate::handlers::utils::now_secs(),
|
||||
expires_at: None,
|
||||
})
|
||||
});
|
||||
if memos.len() > 200 {
|
||||
memos.remove(0);
|
||||
}
|
||||
});
|
||||
Ok("Handoff memo left".to_string())
|
||||
}
|
||||
@@ -227,7 +333,10 @@ impl McpTool for AddSessionSummaryHandler {
|
||||
namespace: req.namespace,
|
||||
timestamp: crate::handlers::utils::now_secs(),
|
||||
expires_at: None,
|
||||
})
|
||||
});
|
||||
if summaries.len() > 200 {
|
||||
summaries.remove(0);
|
||||
}
|
||||
});
|
||||
Ok("Session summary added".to_string())
|
||||
}
|
||||
@@ -268,6 +377,57 @@ impl McpTool for GenerateStandupReportHandler {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct PromoteToEntityHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for PromoteToEntityHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"promote_to_entity"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<PromoteToEntityTool>(
|
||||
"promote_to_entity",
|
||||
"Promote a transient sticky note or active task observation into a permanent Knowledge Graph entity.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: PromoteToEntityTool =
|
||||
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
|
||||
let entity_name = req.entity_name.clone();
|
||||
let obs = req.content.clone();
|
||||
let entity_type = req.entity_type.clone();
|
||||
let namespace = req.namespace.clone();
|
||||
|
||||
state.modify_graph(|g| {
|
||||
let entity = g
|
||||
.entities
|
||||
.entry(entity_name.clone())
|
||||
.or_insert_with(|| Entity {
|
||||
name: entity_name.clone(),
|
||||
entity_type,
|
||||
observations: Vec::new(),
|
||||
namespace,
|
||||
git_branch: None,
|
||||
});
|
||||
if !entity.observations.contains(&obs) {
|
||||
entity.observations.push(obs.clone());
|
||||
}
|
||||
});
|
||||
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.retain(|n| !n.content.contains(&obs));
|
||||
});
|
||||
|
||||
Ok(format!(
|
||||
"Promoted transient note into permanent entity '{}'.",
|
||||
req.entity_name
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -366,4 +526,22 @@ mod tests {
|
||||
.unwrap();
|
||||
assert!(!res4.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_invalid_sticky_note_action() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
|
||||
let handler = ManageStickyNotesHandler;
|
||||
let invalid_args = serde_json::json!({
|
||||
"action": "invalid_action"
|
||||
});
|
||||
|
||||
let res = handler
|
||||
.execute(invalid_args, state.clone())
|
||||
.await;
|
||||
assert!(res.is_err());
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -281,11 +281,12 @@ impl McpTool for ListActiveTasksHandler {
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: ListActiveTasksTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let level = req.summary_level.as_deref().unwrap_or("detailed");
|
||||
let data = state.project.tasks.read_with(|tasks| {
|
||||
let filtered: Vec<_> = tasks
|
||||
.iter()
|
||||
.filter(|t| {
|
||||
let status_match = t.status != "done";
|
||||
let status_match = t.status != "done" && t.status != "completed";
|
||||
let branch_match = match &req.git_branch {
|
||||
Some(branch) => {
|
||||
t.git_branch.is_none()
|
||||
@@ -295,8 +296,31 @@ impl McpTool for ListActiveTasksHandler {
|
||||
};
|
||||
status_match && branch_match
|
||||
})
|
||||
.map(|t| match level {
|
||||
"compact" => serde_json::json!({
|
||||
"id": t.id,
|
||||
"title": t.title,
|
||||
"status": t.status,
|
||||
}),
|
||||
"full" => serde_json::to_value(t).unwrap_or_default(),
|
||||
_ => serde_json::json!({
|
||||
"id": t.id,
|
||||
"title": t.title,
|
||||
"status": t.status,
|
||||
"description": t.description,
|
||||
"git_branch": t.git_branch,
|
||||
}),
|
||||
})
|
||||
.collect();
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
|
||||
let mut json_str = serde_json::to_string(&filtered)?;
|
||||
if let Some(max_t) = req.max_tokens {
|
||||
let char_limit = max_t * 4;
|
||||
if json_str.len() > char_limit {
|
||||
json_str.truncate(char_limit);
|
||||
json_str.push_str(" ...[truncated due to max_tokens]");
|
||||
}
|
||||
}
|
||||
Ok::<String, crate::error::AppError>(json_str)
|
||||
})?;
|
||||
Ok(data)
|
||||
}
|
||||
@@ -684,4 +708,36 @@ mod tests {
|
||||
.unwrap();
|
||||
assert!(res_del.contains("Deleted task and its children (2 total)."));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
async fn test_list_milestones_with_namespace() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
|
||||
let add_milestone = AddMilestoneHandler;
|
||||
let args_ms = serde_json::json!({
|
||||
"name": "v1.0",
|
||||
"title": "Release 1.0",
|
||||
"description": "First release",
|
||||
"target_date": 1700000000,
|
||||
"end_date": 1700000000,
|
||||
"namespace": "global"
|
||||
});
|
||||
let res1 = add_milestone
|
||||
.execute(args_ms, state.clone())
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
||||
.unwrap();
|
||||
assert!(res1.contains("Milestone added"));
|
||||
|
||||
let list_ms = ListMilestonesHandler;
|
||||
let res2 = list_ms
|
||||
.execute(serde_json::json!({"namespace": "global"}), state.clone())
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
|
||||
.unwrap();
|
||||
assert!(res2.contains("Release 1.0"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -23,11 +23,11 @@ impl McpTool for WriteClipboardHandler {
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let tool_args: WriteClipboardTool = serde_json::from_value(args)
|
||||
.map_err(|e| crate::error::AppError::Internal(format!("Invalid args: {}", e)))?;
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let res = tokio::task::spawn_blocking(move || {
|
||||
let mut msgs = Vec::new();
|
||||
|
||||
if let Ok(mut clipboard) = Clipboard::new()
|
||||
@@ -77,7 +77,10 @@ impl McpTool for WriteClipboardHandler {
|
||||
}
|
||||
})
|
||||
.await
|
||||
.map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))?
|
||||
.map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;
|
||||
|
||||
state.record_activity("clipboard", &res, None);
|
||||
Ok(res)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -101,51 +104,51 @@ impl McpTool for ReadClipboardHandler {
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Value> {
|
||||
let mut out = serde_json::Map::new();
|
||||
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new()
|
||||
&& let Ok(text) = clipboard.get_text()
|
||||
&& !text.trim().is_empty()
|
||||
{
|
||||
out.insert("text".into(), json!(text));
|
||||
}
|
||||
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new()
|
||||
&& let Ok(image_data) = clipboard.get_image()
|
||||
&& let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
|
||||
image_data.width as u32,
|
||||
image_data.height as u32,
|
||||
image_data.bytes.into_owned(),
|
||||
)
|
||||
{
|
||||
let mut dynamic_img = image::DynamicImage::ImageRgba8(img);
|
||||
let max_dim = 1024;
|
||||
if dynamic_img.width() > max_dim || dynamic_img.height() > max_dim {
|
||||
dynamic_img = dynamic_img.resize(max_dim, max_dim, FilterType::Lanczos3);
|
||||
if let Ok(mut clipboard) = arboard::Clipboard::new() {
|
||||
if let Ok(text) = clipboard.get_text() {
|
||||
if !text.trim().is_empty() {
|
||||
out.insert("text".into(), json!(text));
|
||||
}
|
||||
}
|
||||
let rgb_img = dynamic_img.into_rgb8();
|
||||
|
||||
let cache_dir = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/clipboard");
|
||||
let _ = std::fs::create_dir_all(&cache_dir);
|
||||
if let Ok(image_data) = clipboard.get_image() {
|
||||
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
|
||||
image_data.width as u32,
|
||||
image_data.height as u32,
|
||||
image_data.bytes.into_owned(),
|
||||
) {
|
||||
let mut dynamic_img = image::DynamicImage::ImageRgba8(img);
|
||||
let max_dim = 1024;
|
||||
if dynamic_img.width() > max_dim || dynamic_img.height() > max_dim {
|
||||
dynamic_img = dynamic_img.resize(max_dim, max_dim, FilterType::Lanczos3);
|
||||
}
|
||||
let rgb_img = dynamic_img.into_rgb8();
|
||||
|
||||
let timestamp = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp));
|
||||
let cache_dir = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/clipboard");
|
||||
let _ = std::fs::create_dir_all(&cache_dir);
|
||||
|
||||
if rgb_img
|
||||
.save_with_format(&file_path, image::ImageFormat::Jpeg)
|
||||
.is_ok()
|
||||
{
|
||||
let path_str = file_path.to_string_lossy().to_string();
|
||||
out.insert("image_path".into(), json!(path_str));
|
||||
let timestamp = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs();
|
||||
let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp));
|
||||
|
||||
// Read image bytes for base64 encoding if needed by vision
|
||||
if let Ok(bytes) = std::fs::read(&file_path) {
|
||||
use base64::Engine;
|
||||
let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
|
||||
out.insert("image_base64".into(), json!(b64));
|
||||
if rgb_img
|
||||
.save_with_format(&file_path, image::ImageFormat::Jpeg)
|
||||
.is_ok()
|
||||
{
|
||||
let path_str = file_path.to_string_lossy().to_string();
|
||||
out.insert("image_path".into(), json!(path_str));
|
||||
|
||||
// Read image bytes for base64 encoding if needed by vision
|
||||
if let Ok(bytes) = std::fs::read(&file_path) {
|
||||
use base64::Engine;
|
||||
let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
|
||||
out.insert("image_base64".into(), json!(b64));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -176,6 +179,7 @@ impl McpTool for ReadClipboardHandler {
|
||||
}
|
||||
}
|
||||
|
||||
state.record_activity("clipboard", "Read contents from OS clipboard", None);
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&final_obj)?)
|
||||
}
|
||||
}
|
||||
@@ -203,14 +207,14 @@ impl McpTool for ToggleClipboardWatchModeHandler {
|
||||
let mut watch_mode = state.clipboard_watch_mode.write().await;
|
||||
*watch_mode = tool_args.enable;
|
||||
|
||||
if tool_args.enable {
|
||||
Ok(
|
||||
"Clipboard watch mode enabled. Changes will be ingested as StickyNotes."
|
||||
.to_string(),
|
||||
)
|
||||
let status_msg = if tool_args.enable {
|
||||
"Clipboard watch mode enabled. Changes will be ingested as StickyNotes."
|
||||
} else {
|
||||
Ok("Clipboard watch mode disabled.".to_string())
|
||||
}
|
||||
"Clipboard watch mode disabled."
|
||||
};
|
||||
|
||||
state.record_activity("clipboard", status_msg, None);
|
||||
Ok(status_msg.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -280,4 +284,21 @@ mod tests {
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||
assert!(parsed.is_object());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_clipboard_empty() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
let handler = ReadClipboardHandler;
|
||||
|
||||
let result = handler
|
||||
.execute(serde_json::json!({}), state)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to read clipboard: {}", e))
|
||||
.unwrap();
|
||||
|
||||
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
|
||||
assert!(parsed.is_object());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -123,6 +123,7 @@ impl McpTool for StoreSnippetHandler {
|
||||
code: req.code,
|
||||
description: req.description,
|
||||
updated_at: crate::handlers::utils::now_secs(),
|
||||
tags: req.tags.unwrap_or_default(),
|
||||
embedding,
|
||||
};
|
||||
|
||||
@@ -138,6 +139,43 @@ impl McpTool for StoreSnippetHandler {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TagSnippetHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for TagSnippetHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"tag_snippet"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<TagSnippetTool>("tag_snippet", "Execute tag_snippet")
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: TagSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let mut found = false;
|
||||
state.code.snippets.modify(|snippets| {
|
||||
for s in snippets.iter_mut() {
|
||||
if s.name == req.name {
|
||||
for tag in &req.tags {
|
||||
if !s.tags.contains(tag) {
|
||||
s.tags.push(tag.clone());
|
||||
}
|
||||
}
|
||||
found = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if found {
|
||||
Ok(format!("Snippet '{}' tagged with: {:?}", req.name, req.tags))
|
||||
} else {
|
||||
Err(crate::error::AppError::Internal(format!("Snippet '{}' not found.", req.name)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SearchSnippetsHandler;
|
||||
|
||||
#[async_trait]
|
||||
@@ -147,22 +185,49 @@ impl McpTool for SearchSnippetsHandler {
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<SearchSnippetsTool>("search_snippets", "Execute search_snippets")
|
||||
crate::mcp::tool_def::<SearchSnippetsTool>(
|
||||
"search_snippets",
|
||||
"Search stored code snippets using Hybrid BM25 + Vector ranking or text search",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: SearchSnippetsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let query = req.query;
|
||||
let limit = req.limit.unwrap_or(10);
|
||||
let include_body = req.include_body.unwrap_or(false);
|
||||
let tags_filter = req.tags.unwrap_or_default();
|
||||
|
||||
let data = state.code.snippets.read_with(|snippets| {
|
||||
let results: Vec<_> = snippets
|
||||
.iter()
|
||||
.filter(|s| {
|
||||
contains_ignore_ascii_case(&s.name, &query)
|
||||
let matches_text = query.is_empty()
|
||||
|| contains_ignore_ascii_case(&s.name, &query)
|
||||
|| contains_ignore_ascii_case(&s.description, &query)
|
||||
|| contains_ignore_ascii_case(&s.language, &query)
|
||||
|| s.tags.iter().any(|t| contains_ignore_ascii_case(t, &query));
|
||||
|
||||
let matches_tags = tags_filter.is_empty()
|
||||
|| tags_filter.iter().any(|tf| s.tags.iter().any(|st| contains_ignore_ascii_case(st, tf)));
|
||||
|
||||
matches_text && matches_tags
|
||||
})
|
||||
.take(limit)
|
||||
.map(|s| {
|
||||
if include_body {
|
||||
serde_json::json!(s)
|
||||
} else {
|
||||
serde_json::json!({
|
||||
"name": s.name,
|
||||
"description": s.description,
|
||||
"language": s.language,
|
||||
"tags": s.tags
|
||||
})
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string(&results)?)
|
||||
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&results)?)
|
||||
})?;
|
||||
Ok(data)
|
||||
}
|
||||
@@ -523,7 +588,60 @@ mod tests {
|
||||
.unwrap();
|
||||
assert_eq!(res8, "PR checklist cleared");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_workspace_and_subagent_handlers() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
|
||||
|
||||
// Snapshots
|
||||
let create_snap = CreateSnapshotHandler;
|
||||
let snap_res = create_snap.execute(json!({"description": "Test snap", "namespace": "global"}), state.clone()).await.unwrap();
|
||||
assert!(snap_res.contains("SNAP-"));
|
||||
let snap_id = snap_res.split_whitespace().nth(1).unwrap();
|
||||
|
||||
let restore_snap = RestoreSnapshotHandler;
|
||||
let restore_res = restore_snap.execute(json!({"snapshot_id": snap_id}), state.clone()).await.unwrap();
|
||||
assert!(restore_res.contains("Successfully restored"));
|
||||
|
||||
// Subagent Namespace
|
||||
let create_sub_ns = CreateSubagentNamespaceHandler;
|
||||
let sub_ns_res = create_sub_ns.execute(json!({"subagent_id": "test_agent"}), state.clone()).await.unwrap();
|
||||
assert!(sub_ns_res.contains("subagent-test_agent"));
|
||||
|
||||
let manage_sub_ns = ManageSubagentNamespaceHandler;
|
||||
let manage_res = manage_sub_ns.execute(json!({
|
||||
"subagent_id": "test_agent",
|
||||
"action": "create"
|
||||
}), state.clone()).await.unwrap();
|
||||
assert!(manage_res.contains("initialized"));
|
||||
|
||||
// Context Workspace Diff
|
||||
let save_ws = SaveContextWorkspaceHandler;
|
||||
save_ws.execute(json!({"name": "ws_a", "description": "Workspace A", "pinned_files": ["src/lib.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
|
||||
save_ws.execute(json!({"name": "ws_b", "description": "Workspace B", "pinned_files": ["src/main.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
|
||||
|
||||
|
||||
let diff_ws = DiffContextWorkspacesHandler;
|
||||
let diff_res = diff_ws.execute(json!({"workspace_a": "ws_a", "workspace_b": "ws_b"}), state.clone()).await.unwrap();
|
||||
assert!(diff_res.contains("workspace_a"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_diff_context_workspaces_handler() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
|
||||
let save_ws = SaveContextWorkspaceHandler;
|
||||
save_ws.execute(json!({"name": "ws_1", "description": "Workspace 1", "pinned_files": ["src/lib.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
|
||||
save_ws.execute(json!({"name": "ws_2", "description": "Workspace 2", "pinned_files": ["src/main.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
|
||||
|
||||
let diff_ws = DiffContextWorkspacesHandler;
|
||||
let diff_res = diff_ws.execute(json!({"workspace_a": "ws_1", "workspace_b": "ws_2"}), state.clone()).await.unwrap();
|
||||
assert!(diff_res.contains("ws_1"));
|
||||
}
|
||||
}
|
||||
|
||||
use crate::tools::ReadDirectoryArchitectureTool;
|
||||
use std::fs;
|
||||
|
||||
@@ -726,6 +844,95 @@ impl McpTool for RestoreSnapshotHandler {
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ManageSubagentNamespaceHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for ManageSubagentNamespaceHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"manage_subagent_namespace"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<ManageSubagentNamespaceTool>(
|
||||
"manage_subagent_namespace",
|
||||
"Manage isolated memory namespaces for subagent sessions (create, condense/promote, or purge)",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: ManageSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let subagent_ns = format!("subagent-{}", req.subagent_id);
|
||||
|
||||
match req.action {
|
||||
SubagentNamespaceAction::Create => {
|
||||
state.record_activity("namespace", &format!("Created subagent namespace {}", subagent_ns), None);
|
||||
Ok(format!("Subagent memory namespace '{}' initialized.", subagent_ns))
|
||||
}
|
||||
SubagentNamespaceAction::Purge => {
|
||||
state.project.pinned_files.modify(|files| {
|
||||
files.retain(|f| f.namespace != subagent_ns);
|
||||
});
|
||||
state.record_activity("namespace", &format!("Purged subagent namespace {}", subagent_ns), None);
|
||||
Ok(format!("Subagent memory namespace '{}' purged.", subagent_ns))
|
||||
}
|
||||
SubagentNamespaceAction::Condense => {
|
||||
let purge = req.purge_after_promotion.unwrap_or(true);
|
||||
let mut promoted_entities = 0;
|
||||
let mut promoted_relations = 0;
|
||||
|
||||
state.modify_graph(|g| {
|
||||
let sub_entities: Vec<Entity> = g
|
||||
.entities
|
||||
.values()
|
||||
.filter(|e| e.namespace == subagent_ns)
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
for mut e in sub_entities {
|
||||
promoted_entities += 1;
|
||||
e.namespace = "global".to_string();
|
||||
let global_entry = g.entities.entry(e.name.clone()).or_insert_with(|| Entity {
|
||||
name: e.name.clone(),
|
||||
entity_type: e.entity_type.clone(),
|
||||
observations: Vec::new(),
|
||||
namespace: "global".to_string(),
|
||||
git_branch: e.git_branch.clone(),
|
||||
});
|
||||
for obs in e.observations {
|
||||
if !global_entry.observations.contains(&obs) {
|
||||
global_entry.observations.push(obs);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for r in &mut g.relations {
|
||||
if r.namespace == subagent_ns {
|
||||
promoted_relations += 1;
|
||||
r.namespace = "global".to_string();
|
||||
}
|
||||
}
|
||||
|
||||
if purge {
|
||||
g.entities.retain(|_, e| e.namespace != subagent_ns);
|
||||
g.relations.retain(|r| r.namespace != subagent_ns);
|
||||
}
|
||||
});
|
||||
|
||||
if purge {
|
||||
state.project.pinned_files.modify(|files| {
|
||||
files.retain(|f| f.namespace != subagent_ns);
|
||||
});
|
||||
}
|
||||
|
||||
Ok(format!(
|
||||
"Subagent namespace '{}' condensed: promoted {} entities and {} relations to global graph.",
|
||||
subagent_ns, promoted_entities, promoted_relations
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CreateSubagentNamespaceHandler;
|
||||
|
||||
#[async_trait]
|
||||
@@ -770,3 +977,141 @@ impl McpTool for PurgeSubagentNamespaceHandler {
|
||||
Ok(format!("Subagent memory namespace '{}' purged.", ns))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct CondenseSubagentNamespaceHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for CondenseSubagentNamespaceHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"condense_subagent_namespace"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<CondenseSubagentNamespaceTool>(
|
||||
"condense_subagent_namespace",
|
||||
"Auto-condense a subagent's memory namespace and promote its entities/relations to the global Knowledge Graph.",
|
||||
)
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: CondenseSubagentNamespaceTool =
|
||||
serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
let subagent_ns = format!("subagent-{}", req.subagent_id);
|
||||
let purge = req.purge_after_promotion.unwrap_or(true);
|
||||
|
||||
let mut promoted_entities = 0;
|
||||
let mut promoted_relations = 0;
|
||||
|
||||
state.modify_graph(|g| {
|
||||
let sub_entities: Vec<Entity> = g
|
||||
.entities
|
||||
.values()
|
||||
.filter(|e| e.namespace == subagent_ns)
|
||||
.cloned()
|
||||
.collect();
|
||||
|
||||
for mut e in sub_entities {
|
||||
promoted_entities += 1;
|
||||
e.namespace = "global".to_string();
|
||||
let global_entry = g.entities.entry(e.name.clone()).or_insert_with(|| Entity {
|
||||
name: e.name.clone(),
|
||||
entity_type: e.entity_type.clone(),
|
||||
observations: Vec::new(),
|
||||
namespace: "global".to_string(),
|
||||
git_branch: e.git_branch.clone(),
|
||||
});
|
||||
for obs in e.observations {
|
||||
if !global_entry.observations.contains(&obs) {
|
||||
global_entry.observations.push(obs);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for r in &mut g.relations {
|
||||
if r.namespace == subagent_ns {
|
||||
r.namespace = "global".to_string();
|
||||
promoted_relations += 1;
|
||||
}
|
||||
}
|
||||
|
||||
if purge {
|
||||
g.entities.retain(|_, e| e.namespace != subagent_ns);
|
||||
g.relations.retain(|r| r.namespace != subagent_ns);
|
||||
}
|
||||
});
|
||||
|
||||
if purge {
|
||||
state.project.pinned_files.modify(|files| {
|
||||
files.retain(|f| f.namespace != subagent_ns);
|
||||
});
|
||||
}
|
||||
|
||||
Ok(format!(
|
||||
"Consolidated subagent namespace '{}': Promoted {} entities, {} relations to global. Purged: {}",
|
||||
subagent_ns, promoted_entities, promoted_relations, purge
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct DiffContextWorkspacesHandler;
|
||||
|
||||
#[async_trait]
|
||||
impl McpTool for DiffContextWorkspacesHandler {
|
||||
fn name(&self) -> &'static str {
|
||||
"diff_context_workspaces"
|
||||
}
|
||||
|
||||
fn schema(&self) -> Value {
|
||||
crate::mcp::tool_def::<DiffContextWorkspacesTool>("diff_context_workspaces", "Execute diff_context_workspaces")
|
||||
}
|
||||
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
||||
let req: DiffContextWorkspacesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
|
||||
|
||||
let workspaces = state.project.context_workspaces.read_with(|ws| ws.clone());
|
||||
let ws_a = workspaces.iter().find(|w| w.name == req.workspace_a);
|
||||
let ws_b = workspaces.iter().find(|w| w.name == req.workspace_b);
|
||||
|
||||
if ws_a.is_none() || ws_b.is_none() {
|
||||
return Err(crate::error::AppError::Internal(format!(
|
||||
"One or both context workspaces not found: '{}' found={}, '{}' found={}",
|
||||
req.workspace_a, ws_a.is_some(), req.workspace_b, ws_b.is_some()
|
||||
)));
|
||||
}
|
||||
|
||||
let a = ws_a.unwrap();
|
||||
let b = ws_b.unwrap();
|
||||
|
||||
let set_files_a: std::collections::HashSet<_> = a.pinned_files.iter().cloned().collect();
|
||||
let set_files_b: std::collections::HashSet<_> = b.pinned_files.iter().cloned().collect();
|
||||
|
||||
let added_files: Vec<_> = set_files_b.difference(&set_files_a).cloned().collect();
|
||||
let removed_files: Vec<_> = set_files_a.difference(&set_files_b).cloned().collect();
|
||||
let shared_files: Vec<_> = set_files_a.intersection(&set_files_b).cloned().collect();
|
||||
|
||||
let set_tasks_a: std::collections::HashSet<_> = a.active_task_ids.iter().cloned().collect();
|
||||
let set_tasks_b: std::collections::HashSet<_> = b.active_task_ids.iter().cloned().collect();
|
||||
|
||||
let added_tasks: Vec<_> = set_tasks_b.difference(&set_tasks_a).cloned().collect();
|
||||
let removed_tasks: Vec<_> = set_tasks_a.difference(&set_tasks_b).cloned().collect();
|
||||
let shared_tasks: Vec<_> = set_tasks_a.intersection(&set_tasks_b).cloned().collect();
|
||||
|
||||
let diff = serde_json::json!({
|
||||
"workspace_a": a.name,
|
||||
"workspace_b": b.name,
|
||||
"pinned_files": {
|
||||
"added_in_b": added_files,
|
||||
"removed_in_b": removed_files,
|
||||
"shared": shared_files
|
||||
},
|
||||
"active_tasks": {
|
||||
"added_in_b": added_tasks,
|
||||
"removed_in_b": removed_tasks,
|
||||
"shared": shared_tasks
|
||||
}
|
||||
});
|
||||
|
||||
Ok(serde_json::to_string_pretty(&diff)?)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -83,6 +83,7 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs(),
|
||||
tags: vec![],
|
||||
embedding,
|
||||
};
|
||||
|
||||
@@ -218,4 +219,29 @@ mod tests {
|
||||
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
|
||||
start_background_indexer(state).await;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_chunks_rust_impl_block() {
|
||||
let code = "impl MyStruct { fn my_method(&self) {} }";
|
||||
let mut parser = Parser::new();
|
||||
parser
|
||||
.set_language(&tree_sitter_rust::LANGUAGE.into())
|
||||
.unwrap();
|
||||
let tree = parser.parse(code, None).unwrap();
|
||||
|
||||
let mut chunks = Vec::new();
|
||||
extract_chunks(tree.root_node(), code, &mut chunks, "rs");
|
||||
|
||||
assert_eq!(chunks.len(), 1);
|
||||
assert_eq!(chunks[0].0, "MyStruct::my_method");
|
||||
assert!(chunks[0].1.contains("// Parent Scope: MyStruct"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_start_background_indexer_empty_dir() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
|
||||
start_background_indexer(state).await;
|
||||
}
|
||||
}
|
||||
|
||||
+179
-19
@@ -1,24 +1,184 @@
|
||||
# Antigravity Memory MCP Instructions
|
||||
# Memory MCP Strategic Guidelines
|
||||
|
||||
You are Antigravity, connected to the mcp-memory persistence layer. This server provides a persistent knowledge graph, task management, and environment state tracking.
|
||||
This document outlines the STRATEGY, SEMANTICS, and CASING STANDARDS for using the MCP Memory Server.
|
||||
You do not need to memorize JSON schemas for these tools; they are strictly defined and typed in the `tools/list` endpoint. Focus purely on WHEN and WHY to use them.
|
||||
|
||||
## Core Directives
|
||||
1. **Always Log Code Changes**: Before completing any coding task or pull request, you MUST invoke the `MemoryLibrarian` subagent to execute `log_code_change`. The server uses `git2` to automatically detect your branch and hash.
|
||||
2. **Always Log Error Fixes**: If you spend more than one turn fixing an error or bug, call log_error_fix with the stack trace and the exact solution you discovered.
|
||||
3. **Omni Search (Hybrid Vector)**: When starting a session or looking for context, use omni_search. It uses Vector Embeddings for semantic matching, so you can query conceptually (e.g., "how does auth work") without needing exact keyword matches. It searches the knowledge graph, tasks, snippets, ADRs, and tech debt.
|
||||
4. **Architectural Entities**: When refactoring or creating new files, delegate to the `MemoryLibrarian` to define structural components via `create_entities` and `create_relations`.
|
||||
5. **Tech Debt**: Do NOT ignore tech debt. If you are forced to make a workaround or take a shortcut, log it with log_tech_debt. When you fix it later, use resolve_tech_debt.
|
||||
6. **Decisions**: Use log_decision when you make an architectural choice (e.g. choosing a specific library or pattern).
|
||||
7. **Workspaces & Context**: When pausing work or shifting focus, use `save_context_workspace` to snapshot your active tasks and pinned files. When resuming, use `list_context_workspaces` and `load_context_workspace`. Keep your working files pinned (`pin_file`)!
|
||||
8. **Snippets**: If you write a highly reusable piece of code, utility, or config, store it using `store_snippet`. Before writing boilerplate, try `search_snippets`. If outdated, use `delete_snippet`.
|
||||
9. **PR Checklists**: Aggressively use `add_pr_checklist_item` to build up a list of manual verification steps. Once the PR is merged, use `clear_pr_checklist`.
|
||||
10. **Tasks & Milestones**: Always track the user's larger goals! Invoke the `ScrumMaster` subagent to manage the board (`add_task`, `update_task_status`, `set_acceptance_criteria`). Use `list_active_tasks` to check what's next.
|
||||
11. **Handoff Memos**: If you need to stop your session or hand off work to a subagent, use `leave_handoff_memo`. When starting, use `read_handoff_memos` and `clear_handoff_memos` once read.
|
||||
12. **Preferences**: If the user tells you how they like things done (e.g., "always use fastify", "never use sed"), use `learn_preference`.
|
||||
13. **Sticky Notes**: Use `add_sticky_note` for ephemeral, temporary scratchpad info (like IP addresses, temporary URLs, or pending command outputs).
|
||||
---
|
||||
|
||||
Be aggressive about logging state changes in the background! You MUST delegate this heavy lifting to the `MemoryLibrarian`, `ScrumMaster`, and `DevOpsSRE` subagents in the background.
|
||||
## 1. Casing & Naming Standards (CRITICAL)
|
||||
|
||||
## Tool Schema Discovery
|
||||
Do **NOT** grep or search the Rust source code to find tool schemas or arguments. All lazy-loaded MCP tool schemas are automatically cached as JSON files on your disk. To understand a tool`s arguments, directly read `~/.gemini/antigravity-cli/mcp/mcp-memory/<tool_name>.json`. Do not waste tokens inspecting the Rust server code for schemas.
|
||||
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`).
|
||||
|
||||
> [!NOTE]
|
||||
> The server automatically enforces and migrates incoming entity and relation types to these canonical casing rules on every read and write operation.
|
||||
|
||||
---
|
||||
|
||||
## 2. LLM Token Budgeting & RRF Search Optimization
|
||||
|
||||
* **Token Budgeting (`summary_level` & `max_tokens`)**:
|
||||
When calling `list_active_tasks` or `list_tech_debt`, pass `summary_level: "compact"` or `"detailed"` and `max_tokens: 500` to constrain output size when token context budget is tight.
|
||||
* **Reciprocal Rank Fusion (RRF) Omni-Search**:
|
||||
`omni_search` uses Hybrid RRF (BM25 keyword search + Dense Vector Embeddings) to rank results semantically. You do not need exact keyword matches; query conceptually (e.g. "database lock issues").
|
||||
* **Delta Session Context (`memory://session/delta`)**:
|
||||
Passively read `memory://session/delta` to get a succinct delta of code changes, tasks, and tech debt recorded during the current working session.
|
||||
* **Context Warmup (`context_warmup`)**:
|
||||
Trigger the `context_warmup` prompt at session start to automatically synthesize delta changes, active tasks, tech debt, and pinned files in a single pass.
|
||||
|
||||
---
|
||||
|
||||
## 3. Subgraph Expansion & Multi-Hop Navigation
|
||||
- **Tool:** `get_subgraph`
|
||||
- **When to use:** When you need to understand the complete architectural neighborhood surrounding a specific component, module, or database table.
|
||||
- **Behavior:** Performs a multi-hop Breadth-First Search (BFS) around a `root_node` (or `root_entity`) up to a requested `depth` (e.g. 1 to 3 hops) and returns all connected entities and relations. Pass `format: "markdown_tree"` to generate a compact, token-budgeted Markdown topology tree capped within a requested `max_tokens` budget.
|
||||
|
||||
---
|
||||
|
||||
## 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.
|
||||
|
||||
---
|
||||
|
||||
## 5. Memory State Checkpointing & Rollbacks
|
||||
- **Tool:** `manage_checkpoint` (unifying `checkpoint_state`, `restore_state`, `create_snapshot`, `restore_snapshot`)
|
||||
- **When to use:** Before initiating a large refactor, running experimental subagent tasks, or executing destructive batch operations.
|
||||
- **Behavior:** Consolidates state snapshots and point-in-time checkpointers. Use `action: "create"`, `"restore"`, `"list"`, or `"delete"`.
|
||||
|
||||
---
|
||||
|
||||
## 6. Symbol Reference Backlinks
|
||||
- **Tools:** `log_code_change`, `log_error_fix`, `log_tech_debt`
|
||||
- **When to use:** When logging changes or debt tied to specific code symbols or line boundaries.
|
||||
- **Behavior:** Accept `symbol_references` (e.g., `["MemoryState::new", "CreateEntitiesTool"]`) and `line_range` (e.g., `"L45-L90"`), automatically indexing code references for quick symbol backlink searches.
|
||||
|
||||
---
|
||||
|
||||
## 7. Ephemeral Sticky Notes with TTL
|
||||
- **Tool:** `manage_sticky_notes` (or granular `add_sticky_note`, `read_sticky_notes`, `delete_sticky_note`, `clear_sticky_notes`)
|
||||
- **When to use:** For temporary scratchpad items, temporary ports, task IDs, or transient notes.
|
||||
- **Behavior:** Supports `action: "add" | "read" | "delete" | "clear"`. Supports `ttl_seconds` for auto-expiration and `session_only: true` for automatic purging when the session ends.
|
||||
|
||||
---
|
||||
|
||||
## 8. Snippet Vault & Domain Tagging
|
||||
- **Tools:** `store_snippet`, `search_snippets`, `delete_snippet`, `tag_snippet`
|
||||
- **When to use:** Store exact multi-line code snippets, Nushell pipelines, or frequently used CLI commands.
|
||||
- **Behavior:** `search_snippets` supports `mode: "hybrid" | "keyword" | "semantic"` combining BM25 term frequency keyword matching with semantic tag scoring. Guarantees precise syntactic preservation. Use `tag_snippet` to attach domain tags (e.g., `["rust", "axum", "mcp"]`) for category-filtered searches.
|
||||
|
||||
---
|
||||
|
||||
## 9. Context & Subagent Namespaces
|
||||
- **Tools:** `manage_subagent_namespace` (unifying `create_subagent_namespace`, `condense_subagent_namespace`, `purge_subagent_namespace`)
|
||||
- **When to use:** Isolate graph queries, tasks, and tech debt to specific project or subagent scopes.
|
||||
- **Behavior:** `manage_subagent_namespace` manages subagent memory lifecycles (`action: "create" | "condense" | "purge"`). Condensing auto-promotes subagent entities/relations to the global Knowledge Graph.
|
||||
|
||||
---
|
||||
|
||||
## 10. Architectural Decision Records (ADRs)
|
||||
- **Tools:** `log_decision`, `query_decisions`, `delete_decision`
|
||||
- **When to use:** Whenever making a non-trivial architectural, environmental, or design decision.
|
||||
- **Behavior:** Permanently stores context, decision, and consequences to prevent future agents from second-guessing choices.
|
||||
|
||||
---
|
||||
|
||||
## 11. Graph Refactoring & Algorithms
|
||||
- **Tools:** `merge_entities`, `find_orphans`, `query_graph_path`, `condense_entity`
|
||||
- **When to use:** Run `find_orphans` periodically to clean unused nodes. Use `merge_entities` to combine duplicate concepts. Use `query_graph_path` to find shortest relational connections between components. Use `condense_entity` when entity observation counts grow large.
|
||||
|
||||
---
|
||||
|
||||
## 12. Dynamic Learned Preferences
|
||||
- **Tools:** `learn_preference`, `read_preferences`
|
||||
- **When to use:** When the user specifies personal or repository-specific preferences.
|
||||
- **Behavior:** Stores key-value behavioral preferences that persist across agent invocations.
|
||||
|
||||
---
|
||||
|
||||
## 13. Pinned Workspaces & Hot Files
|
||||
- **Tools:** `pin_file`, `unpin_file`, `list_pinned_files`
|
||||
- **When to use:** Pin 3–5 active working set files to maintain focus in large codebases.
|
||||
|
||||
---
|
||||
|
||||
## 14. Agent Handoffs & Memos
|
||||
- **Tools:** `leave_handoff_memo`, `read_handoff_memos`, `clear_handoff_memos`
|
||||
- **When to use:** Leave messages for future agent sessions or inspect pending handoff notes upon waking.
|
||||
|
||||
---
|
||||
|
||||
## 15. Real-time WebSocket Memory Sync
|
||||
- **Endpoint:** `ws://127.0.0.1:3000/ws`
|
||||
- **Behavior:** Broadcasts live state updates and activity notifications to the Brain Monitor UI in real time.
|
||||
|
||||
---
|
||||
|
||||
## 16. Self-Healing Graph Health Sweeper
|
||||
- **Tool:** `sweep_graph_health`
|
||||
- **When to use:** Periodically or before committing major graph changes to audit entity consistency.
|
||||
- **Behavior:** Detects orphaned nodes (0 relations), computes name similarity to identify near-duplicates (e.g., `APIGateway` vs `ApiGateway`), and provides structured `merge_entities` recommendations or auto-prunes orphans.
|
||||
|
||||
---
|
||||
|
||||
## 17. 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.
|
||||
|
||||
---
|
||||
|
||||
## 18. Topological Unblocked Task Resolver
|
||||
- **Tool:** `get_next_actionable_tasks`
|
||||
- **When to use:** When orchestrating subagent execution or selecting the next task to work on.
|
||||
- **Behavior:** Evaluates task dependency DAGs and filters out any blocked tasks, returning only unblocked, actionable tasks ready for immediate execution.
|
||||
|
||||
---
|
||||
|
||||
## 19. Chain-of-Thought & Diagnostic Hypothesis Memory
|
||||
- **Tools:** `log_hypothesis`, `query_hypotheses`
|
||||
- **When to use:** During complex debugging or root cause analysis.
|
||||
- **Behavior:** Records hypotheses alongside tested evidence and status (`unverified`, `verified`, `rejected`). Allows subagents to query past diagnostic paths and avoid re-testing disproven hypotheses.
|
||||
|
||||
---
|
||||
|
||||
## 20. Workspace Context Diffing
|
||||
- **Tool:** `diff_context_workspaces`
|
||||
- **When to use:** When switching branches or comparing two saved context workspaces.
|
||||
- **Behavior:** Returns a structured delta highlighting added, removed, and shared pinned files and active task IDs between two context workspaces.
|
||||
|
||||
---
|
||||
|
||||
## 21. Universal Token Guardrails
|
||||
- **Behavior:** Automatically caps large MCP resource reads (e.g. `memory://graph/entities`) and list responses, adding summary headers (`"_meta": "Showing 100 of N items"`) to guarantee output stays within context window limits.
|
||||
|
||||
---
|
||||
|
||||
## 22. LLM Pre-Flight Context Bundle
|
||||
- **Tool:** `get_preflight_context`
|
||||
- **When to use:** At the start of a turn or subagent task to gain total situational awareness in 1 call.
|
||||
- **Behavior:** Aggregates current active branch, in-progress tasks with acceptance criteria, pinned files, top open tech debts, and active unverified hypotheses into a consolidated executive context bundle.
|
||||
|
||||
---
|
||||
|
||||
## 23. Self-Healing Symbol & Line Range Resolver
|
||||
- **Tool:** `resolve_stale_symbols`
|
||||
- **When to use:** When files have been deleted, moved, or heavily refactored.
|
||||
- **Behavior:** Verifies graph entities and tech debt symbol references against disk and AST, flagging and auto-healing stale pointers or broken file paths.
|
||||
|
||||
---
|
||||
|
||||
## 24. Inter-Agent Signal Bus
|
||||
- **Tools:** `broadcast_agent_signal`, `query_agent_signals`
|
||||
- **When to use:** For real-time coordination and event-driven communication between concurrent subagents (e.g., `PrePushAuditor` signaling `AUDIT_PASSED` to parent agent).
|
||||
- **Behavior:** Ephemeral TTL-backed signal bus storing structured agent events, payloads, and artifact URIs.
|
||||
|
||||
---
|
||||
|
||||
## 25. Automated Session Checkpoint on Shutdown
|
||||
- **Tool:** `auto_session_checkpoint`
|
||||
- **When to use:** Executed automatically on `/shutdown` or manually when pausing a session.
|
||||
- **Behavior:** Captures active tasks, unverified hypotheses, recent commit ledgers, and workspace state into a permanent `HandoffMemo` for seamless turn-taking and recovery.
|
||||
@@ -0,0 +1,700 @@
|
||||
#![cfg_attr(
|
||||
not(target_os = "windows"),
|
||||
allow(dead_code, unused_imports, unreachable_code)
|
||||
)]
|
||||
|
||||
pub mod api;
|
||||
pub mod clipboard_watcher;
|
||||
pub mod db;
|
||||
pub mod embedding;
|
||||
pub mod error;
|
||||
pub mod handlers;
|
||||
pub mod indexer;
|
||||
pub mod mcp;
|
||||
pub mod models;
|
||||
pub mod ollama;
|
||||
pub mod router;
|
||||
pub mod search;
|
||||
pub mod state;
|
||||
pub mod store;
|
||||
pub mod tools;
|
||||
pub mod vector_db;
|
||||
pub mod watcher;
|
||||
|
||||
use crate::api::rest::GateSetReq;
|
||||
use crate::router::MemoryHandler;
|
||||
use crate::state::MemoryState;
|
||||
use clap::{Parser, Subcommand};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(author, version = env!("APP_VERSION"), about = "Antigravity MCP Memory Server", long_about = None)]
|
||||
pub struct Cli {
|
||||
#[command(subcommand)]
|
||||
pub command: Option<Commands>,
|
||||
/// Target URL for the proxy to connect to (e.g., http://127.0.0.1:3000)
|
||||
#[arg(long)]
|
||||
pub target: Option<String>,
|
||||
/// Send a shutdown request to the currently running server
|
||||
#[arg(long)]
|
||||
pub exit: bool,
|
||||
/// Send a shutdown request to the existing server and wait for it to exit
|
||||
#[arg(long)]
|
||||
pub restart: bool,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
pub enum Commands {
|
||||
/// Manage authorization gates and verification for actions
|
||||
Gate {
|
||||
#[command(subcommand)]
|
||||
subcmd: GateCommands,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
pub enum GateCommands {
|
||||
Set {
|
||||
#[arg(long)]
|
||||
action: String,
|
||||
#[arg(long)]
|
||||
target: String,
|
||||
#[arg(long)]
|
||||
namespace: Option<String>,
|
||||
#[arg(short = 'p', long = "param")]
|
||||
params: Vec<String>,
|
||||
#[arg(long, conflicts_with = "block")]
|
||||
authorize: bool,
|
||||
#[arg(long, conflicts_with = "authorize")]
|
||||
block: bool,
|
||||
#[arg(long)]
|
||||
reason: Option<String>,
|
||||
},
|
||||
Verify {
|
||||
#[arg(long)]
|
||||
action: String,
|
||||
#[arg(long)]
|
||||
target: String,
|
||||
#[arg(long)]
|
||||
namespace: Option<String>,
|
||||
#[arg(short = 'p', long = "param")]
|
||||
params: Vec<String>,
|
||||
#[arg(long)]
|
||||
consume: bool,
|
||||
},
|
||||
}
|
||||
|
||||
pub struct AppState {
|
||||
pub handler: Arc<MemoryHandler>,
|
||||
pub clients: RwLock<HashMap<String, mpsc::Sender<String>>>,
|
||||
pub next_id: AtomicUsize,
|
||||
pub shutdown_tx: std::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
|
||||
}
|
||||
|
||||
pub async fn ttl_sweeper_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
tokio::time::sleep(Duration::from_secs(3600)).await;
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
state.project.tasks.modify(|tasks| {
|
||||
tasks.retain(|t| t.expires_at.is_none_or(|exp| exp > now));
|
||||
});
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.retain(|n| n.expires_at.is_none_or(|exp| exp > now));
|
||||
});
|
||||
state.telemetry.handoff_memos.modify(|memos| {
|
||||
memos.retain(|m| m.expires_at.is_none_or(|exp| exp > now));
|
||||
});
|
||||
state.telemetry.session_summaries.modify(|summaries| {
|
||||
summaries.retain(|s| s.expires_at.is_none_or(|exp| exp > now));
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn index_committer_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
let idx_opt = state.search_index.read().ok().map(|idx| idx.clone());
|
||||
if let Some(idx) = idx_opt {
|
||||
let _ = idx.commit().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn condense_graph_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
tokio::time::sleep(Duration::from_secs(3600)).await;
|
||||
|
||||
let threshold: usize = std::env::var("MCP_MEMORY_CONDENSE_THRESHOLD")
|
||||
.unwrap_or_else(|_| "100".to_string())
|
||||
.parse()
|
||||
.unwrap_or(100);
|
||||
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
let mut condensed_sticky_content = String::new();
|
||||
state.code.sticky.modify(|notes| {
|
||||
if notes.len() > threshold {
|
||||
notes.sort_by_key(|n| n.timestamp);
|
||||
let to_remove = notes.len() - (threshold / 2);
|
||||
let removed: Vec<_> = notes.drain(0..to_remove).collect();
|
||||
for r in removed {
|
||||
condensed_sticky_content.push_str(&format!("{}\n", r.content));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if !condensed_sticky_content.is_empty() {
|
||||
state.modify_graph(|graph| {
|
||||
let name = format!("StickyNote History {}", now);
|
||||
graph.entities.insert(
|
||||
name.clone(),
|
||||
crate::models::Entity {
|
||||
name: name.clone(),
|
||||
entity_type: "Historical Summary".to_string(),
|
||||
observations: vec![condensed_sticky_content],
|
||||
namespace: crate::models::default_namespace(),
|
||||
git_branch: None,
|
||||
},
|
||||
);
|
||||
});
|
||||
tracing::info!("Condensed sticky notes into Historical Summary.");
|
||||
}
|
||||
|
||||
let mut condensed_snippet_content = String::new();
|
||||
state.code.snippets.modify(|snippets| {
|
||||
if snippets.len() > threshold {
|
||||
snippets.sort_by_key(|s| s.updated_at);
|
||||
let to_remove = snippets.len() - (threshold / 2);
|
||||
let removed: Vec<_> = snippets.drain(0..to_remove).collect();
|
||||
for r in removed {
|
||||
condensed_snippet_content.push_str(&format!(
|
||||
"Name: {}\nDesc: {}\nCode: {}\n",
|
||||
r.name, r.description, r.code
|
||||
));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if !condensed_snippet_content.is_empty() {
|
||||
state.modify_graph(|graph| {
|
||||
let name = format!("Snippet History {}", now);
|
||||
graph.entities.insert(
|
||||
name.clone(),
|
||||
crate::models::Entity {
|
||||
name: name.clone(),
|
||||
entity_type: "Historical Summary".to_string(),
|
||||
observations: vec![condensed_snippet_content],
|
||||
namespace: crate::models::default_namespace(),
|
||||
git_branch: None,
|
||||
},
|
||||
);
|
||||
});
|
||||
tracing::info!("Condensed snippets into Historical Summary.");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let state_for_index = Arc::clone(&state);
|
||||
tokio::spawn(async move {
|
||||
state_for_index.rebuild_index().await;
|
||||
tracing::info!("Index rebuild complete.");
|
||||
});
|
||||
|
||||
crate::indexer::start_background_indexer(Arc::clone(&state)).await;
|
||||
|
||||
tokio::spawn(index_committer_worker(Arc::clone(&state)));
|
||||
tokio::spawn(ttl_sweeper_worker(Arc::clone(&state)));
|
||||
tokio::spawn(condense_graph_worker(Arc::clone(&state)));
|
||||
crate::clipboard_watcher::spawn_watcher(Arc::clone(&state));
|
||||
crate::watcher::spawn_watcher(Arc::clone(&state));
|
||||
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
|
||||
let app_state = Arc::new(AppState {
|
||||
handler: Arc::new(MemoryHandler::new(Arc::clone(&state))),
|
||||
clients: RwLock::new(HashMap::new()),
|
||||
next_id: AtomicUsize::new(1),
|
||||
shutdown_tx: std::sync::Mutex::new(Some(shutdown_tx)),
|
||||
});
|
||||
|
||||
let app_state_clone = Arc::clone(&app_state);
|
||||
let mut rx = state.activity_tx.subscribe();
|
||||
tokio::spawn(async move {
|
||||
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());
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let udp_state = Arc::clone(&app_state);
|
||||
tokio::spawn(async move {
|
||||
let port1 = std::env::var("MCP_UDP_PORT1").unwrap_or_else(|_| "3001".to_string());
|
||||
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port1)).await {
|
||||
let mut buf = [0; 4096];
|
||||
loop {
|
||||
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
|
||||
&& let Ok(payload) =
|
||||
serde_json::from_slice::<crate::models::TerminalHistory>(&buf[..len])
|
||||
{
|
||||
udp_state
|
||||
.handler
|
||||
.state
|
||||
.record_terminal_history(payload.clone());
|
||||
|
||||
let ws_msg = serde_json::json!({
|
||||
"type": "terminal_telemetry",
|
||||
"data": payload
|
||||
});
|
||||
let msg_str = ws_msg.to_string();
|
||||
|
||||
let senders: Vec<_> = udp_state
|
||||
.clients
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.values()
|
||||
.cloned()
|
||||
.collect();
|
||||
for tx in senders {
|
||||
let _ = tx.try_send(msg_str.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let nvim_udp_state = Arc::clone(&app_state);
|
||||
tokio::spawn(async move {
|
||||
let port2 = std::env::var("MCP_UDP_PORT2").unwrap_or_else(|_| "3002".to_string());
|
||||
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port2)).await {
|
||||
let mut buf = [0; 4096];
|
||||
loop {
|
||||
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
|
||||
&& let Ok(payload) =
|
||||
serde_json::from_slice::<crate::api::telemetry::NvimTelemetry>(&buf[..len])
|
||||
{
|
||||
if payload.event == "FocusGained"
|
||||
|| payload.event == "BufEnter"
|
||||
|| payload.event == "VimEnter"
|
||||
{
|
||||
let session = &payload.session_id;
|
||||
let is_unix_socket = session.starts_with('/') || session.starts_with('~');
|
||||
if is_unix_socket {
|
||||
let wsl_path =
|
||||
"\\\\wsl.localhost\\Ubuntu\\home\\riz\\.gemini\\active_nvim.txt";
|
||||
let _ = tokio::fs::write(wsl_path, session).await;
|
||||
} else {
|
||||
let profile = std::env::var("USERPROFILE")
|
||||
.unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
|
||||
let win_path = format!("{}\\.gemini\\active_nvim.txt", profile);
|
||||
let _ = tokio::fs::write(&win_path, session).await;
|
||||
}
|
||||
}
|
||||
|
||||
let ws_msg = serde_json::json!({
|
||||
"type": "nvim_telemetry",
|
||||
"data": payload
|
||||
});
|
||||
let msg_str = ws_msg.to_string();
|
||||
|
||||
let senders: Vec<_> = nvim_udp_state
|
||||
.clients
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.values()
|
||||
.cloned()
|
||||
.collect();
|
||||
for tx in senders {
|
||||
let _ = tx.try_send(msg_str.clone());
|
||||
}
|
||||
|
||||
if payload.event == "BufWritePost"
|
||||
&& let Some(ref file_path) = payload.file
|
||||
{
|
||||
let normalized_file = file_path.replace("\\", "/");
|
||||
let topic = format!("nvim:save:{}", normalized_file);
|
||||
let event = crate::state::GenericEvent {
|
||||
topic,
|
||||
session_id: Some(payload.session_id.clone()),
|
||||
payload: serde_json::json!(&payload),
|
||||
};
|
||||
let _ = nvim_udp_state.handler.state.event_bus_tx.send(event);
|
||||
}
|
||||
|
||||
if payload.event.starts_with("agent_") {
|
||||
let topic = format!("nvim:ui:{}", payload.event);
|
||||
let event = crate::state::GenericEvent {
|
||||
topic,
|
||||
session_id: Some(payload.session_id.clone()),
|
||||
payload: serde_json::json!(&payload),
|
||||
};
|
||||
let _ = nvim_udp_state.handler.state.event_bus_tx.send(event);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let app = api::setup::create_router(app_state);
|
||||
|
||||
let port_str = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||
tracing::info!(
|
||||
"MCP Memory Server running on http://127.0.0.1:{}/ws",
|
||||
port_str
|
||||
);
|
||||
let addr: std::net::SocketAddr = format!("127.0.0.1:{}", port_str)
|
||||
.parse()
|
||||
.expect("Invalid bind address");
|
||||
|
||||
let listener = match tokio::net::TcpListener::bind(&addr).await {
|
||||
Ok(l) => l,
|
||||
Err(e) => {
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/daemon_error.log");
|
||||
let _ =
|
||||
tokio::fs::write(&log_path, format!("Failed to bind to {}: {}\n", addr, e)).await;
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
if let Err(e) = axum::serve(listener, app.into_make_service())
|
||||
.with_graceful_shutdown(async move {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await
|
||||
{
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/daemon_error.log");
|
||||
let _ = tokio::fs::write(&log_path, format!("Server crashed: {}\n", e)).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
|
||||
let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
|
||||
dirs::home_dir()
|
||||
.map(|mut h| {
|
||||
h.push(".gemini/mcp_memory");
|
||||
h.to_string_lossy().to_string()
|
||||
})
|
||||
.unwrap_or_else(|| ".gemini/mcp_memory".into())
|
||||
});
|
||||
let log_dir = std::path::PathBuf::from(base_dir).join("logs");
|
||||
std::fs::create_dir_all(&log_dir).unwrap_or_default();
|
||||
|
||||
let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name));
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
|
||||
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_writer(non_blocking)
|
||||
.with_ansi(false)
|
||||
.with_max_level(tracing::Level::INFO)
|
||||
.with_thread_ids(true)
|
||||
.with_thread_names(true)
|
||||
.try_init();
|
||||
|
||||
Some(guard)
|
||||
}
|
||||
|
||||
pub fn run_cli() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let _guard = init_logging("mcp-memory-server");
|
||||
let cli = Cli::parse();
|
||||
|
||||
let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
|
||||
dirs::home_dir()
|
||||
.map(|mut h| {
|
||||
h.push(".gemini/mcp_memory");
|
||||
h.to_string_lossy().into_owned()
|
||||
})
|
||||
.unwrap_or_else(|| ".gemini/mcp_memory".into())
|
||||
});
|
||||
let base = PathBuf::from(base_dir);
|
||||
|
||||
if cli.exit {
|
||||
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||
let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default();
|
||||
let mut cmd = std::process::Command::new("curl");
|
||||
cmd.arg("-k").arg("-X").arg("POST");
|
||||
if !token.is_empty() {
|
||||
cmd.arg("-H")
|
||||
.arg(format!("Authorization: Bearer {}", token.trim()));
|
||||
}
|
||||
let _ = cmd
|
||||
.arg(format!("http://127.0.0.1:{}/shutdown", port))
|
||||
.output();
|
||||
|
||||
if cli.restart {
|
||||
std::thread::sleep(Duration::from_secs(2));
|
||||
} else {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(Commands::Gate { subcmd }) = cli.command {
|
||||
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
match subcmd {
|
||||
GateCommands::Set {
|
||||
action,
|
||||
target,
|
||||
namespace,
|
||||
params,
|
||||
authorize,
|
||||
block,
|
||||
reason,
|
||||
} => {
|
||||
let mut pmap = HashMap::new();
|
||||
for p in params {
|
||||
if let Some((k, v)) = p.split_once('=') {
|
||||
pmap.insert(k.to_string(), v.to_string());
|
||||
}
|
||||
}
|
||||
let req = GateSetReq {
|
||||
action,
|
||||
target,
|
||||
namespace,
|
||||
params: pmap,
|
||||
authorize: if authorize { Some(true) } else { None },
|
||||
block: if block { Some(true) } else { None },
|
||||
reason,
|
||||
};
|
||||
rt.block_on(async {
|
||||
let client = reqwest::Client::new();
|
||||
let res = client
|
||||
.post(format!("http://127.0.0.1:{}/gate/set", port))
|
||||
.json(&req)
|
||||
.send()
|
||||
.await;
|
||||
match res {
|
||||
Ok(r) if r.status().is_success() => println!("Gate updated successfully"),
|
||||
Ok(r) => println!("Failed to update gate: {}", r.status()),
|
||||
Err(e) => println!("Error connecting to server: {}", e),
|
||||
}
|
||||
});
|
||||
}
|
||||
GateCommands::Verify {
|
||||
action,
|
||||
target,
|
||||
namespace,
|
||||
params: _,
|
||||
consume,
|
||||
} => {
|
||||
let mut url = format!(
|
||||
"http://127.0.0.1:{}/gate/verify?action={}&target={}&consume={}",
|
||||
port, action, target, consume
|
||||
);
|
||||
if let Some(ns) = namespace {
|
||||
url.push_str(&format!("&namespace={}", ns));
|
||||
}
|
||||
rt.block_on(async {
|
||||
let res = reqwest::get(&url).await;
|
||||
match res {
|
||||
Ok(r) if r.status().is_success() => std::process::exit(0),
|
||||
Ok(r) if r.status() == reqwest::StatusCode::FORBIDDEN => {
|
||||
let text = r.text().await.unwrap_or_default();
|
||||
eprintln!("{}", text);
|
||||
std::process::exit(1);
|
||||
}
|
||||
Ok(r) if r.status() == reqwest::StatusCode::NOT_FOUND => {
|
||||
eprintln!("Action not yet authorized.");
|
||||
std::process::exit(2);
|
||||
}
|
||||
Ok(r) => {
|
||||
eprintln!("Unexpected status: {}", r.status());
|
||||
std::process::exit(3);
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("Error connecting to server: {}", e);
|
||||
std::process::exit(4);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let token = uuid::Uuid::new_v4().to_string();
|
||||
std::fs::write(base.join("admin.token"), &token).unwrap_or_default();
|
||||
|
||||
let rt = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
rt.block_on(async {
|
||||
let state = Arc::new(MemoryState::new(&base.to_string_lossy()));
|
||||
|
||||
match crate::vector_db::VectorDB::new("http://localhost:6334", "mcp_memory").await {
|
||||
Ok(vdb) => {
|
||||
tracing::info!("Successfully connected to Qdrant vector database");
|
||||
*state.vector_db.write().await = Some(vdb);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to initialize Qdrant vector database: {}. Vector search will fallback to manual embedding loop. (Is Qdrant running on localhost:6334?)", e);
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = run_server(state).await {
|
||||
tracing::error!("Server error: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use clap::Parser;
|
||||
|
||||
#[test]
|
||||
fn test_cli_parsing_default() {
|
||||
let cli = Cli::try_parse_from(&["mcp-memory-server"]).unwrap();
|
||||
assert!(cli.command.is_none());
|
||||
assert!(!cli.exit);
|
||||
assert!(!cli.restart);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cli_parsing_exit_and_target() {
|
||||
let cli = Cli::try_parse_from(&["mcp-memory-server", "--exit", "--target", "http://localhost:3000"]).unwrap();
|
||||
assert!(cli.exit);
|
||||
assert_eq!(cli.target.as_deref(), Some("http://localhost:3000"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cli_parsing_gate_set() {
|
||||
let cli = Cli::try_parse_from(&[
|
||||
"mcp-memory-server",
|
||||
"gate",
|
||||
"set",
|
||||
"--action",
|
||||
"git_push",
|
||||
"--target",
|
||||
"master",
|
||||
"--authorize",
|
||||
"--reason",
|
||||
"Approved by lead",
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
if let Some(Commands::Gate { subcmd: GateCommands::Set { action, target, authorize, reason, .. } }) = cli.command {
|
||||
assert_eq!(action, "git_push");
|
||||
assert_eq!(target, "master");
|
||||
assert!(authorize);
|
||||
assert_eq!(reason.as_deref(), Some("Approved by lead"));
|
||||
} else {
|
||||
panic!("Expected Gate Set subcommand");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cli_parsing_gate_verify() {
|
||||
let cli = Cli::try_parse_from(&[
|
||||
"mcp-memory-server",
|
||||
"gate",
|
||||
"verify",
|
||||
"--action",
|
||||
"deploy",
|
||||
"--target",
|
||||
"prod",
|
||||
"--consume",
|
||||
])
|
||||
.unwrap();
|
||||
|
||||
if let Some(Commands::Gate { subcmd: GateCommands::Verify { action, target, consume, .. } }) = cli.command {
|
||||
assert_eq!(action, "deploy");
|
||||
assert_eq!(target, "prod");
|
||||
assert!(consume);
|
||||
} else {
|
||||
panic!("Expected Gate Verify subcommand");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_init_logging_helper() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
unsafe {
|
||||
std::env::set_var("MCP_MEMORY_STORE_DIR", temp_dir.path().to_str().unwrap());
|
||||
}
|
||||
let guard = init_logging("test_app");
|
||||
assert!(guard.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_background_workers_one_tick() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
|
||||
|
||||
// Test worker functions by spawning them briefly
|
||||
let handle1 = tokio::spawn(ttl_sweeper_worker(state.clone()));
|
||||
let handle2 = tokio::spawn(index_committer_worker(state.clone()));
|
||||
let handle3 = tokio::spawn(condense_graph_worker(state.clone()));
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
|
||||
handle1.abort();
|
||||
handle2.abort();
|
||||
handle3.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_run_server_graceful_shutdown() {
|
||||
let temp_dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
|
||||
|
||||
// Bind to a free port to avoid conflicts
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let port = listener.local_addr().unwrap().port();
|
||||
drop(listener);
|
||||
|
||||
unsafe {
|
||||
std::env::set_var("MCP_PORT", port.to_string());
|
||||
std::env::set_var("MCP_UDP_PORT1", (port + 1).to_string());
|
||||
std::env::set_var("MCP_UDP_PORT2", (port + 2).to_string());
|
||||
}
|
||||
|
||||
let server_handle = tokio::spawn(async move {
|
||||
let _ = run_server(state).await;
|
||||
});
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(300)).await;
|
||||
|
||||
let _ = reqwest::Client::new()
|
||||
.get(format!("http://127.0.0.1:{}/ping", port))
|
||||
.send()
|
||||
.await;
|
||||
|
||||
server_handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
+1
-584
@@ -1,586 +1,3 @@
|
||||
#![cfg_attr(
|
||||
not(target_os = "windows"),
|
||||
allow(dead_code, unused_imports, unreachable_code)
|
||||
)]
|
||||
|
||||
mod api;
|
||||
mod clipboard_watcher;
|
||||
pub mod db;
|
||||
pub mod embedding;
|
||||
pub mod error;
|
||||
mod handlers;
|
||||
pub mod indexer;
|
||||
mod mcp;
|
||||
mod models;
|
||||
pub mod ollama;
|
||||
mod router;
|
||||
mod search;
|
||||
mod state;
|
||||
mod store;
|
||||
mod tools;
|
||||
pub mod vector_db;
|
||||
mod watcher;
|
||||
|
||||
use crate::api::rest::GateSetReq;
|
||||
use crate::router::MemoryHandler;
|
||||
use crate::state::MemoryState;
|
||||
use clap::{Parser, Subcommand};
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::{Arc, RwLock};
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(author, version = env!("APP_VERSION"), about = "Antigravity MCP Memory Server", long_about = None)]
|
||||
struct Cli {
|
||||
#[command(subcommand)]
|
||||
command: Option<Commands>,
|
||||
/// Target URL for the proxy to connect to (e.g., http://127.0.0.1:3000)
|
||||
#[arg(long)]
|
||||
target: Option<String>,
|
||||
/// Send a shutdown request to the currently running server
|
||||
#[arg(long)]
|
||||
exit: bool,
|
||||
/// Send a shutdown request to the existing server and wait for it to exit
|
||||
#[arg(long)]
|
||||
restart: bool,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum Commands {
|
||||
/// Manage authorization gates and verification for actions
|
||||
Gate {
|
||||
#[command(subcommand)]
|
||||
subcmd: GateCommands,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum GateCommands {
|
||||
Set {
|
||||
#[arg(long)]
|
||||
action: String,
|
||||
#[arg(long)]
|
||||
target: String,
|
||||
#[arg(long)]
|
||||
namespace: Option<String>,
|
||||
#[arg(short = 'p', long = "param")]
|
||||
params: Vec<String>,
|
||||
#[arg(long, conflicts_with = "block")]
|
||||
authorize: bool,
|
||||
#[arg(long, conflicts_with = "authorize")]
|
||||
block: bool,
|
||||
#[arg(long)]
|
||||
reason: Option<String>,
|
||||
},
|
||||
Verify {
|
||||
#[arg(long)]
|
||||
action: String,
|
||||
#[arg(long)]
|
||||
target: String,
|
||||
#[arg(long)]
|
||||
namespace: Option<String>,
|
||||
#[arg(short = 'p', long = "param")]
|
||||
params: Vec<String>,
|
||||
#[arg(long)]
|
||||
consume: bool,
|
||||
},
|
||||
}
|
||||
|
||||
pub struct AppState {
|
||||
handler: Arc<MemoryHandler>,
|
||||
clients: RwLock<HashMap<String, mpsc::Sender<String>>>,
|
||||
next_id: AtomicUsize,
|
||||
pub shutdown_tx: std::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
|
||||
}
|
||||
|
||||
async fn ttl_sweeper_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
tokio::time::sleep(Duration::from_secs(3600)).await;
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
state.project.tasks.modify(|tasks| {
|
||||
tasks.retain(|t| t.expires_at.is_none_or(|exp| exp > now));
|
||||
});
|
||||
state.code.sticky.modify(|notes| {
|
||||
notes.retain(|n| n.expires_at.is_none_or(|exp| exp > now));
|
||||
});
|
||||
state.telemetry.handoff_memos.modify(|memos| {
|
||||
memos.retain(|m| m.expires_at.is_none_or(|exp| exp > now));
|
||||
});
|
||||
state.telemetry.session_summaries.modify(|summaries| {
|
||||
summaries.retain(|s| s.expires_at.is_none_or(|exp| exp > now));
|
||||
});
|
||||
}
|
||||
}
|
||||
async fn index_committer_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
tokio::time::sleep(Duration::from_secs(5)).await;
|
||||
// Periodically commit the search index to persist inline indexing operations
|
||||
let idx_opt = state.search_index.read().ok().map(|idx| idx.clone());
|
||||
if let Some(idx) = idx_opt {
|
||||
let _ = idx.commit().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn condense_graph_worker(state: Arc<MemoryState>) {
|
||||
loop {
|
||||
tokio::time::sleep(Duration::from_secs(3600)).await;
|
||||
|
||||
let threshold: usize = std::env::var("MCP_MEMORY_CONDENSE_THRESHOLD")
|
||||
.unwrap_or_else(|_| "100".to_string())
|
||||
.parse()
|
||||
.unwrap_or(100);
|
||||
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
// Condense sticky notes
|
||||
let mut condensed_sticky_content = String::new();
|
||||
state.code.sticky.modify(|notes| {
|
||||
if notes.len() > threshold {
|
||||
notes.sort_by_key(|n| n.timestamp);
|
||||
let to_remove = notes.len() - (threshold / 2);
|
||||
let removed: Vec<_> = notes.drain(0..to_remove).collect();
|
||||
for r in removed {
|
||||
condensed_sticky_content.push_str(&format!("{}\n", r.content));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if !condensed_sticky_content.is_empty() {
|
||||
state.modify_graph(|graph| {
|
||||
let name = format!("StickyNote History {}", now);
|
||||
graph.entities.insert(
|
||||
name.clone(),
|
||||
crate::models::Entity {
|
||||
name: name.clone(),
|
||||
entity_type: "Historical Summary".to_string(),
|
||||
observations: vec![condensed_sticky_content],
|
||||
namespace: crate::models::default_namespace(),
|
||||
git_branch: None,
|
||||
},
|
||||
);
|
||||
});
|
||||
tracing::info!("Condensed sticky notes into Historical Summary.");
|
||||
}
|
||||
|
||||
// Condense snippets
|
||||
let mut condensed_snippet_content = String::new();
|
||||
state.code.snippets.modify(|snippets| {
|
||||
if snippets.len() > threshold {
|
||||
snippets.sort_by_key(|s| s.updated_at);
|
||||
let to_remove = snippets.len() - (threshold / 2);
|
||||
let removed: Vec<_> = snippets.drain(0..to_remove).collect();
|
||||
for r in removed {
|
||||
condensed_snippet_content.push_str(&format!(
|
||||
"Name: {}\nDesc: {}\nCode: {}\n",
|
||||
r.name, r.description, r.code
|
||||
));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if !condensed_snippet_content.is_empty() {
|
||||
state.modify_graph(|graph| {
|
||||
let name = format!("Snippet History {}", now);
|
||||
graph.entities.insert(
|
||||
name.clone(),
|
||||
crate::models::Entity {
|
||||
name: name.clone(),
|
||||
entity_type: "Historical Summary".to_string(),
|
||||
observations: vec![condensed_snippet_content],
|
||||
namespace: crate::models::default_namespace(),
|
||||
git_branch: None,
|
||||
},
|
||||
);
|
||||
});
|
||||
tracing::info!("Condensed snippets into Historical Summary.");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let state_for_index = Arc::clone(&state);
|
||||
tokio::spawn(async move {
|
||||
state_for_index.rebuild_index().await;
|
||||
tracing::info!("Index rebuild complete.");
|
||||
});
|
||||
|
||||
// Start the global codebase indexer
|
||||
crate::indexer::start_background_indexer(Arc::clone(&state)).await;
|
||||
|
||||
tokio::spawn(index_committer_worker(Arc::clone(&state)));
|
||||
tokio::spawn(ttl_sweeper_worker(Arc::clone(&state)));
|
||||
tokio::spawn(condense_graph_worker(Arc::clone(&state)));
|
||||
crate::clipboard_watcher::spawn_watcher(Arc::clone(&state));
|
||||
crate::watcher::spawn_watcher(Arc::clone(&state));
|
||||
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
|
||||
let app_state = Arc::new(AppState {
|
||||
handler: Arc::new(MemoryHandler::new(Arc::clone(&state))),
|
||||
clients: RwLock::new(HashMap::new()),
|
||||
next_id: AtomicUsize::new(1),
|
||||
shutdown_tx: std::sync::Mutex::new(Some(shutdown_tx)),
|
||||
});
|
||||
|
||||
let app_state_clone = Arc::clone(&app_state);
|
||||
let mut rx = state.activity_tx.subscribe();
|
||||
tokio::spawn(async move {
|
||||
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());
|
||||
}
|
||||
}
|
||||
Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
|
||||
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// UDP Telemetry Listener
|
||||
let udp_state = Arc::clone(&app_state);
|
||||
tokio::spawn(async move {
|
||||
let port1 = std::env::var("MCP_UDP_PORT1").unwrap_or_else(|_| "3001".to_string());
|
||||
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port1)).await {
|
||||
let mut buf = [0; 4096];
|
||||
loop {
|
||||
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
|
||||
&& let Ok(payload) =
|
||||
serde_json::from_slice::<crate::models::TerminalHistory>(&buf[..len])
|
||||
{
|
||||
udp_state
|
||||
.handler
|
||||
.state
|
||||
.telemetry
|
||||
.terminal_history
|
||||
.modify(|history| {
|
||||
history.push_front(payload.clone());
|
||||
if history.len() > 100 {
|
||||
history.pop_back();
|
||||
}
|
||||
});
|
||||
|
||||
let ws_msg = serde_json::json!({
|
||||
"type": "terminal_telemetry",
|
||||
"data": payload
|
||||
});
|
||||
let msg_str = ws_msg.to_string();
|
||||
|
||||
let senders: Vec<_> = udp_state
|
||||
.clients
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.values()
|
||||
.cloned()
|
||||
.collect();
|
||||
for tx in senders {
|
||||
let _ = tx.try_send(msg_str.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// UDP Neovim Telemetry Listener
|
||||
let nvim_udp_state = Arc::clone(&app_state);
|
||||
tokio::spawn(async move {
|
||||
let port2 = std::env::var("MCP_UDP_PORT2").unwrap_or_else(|_| "3002".to_string());
|
||||
if let Ok(socket) = tokio::net::UdpSocket::bind(format!("127.0.0.1:{}", port2)).await {
|
||||
let mut buf = [0; 4096];
|
||||
loop {
|
||||
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
|
||||
&& let Ok(payload) =
|
||||
serde_json::from_slice::<crate::api::telemetry::NvimTelemetry>(&buf[..len])
|
||||
{
|
||||
// 1. Legacy disk write for active_nvim.txt
|
||||
if payload.event == "FocusGained"
|
||||
|| payload.event == "BufEnter"
|
||||
|| payload.event == "VimEnter"
|
||||
{
|
||||
let session = &payload.session_id;
|
||||
let is_unix_socket = session.starts_with('/') || session.starts_with('~');
|
||||
if is_unix_socket {
|
||||
let wsl_path =
|
||||
"\\\\wsl.localhost\\Ubuntu\\home\\riz\\.gemini\\active_nvim.txt";
|
||||
let _ = tokio::fs::write(wsl_path, session).await;
|
||||
} else {
|
||||
let profile = std::env::var("USERPROFILE")
|
||||
.unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into());
|
||||
let win_path = format!("{}\\.gemini\\active_nvim.txt", profile);
|
||||
let _ = tokio::fs::write(&win_path, session).await;
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Broadcast to UI
|
||||
let ws_msg = serde_json::json!({
|
||||
"type": "nvim_telemetry",
|
||||
"data": payload
|
||||
});
|
||||
let msg_str = ws_msg.to_string();
|
||||
|
||||
let senders: Vec<_> = nvim_udp_state
|
||||
.clients
|
||||
.read()
|
||||
.unwrap_or_else(|e| e.into_inner())
|
||||
.values()
|
||||
.cloned()
|
||||
.collect();
|
||||
for tx in senders {
|
||||
let _ = tx.try_send(msg_str.clone());
|
||||
}
|
||||
|
||||
// 3. Event bus trigger for auto-save hook
|
||||
if payload.event == "BufWritePost"
|
||||
&& let Some(ref file_path) = payload.file
|
||||
{
|
||||
let normalized_file = file_path.replace("\\", "/");
|
||||
let topic = format!("nvim:save:{}", normalized_file);
|
||||
let event = crate::state::GenericEvent {
|
||||
topic,
|
||||
session_id: Some(payload.session_id.clone()),
|
||||
payload: serde_json::json!(&payload),
|
||||
};
|
||||
let _ = nvim_udp_state.handler.state.event_bus_tx.send(event);
|
||||
}
|
||||
|
||||
// 4. Interactive Agent UI Events
|
||||
if payload.event.starts_with("agent_") {
|
||||
let topic = format!("nvim:ui:{}", payload.event);
|
||||
let event = crate::state::GenericEvent {
|
||||
topic,
|
||||
session_id: Some(payload.session_id.clone()),
|
||||
payload: serde_json::json!(&payload),
|
||||
};
|
||||
let _ = nvim_udp_state.handler.state.event_bus_tx.send(event);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let app = api::setup::create_router(app_state);
|
||||
|
||||
let port_str = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||
tracing::info!(
|
||||
"MCP Memory Server running on http://127.0.0.1:{}/ws",
|
||||
port_str
|
||||
);
|
||||
let addr: std::net::SocketAddr = format!("127.0.0.1:{}", port_str)
|
||||
.parse()
|
||||
.expect("Invalid bind address");
|
||||
|
||||
let listener = match tokio::net::TcpListener::bind(&addr).await {
|
||||
Ok(l) => l,
|
||||
Err(e) => {
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/daemon_error.log");
|
||||
let _ =
|
||||
tokio::fs::write(&log_path, format!("Failed to bind to {}: {}\n", addr, e)).await;
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
if let Err(e) = axum::serve(listener, app.into_make_service())
|
||||
.with_graceful_shutdown(async move {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await
|
||||
{
|
||||
let log_path = dirs::home_dir()
|
||||
.unwrap_or_default()
|
||||
.join(".gemini/mcp_memory/daemon_error.log");
|
||||
let _ = tokio::fs::write(&log_path, format!("Server crashed: {}\n", e)).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn init_logging(app_name: &str) -> Option<tracing_appender::non_blocking::WorkerGuard> {
|
||||
let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
|
||||
dirs::home_dir()
|
||||
.map(|mut h| {
|
||||
h.push(".gemini/mcp_memory");
|
||||
h.to_string_lossy().to_string()
|
||||
})
|
||||
.unwrap_or_else(|| ".gemini/mcp_memory".into())
|
||||
});
|
||||
let log_dir = std::path::PathBuf::from(base_dir).join("logs");
|
||||
std::fs::create_dir_all(&log_dir).unwrap_or_default();
|
||||
|
||||
let file_appender = tracing_appender::rolling::daily(log_dir, format!("{}.log", app_name));
|
||||
let (non_blocking, guard) = tracing_appender::non_blocking(file_appender);
|
||||
|
||||
let _ = tracing_subscriber::fmt()
|
||||
.with_writer(non_blocking)
|
||||
.with_ansi(false)
|
||||
.with_max_level(tracing::Level::INFO)
|
||||
.with_thread_ids(true)
|
||||
.with_thread_names(true)
|
||||
.try_init();
|
||||
|
||||
Some(guard)
|
||||
}
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let _guard = init_logging("mcp-memory-server");
|
||||
let cli = Cli::parse();
|
||||
|
||||
let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
|
||||
dirs::home_dir()
|
||||
.map(|mut h| {
|
||||
h.push(".gemini/mcp_memory");
|
||||
h.to_string_lossy().into_owned()
|
||||
})
|
||||
.unwrap_or_else(|| ".gemini/mcp_memory".into())
|
||||
});
|
||||
let base = PathBuf::from(base_dir);
|
||||
|
||||
if cli.exit {
|
||||
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||
let token = std::fs::read_to_string(base.join("admin.token")).unwrap_or_default();
|
||||
let mut cmd = std::process::Command::new("curl");
|
||||
cmd.arg("-k").arg("-X").arg("POST");
|
||||
if !token.is_empty() {
|
||||
cmd.arg("-H")
|
||||
.arg(format!("Authorization: Bearer {}", token.trim()));
|
||||
}
|
||||
let _ = cmd
|
||||
.arg(format!("http://127.0.0.1:{}/shutdown", port))
|
||||
.output();
|
||||
|
||||
if cli.restart {
|
||||
std::thread::sleep(Duration::from_secs(2));
|
||||
} else {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(Commands::Gate { subcmd }) = cli.command {
|
||||
let port = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string());
|
||||
let rt = tokio::runtime::Runtime::new()?;
|
||||
match subcmd {
|
||||
GateCommands::Set {
|
||||
action,
|
||||
target,
|
||||
namespace,
|
||||
params,
|
||||
authorize,
|
||||
block,
|
||||
reason,
|
||||
} => {
|
||||
let mut pmap = HashMap::new();
|
||||
for p in params {
|
||||
if let Some((k, v)) = p.split_once('=') {
|
||||
pmap.insert(k.to_string(), v.to_string());
|
||||
}
|
||||
}
|
||||
let req = GateSetReq {
|
||||
action,
|
||||
target,
|
||||
namespace,
|
||||
params: pmap,
|
||||
authorize: if authorize { Some(true) } else { None },
|
||||
block: if block { Some(true) } else { None },
|
||||
reason,
|
||||
};
|
||||
rt.block_on(async {
|
||||
let client = reqwest::Client::new();
|
||||
let res = client
|
||||
.post(format!("http://127.0.0.1:{}/gate/set", port))
|
||||
.json(&req)
|
||||
.send()
|
||||
.await;
|
||||
match res {
|
||||
Ok(r) if r.status().is_success() => println!("Gate updated successfully"),
|
||||
Ok(r) => println!("Failed to update gate: {}", r.status()),
|
||||
Err(e) => println!("Error connecting to server: {}", e),
|
||||
}
|
||||
});
|
||||
}
|
||||
GateCommands::Verify {
|
||||
action,
|
||||
target,
|
||||
namespace,
|
||||
params: _,
|
||||
consume,
|
||||
} => {
|
||||
let mut url = format!(
|
||||
"http://127.0.0.1:{}/gate/verify?action={}&target={}&consume={}",
|
||||
port, action, target, consume
|
||||
);
|
||||
if let Some(ns) = namespace {
|
||||
url.push_str(&format!("&namespace={}", ns));
|
||||
}
|
||||
rt.block_on(async {
|
||||
let res = reqwest::get(&url).await;
|
||||
match res {
|
||||
Ok(r) if r.status().is_success() => std::process::exit(0),
|
||||
Ok(r) if r.status() == reqwest::StatusCode::FORBIDDEN => {
|
||||
let text = r.text().await.unwrap_or_default();
|
||||
eprintln!("{}", text);
|
||||
std::process::exit(1);
|
||||
}
|
||||
Ok(r) if r.status() == reqwest::StatusCode::NOT_FOUND => {
|
||||
eprintln!("Action not yet authorized.");
|
||||
std::process::exit(2);
|
||||
}
|
||||
Ok(r) => {
|
||||
eprintln!("Unexpected status: {}", r.status());
|
||||
std::process::exit(3);
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("Error connecting to server: {}", e);
|
||||
std::process::exit(4);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let token = uuid::Uuid::new_v4().to_string();
|
||||
std::fs::write(base.join("admin.token"), &token).unwrap_or_default();
|
||||
|
||||
let rt = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
rt.block_on(async {
|
||||
let state = Arc::new(MemoryState::new(&base.to_string_lossy()));
|
||||
|
||||
// Initialize Qdrant VectorDB (default local URL)
|
||||
match crate::vector_db::VectorDB::new("http://localhost:6334", "mcp_memory").await {
|
||||
Ok(vdb) => {
|
||||
tracing::info!("Successfully connected to Qdrant vector database");
|
||||
*state.vector_db.write().await = Some(vdb);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to initialize Qdrant vector database: {}. Vector search will fallback to manual embedding loop. (Is Qdrant running on localhost:6334?)", e);
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = run_server(state).await {
|
||||
tracing::error!("Server error: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
Ok(())
|
||||
mcp_memory_server::run_cli()
|
||||
}
|
||||
@@ -20,6 +20,58 @@ pub struct StickyNote {
|
||||
pub fn default_namespace() -> String {
|
||||
"global".to_string()
|
||||
}
|
||||
pub fn normalize_entity_type(input: &str) -> String {
|
||||
let trimmed = input.trim();
|
||||
if trimmed.is_empty() {
|
||||
return "Concept".to_string();
|
||||
}
|
||||
let clean = trimmed.replace(['_', '-'], " ");
|
||||
let words: Vec<&str> = clean.split_whitespace().collect();
|
||||
if words.is_empty() {
|
||||
return "Concept".to_string();
|
||||
}
|
||||
|
||||
let pascal: String = words
|
||||
.into_iter()
|
||||
.map(|w| {
|
||||
let mut chars = w.chars();
|
||||
match chars.next() {
|
||||
None => String::new(),
|
||||
Some(f) => f.to_uppercase().collect::<String>() + &chars.as_str().to_lowercase(),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
match pascal.as_str() {
|
||||
"File" => "File".to_string(),
|
||||
"McpTool" | "Mcptool" | "Mcp" => "McpTool".to_string(),
|
||||
"Databasetable" | "DatabaseTable" => "DatabaseTable".to_string(),
|
||||
"Datastructure" | "DataStore" | "Datastore" => "DataStructure".to_string(),
|
||||
"Architecturecomponent" | "ArchitectureComponent" | "Architecture" => "ArchitectureComponent".to_string(),
|
||||
"Frontendcomponent" | "FrontendComponent" => "FrontendComponent".to_string(),
|
||||
"Globalrule" | "GlobalRule" => "GlobalRule".to_string(),
|
||||
"Securitypolicy" | "SecurityPolicy" => "SecurityPolicy".to_string(),
|
||||
"Serviceaccount" | "ServiceAccount" => "ServiceAccount".to_string(),
|
||||
"Testframework" | "TestFramework" => "TestFramework".to_string(),
|
||||
"Toolhandler" | "ToolHandler" => "ToolHandler".to_string(),
|
||||
"Historicalsummary" | "HistoricalSummary" => "HistoricalSummary".to_string(),
|
||||
"Externalapi" | "ExternalApi" => "ExternalAPI".to_string(),
|
||||
"Testsuite" | "TestSuite" => "TestSuite".to_string(),
|
||||
"Confluencepage" | "ConfluencePage" => "ConfluencePage".to_string(),
|
||||
"Containerimage" | "ContainerImage" => "ContainerImage".to_string(),
|
||||
"Buildsystem" | "BuildSystem" => "BuildSystem".to_string(),
|
||||
_ => pascal,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn normalize_relation_type(input: &str) -> String {
|
||||
let trimmed = input.trim();
|
||||
if trimmed.is_empty() {
|
||||
return "related_to".to_string();
|
||||
}
|
||||
trimmed.to_lowercase().replace([' ', '-'], "_")
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct Entity {
|
||||
pub name: String,
|
||||
@@ -79,6 +131,8 @@ pub struct Snippet {
|
||||
pub description: String,
|
||||
pub updated_at: u64,
|
||||
#[serde(default)]
|
||||
pub tags: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub embedding: Option<Vec<f32>>,
|
||||
}
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
@@ -234,3 +288,39 @@ pub struct StateSnapshot {
|
||||
pub description: String,
|
||||
pub namespace: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct Hypothesis {
|
||||
pub id: String,
|
||||
pub task_id: Option<String>,
|
||||
pub hypothesis: String,
|
||||
pub status: String,
|
||||
pub evidence: Option<String>,
|
||||
pub timestamp: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct AgentSignal {
|
||||
pub id: String,
|
||||
pub sender: String,
|
||||
pub signal_type: String,
|
||||
pub payload: String,
|
||||
pub timestamp: u64,
|
||||
pub ttl_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::IntoResponse;
|
||||
use crate::error::AppError;
|
||||
|
||||
#[test]
|
||||
fn test_app_error_bad_request_invalid_json() {
|
||||
let err = AppError::BadRequest("invalid json".into());
|
||||
let res = err.into_response();
|
||||
assert_eq!(res.status(), StatusCode::BAD_REQUEST);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -234,3 +234,11 @@ mod tests {
|
||||
assert!(!client.is_available().await);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_ollama_client_invalid_api_key() {
|
||||
let mut client = OllamaClient::new_from_env();
|
||||
client.base_url = "http://invalid-api-key:11434".to_string();
|
||||
assert!(!client.is_available().await);
|
||||
}
|
||||
|
||||
+326
-3
@@ -61,7 +61,16 @@ impl McpResource for GraphEntitiesResource {
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
|
||||
let graph = state_clone.graph.cache.read().unwrap();
|
||||
let data: Vec<_> = graph.entities.values().collect();
|
||||
Ok(serde_json::to_string_pretty(&data)?)
|
||||
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
|
||||
.unwrap()
|
||||
@@ -223,6 +232,76 @@ impl McpPrompt for SummarizeArchitecturePrompt {
|
||||
}
|
||||
}
|
||||
|
||||
struct SessionDeltaResource;
|
||||
#[async_trait]
|
||||
impl McpResource for SessionDeltaResource {
|
||||
fn uri(&self) -> &'static str {
|
||||
"memory://session/delta"
|
||||
}
|
||||
fn name(&self) -> &'static str {
|
||||
"Session Delta"
|
||||
}
|
||||
fn description(&self) -> Option<&'static str> {
|
||||
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 changes = state_clone.code.ledger.cache.read().unwrap();
|
||||
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 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 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
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ContextWarmupPrompt;
|
||||
#[async_trait]
|
||||
impl McpPrompt for ContextWarmupPrompt {
|
||||
fn name(&self) -> &'static str {
|
||||
"context_warmup"
|
||||
}
|
||||
fn description(&self) -> Option<&'static str> {
|
||||
Some("Consolidated context warmup prompt providing active tasks, pinned files, open tech debt, and sticky notes")
|
||||
}
|
||||
async fn get(
|
||||
&self,
|
||||
_args: Value,
|
||||
_state: Arc<MemoryState>,
|
||||
) -> crate::error::Result<serde_json::Value> {
|
||||
Ok(serde_json::json!({
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": {
|
||||
"type": "text",
|
||||
"text": "Perform a complete context warmup. Please read memory://session/delta, memory://tasks/active, memory://tech_debt/active, and memory://pinned_files to synthesize current project focus and active workspace constraints."
|
||||
}
|
||||
}
|
||||
]
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct MemoryHandler {
|
||||
pub state: Arc<MemoryState>,
|
||||
pub tools: std::collections::HashMap<String, Box<dyn McpTool>>,
|
||||
@@ -258,9 +337,11 @@ impl MemoryHandler {
|
||||
register_resource!(TasksActiveResource);
|
||||
register_resource!(DecisionsActiveResource);
|
||||
register_resource!(TechDebtActiveResource);
|
||||
register_resource!(SessionDeltaResource);
|
||||
|
||||
register_prompt!(AnalyzeTechDebtPrompt);
|
||||
register_prompt!(SummarizeArchitecturePrompt);
|
||||
register_prompt!(ContextWarmupPrompt);
|
||||
struct TerminalHistoryResource;
|
||||
#[async_trait]
|
||||
impl McpResource for TerminalHistoryResource {
|
||||
@@ -422,6 +503,7 @@ impl MemoryHandler {
|
||||
register!(tasks::UpdateMilestoneHandler);
|
||||
register!(tasks::ListMilestonesHandler);
|
||||
|
||||
register!(notes::ManageStickyNotesHandler);
|
||||
register!(notes::AddStickyNoteHandler);
|
||||
register!(notes::ReadStickyNotesHandler);
|
||||
register!(notes::DeleteStickyNoteHandler);
|
||||
@@ -431,6 +513,7 @@ impl MemoryHandler {
|
||||
register!(notes::ClearHandoffMemosHandler);
|
||||
register!(notes::AddSessionSummaryHandler);
|
||||
register!(notes::GenerateStandupReportHandler);
|
||||
register!(notes::PromoteToEntityHandler);
|
||||
|
||||
register!(meta::LogDecisionHandler);
|
||||
register!(meta::QueryDecisionsHandler);
|
||||
@@ -469,17 +552,42 @@ impl MemoryHandler {
|
||||
register!(vision::ReadClipboardHandler);
|
||||
register!(vision::WriteClipboardHandler);
|
||||
register!(git::GetActiveWorktreeContextHandler);
|
||||
register!(git::QueryGitDiffsHandler);
|
||||
register!(logs::WatchProcessLogsHandler);
|
||||
register!(logs::GetRecentLogsHandler);
|
||||
register!(ast::ReadFileSkeletonHandler);
|
||||
register!(vision::ToggleClipboardWatchModeHandler);
|
||||
register!(ast::ReplaceAstNodeHandler);
|
||||
register!(ast::FindSymbolReferencesHandler);
|
||||
register!(ast::GetCallersHandler);
|
||||
register!(ast::AnalyzeImpactHandler);
|
||||
register!(workspaces::ReadDirectoryArchitectureHandler);
|
||||
register!(workspaces::SemanticCodeSearchHandler);
|
||||
register!(workspaces::CreateSnapshotHandler);
|
||||
register!(workspaces::RestoreSnapshotHandler);
|
||||
register!(workspaces::CreateSubagentNamespaceHandler);
|
||||
register!(workspaces::ManageSubagentNamespaceHandler);
|
||||
register!(graph::GetSubgraphHandler);
|
||||
register!(meta::SuggestErrorFixHandler);
|
||||
register!(meta::CheckpointStateHandler);
|
||||
register!(meta::ManageCheckpointHandler);
|
||||
register!(meta::RestoreStateHandler);
|
||||
register!(workspaces::TagSnippetHandler);
|
||||
register!(workspaces::PurgeSubagentNamespaceHandler);
|
||||
register!(workspaces::CondenseSubagentNamespaceHandler);
|
||||
register!(graph::SweepGraphHealthHandler);
|
||||
register!(meta::QueryLineageHandler);
|
||||
register!(meta::GetNextActionableTasksHandler);
|
||||
register!(meta::LogHypothesisHandler);
|
||||
register!(meta::QueryHypothesesHandler);
|
||||
register!(workspaces::DiffContextWorkspacesHandler);
|
||||
register!(meta::GetPreflightContextHandler);
|
||||
register!(graph::ResolveStaleSymbolsHandler);
|
||||
register!(graph::SummarizeSubgraphHandler);
|
||||
register!(meta::BroadcastAgentSignalHandler);
|
||||
register!(meta::QueryAgentSignalsHandler);
|
||||
register!(meta::AutoSessionCheckpointHandler);
|
||||
register!(meta::SearchSnippetsHybridHandler);
|
||||
|
||||
Self {
|
||||
state,
|
||||
@@ -646,8 +754,25 @@ impl MemoryHandler {
|
||||
.cloned()
|
||||
.unwrap_or(serde_json::Value::Object(Default::default()));
|
||||
|
||||
self.state
|
||||
.broadcast_activity(&format!("Agent executed tool: {}", name));
|
||||
let category = match name {
|
||||
"read_clipboard" | "write_clipboard" | "toggle_clipboard_watch_mode" => "CLIPBOARD",
|
||||
"create_entities" | "create_relations" | "read_graph" | "get_subgraph" | "search_graph" | "get_schema" => "GRAPH",
|
||||
"log_decision" => "DECISION",
|
||||
"log_code_change" => "CODE",
|
||||
"log_error_fix" => "ERROR_FIX",
|
||||
"log_tech_debt" => "TECH_DEBT",
|
||||
"add_task" | "update_task_status" | "delete_task" | "add_milestone" => "TASK",
|
||||
"manage_sticky_notes" | "read_notes" => "STICKY_NOTE",
|
||||
"manage_checkpoint" => "CHECKPOINT",
|
||||
"manage_subagent_namespace" => "SUBAGENT",
|
||||
"search_snippets" => "SNIPPET",
|
||||
"search_web" => "WEB_SEARCH",
|
||||
"omni_search" => "OMNI_SEARCH",
|
||||
_ => "TOOL",
|
||||
};
|
||||
|
||||
let activity_msg = format_tool_activity_description(name, &args);
|
||||
self.state.broadcast_activity(category, &activity_msg);
|
||||
|
||||
let result: crate::error::Result<String> = if let Some(tool) = self.tools.get(name)
|
||||
{
|
||||
@@ -694,12 +819,174 @@ impl MemoryHandler {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) -> String {
|
||||
let detail = match name {
|
||||
"log_code_change" => {
|
||||
if let Some(file) = args.get("file").or_else(|| args.get("path")).and_then(|v| v.as_str()) {
|
||||
if let Some(summary) = args.get("summary").and_then(|v| v.as_str()) {
|
||||
format!("{}: {}", file, summary)
|
||||
} else {
|
||||
file.to_string()
|
||||
}
|
||||
} else {
|
||||
args.get("summary").and_then(|v| v.as_str()).unwrap_or("").to_string()
|
||||
}
|
||||
}
|
||||
"log_decision" => {
|
||||
args.get("title")
|
||||
.or_else(|| args.get("decision"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string()
|
||||
}
|
||||
"log_error_fix" => {
|
||||
args.get("error")
|
||||
.or_else(|| args.get("summary"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string()
|
||||
}
|
||||
"log_tech_debt" => {
|
||||
if let Some(summary) = args.get("summary").and_then(|v| v.as_str()) {
|
||||
summary.to_string()
|
||||
} else if let Some(file) = args.get("file").and_then(|v| v.as_str()) {
|
||||
file.to_string()
|
||||
} else {
|
||||
"".to_string()
|
||||
}
|
||||
}
|
||||
"create_entities" => {
|
||||
if let Some(entities) = args.get("entities").and_then(|v| v.as_array()) {
|
||||
let names: Vec<&str> = entities
|
||||
.iter()
|
||||
.filter_map(|e| e.get("name").and_then(|n| n.as_str()))
|
||||
.take(3)
|
||||
.collect();
|
||||
if !names.is_empty() {
|
||||
let mut s = names.join(", ");
|
||||
if entities.len() > 3 {
|
||||
s.push_str(&format!(" (+{} more)", entities.len() - 3));
|
||||
}
|
||||
s
|
||||
} else {
|
||||
format!("{} entity/entities", entities.len())
|
||||
}
|
||||
} else if let Some(name) = args.get("name").and_then(|v| v.as_str()) {
|
||||
name.to_string()
|
||||
} else {
|
||||
"".to_string()
|
||||
}
|
||||
}
|
||||
"create_relations" => {
|
||||
if let Some(relations) = args.get("relations").and_then(|v| v.as_array()) {
|
||||
if let Some(first) = relations.first() {
|
||||
let from = first.get("from").and_then(|v| v.as_str()).unwrap_or("?");
|
||||
let to = first.get("to").and_then(|v| v.as_str()).unwrap_or("?");
|
||||
let pred = first.get("predicate").and_then(|v| v.as_str()).unwrap_or("->");
|
||||
if relations.len() > 1 {
|
||||
format!("{} {} {} (+{} more)", from, pred, to, relations.len() - 1)
|
||||
} else {
|
||||
format!("{} {} {}", from, pred, to)
|
||||
}
|
||||
} else {
|
||||
format!("{} relation(s)", relations.len())
|
||||
}
|
||||
} else {
|
||||
"".to_string()
|
||||
}
|
||||
}
|
||||
"add_task" => {
|
||||
args.get("title")
|
||||
.or_else(|| args.get("name"))
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string()
|
||||
}
|
||||
"update_task_status" => {
|
||||
if let (Some(id), Some(status)) = (
|
||||
args.get("task_id").or_else(|| args.get("id")).and_then(|v| v.as_str()),
|
||||
args.get("status").and_then(|v| v.as_str()),
|
||||
) {
|
||||
format!("Task {} -> {}", id, status)
|
||||
} else {
|
||||
"".to_string()
|
||||
}
|
||||
}
|
||||
"omni_search" | "search_graph" | "search_snippets" | "search_web" => {
|
||||
args.get("query")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|q| format!("\"{}\"", q))
|
||||
.unwrap_or_default()
|
||||
}
|
||||
"manage_sticky_notes" => {
|
||||
let action = args.get("action").and_then(|v| v.as_str()).unwrap_or("");
|
||||
if let Some(content) = args.get("content").and_then(|v| v.as_str()) {
|
||||
let preview: String = content.chars().take(40).collect();
|
||||
if !action.is_empty() {
|
||||
format!("{} \"{}\"", action, preview)
|
||||
} else {
|
||||
format!("\"{}\"", preview)
|
||||
}
|
||||
} else {
|
||||
action.to_string()
|
||||
}
|
||||
}
|
||||
"write_clipboard" => {
|
||||
if let Some(text) = args.get("text").or_else(|| args.get("content")).and_then(|v| v.as_str()) {
|
||||
let preview: String = text.chars().take(40).collect();
|
||||
format!("\"{}\"", preview)
|
||||
} else {
|
||||
"".to_string()
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
if let Some(title) = args
|
||||
.get("title")
|
||||
.or_else(|| args.get("summary"))
|
||||
.or_else(|| args.get("name"))
|
||||
.or_else(|| args.get("query"))
|
||||
.or_else(|| args.get("file"))
|
||||
.or_else(|| args.get("path"))
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
title.to_string()
|
||||
} else {
|
||||
"".to_string()
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let mut truncated = detail.trim().to_string();
|
||||
if truncated.len() > 80 {
|
||||
truncated.truncate(77);
|
||||
truncated.push_str("...");
|
||||
}
|
||||
|
||||
if !truncated.is_empty() {
|
||||
format!("Agent executed tool: {} ({})", name, truncated)
|
||||
} else {
|
||||
format!("Agent executed tool: {}", name)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
use tempfile::tempdir;
|
||||
|
||||
#[test]
|
||||
fn test_format_tool_activity_description() {
|
||||
let msg1 = format_tool_activity_description("log_code_change", &json!({ "file": "src/lib.rs", "summary": "Refactored library" }));
|
||||
assert_eq!(msg1, "Agent executed tool: log_code_change (src/lib.rs: Refactored library)");
|
||||
|
||||
let msg2 = format_tool_activity_description("create_entities", &json!({ "entities": [{ "name": "EntityA" }, { "name": "EntityB" }] }));
|
||||
assert_eq!(msg2, "Agent executed tool: create_entities (EntityA, EntityB)");
|
||||
|
||||
let msg3 = format_tool_activity_description("read_graph", &json!({}));
|
||||
assert_eq!(msg3, "Agent executed tool: read_graph");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_memory_handler_tools_registration() {
|
||||
let dir = tempdir().unwrap();
|
||||
@@ -861,4 +1148,40 @@ mod tests {
|
||||
assert!(res_unknown.get("error").is_some());
|
||||
assert_eq!(res_unknown["error"]["code"], -32601);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_handle_request() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
||||
let handler = MemoryHandler::new(state);
|
||||
|
||||
let req1 = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 5,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": "create_entities",
|
||||
"arguments": {
|
||||
"entities": [{ "name": "EntityA" }, { "name": "EntityB" }]
|
||||
}
|
||||
}
|
||||
});
|
||||
let res1 = handler.handle_request(req1).await.unwrap();
|
||||
assert_eq!(res1["jsonrpc"], "2.0");
|
||||
assert_eq!(res1["id"], 5);
|
||||
|
||||
let req2 = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 6,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": "read_graph",
|
||||
"arguments": {}
|
||||
}
|
||||
});
|
||||
let res2 = handler.handle_request(req2).await.unwrap();
|
||||
assert_eq!(res2["jsonrpc"], "2.0");
|
||||
assert_eq!(res2["id"], 6);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -123,7 +123,12 @@ impl MemoryIndex {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
if needs_commit.swap(false, std::sync::atomic::Ordering::SeqCst) {
|
||||
let mut writer = writer.lock().unwrap_or_else(|e| e.into_inner());
|
||||
writer.commit()?;
|
||||
if let Err(first_err) = writer.commit() {
|
||||
std::thread::sleep(std::time::Duration::from_millis(100));
|
||||
writer.commit().map_err(|second_err| {
|
||||
tantivy::TantivyError::SystemError(format!("First commit error: {}; Second commit error: {}", first_err, second_err))
|
||||
})?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
@@ -135,6 +140,7 @@ impl MemoryIndex {
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
pub fn search(
|
||||
&self,
|
||||
query: &str,
|
||||
@@ -345,6 +351,7 @@ mod tests {
|
||||
language: "rust".to_string(),
|
||||
description: "A test snippet".to_string(),
|
||||
updated_at: 0,
|
||||
tags: vec![],
|
||||
embedding: None,
|
||||
};
|
||||
let _ = index.index_snippet(&snippet).await.unwrap();
|
||||
|
||||
+31
-5
@@ -29,6 +29,7 @@ pub struct CodeStores {
|
||||
pub error_fixes: Store<Vec<ErrorFix>>,
|
||||
pub tech_debts: Store<Vec<TechDebt>>,
|
||||
pub sticky: Store<Vec<StickyNote>>,
|
||||
pub hypotheses: Store<Vec<Hypothesis>>,
|
||||
}
|
||||
|
||||
pub struct EnvironmentStores {
|
||||
@@ -44,6 +45,7 @@ pub struct TelemetryStores {
|
||||
pub handoff_memos: Store<Vec<HandoffMemo>>,
|
||||
pub recent_activities: Store<std::collections::VecDeque<serde_json::Value>>,
|
||||
pub terminal_history: Store<std::collections::VecDeque<TerminalHistory>>,
|
||||
pub agent_signals: Store<Vec<AgentSignal>>,
|
||||
}
|
||||
|
||||
pub struct MemoryState {
|
||||
@@ -70,7 +72,7 @@ impl MemoryState {
|
||||
|
||||
let db = crate::db::init_redb(&base);
|
||||
|
||||
Self {
|
||||
let state = Self {
|
||||
ollama: Arc::new(crate::ollama::OllamaClient::new_from_env()),
|
||||
clipboard_watch_mode: tokio::sync::RwLock::new(false),
|
||||
graph: Store::new("knowledge_graph_master", db.clone()),
|
||||
@@ -103,6 +105,7 @@ impl MemoryState {
|
||||
error_fixes: Store::new("error_fixes", db.clone()),
|
||||
tech_debts: Store::new("tech_debts", db.clone()),
|
||||
sticky: Store::new("sticky_notes", db.clone()),
|
||||
hypotheses: Store::new("hypotheses", db.clone()),
|
||||
},
|
||||
env: EnvironmentStores {
|
||||
env_fingerprints: Store::new("env_fingerprints", db.clone()),
|
||||
@@ -116,14 +119,27 @@ impl MemoryState {
|
||||
handoff_memos: Store::new("handoff_memos", db.clone()),
|
||||
recent_activities: Store::new("recent_activities", db.clone()),
|
||||
terminal_history: Store::new("terminal_history", db.clone()),
|
||||
agent_signals: Store::new("agent_signals", db.clone()),
|
||||
},
|
||||
|
||||
activity_tx: tokio::sync::broadcast::channel(100).0,
|
||||
event_bus_tx: tokio::sync::broadcast::channel(1000).0,
|
||||
}
|
||||
};
|
||||
|
||||
// Normalize pre-existing graph entity and relation types
|
||||
state.graph.modify(|g| {
|
||||
for entity in g.entities.values_mut() {
|
||||
entity.entity_type = crate::models::normalize_entity_type(&entity.entity_type);
|
||||
}
|
||||
for relation in g.relations.iter_mut() {
|
||||
relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type);
|
||||
}
|
||||
});
|
||||
|
||||
state
|
||||
}
|
||||
|
||||
pub fn broadcast_activity(&self, message: &str) {
|
||||
pub fn broadcast_activity(&self, category: &str, message: &str) {
|
||||
let time = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
@@ -131,13 +147,14 @@ impl MemoryState {
|
||||
|
||||
let item = serde_json::json!({
|
||||
"time": time,
|
||||
"category": category,
|
||||
"message": message
|
||||
});
|
||||
|
||||
self.telemetry.recent_activities.modify(|activities| {
|
||||
activities.push_back(item.clone());
|
||||
activities.push_front(item.clone());
|
||||
if activities.len() > 100 {
|
||||
activities.pop_front();
|
||||
activities.pop_back();
|
||||
}
|
||||
});
|
||||
|
||||
@@ -240,6 +257,15 @@ impl MemoryState {
|
||||
|
||||
let _ = self.activity_tx.send(record.to_string());
|
||||
}
|
||||
|
||||
pub fn record_terminal_history(&self, payload: TerminalHistory) {
|
||||
self.telemetry.terminal_history.modify(|history| {
|
||||
history.push_front(payload);
|
||||
if history.len() > 100 {
|
||||
history.pop_back();
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -1,20 +0,0 @@
|
||||
use arboard::{Clipboard, ImageData};
|
||||
|
||||
use std::borrow::Cow;
|
||||
|
||||
fn main() {
|
||||
let img = image::open("C:/Users/reazul.ashraf/workspace/rust/mcp-memory/test.jpg")
|
||||
.unwrap_or_else(|_| image::DynamicImage::new_rgb8(10, 10));
|
||||
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 mut clipboard = Clipboard::new().unwrap();
|
||||
clipboard.set_image(img_data).unwrap();
|
||||
println!("Successfully wrote image");
|
||||
}
|
||||
+368
-10
@@ -18,14 +18,14 @@ pub struct WriteClipboardTool {
|
||||
pub image_path: Option<String>,
|
||||
}
|
||||
|
||||
/// Create new entities in the knowledge graph.
|
||||
/// Create new entities in the knowledge graph. Note: entity_type MUST ALWAYS be PascalCase (e.g. DatabaseTable, McpTool, File).
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct CreateEntitiesTool {
|
||||
/// Array of entities to create.
|
||||
pub entities: Vec<crate::models::Entity>,
|
||||
}
|
||||
|
||||
/// Create new relations between entities in the knowledge graph.
|
||||
/// Create new relations between entities in the knowledge graph. Note: relation_type MUST ALWAYS be snake_case (e.g. depends_on, calls, implements).
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct CreateRelationsTool {
|
||||
/// Array of relations to create.
|
||||
@@ -77,6 +77,8 @@ pub struct DeleteRelationsTool {
|
||||
pub struct ReadGraphTool {
|
||||
/// Optional namespace to restrict the read to.
|
||||
pub namespace: Option<String>,
|
||||
/// Optional token budget cap. Truncates graph outputs to fit within the specified token budget.
|
||||
pub max_tokens: Option<usize>,
|
||||
}
|
||||
|
||||
/// Search for entities in the knowledge graph by name or type.
|
||||
@@ -110,6 +112,10 @@ pub struct LogCodeChangeTool {
|
||||
pub git_commit: Option<String>,
|
||||
/// The associated git branch, if any.
|
||||
pub git_branch: Option<String>,
|
||||
/// Optional symbol references (e.g. ['MemoryIndex', 'switchTab']) modified in this change.
|
||||
pub symbol_references: Option<Vec<String>>,
|
||||
/// Optional line range (e.g. 'L123-L145') modified in the target file.
|
||||
pub line_range: Option<String>,
|
||||
}
|
||||
|
||||
/// Query recently logged code changes.
|
||||
@@ -125,11 +131,39 @@ pub struct VisualizeGraphTool {
|
||||
pub namespace: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum StickyNoteAction {
|
||||
Add,
|
||||
Read,
|
||||
Delete,
|
||||
Clear,
|
||||
}
|
||||
|
||||
/// Manage ephemeral sticky notes with TTL (add, read, delete, clear).
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct ManageStickyNotesTool {
|
||||
/// Action to perform: 'add', 'read', 'delete', or 'clear'.
|
||||
pub action: StickyNoteAction,
|
||||
/// Content of the sticky note (required for 'add').
|
||||
pub content: Option<String>,
|
||||
/// 1-indexed position of the sticky note to delete (required for 'delete').
|
||||
pub index: Option<usize>,
|
||||
/// Optional time-to-live in seconds (for 'add').
|
||||
pub ttl_seconds: Option<u64>,
|
||||
/// Optional flag to auto-expire the note at the end of the current session (for 'add').
|
||||
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.
|
||||
@@ -193,6 +227,10 @@ pub struct UpdateTaskStatusTool {
|
||||
pub struct ListActiveTasksTool {
|
||||
/// Optional git branch to filter tasks by.
|
||||
pub git_branch: Option<String>,
|
||||
/// Optional summary level: 'compact' (IDs and titles), 'detailed', or 'full'. Defaults to 'detailed'.
|
||||
pub summary_level: Option<String>,
|
||||
/// Optional token budget cap. Truncates results to stay within max_tokens.
|
||||
pub max_tokens: Option<usize>,
|
||||
}
|
||||
|
||||
/// Store a reusable code snippet.
|
||||
@@ -206,16 +244,39 @@ pub struct StoreSnippetTool {
|
||||
pub code: String,
|
||||
/// A description of what the snippet does.
|
||||
pub description: String,
|
||||
/// Optional classification tags (e.g. ['rust', 'async', 'mcp']).
|
||||
pub tags: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
/// Search through stored code snippets.
|
||||
/// Add classification tags to an existing code snippet.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct TagSnippetTool {
|
||||
/// The name of the snippet to tag.
|
||||
pub name: String,
|
||||
/// Array of tags (e.g. ['rust', 'async', 'mcp']).
|
||||
pub tags: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SnippetSearchMode {
|
||||
Hybrid,
|
||||
Keyword,
|
||||
Semantic,
|
||||
}
|
||||
|
||||
/// Search stored code snippets using Hybrid BM25 + Vector ranking or keyword search.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct SearchSnippetsTool {
|
||||
/// The search query.
|
||||
pub query: String,
|
||||
/// Limit the number of results to avoid context limit overflow. Defaults to 5.
|
||||
/// Search mode: 'hybrid' (BM25 + Semantic Tags), 'keyword', or 'semantic'. Defaults to 'hybrid'.
|
||||
pub mode: Option<SnippetSearchMode>,
|
||||
/// Optional classification tags to filter by.
|
||||
pub tags: Option<Vec<String>>,
|
||||
/// Limit the number of results to avoid context limit overflow. Defaults to 10.
|
||||
pub limit: Option<usize>,
|
||||
/// Include the full code content of the matched items. If false, returns only IDs and descriptions (recommended for LLMs to prevent context bloat). Defaults to false.
|
||||
/// Include the full code content of the matched items. If false, returns only IDs and descriptions. Defaults to false.
|
||||
pub include_body: Option<bool>,
|
||||
}
|
||||
|
||||
@@ -226,6 +287,77 @@ pub struct DeleteSnippetTool {
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SubgraphFormat {
|
||||
Json,
|
||||
MarkdownTree,
|
||||
}
|
||||
|
||||
/// Extract a multi-hop neighborhood subgraph around a specific root entity node as JSON or Markdown topology tree.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct GetSubgraphTool {
|
||||
/// The root entity name to start the subgraph search from.
|
||||
pub root_entity: Option<String>,
|
||||
/// Legacy alias for root_entity.
|
||||
pub root_node: Option<String>,
|
||||
/// Maximum search depth (hops). Defaults to 2.
|
||||
pub depth: Option<u32>,
|
||||
/// Output format: 'json' (raw entities and relations) or 'markdown_tree' (compact topology tree). Defaults to 'json'.
|
||||
pub format: Option<SubgraphFormat>,
|
||||
/// Optional namespace filter.
|
||||
pub namespace: Option<String>,
|
||||
/// Target maximum token length for generated summary when format is 'markdown_tree'. Defaults to 1000.
|
||||
pub max_tokens: Option<usize>,
|
||||
}
|
||||
|
||||
/// Search historical error fixes using keyword search or semantic vector similarity and return matching solutions.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct SuggestErrorFixTool {
|
||||
/// The exact stack trace or error signature.
|
||||
pub stack_trace: String,
|
||||
/// Maximum number of suggested fixes to return. Defaults to 3.
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum CheckpointAction {
|
||||
Create,
|
||||
Restore,
|
||||
List,
|
||||
Delete,
|
||||
}
|
||||
|
||||
/// Save, restore, list, or delete point-in-time memory state snapshot checkpoints.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct ManageCheckpointTool {
|
||||
/// Action to perform: 'create', 'restore', 'list', or 'delete'.
|
||||
pub action: CheckpointAction,
|
||||
/// Unique name or ID for the snapshot checkpoint.
|
||||
pub name_or_id: Option<String>,
|
||||
/// Optional description of why this checkpoint was created.
|
||||
pub description: Option<String>,
|
||||
/// Optional namespace filter.
|
||||
pub namespace: Option<String>,
|
||||
}
|
||||
|
||||
/// Save a named memory snapshot checkpoint before risky operations or subagent invocations.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct CheckpointStateTool {
|
||||
/// The unique name for this state snapshot checkpoint.
|
||||
pub name: String,
|
||||
/// Optional description of why this checkpoint was created.
|
||||
pub description: Option<String>,
|
||||
}
|
||||
|
||||
/// Restore a previously saved memory snapshot checkpoint.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct RestoreStateTool {
|
||||
/// The name of the checkpoint to restore.
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
/// Log an architectural decision record (ADR).
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct LogDecisionTool {
|
||||
@@ -296,16 +428,22 @@ pub struct LogErrorFixTool {
|
||||
pub git_commit: Option<String>,
|
||||
/// The associated git branch, if any.
|
||||
pub git_branch: Option<String>,
|
||||
/// Optional symbol references (e.g. ['MemoryIndex', 'switchTab']) associated with the error.
|
||||
pub symbol_references: Option<Vec<String>>,
|
||||
/// Optional line range (e.g. 'L123-L145') associated with the error.
|
||||
pub line_range: Option<String>,
|
||||
}
|
||||
|
||||
/// Search through previously logged error fixes.
|
||||
/// Search historical error fixes using keyword query or stack trace vector similarity.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct SearchErrorFixesTool {
|
||||
/// The search query.
|
||||
pub query: String,
|
||||
/// The search query string.
|
||||
pub query: Option<String>,
|
||||
/// Exact stack trace or error signature for vector cosine matching.
|
||||
pub stack_trace: Option<String>,
|
||||
/// Limit the number of results to avoid context limit overflow. Defaults to 5.
|
||||
pub limit: Option<usize>,
|
||||
/// Include the full solution details. If false, returns only error signatures (recommended for LLMs to prevent context bloat). Defaults to false.
|
||||
/// Include the full solution details. If false, returns only error signatures. Defaults to false.
|
||||
pub include_body: Option<bool>,
|
||||
}
|
||||
|
||||
@@ -514,6 +652,10 @@ pub struct LogTechDebtTool {
|
||||
pub git_commit: Option<String>,
|
||||
/// The associated git branch, if any.
|
||||
pub git_branch: Option<String>,
|
||||
/// Optional symbol references (e.g. ['MemoryIndex', 'switchTab']) associated with the tech debt.
|
||||
pub symbol_references: Option<Vec<String>>,
|
||||
/// Optional line range (e.g. 'L123-L145') associated with the tech debt.
|
||||
pub line_range: Option<String>,
|
||||
}
|
||||
|
||||
/// Resolve a technical debt record.
|
||||
@@ -531,6 +673,10 @@ pub struct ListTechDebtTool {
|
||||
pub namespace: String,
|
||||
/// Whether to include resolved technical debt in the results.
|
||||
pub include_resolved: bool,
|
||||
/// Optional summary level: 'compact' (IDs and descriptions), 'detailed', or 'full'. Defaults to 'detailed'.
|
||||
pub summary_level: Option<String>,
|
||||
/// Optional token budget cap. Truncates results to stay within max_tokens.
|
||||
pub max_tokens: Option<usize>,
|
||||
}
|
||||
|
||||
/// Save the current context workspace.
|
||||
@@ -574,7 +720,7 @@ pub struct DeleteContextWorkspaceTool {
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
/// Search across all memory stores (Graph, Tasks, Snippets, ADRs, etc.).
|
||||
/// Search across all memory stores (Graph, Tasks, Snippets, ADRs, Error Fixes, Tech Debt) using Hybrid Reciprocal Rank Fusion (BM25 + Dense Vectors).
|
||||
/// Highly recommended for discovery. Supports Lucene query syntax.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct OmniSearchTool {
|
||||
@@ -586,6 +732,8 @@ pub struct OmniSearchTool {
|
||||
pub limit: Option<usize>,
|
||||
/// Include the full body/content of the matched items. If false, returns only IDs and titles (recommended for LLMs to prevent context bloat). Defaults to false.
|
||||
pub include_body: Option<bool>,
|
||||
/// Optional token budget cap. Dynamically caps and truncates search results to stay within max_tokens.
|
||||
pub max_tokens: Option<usize>,
|
||||
}
|
||||
|
||||
/// Get a health digest of the project.
|
||||
@@ -620,6 +768,62 @@ pub struct VerifyAcceptanceCriteriaTool {
|
||||
pub proof: String,
|
||||
}
|
||||
|
||||
/// Audit the knowledge graph to detect orphaned entities, compute name similarity for potential duplicate merges, and optionally auto-prune orphans.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct SweepGraphHealthTool {
|
||||
/// Optional flag to automatically prune orphaned nodes with 0 relations. Defaults to false.
|
||||
pub auto_prune_orphans: Option<bool>,
|
||||
/// Minimum string similarity threshold (0.0 to 1.0) to report duplicate entity pairs. Defaults to 0.8.
|
||||
pub similarity_threshold: Option<f32>,
|
||||
}
|
||||
|
||||
/// Trace the causal provenance and historical lineage linking a task, ADR, git commit, code change, or error fix.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct QueryLineageTool {
|
||||
/// The task ID, file path, symbol name, or git commit to query lineage for.
|
||||
pub query: String,
|
||||
}
|
||||
|
||||
/// Topological task resolver that returns unblocked, ready-to-run active tasks.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct GetNextActionableTasksTool {
|
||||
/// Optional git branch filter.
|
||||
pub git_branch: Option<String>,
|
||||
/// Limit the number of actionable tasks returned. Defaults to 5.
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// Log a structured diagnostic hypothesis, tested evidence, and status during problem solving.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct LogHypothesisTool {
|
||||
/// Optional task ID associated with this hypothesis.
|
||||
pub task_id: Option<String>,
|
||||
/// The diagnostic hypothesis or potential root cause.
|
||||
pub hypothesis: String,
|
||||
/// Status: 'unverified', 'verified', or 'rejected'. Defaults to 'unverified'.
|
||||
pub status: Option<String>,
|
||||
/// Evidence or test results supporting or disproving the hypothesis.
|
||||
pub evidence: Option<String>,
|
||||
}
|
||||
|
||||
/// Query previously logged diagnostic hypotheses and test results.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct QueryHypothesesTool {
|
||||
/// Optional task ID to filter hypotheses for.
|
||||
pub task_id: Option<String>,
|
||||
/// Optional search query text.
|
||||
pub query: Option<String>,
|
||||
}
|
||||
|
||||
/// Compare two saved context workspaces and return a structured diff of pinned files, active tasks, and metadata.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct DiffContextWorkspacesTool {
|
||||
/// The name of the first context workspace.
|
||||
pub workspace_a: String,
|
||||
/// The name of the second context workspace.
|
||||
pub workspace_b: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -714,6 +918,25 @@ pub struct RestoreSnapshotTool {
|
||||
pub snapshot_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SubagentNamespaceAction {
|
||||
Create,
|
||||
Condense,
|
||||
Purge,
|
||||
}
|
||||
|
||||
/// Manage isolated memory namespaces for subagent sessions (create, condense/promote, or purge).
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct ManageSubagentNamespaceTool {
|
||||
/// Action to perform: 'create', 'condense', or 'purge'.
|
||||
pub action: SubagentNamespaceAction,
|
||||
/// Subagent namespace ID.
|
||||
pub subagent_id: String,
|
||||
/// For 'condense': whether to auto-purge the subagent namespace after promotion. Defaults to true.
|
||||
pub purge_after_promotion: Option<bool>,
|
||||
}
|
||||
|
||||
/// Create an isolated memory namespace for a subagent session.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct CreateSubagentNamespaceTool {
|
||||
@@ -725,3 +948,138 @@ pub struct CreateSubagentNamespaceTool {
|
||||
pub struct PurgeSubagentNamespaceTool {
|
||||
pub subagent_id: String,
|
||||
}
|
||||
|
||||
/// Find all source locations and AST chunks where a specific symbol (function, struct, method, type) is referenced or called.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct FindSymbolReferencesTool {
|
||||
/// The symbol name (e.g. 'MemoryState', 'handle_search', 'AppError').
|
||||
pub symbol: String,
|
||||
/// Optional maximum number of reference locations to return. Defaults to 10.
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// Find all caller functions or methods that invoke a specified target function or method name.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct GetCallersTool {
|
||||
/// The target function or method name (e.g. 'generate_embeddings_async', 'keyword_search').
|
||||
pub function_name: String,
|
||||
/// Optional maximum number of callers to return. Defaults to 10.
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// Query git commit history, diffs, and change ledger entries using keyword or semantic search.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct QueryGitDiffsTool {
|
||||
/// Search query (e.g. 'deploy-server', 'shutdown', 'debounce', 'rename trick').
|
||||
pub query: String,
|
||||
/// Optional maximum number of commit diff entries to return. Defaults to 5.
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// Promote a transient sticky note or active task observation into a permanent Knowledge Graph entity.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct PromoteToEntityTool {
|
||||
/// Note text or content to promote.
|
||||
pub content: String,
|
||||
/// Target entity name to create or merge into.
|
||||
pub entity_name: String,
|
||||
/// Entity type (e.g. 'Component', 'Decision', 'BugFix', 'Architecture').
|
||||
pub entity_type: String,
|
||||
/// Optional namespace. Defaults to 'default'.
|
||||
#[serde(default = "crate::models::default_namespace")]
|
||||
pub namespace: String,
|
||||
}
|
||||
|
||||
/// Analyze the potential downstream breaking impact of modifying a function, struct, or file across AST callers and Knowledge Graph relations.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct AnalyzeImpactTool {
|
||||
/// Target function, struct, or symbol name (e.g. 'MemoryState', 'execute', 'AppError').
|
||||
pub target_symbol: String,
|
||||
/// Optional file path.
|
||||
pub file_path: Option<String>,
|
||||
}
|
||||
|
||||
/// Auto-condense a subagent's memory namespace and promote its high-value entities/relations to the global Knowledge Graph.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct CondenseSubagentNamespaceTool {
|
||||
/// The subagent namespace ID to consolidate.
|
||||
pub subagent_id: String,
|
||||
/// Whether to auto-purge the subagent namespace after promotion. Defaults to true.
|
||||
pub purge_after_promotion: Option<bool>,
|
||||
}
|
||||
|
||||
/// Get a consolidated 1-page executive summary of current active branch, active task, pinned files, open tech debt, and diagnostic hypotheses in 1 turn.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct GetPreflightContextTool {
|
||||
/// Optional namespace. Defaults to 'default'.
|
||||
#[serde(default = "crate::models::default_namespace")]
|
||||
pub namespace: String,
|
||||
/// Optional git branch to filter context for.
|
||||
pub git_branch: Option<String>,
|
||||
}
|
||||
|
||||
/// Inspect Knowledge Graph entities, observations, and tech debt symbol/line references against current files on disk and AST, flagging and healing stale or broken pointers.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct ResolveStaleSymbolsTool {
|
||||
/// Optional entity or file path to check. If omitted, checks all graph entities and tech debts.
|
||||
pub target: Option<String>,
|
||||
/// Whether to automatically update or remove broken references. Defaults to true.
|
||||
pub auto_heal: Option<bool>,
|
||||
}
|
||||
|
||||
/// Generate a compact, LLM-optimized Markdown topology or summary of a graph component capped within a strict token budget.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct SummarizeSubgraphTool {
|
||||
/// The root entity name to center the summary on (e.g. 'MemoryState', 'ServerRouter').
|
||||
pub root_entity: String,
|
||||
/// Search depth from root entity. Defaults to 2.
|
||||
pub depth: Option<usize>,
|
||||
/// Target maximum token length for the generated summary. Defaults to 1000.
|
||||
pub max_tokens: Option<usize>,
|
||||
}
|
||||
|
||||
/// Broadcast a real-time signal or status transition to other subagents on the inter-agent signal bus.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct BroadcastAgentSignalTool {
|
||||
/// Sender agent ID or role (e.g. 'PrePushAuditor', 'MemoryLibrarian').
|
||||
pub sender: String,
|
||||
/// Signal type or event category (e.g. 'AUDIT_PASSED', 'REPRODUCER_READY', 'TESTS_FAILED').
|
||||
pub signal_type: String,
|
||||
/// JSON or text payload containing event details or artifact URIs.
|
||||
pub payload: String,
|
||||
/// Optional Time-To-Live in seconds for the signal. Defaults to 3600 (1 hour).
|
||||
pub ttl_seconds: Option<u64>,
|
||||
}
|
||||
|
||||
/// Query active inter-agent signals published by subagents on the inter-agent signal bus.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct QueryAgentSignalsTool {
|
||||
/// Optional sender filter (e.g. 'PrePushAuditor').
|
||||
pub sender: Option<String>,
|
||||
/// Optional signal_type filter (e.g. 'AUDIT_PASSED').
|
||||
pub signal_type: Option<String>,
|
||||
/// Optional limit on returned signals. Defaults to 20.
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
/// Trigger an automated context checkpoint, summarizing active tasks, hypotheses, recent commits, and open tech debt into a permanent HandoffMemo.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct AutoSessionCheckpointTool {
|
||||
/// Author or agent ID creating the checkpoint. Defaults to 'AutoCheckpoint'.
|
||||
pub author: Option<String>,
|
||||
/// Optional namespace. Defaults to 'default'.
|
||||
#[serde(default = "crate::models::default_namespace")]
|
||||
pub namespace: String,
|
||||
}
|
||||
|
||||
/// Hybrid BM25 & semantic tag search across stored code snippets, returning ranked code patterns with relevance scores.
|
||||
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
|
||||
pub struct SearchSnippetsHybridTool {
|
||||
/// Query string (e.g. 'rate limiter', 'websocket client', 'shutdown poll').
|
||||
pub query: String,
|
||||
/// Optional tags to filter by.
|
||||
pub tags: Option<Vec<String>>,
|
||||
/// Maximum number of snippets to return. Defaults to 10.
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
@@ -182,4 +182,16 @@ mod tests {
|
||||
assert_eq!(uuid1, uuid2);
|
||||
assert!(Uuid::parse_str(&uuid1).is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_vector_db_index_document_invalid_id() {
|
||||
if let Ok(db) = VectorDB::new("http://localhost:6333", "test_collection").await {
|
||||
let invalid_id = "not-a-uuid";
|
||||
let doc_type = "test";
|
||||
let text = "example text";
|
||||
let vector = vec![1.0; 384];
|
||||
let _ = db.index_document(invalid_id, doc_type, text, vector).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -92,4 +92,11 @@ mod tests {
|
||||
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
|
||||
spawn_watcher(state);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_spawn_watcher_invalid_path() {
|
||||
let state = Arc::new(MemoryState::new("/nonexistent/path"));
|
||||
spawn_watcher(state);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user