feat(server): refactor handlers, router, state management, and memory tools

This commit is contained in:
Riz Ashraf committed 2026-10-02 07:27:37 +01:00
1 parent 87ddb01063
commit a083719cf1
36 files changed
+1899 -597

No files matched your search

+13
View File
@@ -0,0 +1,13 @@
[build]
rustc-wrapper = "sccache"
[target.x86_64-pc-windows-msvc]
linker = "rust-lld"
[profile.dev]
debug = 1 # Line tables only: cuts object file & PDB size by ~50%
codegen-units = 16 # Maximizes parallel CPU compilation threads
incremental = false # Enables 100% sccache object file caching across builds
[http]
check-revoke = false # Prevents SChannel revocation check errors on corporate VPNs
Generated
+28 -6
View File
@@ -397,6 +397,19 @@ dependencies = [
"no_std_io2", "no_std_io2",
] ]
[[package]]
name = "blake3"
version = "1.8.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6d9e454fc11f76977dc803893aff6304ed33d6a26efae8696573bea74baa27ae"
dependencies = [
"arrayvec",
"cc",
"cfg-if",
"constant_time_eq",
"cpufeatures 0.3.1",
]
[[package]] [[package]]
name = "block-buffer" name = "block-buffer"
version = "0.10.4" version = "0.10.4"
@@ -670,6 +683,12 @@ version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c"
[[package]]
name = "constant_time_eq"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b"
[[package]] [[package]]
name = "cookie" name = "cookie"
version = "0.18.2" version = "0.18.2"
@@ -2337,7 +2356,7 @@ dependencies = [
[[package]] [[package]]
name = "mcp-memory-linux-nvim" name = "mcp-memory-linux-nvim"
version = "0.1.0" version = "1.0.0"
dependencies = [ dependencies = [
"nvim-core", "nvim-core",
"serde_json", "serde_json",
@@ -2346,12 +2365,13 @@ dependencies = [
[[package]] [[package]]
name = "mcp-memory-server" name = "mcp-memory-server"
version = "0.1.0" version = "1.0.0"
dependencies = [ dependencies = [
"arboard", "arboard",
"async-trait", "async-trait",
"axum", "axum",
"base64 0.23.1", "base64 0.23.1",
"blake3",
"chrono", "chrono",
"clap", "clap",
"clipboard-win", "clipboard-win",
@@ -2373,6 +2393,7 @@ dependencies = [
"tempfile", "tempfile",
"thiserror", "thiserror",
"tokio", "tokio",
"tokio-stream",
"tower", "tower",
"tracing", "tracing",
"tracing-appender", "tracing-appender",
@@ -2391,7 +2412,7 @@ dependencies = [
[[package]] [[package]]
name = "mcp-memory-stub" name = "mcp-memory-stub"
version = "0.1.0" version = "1.0.0"
dependencies = [ dependencies = [
"async-channel", "async-channel",
"clap", "clap",
@@ -2410,7 +2431,7 @@ dependencies = [
[[package]] [[package]]
name = "mcp-memory-win-nvim" name = "mcp-memory-win-nvim"
version = "0.1.0" version = "1.0.0"
dependencies = [ dependencies = [
"nvim-core", "nvim-core",
"serde_json", "serde_json",
@@ -2419,7 +2440,7 @@ dependencies = [
[[package]] [[package]]
name = "mcp-stdio" name = "mcp-stdio"
version = "0.1.0" version = "1.0.0"
dependencies = [ dependencies = [
"tokio", "tokio",
] ]
@@ -2713,7 +2734,7 @@ dependencies = [
[[package]] [[package]]
name = "nvim-core" name = "nvim-core"
version = "0.1.0" version = "1.0.0"
dependencies = [ dependencies = [
"bytes", "bytes",
"dashmap", "dashmap",
@@ -4559,6 +4580,7 @@ dependencies = [
"futures-core", "futures-core",
"pin-project-lite", "pin-project-lite",
"tokio", "tokio",
"tokio-util",
] ]
[[package]] [[package]]
+1 -1
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "mcp-memory-linux-nvim" name = "mcp-memory-linux-nvim"
version = "0.1.0" version = "1.0.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]
+1 -1
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "mcp-stdio" name = "mcp-stdio"
version = "0.1.0" version = "1.0.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]
+2 -2
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "nvim-core" name = "nvim-core"
version = "0.1.0" version = "1.0.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]
@@ -13,7 +13,7 @@ tracing-appender = "0.2.5"
tracing-subscriber = "0.3.23" tracing-subscriber = "0.3.23"
dirs = "7.0.0" dirs = "7.0.0"
rmcp = { version = "3.4.0", features = ["server"] } rmcp = { version = "3.4.0", features = ["server"] }
mcp-stdio = { version = "0.1.0", path = "../mcp-stdio" } mcp-stdio = { version = "1.0.0", path = "../mcp-stdio" }
bytes = "1.12.1" bytes = "1.12.1"
dashmap = "6.2.1" dashmap = "6.2.1"
+11 -9
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "mcp-memory-server" name = "mcp-memory-server"
version = "0.1.0" version = "1.0.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]
@@ -9,12 +9,12 @@ axum = { version = "0.8", features = ["ws"] }
clap = { version = "4.6.6", features = ["derive"] } clap = { version = "4.6.6", features = ["derive"] }
dirs = "7.0.0" dirs = "7.0.0"
futures-util = "0.3.34" futures-util = "0.3.34"
redb = "4.2.0" redb = "4.3.0"
reqwest = { version = "0.13", default-features = false, features = ["stream", "json"] } reqwest = { version = "0.13", default-features = false, features = ["stream", "json"] }
schemars = "1.2.2" schemars = "1.2.2"
serde = { version = "1.0.229", features = ["derive"] } serde = { version = "1.0.229", features = ["derive"] }
serde_json = "1.0.151" serde_json = "1.0.151"
tantivy = "0.26.1" tantivy = "0.26.2"
tokio = { version = "1.53.1", features = ["full"] } tokio = { version = "1.53.1", features = ["full"] }
tracing = "0.1.44" tracing = "0.1.44"
tracing-subscriber = "0.3.23" tracing-subscriber = "0.3.23"
@@ -24,15 +24,15 @@ rmcp = { version = "3.4.0", features = ["server"] }
thiserror = "2.0.20" thiserror = "2.0.20"
tower = { version = "0.5.3", features = ["util"] } tower = { version = "0.5.3", features = ["util"] }
fastembed = "7.1.0" fastembed = "7.1.0"
arboard = "3.4.0" arboard = "3.6.1"
image = "0.25.0" image = "0.25.10"
base64 = "0.23.1" base64 = "0.23.1"
clipboard-win = "5.4.1" clipboard-win = "5.4.1"
git2 = "0.19.0" git2 = "0.19.0"
tree-sitter = "0.23.0" tree-sitter = "0.23.2"
tree-sitter-rust = "0.23.0" tree-sitter-rust = "0.23.3"
tree-sitter-typescript = "0.23.0" tree-sitter-typescript = "0.23.2"
tree-sitter-python = "0.23.0" tree-sitter-python = "0.23.6"
tree-sitter-java = "0.23.5" tree-sitter-java = "0.23.5"
tree-sitter-c = "0.24.2" tree-sitter-c = "0.24.2"
tree-sitter-cpp = "0.23.4" tree-sitter-cpp = "0.23.4"
@@ -40,6 +40,8 @@ tree-sitter-go = "0.25.0"
ignore = "0.4.33" ignore = "0.4.33"
notify = "7" notify = "7"
qdrant-client = "1.19.0" qdrant-client = "1.19.0"
blake3 = "1.8"
tokio-stream = { version = "0.1", features = ["sync"] }
[build-dependencies] [build-dependencies]
chrono = "0.4.45" chrono = "0.4.45"
+68 -10
View File
@@ -9,6 +9,7 @@ use axum::{
routing::{get, post}, routing::{get, post},
}; };
use std::sync::Arc; use std::sync::Arc;
use tokio_stream::StreamExt;
pub fn create_router(app_state: Arc<AppState>) -> Router { pub fn create_router(app_state: Arc<AppState>) -> Router {
Router::new() Router::new()
@@ -24,11 +25,78 @@ pub fn create_router(app_state: Arc<AppState>) -> Router {
.route("/ws", get(ws_handler)) .route("/ws", get(ws_handler))
.route("/health", get(health_handler)) .route("/health", get(health_handler))
.route("/ping", get(|| async move { (axum::http::StatusCode::OK, "pong") })) .route("/ping", get(|| async move { (axum::http::StatusCode::OK, "pong") }))
.route(
"/metrics",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let graph = state_clone.graph.cache.read().unwrap();
let entity_count = graph.entities.len();
let relation_count = graph.relations.len();
let tasks = state_clone.project.tasks.cache.read().unwrap();
let active_tasks = tasks.iter().filter(|t| t.status != "completed" && t.status != "done").count();
let adrs = state_clone.code.adrs.cache.read().unwrap();
let adr_count = adrs.len();
let tech_debts = state_clone.code.tech_debts.cache.read().unwrap();
let debt_count = tech_debts.iter().filter(|d| !d.is_resolved).count();
let metrics = format!(
"# HELP mcp_memory_entities_total Total entities in knowledge graph\n\
# TYPE mcp_memory_entities_total gauge\n\
mcp_memory_entities_total {}\n\
# HELP mcp_memory_relations_total Total relations in knowledge graph\n\
# TYPE mcp_memory_relations_total gauge\n\
mcp_memory_relations_total {}\n\
# HELP mcp_memory_active_tasks Active tasks count\n\
# TYPE mcp_memory_active_tasks gauge\n\
mcp_memory_active_tasks {}\n\
# HELP mcp_memory_adrs_total Total ADR decisions\n\
# TYPE mcp_memory_adrs_total gauge\n\
mcp_memory_adrs_total {}\n\
# HELP mcp_memory_unresolved_tech_debt Unresolved tech debt count\n\
# TYPE mcp_memory_unresolved_tech_debt gauge\n\
mcp_memory_unresolved_tech_debt {}\n",
entity_count, relation_count, active_tasks, adr_count, debt_count
);
([(axum::http::header::CONTENT_TYPE, "text/plain; version=0.0.4")], metrics)
}
}),
)
.route("/nvim/telemetry", post(nvim_telemetry_handler)) .route("/nvim/telemetry", post(nvim_telemetry_handler))
.route("/api/terminal/history", get(crate::api::telemetry::get_terminal_history_handler)) .route("/api/terminal/history", get(crate::api::telemetry::get_terminal_history_handler))
.route("/terminal/telemetry", post(crate::api::telemetry::terminal_telemetry_handler)) .route("/terminal/telemetry", post(crate::api::telemetry::terminal_telemetry_handler))
.route("/events/wait", get(crate::api::events::wait_for_event_handler)) .route("/events/wait", get(crate::api::events::wait_for_event_handler))
.route("/events", post(crate::api::events::post_event_handler)) .route("/events", post(crate::api::events::post_event_handler))
.route(
"/api/activity",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let activity_json = state_clone
.telemetry
.recent_activities
.read_with(|act| serde_json::to_string(act).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], activity_json)
}
}),
)
.route(
"/api/activity/stream",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let rx = state_clone.activity_tx.subscribe();
let stream = tokio_stream::wrappers::BroadcastStream::new(rx)
.filter_map(|msg| match msg {
Ok(data) => Some(Ok::<_, std::convert::Infallible>(
axum::response::sse::Event::default().data(data),
)),
Err(_) => None,
});
axum::response::sse::Sse::new(stream)
}
}),
)
.route("/gate/verify", get(gate_verify_handler)) .route("/gate/verify", get(gate_verify_handler))
.route("/gate/set", post(gate_set_handler)) .route("/gate/set", post(gate_set_handler))
.route( .route(
@@ -140,16 +208,6 @@ pub fn create_router(app_state: Arc<AppState>) -> Router {
} }
}), }),
) )
.route(
"/api/activity",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let activities_json = state_clone.telemetry.recent_activities.read_with(|a| serde_json::to_string(a).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], activities_json)
}
}),
)
.route( .route(
"/api/tech_debts", "/api/tech_debts",
get({ get({
+37 -15
View File
@@ -71,6 +71,10 @@ pub async fn nvim_telemetry_handler(
let _ = state.handler.state.event_bus_tx.send(event); let _ = state.handler.state.event_bus_tx.send(event);
} }
if let Some(ref f) = payload.file {
state.handler.state.record_activity("editor", &format!("Neovim {}: {}", payload.event, f), None);
}
axum::Json(serde_json::json!({"status": "ok"})) axum::Json(serde_json::json!({"status": "ok"}))
} }
@@ -79,20 +83,35 @@ use crate::models::TerminalHistory;
pub async fn get_terminal_history_handler( pub async fn get_terminal_history_handler(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
) -> impl axum::response::IntoResponse { ) -> impl axum::response::IntoResponse {
let history_json = state.handler.state.telemetry.terminal_history.read_with(|h| serde_json::to_string(h).unwrap_or_else(|_| "[]".to_string())); let history_json = state
([(axum::http::header::CONTENT_TYPE, "application/json")], history_json) .handler
.state
.telemetry
.terminal_history
.read_with(|h| serde_json::to_string(h).unwrap_or_else(|_| "[]".to_string()));
(
[(axum::http::header::CONTENT_TYPE, "application/json")],
history_json,
)
} }
pub async fn terminal_telemetry_handler( pub async fn terminal_telemetry_handler(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
axum::Json(payload): axum::Json<TerminalHistory>, axum::Json(payload): axum::Json<TerminalHistory>,
) -> impl axum::response::IntoResponse { ) -> impl axum::response::IntoResponse {
state.handler.state.telemetry.terminal_history.modify(|history| { state.handler.state.record_activity("terminal", &format!("Ran command: {}", payload.command), Some(&format!("exit_code: {}, cwd: {}", payload.exit_code, payload.cwd)));
history.push_front(payload.clone());
if history.len() > 100 { state
history.pop_back(); .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!({ let ws_msg = serde_json::json!({
"type": "terminal_telemetry", "type": "terminal_telemetry",
@@ -120,12 +139,12 @@ mod tests {
use crate::router::MemoryHandler; use crate::router::MemoryHandler;
use crate::state::MemoryState; use crate::state::MemoryState;
use axum::http::Request; use axum::http::Request;
use tower::ServiceExt; use axum::routing::get;
use axum::routing::post;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::atomic::AtomicUsize; use std::sync::atomic::AtomicUsize;
use tempfile::tempdir; use tempfile::tempdir;
use axum::routing::post; use tower::ServiceExt;
use axum::routing::get;
#[tokio::test] #[tokio::test]
async fn test_terminal_history() { async fn test_terminal_history() {
@@ -158,7 +177,9 @@ mod tests {
.method("POST") .method("POST")
.uri("/terminal/telemetry") .uri("/terminal/telemetry")
.header("content-type", "application/json") .header("content-type", "application/json")
.body(axum::body::Body::from(serde_json::to_string(&payload).unwrap())) .body(axum::body::Body::from(
serde_json::to_string(&payload).unwrap(),
))
.unwrap(); .unwrap();
let response = app.clone().oneshot(req).await.unwrap(); let response = app.clone().oneshot(req).await.unwrap();
@@ -173,11 +194,12 @@ mod tests {
let response2 = app.oneshot(req2).await.unwrap(); let response2 = app.oneshot(req2).await.unwrap();
assert_eq!(response2.status(), 200); assert_eq!(response2.status(), 200);
let body_bytes = axum::body::to_bytes(response2.into_body(), usize::MAX).await.unwrap(); let body_bytes = axum::body::to_bytes(response2.into_body(), usize::MAX)
.await
.unwrap();
let history: Vec<TerminalHistory> = serde_json::from_slice(&body_bytes).unwrap(); let history: Vec<TerminalHistory> = serde_json::from_slice(&body_bytes).unwrap();
assert_eq!(history.len(), 1); assert_eq!(history.len(), 1);
assert_eq!(history[0].command, "echo test"); assert_eq!(history[0].command, "echo test");
} }
} }
+25 -25
View File
@@ -1,8 +1,8 @@
use crate::state::MemoryState;
use crate::models::StickyNote; use crate::models::StickyNote;
use std::sync::Arc; use crate::state::MemoryState;
use tokio::time::{sleep, Duration};
use arboard::Clipboard; use arboard::Clipboard;
use std::sync::Arc;
use tokio::time::{Duration, sleep};
pub fn spawn_watcher(state: Arc<MemoryState>) { pub fn spawn_watcher(state: Arc<MemoryState>) {
tokio::spawn(async move { tokio::spawn(async move {
let mut last_text = String::new(); let mut last_text = String::new();
@@ -18,28 +18,28 @@ pub fn spawn_watcher(state: Arc<MemoryState>) {
continue; continue;
} }
if let Ok(mut clipboard) = Clipboard::new() { if let Ok(mut clipboard) = Clipboard::new()
if let Ok(text) = clipboard.get_text() { && let Ok(text) = clipboard.get_text()
if text != last_text && !text.trim().is_empty() { && text != last_text
last_text = text.clone(); && !text.trim().is_empty()
{
let note = StickyNote { last_text = text.clone();
timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH) let note = StickyNote {
.unwrap() timestamp: std::time::SystemTime::now()
.as_secs(), .duration_since(std::time::UNIX_EPOCH)
content: text, .unwrap()
expires_at: None, .as_secs(),
}; content: text,
expires_at: None,
state.code.sticky.modify(|notes| { };
notes.push(note.clone());
}); state.code.sticky.modify(|notes| {
notes.push(note.clone());
// We use rebuild_index to index the new sticky note });
state.rebuild_index().await;
} // We use rebuild_index to index the new sticky note
} state.rebuild_index().await;
} }
} }
}); });
+4 -1
View File
@@ -8,7 +8,10 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
let db = match redb::Database::create(&redb_path) { let db = match redb::Database::create(&redb_path) {
Ok(db) => Arc::new(db), Ok(db) => Arc::new(db),
Err(e) => { Err(e) => {
let err_msg = format!("CRITICAL FATAL ERROR: Failed to open or create redb database at {:?}. Error: {}. This almost always means another instance of mcp-memory-server is already running and holds the exclusive file lock on the database.", redb_path, e); let err_msg = format!(
"CRITICAL FATAL ERROR: Failed to open or create redb database at {:?}. Error: {}. This almost always means another instance of mcp-memory-server is already running and holds the exclusive file lock on the database.",
redb_path, e
);
eprintln!("{}", err_msg); eprintln!("{}", err_msg);
tracing::error!("{}", err_msg); tracing::error!("{}", err_msg);
std::process::exit(1); std::process::exit(1);
+34 -5
View File
@@ -1,7 +1,7 @@
#[allow(deprecated)] #[allow(deprecated)]
use fastembed::{TextEmbedding, EmbeddingModel}; use fastembed::{EmbeddingModel, TextEmbedding};
use std::sync::OnceLock;
use std::sync::Mutex; use std::sync::Mutex;
use std::sync::OnceLock;
static EMBEDDING_MODEL: OnceLock<Mutex<TextEmbedding>> = OnceLock::new(); static EMBEDDING_MODEL: OnceLock<Mutex<TextEmbedding>> = OnceLock::new();
@@ -12,7 +12,7 @@ pub fn get_embedding_model() -> Result<&'static Mutex<TextEmbedding>, String> {
} }
let options = fastembed::InitOptions::new(EmbeddingModel::AllMiniLML6V2) let options = fastembed::InitOptions::new(EmbeddingModel::AllMiniLML6V2)
.with_show_download_progress(true); .with_show_download_progress(true);
let model = TextEmbedding::try_new(options).map_err(|e| e.to_string())?; let model = TextEmbedding::try_new(options).map_err(|e| e.to_string())?;
let _ = EMBEDDING_MODEL.set(Mutex::new(model)); let _ = EMBEDDING_MODEL.set(Mutex::new(model));
Ok(EMBEDDING_MODEL.get().unwrap()) Ok(EMBEDDING_MODEL.get().unwrap())
@@ -24,7 +24,9 @@ pub async fn generate_embedding_async(text: String) -> Result<Vec<f32>, String>
let mut model = model_mutex.lock().map_err(|e| e.to_string())?; let mut model = model_mutex.lock().map_err(|e| e.to_string())?;
let embeddings = model.embed(vec![text], None).map_err(|e| e.to_string())?; let embeddings = model.embed(vec![text], None).map_err(|e| e.to_string())?;
Ok(embeddings.into_iter().next().unwrap_or_default()) Ok(embeddings.into_iter().next().unwrap_or_default())
}).await.map_err(|e| e.to_string())? })
.await
.map_err(|e| e.to_string())?
} }
pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
@@ -49,5 +51,32 @@ pub async fn generate_embeddings_async(texts: Vec<String>) -> Result<Vec<Vec<f32
let mut model = model_mutex.lock().map_err(|e| e.to_string())?; let mut model = model_mutex.lock().map_err(|e| e.to_string())?;
let embeddings = model.embed(texts, None).map_err(|e| e.to_string())?; let embeddings = model.embed(texts, None).map_err(|e| e.to_string())?;
Ok(embeddings) Ok(embeddings)
}).await.map_err(|e| e.to_string())? })
.await
.map_err(|e| e.to_string())?
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cosine_similarity_edge_cases() {
assert_eq!(cosine_similarity(&[], &[]), 0.0);
assert_eq!(cosine_similarity(&[1.0, 2.0], &[1.0]), 0.0);
assert_eq!(cosine_similarity(&[0.0, 0.0], &[0.0, 0.0]), 0.0);
let v1 = vec![1.0, 0.0, 0.0];
let v2 = vec![1.0, 0.0, 0.0];
assert!((cosine_similarity(&v1, &v2) - 1.0).abs() < 1e-5);
let v3 = vec![0.0, 1.0, 0.0];
assert!((cosine_similarity(&v1, &v3) - 0.0).abs() < 1e-5);
}
#[tokio::test]
async fn test_generate_embeddings_async_empty() {
let res = generate_embeddings_async(vec![]).await.unwrap();
assert!(res.is_empty());
}
} }
+16 -2
View File
@@ -42,8 +42,14 @@ impl IntoResponse for AppError {
AppError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg.clone()), AppError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg.clone()),
AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()), AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()),
AppError::Internal(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg.clone()), AppError::Internal(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg.clone()),
AppError::JsonError(e) => (StatusCode::INTERNAL_SERVER_ERROR, format!("JSON Error: {}", e)), AppError::JsonError(e) => (
AppError::IoError(e) => (StatusCode::INTERNAL_SERVER_ERROR, format!("I/O Error: {}", e)), StatusCode::INTERNAL_SERVER_ERROR,
format!("JSON Error: {}", e),
),
AppError::IoError(e) => (
StatusCode::INTERNAL_SERVER_ERROR,
format!("I/O Error: {}", e),
),
}; };
let body = Json(json!({ let body = Json(json!({
@@ -115,3 +121,11 @@ impl From<&str> for AppError {
AppError::Internal(msg.to_string()) AppError::Internal(msg.to_string())
} }
} }
// Generated by improve-test (Attempt 1)
#[test]
fn test_app_error_unknown_error() {
let err = AppError::Internal("unknown error".into());
let res = err.into_response();
assert_eq!(res.status(), StatusCode::INTERNAL_SERVER_ERROR);
}
+92 -42
View File
@@ -4,7 +4,7 @@ use crate::tools::ReadFileSkeletonTool;
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::Value; use serde_json::Value;
use std::sync::Arc; use std::sync::Arc;
use tree_sitter::{Parser, Node}; use tree_sitter::{Node, Parser};
pub struct ReadFileSkeletonHandler; pub struct ReadFileSkeletonHandler;
@@ -22,20 +22,22 @@ impl McpTool for ReadFileSkeletonHandler {
} }
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: ReadFileSkeletonTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let tool_args: ReadFileSkeletonTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let file_path = tool_args.file_path.clone(); let file_path = tool_args.file_path.clone();
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> { let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let code = std::fs::read_to_string(&file_path) let code = std::fs::read_to_string(&file_path).map_err(|e| {
.map_err(|e| crate::error::AppError::Internal(format!("Failed to read file: {}", e)))?; crate::error::AppError::Internal(format!("Failed to read file: {}", e))
})?;
let mut parser = Parser::new(); let mut parser = Parser::new();
let ext = std::path::Path::new(&file_path) let ext = std::path::Path::new(&file_path)
.extension() .extension()
.and_then(|s| s.to_str()) .and_then(|s| s.to_str())
.unwrap_or(""); .unwrap_or("");
let language = match ext { let language = match ext {
"rs" => tree_sitter_rust::LANGUAGE, "rs" => tree_sitter_rust::LANGUAGE,
"ts" | "tsx" | "js" | "jsx" => tree_sitter_typescript::LANGUAGE_TYPESCRIPT, "ts" | "tsx" | "js" | "jsx" => tree_sitter_typescript::LANGUAGE_TYPESCRIPT,
@@ -46,20 +48,41 @@ impl McpTool for ReadFileSkeletonHandler {
"go" => tree_sitter_go::LANGUAGE, "go" => tree_sitter_go::LANGUAGE,
_ => return Ok(code), _ => return Ok(code),
}; };
parser.set_language(&language.into()).map_err(|e| e.to_string())?; parser
.set_language(&language.into())
let tree = parser.parse(&code, None).ok_or_else(|| crate::error::AppError::Internal("Failed to parse code".to_string()))?; .map_err(|e| e.to_string())?;
let tree = parser.parse(&code, None).ok_or_else(|| {
crate::error::AppError::Internal("Failed to parse code".to_string())
})?;
let mut result_skeleton = String::new(); let mut result_skeleton = String::new();
fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) { fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) {
let kind = node.kind(); let kind = node.kind();
let is_structural = matches!(kind, "use_declaration" | "import_statement" | "import_from_statement" | "struct_item" | "enum_item" | "trait_item" | "impl_item" | "function_item" | "function_declaration" | "function_definition" | "method_definition" | "interface_declaration" | "type_alias_declaration" | "class_declaration" | "class_definition"); let is_structural = matches!(
kind,
"use_declaration"
| "import_statement"
| "import_from_statement"
| "struct_item"
| "enum_item"
| "trait_item"
| "impl_item"
| "function_item"
| "function_declaration"
| "function_definition"
| "method_definition"
| "interface_declaration"
| "type_alias_declaration"
| "class_declaration"
| "class_definition"
);
if is_structural { if is_structural {
let indent = " ".repeat(depth); let indent = " ".repeat(depth);
let node_text = node.utf8_text(code.as_bytes()).unwrap_or(""); let node_text = node.utf8_text(code.as_bytes()).unwrap_or("");
let mut signature = String::new(); let mut signature = String::new();
for line in node_text.lines() { for line in node_text.lines() {
@@ -76,7 +99,7 @@ impl McpTool for ReadFileSkeletonHandler {
if signature.is_empty() { if signature.is_empty() {
signature = node_text.to_string(); signature = node_text.to_string();
} }
out.push_str(&indent); out.push_str(&indent);
out.push_str(signature.trim()); out.push_str(signature.trim());
out.push('\n'); out.push('\n');
@@ -87,9 +110,9 @@ impl McpTool for ReadFileSkeletonHandler {
} }
} }
} }
extract_skeleton(tree.root_node(), &code, &mut result_skeleton, 0); extract_skeleton(tree.root_node(), &code, &mut result_skeleton, 0);
if result_skeleton.is_empty() { if result_skeleton.is_empty() {
Ok(code) Ok(code)
} else { } else {
@@ -120,20 +143,22 @@ impl McpTool for ReplaceAstNodeHandler {
} }
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: ReplaceAstNodeTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let tool_args: ReplaceAstNodeTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let file_path = tool_args.file_path.clone(); let file_path = tool_args.file_path.clone();
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> { let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let code = std::fs::read_to_string(&file_path) let code = std::fs::read_to_string(&file_path).map_err(|e| {
.map_err(|e| crate::error::AppError::Internal(format!("Failed to read file: {}", e)))?; crate::error::AppError::Internal(format!("Failed to read file: {}", e))
})?;
let mut parser = Parser::new(); let mut parser = Parser::new();
let ext = std::path::Path::new(&file_path) let ext = std::path::Path::new(&file_path)
.extension() .extension()
.and_then(|s| s.to_str()) .and_then(|s| s.to_str())
.unwrap_or(""); .unwrap_or("");
let language = match ext { let language = match ext {
"rs" => tree_sitter_rust::LANGUAGE, "rs" => tree_sitter_rust::LANGUAGE,
"ts" | "tsx" | "js" | "jsx" => tree_sitter_typescript::LANGUAGE_TYPESCRIPT, "ts" | "tsx" | "js" | "jsx" => tree_sitter_typescript::LANGUAGE_TYPESCRIPT,
@@ -142,14 +167,28 @@ impl McpTool for ReplaceAstNodeHandler {
"c" | "h" => tree_sitter_c::LANGUAGE, "c" | "h" => tree_sitter_c::LANGUAGE,
"cpp" | "cc" | "cxx" | "hpp" | "hxx" => tree_sitter_cpp::LANGUAGE, "cpp" | "cc" | "cxx" | "hpp" | "hxx" => tree_sitter_cpp::LANGUAGE,
"go" => tree_sitter_go::LANGUAGE, "go" => tree_sitter_go::LANGUAGE,
_ => return Err(crate::error::AppError::Internal(format!("Unsupported language for AST replacement: {}", ext))), _ => {
return Err(crate::error::AppError::Internal(format!(
"Unsupported language for AST replacement: {}",
ext
)));
}
}; };
parser.set_language(&language.into()).map_err(|e| e.to_string())?; parser
let tree = parser.parse(&code, None).ok_or_else(|| crate::error::AppError::Internal("Failed to parse code".to_string()))?; .set_language(&language.into())
.map_err(|e| e.to_string())?;
let tree = parser.parse(&code, None).ok_or_else(|| {
crate::error::AppError::Internal("Failed to parse code".to_string())
})?;
// Search for the node // Search for the node
fn find_node<'a>(node: Node<'a>, code: &str, target_type: &str, target_name: &str) -> Option<Node<'a>> { fn find_node<'a>(
node: Node<'a>,
code: &str,
target_type: &str,
target_name: &str,
) -> Option<Node<'a>> {
if node.kind() == target_type { if node.kind() == target_type {
// Try to find the name/identifier // Try to find the name/identifier
let mut cursor = node.walk(); let mut cursor = node.walk();
@@ -163,7 +202,7 @@ impl McpTool for ReplaceAstNodeHandler {
} }
} }
} }
let mut cursor = node.walk(); let mut cursor = node.walk();
for child in node.children(&mut cursor) { for child in node.children(&mut cursor) {
if let Some(found) = find_node(child, code, target_type, target_name) { if let Some(found) = find_node(child, code, target_type, target_name) {
@@ -172,22 +211,33 @@ impl McpTool for ReplaceAstNodeHandler {
} }
None None
} }
let target_node = find_node(tree.root_node(), &code, &tool_args.node_type, &tool_args.node_name); let target_node = find_node(
tree.root_node(),
&code,
&tool_args.node_type,
&tool_args.node_name,
);
if let Some(node) = target_node { if let Some(node) = target_node {
let start_byte = node.start_byte(); let start_byte = node.start_byte();
let end_byte = node.end_byte(); let end_byte = node.end_byte();
let mut new_file_content = String::new(); let mut new_file_content = String::new();
new_file_content.push_str(&code[..start_byte]); new_file_content.push_str(&code[..start_byte]);
new_file_content.push_str(&tool_args.new_content); new_file_content.push_str(&tool_args.new_content);
new_file_content.push_str(&code[end_byte..]); new_file_content.push_str(&code[end_byte..]);
std::fs::write(&file_path, new_file_content).map_err(|e| e.to_string())?; std::fs::write(&file_path, new_file_content).map_err(|e| e.to_string())?;
Ok(format!("Successfully replaced node {} of type {} in {}", tool_args.node_name, tool_args.node_type, file_path)) Ok(format!(
"Successfully replaced node {} of type {} in {}",
tool_args.node_name, tool_args.node_type, file_path
))
} else { } else {
Err(crate::error::AppError::Internal(format!("Could not find node {} of type {}", tool_args.node_name, tool_args.node_type))) Err(crate::error::AppError::Internal(format!(
"Could not find node {} of type {}",
tool_args.node_name, tool_args.node_type
)))
} }
}) })
.await .await
@@ -208,7 +258,7 @@ mod tests {
let dir = tempdir().unwrap(); let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap())); let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let file_path = dir.path().join("test_skeleton.rs"); let file_path = dir.path().join("test_skeleton.rs");
let code = "fn my_func() {\n let x = 1;\n}\n\nstruct MyStruct {\n val: i32\n}"; let code = "fn my_func() {\n let x = 1;\n}\n\nstruct MyStruct {\n val: i32\n}";
std::fs::write(&file_path, code).unwrap(); std::fs::write(&file_path, code).unwrap();
@@ -227,7 +277,7 @@ mod tests {
let dir = tempdir().unwrap(); let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap())); let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let file_path = dir.path().join("test_replace.rs"); let file_path = dir.path().join("test_replace.rs");
let code = "fn my_func() {\n let x = 1;\n}\n\nstruct MyStruct {\n val: i32\n}"; let code = "fn my_func() {\n let x = 1;\n}\n\nstruct MyStruct {\n val: i32\n}";
std::fs::write(&file_path, code).unwrap(); std::fs::write(&file_path, code).unwrap();
+26 -7
View File
@@ -58,7 +58,8 @@ impl McpTool for ReadEnvFingerprintHandler {
let req: ReadEnvFingerprintTool = let req: ReadEnvFingerprintTool =
serde_json::from_value(args).map_err(|e| e.to_string())?; serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state let data = state
.env.env_fingerprints .env
.env_fingerprints
.read_with(|fps| fps.get(&req.namespace).cloned()); .read_with(|fps| fps.get(&req.namespace).cloned());
if let Some(fp) = data { if let Some(fp) = data {
let data = Ok::<String, crate::error::AppError>(serde_json::to_string(&fp)?)?; let data = Ok::<String, crate::error::AppError>(serde_json::to_string(&fp)?)?;
@@ -180,13 +181,19 @@ mod tests {
} }
}); });
let res = update_handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res = update_handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res, "Env fingerprint updated"); assert_eq!(res, "Env fingerprint updated");
let read_handler = ReadEnvFingerprintHandler; let read_handler = ReadEnvFingerprintHandler;
let res2 = read_handler let res2 = read_handler
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res2.contains("rustc")); assert!(res2.contains("rustc"));
assert!(res2.contains("1.70.0")); assert!(res2.contains("1.70.0"));
} }
@@ -211,7 +218,9 @@ mod tests {
let handler = GetEnvironmentDetailsHandler; let handler = GetEnvironmentDetailsHandler;
let res = handler let res = handler
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res.contains("global")); assert!(res.contains("global"));
} }
@@ -230,7 +239,11 @@ mod tests {
"context": "For database access", "context": "For database access",
"namespace": "global" "namespace": "global"
}); });
let res1 = req_handler.execute(args_req, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res1 = req_handler
.execute(args_req, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res1, "Env requirement logged"); assert_eq!(res1, "Env requirement logged");
let reg_handler = RegisterEnvironmentHandler; let reg_handler = RegisterEnvironmentHandler;
@@ -241,13 +254,19 @@ mod tests {
"requires_vpn": true, "requires_vpn": true,
"namespace": "global" "namespace": "global"
}); });
let res2 = reg_handler.execute(args_reg, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res2 = reg_handler
.execute(args_reg, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res2, "Environment registered"); assert_eq!(res2, "Environment registered");
let get_handler = GetEnvironmentDetailsHandler; let get_handler = GetEnvironmentDetailsHandler;
let res3 = get_handler let res3 = get_handler
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res3.contains("prod.local")); assert!(res3.contains("prod.local"));
assert!(!res3.is_empty()); assert!(!res3.is_empty());
} }
+56 -45
View File
@@ -2,9 +2,9 @@ use crate::router::McpTool;
use crate::state::MemoryState; use crate::state::MemoryState;
use crate::tools::GetActiveWorktreeContextTool; use crate::tools::GetActiveWorktreeContextTool;
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::{json, Value}; use serde_json::{Value, json};
use std::sync::Arc;
use std::env; use std::env;
use std::sync::Arc;
pub struct GetActiveWorktreeContextHandler; pub struct GetActiveWorktreeContextHandler;
@@ -21,56 +21,66 @@ impl McpTool for GetActiveWorktreeContextHandler {
) )
} }
async fn execute(&self, _args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Value> { &self,
let cwd = env::current_dir().map_err(|e| e.to_string())?; _args: Value,
let repo = git2::Repository::discover(&cwd).map_err(|e| crate::error::AppError::Internal(format!("Not in a git repository: {}", e)))?; _state: Arc<MemoryState>,
) -> crate::error::Result<String> {
let mut branch_name = String::new(); let result =
if let Ok(head) = repo.head() tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Value> {
&& let Some(name) = head.shorthand() { let cwd = env::current_dir().map_err(|e| e.to_string())?;
let repo = git2::Repository::discover(&cwd).map_err(|e| {
crate::error::AppError::Internal(format!("Not in a git repository: {}", e))
})?;
let mut branch_name = String::new();
if let Ok(head) = repo.head()
&& let Some(name) = head.shorthand()
{
branch_name = name.to_string(); branch_name = name.to_string();
} }
let mut opts = git2::DiffOptions::new(); let mut opts = git2::DiffOptions::new();
let mut diff = None; let mut diff = None;
// Try to diff against HEAD
if let Ok(tree) = repo.head().and_then(|h| h.peel_to_tree()) {
diff = repo.diff_tree_to_workdir_with_index(Some(&tree), Some(&mut opts)).ok();
}
let mut files = Vec::new(); // Try to diff against HEAD
let mut diff_output = String::new(); if let Ok(tree) = repo.head().and_then(|h| h.peel_to_tree()) {
diff = repo
.diff_tree_to_workdir_with_index(Some(&tree), Some(&mut opts))
.ok();
}
if let Some(diff) = diff { let mut files = Vec::new();
let _ = diff.print(git2::DiffFormat::Patch, |_delta, _hunk, line| { let mut diff_output = String::new();
let content = std::str::from_utf8(line.content()).unwrap_or("");
diff_output.push_str(content);
true
});
for delta in diff.deltas() { if let Some(diff) = diff {
if let Some(path) = delta.new_file().path() { let _ = diff.print(git2::DiffFormat::Patch, |_delta, _hunk, line| {
files.push(path.to_string_lossy().into_owned()); let content = std::str::from_utf8(line.content()).unwrap_or("");
diff_output.push_str(content);
true
});
for delta in diff.deltas() {
if let Some(path) = delta.new_file().path() {
files.push(path.to_string_lossy().into_owned());
}
} }
} }
}
// Truncate diff output if it's too large to save tokens // Truncate diff output if it's too large to save tokens
if diff_output.len() > 10000 { if diff_output.len() > 10000 {
diff_output.truncate(10000); diff_output.truncate(10000);
diff_output.push_str("\n... [Diff truncated due to size]"); diff_output.push_str("\n... [Diff truncated due to size]");
} }
Ok(json!({ Ok(json!({
"branch": branch_name, "branch": branch_name,
"modified_files": files, "modified_files": files,
"diff": diff_output "diff": diff_output
})) }))
}) })
.await .await
.map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??; .map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&result)?) Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&result)?)
} }
@@ -79,9 +89,9 @@ impl McpTool for GetActiveWorktreeContextHandler {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use tempfile::tempdir;
use std::sync::Arc;
use serde_json::json; use serde_json::json;
use std::sync::Arc;
use tempfile::tempdir;
#[tokio::test] #[tokio::test]
async fn test_get_active_worktree_context() { async fn test_get_active_worktree_context() {
@@ -89,7 +99,8 @@ mod tests {
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap())); let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let handler = GetActiveWorktreeContextHandler; let handler = GetActiveWorktreeContextHandler;
let result = handler.execute(json!({}), state) let result = handler
.execute(json!({}), state)
.await .await
.map_err(|e| format!("Failed to get worktree context: {}", e)) .map_err(|e| format!("Failed to get worktree context: {}", e))
.unwrap(); .unwrap();
+65 -19
View File
@@ -186,7 +186,8 @@ impl McpTool for CreateRelationsHandler {
let missing: Vec<_> = missing_nodes.into_iter().collect(); let missing: Vec<_> = missing_nodes.into_iter().collect();
return Err(crate::error::AppError::Internal(format!( return Err(crate::error::AppError::Internal(format!(
"Error: Relations dropped due to missing entities: {}", "Error: Relations dropped due to missing entities: {}",
missing.join(", ")))); missing.join(", ")
)));
} }
Ok("Relations created".to_string()) Ok("Relations created".to_string())
} }
@@ -219,7 +220,8 @@ impl McpTool for AddObservationsHandler {
if !missing_entities.is_empty() { if !missing_entities.is_empty() {
return Err(crate::error::AppError::Internal(format!( return Err(crate::error::AppError::Internal(format!(
"Error: Observations dropped for missing entities: {}", "Error: Observations dropped for missing entities: {}",
missing_entities.join(", ")))); missing_entities.join(", ")
)));
} }
Ok("Observations added".to_string()) Ok("Observations added".to_string())
} }
@@ -255,7 +257,8 @@ impl McpTool for DeleteEntitiesHandler {
if !missing.is_empty() { if !missing.is_empty() {
return Err(crate::error::AppError::Internal(format!( return Err(crate::error::AppError::Internal(format!(
"Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.", "Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
missing.join(", ")))); missing.join(", ")
)));
} }
let idx = state.get_search_index(); let idx = state.get_search_index();
@@ -298,7 +301,8 @@ impl McpTool for DeleteObservationsHandler {
if !missing.is_empty() { if !missing.is_empty() {
return Err(crate::error::AppError::Internal(format!( return Err(crate::error::AppError::Internal(format!(
"Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.", "Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
missing.join(", ")))); missing.join(", ")
)));
} }
Ok("Observations deleted".to_string()) Ok("Observations deleted".to_string())
} }
@@ -671,7 +675,11 @@ mod tests {
] ]
}); });
let res = create_handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res = create_handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res, "Entities created"); assert_eq!(res, "Entities created");
// Ensure graph contains the entity // Ensure graph contains the entity
@@ -716,7 +724,11 @@ mod tests {
{"from": "A", "to": "B", "relation_type": "knows"} {"from": "A", "to": "B", "relation_type": "knows"}
] ]
}); });
let res = handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res = handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res, "Relations created"); assert_eq!(res, "Relations created");
// Test semantic LLM schema feedback (User request) // Test semantic LLM schema feedback (User request)
@@ -725,7 +737,11 @@ mod tests {
{"source": "A", "target": "B", "relationType": "knows"} {"source": "A", "target": "B", "relationType": "knows"}
] ]
}); });
let err_res = handler.execute(bad_args, state.clone()).await.unwrap_err().to_string(); let err_res = handler
.execute(bad_args, state.clone())
.await
.unwrap_err()
.to_string();
assert!(err_res.contains("Schema error:")); assert!(err_res.contains("Schema error:"));
assert!(err_res.contains("strictly uses 'from', 'to', and 'relation_type'")); assert!(err_res.contains("strictly uses 'from', 'to', and 'relation_type'"));
} }
@@ -755,25 +771,35 @@ mod tests {
{"entity_name": "A", "contents": ["Obs 1", "Obs 2"]} {"entity_name": "A", "contents": ["Obs 1", "Obs 2"]}
] ]
}); });
let res1 = add_obs.execute(args_obs, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res1 = add_obs
.execute(args_obs, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res1, "Observations added"); assert_eq!(res1, "Observations added");
let read_graph = ReadGraphHandler; let read_graph = ReadGraphHandler;
let res2 = read_graph let res2 = read_graph
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res2.contains("Obs 1")); assert!(res2.contains("Obs 1"));
assert!(res2.contains("Obs 2")); assert!(res2.contains("Obs 2"));
let del_entity = DeleteEntitiesHandler; let del_entity = DeleteEntitiesHandler;
let res4 = del_entity let res4 = del_entity
.execute(json!({"entity_names": ["A"]}), state.clone()) .execute(json!({"entity_names": ["A"]}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res4, "Entities deleted"); assert_eq!(res4, "Entities deleted");
let res5 = read_graph let res5 = read_graph
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(!res5.contains("A")); assert!(!res5.contains("A"));
} }
@@ -791,7 +817,9 @@ mod tests {
}); });
create_handler create_handler
.execute(args_ent, state.clone()) .execute(args_ent, state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let rel_handler = CreateRelationsHandler; let rel_handler = CreateRelationsHandler;
let args_rel = json!({ let args_rel = json!({
@@ -799,25 +827,35 @@ mod tests {
{"from": "X", "to": "Y", "relation_type": "depends_on", "namespace": "global"} {"from": "X", "to": "Y", "relation_type": "depends_on", "namespace": "global"}
] ]
}); });
rel_handler.execute(args_rel, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); rel_handler
.execute(args_rel, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let read_handler = ReadGraphHandler; let read_handler = ReadGraphHandler;
let res_read = read_handler let res_read = read_handler
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res_read.contains("X")); assert!(res_read.contains("X"));
assert!(res_read.contains("depends_on")); assert!(res_read.contains("depends_on"));
let open_handler = OpenNodesHandler; let open_handler = OpenNodesHandler;
let res_open = open_handler let res_open = open_handler
.execute(json!({"names": ["X"]}), state.clone()) .execute(json!({"names": ["X"]}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res_open.contains("Y")); assert!(res_open.contains("Y"));
let viz_handler = VisualizeGraphHandler; let viz_handler = VisualizeGraphHandler;
let res_viz = viz_handler let res_viz = viz_handler
.execute(json!({"query": "X"}), state.clone()) .execute(json!({"query": "X"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(!res_viz.is_empty()); assert!(!res_viz.is_empty());
let condense = CondenseEntityHandler; let condense = CondenseEntityHandler;
@@ -826,7 +864,9 @@ mod tests {
json!({"entity_name": "X", "summarized_observations": ["X condensed"]}), json!({"entity_name": "X", "summarized_observations": ["X condensed"]}),
state.clone(), state.clone(),
) )
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res_cond, "Entity condensed"); assert_eq!(res_cond, "Entity condensed");
let merge = MergeEntitiesHandler; let merge = MergeEntitiesHandler;
@@ -835,11 +875,17 @@ mod tests {
json!({"source_entity": "X", "target_entity": "Y"}), json!({"source_entity": "X", "target_entity": "Y"}),
state.clone(), state.clone(),
) )
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res_merge, "Entities merged"); assert_eq!(res_merge, "Entities merged");
let orphans = FindOrphansHandler; let orphans = FindOrphansHandler;
let res_orphans = orphans.execute(json!({}), state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res_orphans = orphans
.execute(json!({}), state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(!res_orphans.contains("Y")); assert!(!res_orphans.contains("Y"));
} }
} }
+29 -18
View File
@@ -1,11 +1,11 @@
use crate::router::McpTool; use crate::router::McpTool;
use crate::state::MemoryState; use crate::state::MemoryState;
use crate::tools::{WatchProcessLogsTool, GetRecentLogsTool}; use crate::tools::{GetRecentLogsTool, WatchProcessLogsTool};
use async_trait::async_trait; use async_trait::async_trait;
use serde_json::Value; use serde_json::Value;
use std::sync::Arc;
use std::io::{Read, Seek, SeekFrom};
use std::fs::File; use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::sync::Arc;
pub struct WatchProcessLogsHandler; pub struct WatchProcessLogsHandler;
@@ -23,9 +23,13 @@ impl McpTool for WatchProcessLogsHandler {
} }
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: WatchProcessLogsTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let tool_args: WatchProcessLogsTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
if !std::path::Path::new(&tool_args.file_path).exists() { if !std::path::Path::new(&tool_args.file_path).exists() {
return Err(crate::error::AppError::Internal(format!("File does not exist: {}", tool_args.file_path))); return Err(crate::error::AppError::Internal(format!(
"File does not exist: {}",
tool_args.file_path
)));
} }
Ok(format!("Started watching logs for {}", tool_args.file_path)) Ok(format!("Started watching logs for {}", tool_args.file_path))
} }
@@ -47,25 +51,30 @@ impl McpTool for GetRecentLogsHandler {
} }
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: GetRecentLogsTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let tool_args: GetRecentLogsTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> { let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let mut file = File::open(&tool_args.file_path).map_err(|e| crate::error::AppError::Internal(format!("Failed to open file: {}", e)))?; let mut file = File::open(&tool_args.file_path).map_err(|e| {
crate::error::AppError::Internal(format!("Failed to open file: {}", e))
})?;
let len = file.metadata().map_err(|e| e.to_string())?.len(); let len = file.metadata().map_err(|e| e.to_string())?.len();
let read_size = std::cmp::min(16 * 1024, len); let read_size = std::cmp::min(16 * 1024, len);
file.seek(SeekFrom::End(-(read_size as i64))).map_err(|e| e.to_string())?; file.seek(SeekFrom::End(-(read_size as i64)))
.map_err(|e| e.to_string())?;
let mut buffer = String::new(); let mut buffer = String::new();
file.read_to_string(&mut buffer).map_err(|e| e.to_string())?; file.read_to_string(&mut buffer)
.map_err(|e| e.to_string())?;
let lines: Vec<&str> = buffer.lines().collect(); let lines: Vec<&str> = buffer.lines().collect();
let recent_lines = if lines.len() > 100 { let recent_lines = if lines.len() > 100 {
lines[lines.len() - 100..].join("\n") lines[lines.len() - 100..].join("\n")
} else { } else {
buffer buffer
}; };
Ok(recent_lines) Ok(recent_lines)
}) })
.await .await
@@ -78,9 +87,9 @@ impl McpTool for GetRecentLogsHandler {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use tempfile::tempdir;
use std::sync::Arc;
use serde_json::json; use serde_json::json;
use std::sync::Arc;
use tempfile::tempdir;
#[tokio::test] #[tokio::test]
async fn test_watch_process_logs() { async fn test_watch_process_logs() {
@@ -95,7 +104,8 @@ mod tests {
"file_path": log_file.to_str().unwrap() "file_path": log_file.to_str().unwrap()
}); });
let result = handler.execute(args, state) let result = handler
.execute(args, state)
.await .await
.map_err(|e| format!("Failed to watch logs: {}", e)) .map_err(|e| format!("Failed to watch logs: {}", e))
.unwrap(); .unwrap();
@@ -115,7 +125,8 @@ mod tests {
"file_path": log_file.to_str().unwrap() "file_path": log_file.to_str().unwrap()
}); });
let result = handler.execute(args, state) let result = handler
.execute(args, state)
.await .await
.map_err(|e| format!("Failed to get recent logs: {}", e)) .map_err(|e| format!("Failed to get recent logs: {}", e))
.unwrap(); .unwrap();
+208 -88
View File
@@ -39,9 +39,9 @@ impl McpTool for LogDecisionHandler {
final_id = format!("ADR-{:04}", adrs.len() + 1); final_id = format!("ADR-{:04}", adrs.len() + 1);
let a = Adr { let a = Adr {
id: final_id.clone(), id: final_id.clone(),
title: req.title, title: req.title.clone(),
context: req.context, context: req.context,
decision: req.decision, decision: req.decision.clone(),
consequence: req.consequence, consequence: req.consequence,
status: "accepted".to_string(), status: "accepted".to_string(),
supersedes: req.supersedes, supersedes: req.supersedes,
@@ -52,7 +52,8 @@ impl McpTool for LogDecisionHandler {
adrs.push(a); adrs.push(a);
}); });
Ok(format!("Decision logged as {}", final_id)) state.record_activity("decision", &format!("Logged {}: {}", final_id, req.title), Some(&req.decision));
Ok(format!("Logged decision {}: {}", final_id, req.title))
} }
} }
@@ -119,7 +120,9 @@ impl McpTool for DeleteDecisionHandler {
state.rebuild_index().await; state.rebuild_index().await;
Ok("Decision deleted successfully".to_string()) Ok("Decision deleted successfully".to_string())
} else { } else {
Err(crate::error::AppError::Internal("Decision not found".to_string())) Err(crate::error::AppError::Internal(
"Decision not found".to_string(),
))
} }
} }
} }
@@ -139,18 +142,40 @@ impl McpTool for LogErrorFixHandler {
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 req: LogErrorFixTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: LogErrorFixTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let text_to_embed = format!("Signature: {}\nSolution: {}", req.signature, req.solution); let text_to_embed = format!("Signature: {}\nSolution: {}", req.signature, req.solution);
let embedding = crate::embedding::generate_embedding_async(text_to_embed).await.ok(); let mut solution = req.solution;
if state.ollama.is_available().await {
let prompt = format!(
"Analyze this error signature and solution. Output 1 sentence summarizing the root cause and fix:\nSignature: {}\nSolution: {}",
req.signature, solution
);
if let Ok(summary) = state
.ollama
.generate(&prompt, Some(&state.ollama.reasoning_model), None)
.await
{
let clean = summary.trim();
if !clean.is_empty() {
solution = format!("{} (AI Analysis: {})", solution, clean);
}
}
}
let embedding = crate::embedding::generate_embedding_async(text_to_embed)
.await
.ok();
state.code.error_fixes.modify(|fixes| { state.code.error_fixes.modify(|fixes| {
fixes.push(crate::models::ErrorFix { fixes.push(crate::models::ErrorFix {
signature: req.signature, signature: req.signature.clone(),
solution: req.solution, solution: solution.clone(),
timestamp: crate::handlers::utils::now_secs(), timestamp: crate::handlers::utils::now_secs(),
git_commit: req.git_commit, git_commit: req.git_commit,
git_branch: req.git_branch, git_branch: req.git_branch,
embedding, embedding,
}) })
}); });
Ok("Error fix logged".to_string()) state.record_activity("error_fix", &format!("Fixed error: {}", req.signature), Some(&solution));
Ok(format!("Logged error fix for {}: {}", req.signature, solution))
} }
} }
@@ -200,16 +225,32 @@ impl McpTool for LogCodeChangeHandler {
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 req: LogCodeChangeTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: LogCodeChangeTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut description = req.description;
if state.ollama.is_available().await {
let prompt = format!(
"Summarize in 1 concise sentence the architectural impact of changing file '{}': {}",
req.file_path, description
);
if let Ok(summary) = state.ollama.generate(&prompt, None, None).await {
let clean = summary.trim();
if !clean.is_empty() {
description = format!("{} (AI Summary: {})", description, clean);
}
}
}
state.code.ledger.modify(|ledger| { state.code.ledger.modify(|ledger| {
ledger.push(CodeChange { ledger.push(CodeChange {
timestamp: crate::handlers::utils::now_secs(), timestamp: crate::handlers::utils::now_secs(),
file_path: req.file_path, file_path: req.file_path.clone(),
description: req.description, description: description.clone(),
git_commit: req.git_commit, git_commit: req.git_commit,
git_branch: req.git_branch, git_branch: req.git_branch,
}); });
}); });
Ok("Code change logged".to_string()) state.record_activity("code_change", &format!("Modified {}", req.file_path), Some(&description));
Ok(format!("Logged code change for {}: {}", req.file_path, description))
} }
} }
@@ -230,7 +271,8 @@ impl McpTool for QueryRecentChangesHandler {
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 data = state let data = state
.code.ledger .code
.ledger
.read_with(|l| Ok::<String, crate::error::AppError>(serde_json::to_string(l)?))?; .read_with(|l| Ok::<String, crate::error::AppError>(serde_json::to_string(l)?))?;
Ok(data) Ok(data)
} }
@@ -278,7 +320,8 @@ impl McpTool for ReadPreferencesHandler {
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> {
state state
.env.prefs .env
.prefs
.read_with(|prefs| Ok::<String, crate::error::AppError>(serde_json::to_string(prefs)?)) .read_with(|prefs| Ok::<String, crate::error::AppError>(serde_json::to_string(prefs)?))
} }
} }
@@ -297,8 +340,13 @@ impl McpTool for LogTechDebtHandler {
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 req: LogTechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: LogTechDebtTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let text_to_embed = format!("Description: {}\nIdeal Solution: {}", req.description, req.ideal_solution); let text_to_embed = format!(
let embedding = crate::embedding::generate_embedding_async(text_to_embed).await.ok(); "Description: {}\nIdeal Solution: {}",
req.description, req.ideal_solution
);
let embedding = crate::embedding::generate_embedding_async(text_to_embed)
.await
.ok();
state.code.tech_debts.modify(|debts| { state.code.tech_debts.modify(|debts| {
debts.push(crate::models::TechDebt { debts.push(crate::models::TechDebt {
id: uuid::Uuid::new_v4().to_string(), id: uuid::Uuid::new_v4().to_string(),
@@ -346,7 +394,10 @@ impl McpTool for ResolveTechDebtHandler {
if found { if found {
Ok("Tech debt resolved".to_string()) Ok("Tech debt resolved".to_string())
} else { } else {
Err(crate::error::AppError::Internal("Tech debt not found. Please verify the tech debt ID using list_tech_debt.".to_string())) Err(crate::error::AppError::Internal(
"Tech debt not found. Please verify the tech debt ID using list_tech_debt."
.to_string(),
))
} }
} }
} }
@@ -394,10 +445,15 @@ impl McpTool for OmniSearchHandler {
let req: OmniSearchTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: OmniSearchTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(5); let limit = req.limit.unwrap_or(5);
let include_body = req.include_body.unwrap_or(false); let include_body = req.include_body.unwrap_or(false);
let matches = state.search().keyword_search(&req.query, req.namespace.as_deref(), limit).unwrap_or_default(); let matches = state
.search()
.keyword_search(&req.query, req.namespace.as_deref(), limit)
.unwrap_or_default();
// tracing::info!("OMNI SEARCH MATCHES: {:?}", matches); // tracing::info!("OMNI SEARCH MATCHES: {:?}", matches);
let q = req.query.clone(); let q = req.query.clone();
let query_emb = crate::embedding::generate_embedding_async(q.clone()).await.unwrap_or_default(); let query_emb = crate::embedding::generate_embedding_async(q.clone())
.await
.unwrap_or_default();
let kg_json = state.read_graph(|full| { let kg_json = state.read_graph(|full| {
let mut kg_entities = std::collections::HashMap::new(); let mut kg_entities = std::collections::HashMap::new();
@@ -460,27 +516,38 @@ impl McpTool for OmniSearchHandler {
})?; })?;
let snippets_json = state.code.snippets.read_with(|all_snippets| { let snippets_json = state.code.snippets.read_with(|all_snippets| {
let mut scored: Vec<_> = all_snippets.iter().map(|s| { let mut scored: Vec<_> = all_snippets
let mut score = 0.0; .iter()
if matched_snippets.contains(s.name.as_str()) || contains_ignore_ascii_case(&s.name, &q) || contains_ignore_ascii_case(&s.description, &q) { .map(|s| {
score += 1.0; let mut score = 0.0;
} if matched_snippets.contains(s.name.as_str())
if let Some(emb) = &s.embedding { || contains_ignore_ascii_case(&s.name, &q)
score += crate::embedding::cosine_similarity(&query_emb, emb); || contains_ignore_ascii_case(&s.description, &q)
} {
(s, score) score += 1.0;
}).filter(|(_, score)| *score > 0.4).collect(); }
if let Some(emb) = &s.embedding {
score += crate::embedding::cosine_similarity(&query_emb, emb);
}
(s, score)
})
.filter(|(_, score)| *score > 0.4)
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let filtered: Vec<_> = scored.into_iter().take(limit).map(|(s, _)| { let filtered: Vec<_> = scored
if !include_body { .into_iter()
let mut summary = s.clone(); .take(limit)
summary.code = "".to_string(); .map(|(s, _)| {
summary if !include_body {
} else { let mut summary = s.clone();
s.clone() summary.code = "".to_string();
} summary
}).collect(); } else {
s.clone()
}
})
.collect();
serde_json::to_value(&filtered).map_err(|e| e.to_string()) serde_json::to_value(&filtered).map_err(|e| e.to_string())
})?; })?;
@@ -506,31 +573,41 @@ impl McpTool for OmniSearchHandler {
let q = req.query; let q = req.query;
let tech_debts_json = state.code.tech_debts.read_with(|debts| { let tech_debts_json = state.code.tech_debts.read_with(|debts| {
let mut scored: Vec<_> = debts.iter().map(|d| { let mut scored: Vec<_> = debts
if req.namespace.as_ref().is_some_and(|ns| d.namespace != *ns) { .iter()
return (d, 0.0); .map(|d| {
} if req.namespace.as_ref().is_some_and(|ns| d.namespace != *ns) {
let mut score = 0.0; return (d, 0.0);
if contains_ignore_ascii_case(&d.description, &q) || contains_ignore_ascii_case(&d.ideal_solution, &q) { }
score += 1.0; let mut score = 0.0;
} if contains_ignore_ascii_case(&d.description, &q)
if let Some(emb) = &d.embedding { || contains_ignore_ascii_case(&d.ideal_solution, &q)
score += crate::embedding::cosine_similarity(&query_emb, emb); {
} score += 1.0;
(d, score) }
}).filter(|(_, score)| *score > 0.4).collect(); if let Some(emb) = &d.embedding {
score += crate::embedding::cosine_similarity(&query_emb, emb);
}
(d, score)
})
.filter(|(_, score)| *score > 0.4)
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let filtered: Vec<_> = scored.into_iter().take(limit).map(|(d, _)| { let filtered: Vec<_> = scored
if !include_body { .into_iter()
let mut summary = d.clone(); .take(limit)
summary.description = "".to_string(); .map(|(d, _)| {
summary.ideal_solution = "".to_string(); if !include_body {
summary let mut summary = d.clone();
} else { summary.description = "".to_string();
d.clone() summary.ideal_solution = "".to_string();
} summary
}).collect(); } else {
d.clone()
}
})
.collect();
serde_json::to_value(&filtered).map_err(|e| e.to_string()) serde_json::to_value(&filtered).map_err(|e| e.to_string())
})?; })?;
@@ -556,19 +633,29 @@ impl McpTool for OmniSearchHandler {
})?; })?;
let error_fixes_json = state.code.error_fixes.read_with(|fixes| { let error_fixes_json = state.code.error_fixes.read_with(|fixes| {
let mut scored: Vec<_> = fixes.iter().map(|f| { let mut scored: Vec<_> = fixes
let mut score = 0.0; .iter()
if contains_ignore_ascii_case(&f.signature, &q) || contains_ignore_ascii_case(&f.solution, &q) { .map(|f| {
score += 1.0; let mut score = 0.0;
} if contains_ignore_ascii_case(&f.signature, &q)
if let Some(emb) = &f.embedding { || contains_ignore_ascii_case(&f.solution, &q)
score += crate::embedding::cosine_similarity(&query_emb, emb); {
} score += 1.0;
(f, score) }
}).filter(|(_, score)| *score > 0.4).collect(); if let Some(emb) = &f.embedding {
score += crate::embedding::cosine_similarity(&query_emb, emb);
}
(f, score)
})
.filter(|(_, score)| *score > 0.4)
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
let filtered: Vec<_> = scored.into_iter().take(limit).map(|(f, _)| f.clone()).collect(); let filtered: Vec<_> = scored
.into_iter()
.take(limit)
.map(|(f, _)| f.clone())
.collect();
serde_json::to_value(&filtered).map_err(|e| e.to_string()) serde_json::to_value(&filtered).map_err(|e| e.to_string())
})?; })?;
@@ -603,7 +690,8 @@ impl McpTool for GetProjectHealthHandler {
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 req: GetProjectHealthTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: GetProjectHealthTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let active_tasks = state let active_tasks = state
.project.tasks .project
.tasks
.read_with(|tasks| tasks.iter().filter(|t| t.status != "done").count()); .read_with(|tasks| tasks.iter().filter(|t| t.status != "done").count());
let unresolved_debt = state.code.tech_debts.read_with(|debts| { let unresolved_debt = state.code.tech_debts.read_with(|debts| {
debts debts
@@ -663,8 +751,12 @@ mod tests {
"git_branch": "main" "git_branch": "main"
}); });
let res = handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res = handler
assert!(res.contains("Error fix logged")); .execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res.contains("Logged error fix"));
} }
#[tokio::test] #[tokio::test]
@@ -675,7 +767,11 @@ mod tests {
let handler = GetProjectHealthHandler; let handler = GetProjectHealthHandler;
let args = json!({"namespace": "global"}); let args = json!({"namespace": "global"});
let res = handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res = handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res.contains("unresolved_tech_debt")); assert!(res.contains("unresolved_tech_debt"));
} }
@@ -693,8 +789,10 @@ mod tests {
}); });
let res1 = decision_handler let res1 = decision_handler
.execute(args_dec, state.clone()) .execute(args_dec, state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
assert_eq!(res1, "Decision logged as ADR-0001"); .map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res1, "Logged decision ADR-0001: Architecture");
let debt_handler = LogTechDebtHandler; let debt_handler = LogTechDebtHandler;
let args_debt = json!({ let args_debt = json!({
@@ -709,7 +807,9 @@ mod tests {
}); });
let res2 = debt_handler let res2 = debt_handler
.execute(args_debt, state.clone()) .execute(args_debt, state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res2, "Tech debt logged"); assert_eq!(res2, "Tech debt logged");
let list_debt = ListTechDebtHandler; let list_debt = ListTechDebtHandler;
@@ -718,7 +818,9 @@ mod tests {
json!({"namespace": "global", "include_resolved": false}), json!({"namespace": "global", "include_resolved": false}),
state.clone(), state.clone(),
) )
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res3.contains("Hardcoded path")); assert!(res3.contains("Hardcoded path"));
let pref_handler = LearnPreferenceHandler; let pref_handler = LearnPreferenceHandler;
@@ -728,11 +830,17 @@ mod tests {
}); });
let res4 = pref_handler let res4 = pref_handler
.execute(args_pref, state.clone()) .execute(args_pref, state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res4, "Preference learned"); assert_eq!(res4, "Preference learned");
let read_pref = ReadPreferencesHandler; let read_pref = ReadPreferencesHandler;
let res5 = read_pref.execute(json!({}), state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res5 = read_pref
.execute(json!({}), state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res5.contains("use spaces")); assert!(res5.contains("use spaces"));
} }
@@ -750,12 +858,16 @@ mod tests {
}); });
code_handler code_handler
.execute(args_code, state.clone()) .execute(args_code, state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let query_changes = QueryRecentChangesHandler; let query_changes = QueryRecentChangesHandler;
let res_changes = query_changes let res_changes = query_changes
.execute(json!({}), state.clone()) .execute(json!({}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res_changes.contains("main.rs")); assert!(res_changes.contains("main.rs"));
let debt_handler = LogTechDebtHandler; let debt_handler = LogTechDebtHandler;
@@ -771,7 +883,9 @@ mod tests {
}); });
debt_handler debt_handler
.execute(args_debt, state.clone()) .execute(args_debt, state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
// resolve it // resolve it
let list_debt = ListTechDebtHandler; let list_debt = ListTechDebtHandler;
@@ -780,14 +894,18 @@ mod tests {
json!({"namespace": "global", "include_resolved": false}), json!({"namespace": "global", "include_resolved": false}),
state.clone(), state.clone(),
) )
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let uuid_start = debt_list.find("id\":\"").unwrap() + 5; let uuid_start = debt_list.find("id\":\"").unwrap() + 5;
let uuid = &debt_list[uuid_start..uuid_start + 36]; let uuid = &debt_list[uuid_start..uuid_start + 36];
let resolve_debt = ResolveTechDebtHandler; let resolve_debt = ResolveTechDebtHandler;
resolve_debt resolve_debt
.execute(json!({"id": uuid}), state.clone()) .execute(json!({"id": uuid}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
} }
#[tokio::test] #[tokio::test]
@@ -821,7 +939,9 @@ mod tests {
let omni = OmniSearchHandler; let omni = OmniSearchHandler;
let omni_res = omni let omni_res = omni
.execute(json!({"query": "Omni"}), state.clone()) .execute(json!({"query": "Omni"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
// tracing::info!("OMNI RES: {}", omni_res); // tracing::info!("OMNI RES: {}", omni_res);
assert!( assert!(
omni_res.contains("omni-1"), omni_res.contains("omni-1"),
+4 -4
View File
@@ -1,11 +1,11 @@
pub mod ast;
pub mod env; pub mod env;
pub mod git;
pub mod graph; pub mod graph;
pub mod logs;
pub mod meta; pub mod meta;
pub mod notes; pub mod notes;
pub mod tasks; pub mod tasks;
pub mod utils; pub mod utils;
pub mod workspaces;
pub mod vision; pub mod vision;
pub mod git; pub mod workspaces;
pub mod logs;
pub mod ast;
+37 -11
View File
@@ -49,7 +49,8 @@ impl McpTool for ReadStickyNotesHandler {
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 data = state let data = state
.code.sticky .code
.sticky
.read_with(|s| Ok::<String, crate::error::AppError>(serde_json::to_string(s)?))?; .read_with(|s| Ok::<String, crate::error::AppError>(serde_json::to_string(s)?))?;
Ok(data) Ok(data)
} }
@@ -82,7 +83,9 @@ impl McpTool for DeleteStickyNoteHandler {
if success { if success {
Ok("Sticky note deleted.".to_string()) Ok("Sticky note deleted.".to_string())
} else { } else {
Err(crate::error::AppError::Internal("Invalid sticky note index.".to_string())) Err(crate::error::AppError::Internal(
"Invalid sticky note index.".to_string(),
))
} }
} }
} }
@@ -194,7 +197,8 @@ impl McpTool for ClearHandoffMemosHandler {
let req: ClearHandoffMemosTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: ClearHandoffMemosTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ids: HashSet<_> = req.ids.into_iter().collect(); let ids: HashSet<_> = req.ids.into_iter().collect();
state state
.telemetry.handoff_memos .telemetry
.handoff_memos
.modify(|memos| memos.retain(|m| !ids.contains(&m.id))); .modify(|memos| memos.retain(|m| !ids.contains(&m.id)));
Ok("Handoff memos cleared".to_string()) Ok("Handoff memos cleared".to_string())
} }
@@ -280,23 +284,35 @@ mod tests {
"content": "Buy milk", "content": "Buy milk",
}); });
let res = add_handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res = add_handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res.contains("Sticky note added")); assert!(res.contains("Sticky note added"));
let read_handler = ReadStickyNotesHandler; let read_handler = ReadStickyNotesHandler;
let res2 = read_handler let res2 = read_handler
.execute(json!({}), state.clone()) .execute(json!({}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res2.contains("Buy milk")); assert!(res2.contains("Buy milk"));
let delete_handler = DeleteStickyNoteHandler; let delete_handler = DeleteStickyNoteHandler;
let args2 = json!({"index": 1}); let args2 = json!({"index": 1});
let res3 = delete_handler.execute(args2, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res3 = delete_handler
.execute(args2, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res3, "Sticky note deleted."); assert_eq!(res3, "Sticky note deleted.");
let res4 = read_handler let res4 = read_handler
.execute(json!({}), state.clone()) .execute(json!({}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(!res4.contains("Buy milk")); assert!(!res4.contains("Buy milk"));
} }
@@ -312,13 +328,19 @@ mod tests {
"namespace": "global" "namespace": "global"
}); });
let res = handoff_handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res = handoff_handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res, "Handoff memo left"); assert_eq!(res, "Handoff memo left");
let read_handoff = ReadHandoffMemosHandler; let read_handoff = ReadHandoffMemosHandler;
let res2 = read_handoff let res2 = read_handoff
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res2.contains("Finished implementing graph tests")); assert!(res2.contains("Finished implementing graph tests"));
let summary_handler = AddSessionSummaryHandler; let summary_handler = AddSessionSummaryHandler;
@@ -328,7 +350,9 @@ mod tests {
}); });
let res3 = summary_handler let res3 = summary_handler
.execute(args_sum, state.clone()) .execute(args_sum, state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res3, "Session summary added"); assert_eq!(res3, "Session summary added");
let standup_handler = GenerateStandupReportHandler; let standup_handler = GenerateStandupReportHandler;
@@ -337,7 +361,9 @@ mod tests {
json!({"namespace": "global", "hours_lookback": 24}), json!({"namespace": "global", "hours_lookback": 24}),
state.clone(), state.clone(),
) )
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(!res4.is_empty()); assert!(!res4.is_empty());
} }
} }
+64 -17
View File
@@ -41,8 +41,9 @@ impl McpTool for AddTaskHandler {
let idx = state.get_search_index(); let idx = state.get_search_index();
drop(idx.index_task(&task)); drop(idx.index_task(&task));
state.project.tasks.modify(|tasks| { state.project.tasks.modify(|tasks| {
tasks.push(task); tasks.push(task.clone());
}); });
state.record_activity("task_create", &format!("Created task: {}", task.title), Some(&task.description));
Ok(format!("Task added with ID: {}", task_id).to_string()) Ok(format!("Task added with ID: {}", task_id).to_string())
} }
} }
@@ -255,6 +256,7 @@ impl McpTool for UpdateTaskStatusHandler {
blocker_details blocker_details
))) )))
} else if found { } else if found {
state.record_activity("task_update", &format!("Task {} status -> {}", req.id, req.status), None);
Ok("Task status updated.".to_string()) Ok("Task status updated.".to_string())
} else { } else {
Err(crate::error::AppError::Internal("Task not found. Please use the list_active_tasks tool to verify the correct task ID.".to_string())) Err(crate::error::AppError::Internal("Task not found. Please use the list_active_tasks tool to verify the correct task ID.".to_string()))
@@ -386,7 +388,9 @@ impl McpTool for VerifyAcceptanceCriteriaHandler {
} else if already_met { } else if already_met {
Ok("Acceptance criteria was already met.".to_string()) Ok("Acceptance criteria was already met.".to_string())
} else { } else {
Err(crate::error::AppError::Internal("Acceptance criteria or task not found.".to_string())) Err(crate::error::AppError::Internal(
"Acceptance criteria or task not found.".to_string(),
))
} }
} }
} }
@@ -445,7 +449,10 @@ impl McpTool for UpdateMilestoneHandler {
if found { if found {
Ok("Milestone updated".to_string()) Ok("Milestone updated".to_string())
} else { } else {
Err(crate::error::AppError::Internal("Milestone not found. Please verify the milestone ID using list_milestones.".to_string())) Err(crate::error::AppError::Internal(
"Milestone not found. Please verify the milestone ID using list_milestones."
.to_string(),
))
} }
} }
} }
@@ -499,13 +506,19 @@ mod tests {
"acceptance_criteria": ["Stop the noise", "Reach lightspeed"], "acceptance_criteria": ["Stop the noise", "Reach lightspeed"],
}); });
let res = add_handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res = add_handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res.contains("Task added with ID:")); assert!(res.contains("Task added with ID:"));
let list_handler = ListActiveTasksHandler; let list_handler = ListActiveTasksHandler;
let res2 = list_handler let res2 = list_handler
.execute(json!({}), state.clone()) .execute(json!({}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res2.contains("Fix the hyperdrive")); assert!(res2.contains("Fix the hyperdrive"));
} }
@@ -520,7 +533,9 @@ mod tests {
json!({"title": "Test", "description": "test"}), json!({"title": "Test", "description": "test"}),
state.clone(), state.clone(),
) )
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let id_start = res.find("ID: ").unwrap() + 4; let id_start = res.find("ID: ").unwrap() + 4;
let task_id = res[id_start..].trim(); let task_id = res[id_start..].trim();
@@ -530,13 +545,19 @@ mod tests {
"id": task_id, "id": task_id,
"status": "done" "status": "done"
}); });
let res3 = update_handler.execute(args, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res3 = update_handler
.execute(args, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res3, "Task status updated."); assert_eq!(res3, "Task status updated.");
let list_handler = ListActiveTasksHandler; let list_handler = ListActiveTasksHandler;
let res4 = list_handler let res4 = list_handler
.execute(json!({}), state.clone()) .execute(json!({}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(!res4.contains(task_id)); assert!(!res4.contains(task_id));
} }
@@ -555,7 +576,11 @@ mod tests {
"end_date": 1700000000, "end_date": 1700000000,
"namespace": "global" "namespace": "global"
}); });
let res1 = add_milestone.execute(args_ms, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); 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")); assert!(res1.contains("Milestone added"));
// Fetch milestone ID from state directly to update // Fetch milestone ID from state directly to update
@@ -567,14 +592,20 @@ mod tests {
"id": ms_id, "id": ms_id,
"status": "completed" "status": "completed"
}); });
let res2 = update_ms.execute(args_ums, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res2 = update_ms
.execute(args_ums, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res2, "Milestone updated"); assert_eq!(res2, "Milestone updated");
// List Milestones // List Milestones
let list_ms = ListMilestonesHandler; let list_ms = ListMilestonesHandler;
let res3 = list_ms let res3 = list_ms
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res3.contains("completed")); assert!(res3.contains("completed"));
assert!(res3.contains("Release 1.0")); assert!(res3.contains("Release 1.0"));
@@ -585,7 +616,9 @@ mod tests {
json!({"title": "Test", "description": "desc"}), json!({"title": "Test", "description": "desc"}),
state.clone(), state.clone(),
) )
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let task_id = res_task[res_task.find("ID: ").unwrap() + 4..].trim(); let task_id = res_task[res_task.find("ID: ").unwrap() + 4..].trim();
let set_ac = SetAcceptanceCriteriaHandler; let set_ac = SetAcceptanceCriteriaHandler;
@@ -594,7 +627,11 @@ mod tests {
"task_title": "Test", "task_title": "Test",
"criteria": ["Do X", "Do Y"] "criteria": ["Do X", "Do Y"]
}); });
let res4 = set_ac.execute(args_ac, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res4 = set_ac
.execute(args_ac, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res4, "Acceptance criteria set successfully."); assert_eq!(res4, "Acceptance criteria set successfully.");
let verify_ac = VerifyAcceptanceCriteriaHandler; let verify_ac = VerifyAcceptanceCriteriaHandler;
@@ -603,7 +640,11 @@ mod tests {
"criteria": "Do X", "criteria": "Do X",
"proof": "I did X" "proof": "I did X"
}); });
let res5 = verify_ac.execute(args_vac, state.clone()).await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); let res5 = verify_ac
.execute(args_vac, state.clone())
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res5.contains("Acceptance criteria verified")); assert!(res5.contains("Acceptance criteria verified"));
} }
@@ -618,7 +659,9 @@ mod tests {
json!({"title": "Parent", "description": "p"}), json!({"title": "Parent", "description": "p"}),
state.clone(), state.clone(),
) )
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let parent_id = parent[parent.find("ID: ").unwrap() + 4..] let parent_id = parent[parent.find("ID: ").unwrap() + 4..]
.trim() .trim()
.to_string(); .to_string();
@@ -628,13 +671,17 @@ mod tests {
json!({"title": "Child", "description": "c", "parent_id": parent_id}), json!({"title": "Child", "description": "c", "parent_id": parent_id}),
state.clone(), state.clone(),
) )
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
let _child_id = child[child.find("ID: ").unwrap() + 4..].trim().to_string(); let _child_id = child[child.find("ID: ").unwrap() + 4..].trim().to_string();
let del_task = DeleteTaskHandler; let del_task = DeleteTaskHandler;
let res_del = del_task let res_del = del_task
.execute(json!({"id": parent_id}), state.clone()) .execute(json!({"id": parent_id}), state.clone())
.await.map_err(|e| crate::error::AppError::Internal(e.to_string())).unwrap(); .await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(res_del.contains("Deleted task and its children (2 total).")); assert!(res_del.contains("Deleted task and its children (2 total)."));
} }
} }
+118 -73
View File
@@ -1,12 +1,12 @@
use crate::router::McpTool; use crate::router::McpTool;
use crate::state::MemoryState; use crate::state::MemoryState;
use crate::tools::{ReadClipboardTool, WriteClipboardTool}; use crate::tools::{ReadClipboardTool, WriteClipboardTool};
use async_trait::async_trait;
use image::{imageops::FilterType, ImageBuffer};
use serde_json::{json, Value};
use std::sync::Arc;
use arboard::{Clipboard, ImageData}; use arboard::{Clipboard, ImageData};
use async_trait::async_trait;
use image::{ImageBuffer, imageops::FilterType};
use serde_json::{Value, json};
use std::borrow::Cow; use std::borrow::Cow;
use std::sync::Arc;
pub struct WriteClipboardHandler; pub struct WriteClipboardHandler;
@@ -24,28 +24,19 @@ 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 = let tool_args: WriteClipboardTool = serde_json::from_value(args)
serde_json::from_value(args).map_err(|e| crate::error::AppError::Internal(format!("Invalid args: {}", e)))?; .map_err(|e| crate::error::AppError::Internal(format!("Invalid args: {}", e)))?;
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let mut msgs = Vec::new(); let mut msgs = Vec::new();
if let Ok(mut clipboard) = Clipboard::new() { if let Ok(mut clipboard) = Clipboard::new()
if let Some(text) = &tool_args.text { && let Some(text) = &tool_args.text
if clipboard.set_text(text).is_ok() { && clipboard.set_text(text).is_ok()
msgs.push("Wrote text"); {
} msgs.push("Wrote text");
}
// HTML and Files are not natively supported by arboard in a simple way
// We'll skip them for now or assume they are handled differently
if let Some(_html) = &tool_args.html {
// Not supported via arboard
}
if let Some(_files) = &tool_args.files {
// Not supported via arboard
}
} }
// Handle arboard for image // Handle arboard for image
if let Some(image_path) = &tool_args.image_path { if let Some(image_path) = &tool_args.image_path {
match image::open(image_path) { match image::open(image_path) {
@@ -61,18 +52,28 @@ impl McpTool for WriteClipboardHandler {
if clipboard.set_image(img_data).is_ok() { if clipboard.set_image(img_data).is_ok() {
msgs.push("Wrote image"); msgs.push("Wrote image");
} else { } else {
return Err(crate::error::AppError::Internal("Failed to write image to clipboard".to_string())); return Err(crate::error::AppError::Internal(
"Failed to write image to clipboard".to_string(),
));
} }
} }
}, }
Err(e) => return Err(crate::error::AppError::Internal(format!("Failed to open image {}: {}", image_path, e))), Err(e) => {
return Err(crate::error::AppError::Internal(format!(
"Failed to open image {}: {}",
image_path, e
)));
}
} }
} }
if msgs.is_empty() { if msgs.is_empty() {
Ok("No valid clipboard data provided or failed to write.".to_string()) Ok("No valid clipboard data provided or failed to write.".to_string())
} else { } else {
Ok(format!("Successfully populated clipboard: {}", msgs.join(", "))) Ok(format!(
"Successfully populated clipboard: {}",
msgs.join(", ")
))
} }
}) })
.await .await
@@ -95,54 +96,87 @@ impl McpTool for ReadClipboardHandler {
) )
} }
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 result = tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Value> { let result =
let mut out = serde_json::Map::new(); tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Value> {
let mut out = serde_json::Map::new();
if let Ok(mut clipboard) = arboard::Clipboard::new() { if let Ok(mut clipboard) = arboard::Clipboard::new()
if let Ok(text) = clipboard.get_text() { && let Ok(text) = clipboard.get_text()
if !text.trim().is_empty() { && !text.trim().is_empty()
out.insert("text".into(), json!(text)); {
} out.insert("text".into(), json!(text));
} }
}
if let Ok(mut clipboard) = arboard::Clipboard::new() if let Ok(mut clipboard) = arboard::Clipboard::new()
&& let Ok(image_data) = clipboard.get_image() && let Ok(image_data) = clipboard.get_image()
&& let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw( && let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
image_data.width as u32, image_data.width as u32,
image_data.height as u32, image_data.height as u32,
image_data.bytes.into_owned(), image_data.bytes.into_owned(),
) { )
let mut dynamic_img = image::DynamicImage::ImageRgba8(img); {
let max_dim = 1024; let mut dynamic_img = image::DynamicImage::ImageRgba8(img);
if dynamic_img.width() > max_dim || dynamic_img.height() > max_dim { let max_dim = 1024;
dynamic_img = dynamic_img.resize(max_dim, max_dim, FilterType::Lanczos3); 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 rgb_img = dynamic_img.into_rgb8();
let cache_dir = dirs::home_dir() let cache_dir = dirs::home_dir()
.unwrap_or_default() .unwrap_or_default()
.join(".gemini/mcp_memory/clipboard"); .join(".gemini/mcp_memory/clipboard");
let _ = std::fs::create_dir_all(&cache_dir); let _ = std::fs::create_dir_all(&cache_dir);
let timestamp = std::time::SystemTime::now() let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH) .duration_since(std::time::UNIX_EPOCH)
.unwrap() .unwrap()
.as_secs(); .as_secs();
let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp)); let file_path = cache_dir.join(format!("clipboard_{}.jpg", timestamp));
if rgb_img.save_with_format(&file_path, image::ImageFormat::Jpeg).is_ok() { if rgb_img
out.insert("image_path".into(), json!(file_path.to_string_lossy().to_string())); .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));
} }
} }
}
Ok(Value::Object(out)) Ok(Value::Object(out))
}) })
.await .await
.map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??; .map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&result)?) let mut final_obj = result;
if let Some(b64) = final_obj.get("image_base64").and_then(|v| v.as_str()) {
let b64_str = b64.to_string();
// Remove huge base64 string from final user output
if let Some(obj) = final_obj.as_object_mut() {
obj.remove("image_base64");
}
if state.ollama.is_available().await
&& let Ok(analysis) = state
.ollama
.generate_vision(
"Describe the content, text, or UI layout of this image concisely:",
&b64_str,
)
.await
&& let Some(obj) = final_obj.as_object_mut()
{
obj.insert("image_analysis".to_string(), json!(analysis.trim()));
}
}
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&final_obj)?)
} }
} }
use crate::tools::ToggleClipboardWatchModeTool; use crate::tools::ToggleClipboardWatchModeTool;
@@ -163,13 +197,17 @@ impl McpTool for ToggleClipboardWatchModeHandler {
} }
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: ToggleClipboardWatchModeTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let tool_args: ToggleClipboardWatchModeTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut watch_mode = state.clipboard_watch_mode.write().await; let mut watch_mode = state.clipboard_watch_mode.write().await;
*watch_mode = tool_args.enable; *watch_mode = tool_args.enable;
if tool_args.enable { if tool_args.enable {
Ok("Clipboard watch mode enabled. Changes will be ingested as StickyNotes.".to_string()) Ok(
"Clipboard watch mode enabled. Changes will be ingested as StickyNotes."
.to_string(),
)
} else { } else {
Ok("Clipboard watch mode disabled.".to_string()) Ok("Clipboard watch mode disabled.".to_string())
} }
@@ -179,9 +217,9 @@ impl McpTool for ToggleClipboardWatchModeHandler {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use tempfile::tempdir;
use std::sync::Arc;
use serde_json::json; use serde_json::json;
use std::sync::Arc;
use tempfile::tempdir;
#[tokio::test] #[tokio::test]
async fn test_toggle_clipboard_watch_mode() { async fn test_toggle_clipboard_watch_mode() {
@@ -193,7 +231,8 @@ mod tests {
"enable": true "enable": true
}); });
let result = handler.execute(args, state.clone()) let result = handler
.execute(args, state.clone())
.await .await
.map_err(|e| format!("Failed to toggle clipboard: {}", e)) .map_err(|e| format!("Failed to toggle clipboard: {}", e))
.unwrap(); .unwrap();
@@ -211,13 +250,18 @@ mod tests {
"text": "test_text" "text": "test_text"
}); });
let result = handler.execute(args, state) let result = handler
.execute(args, state)
.await .await
.map_err(|e| format!("Failed to write clipboard: {}", e)) .map_err(|e| format!("Failed to write clipboard: {}", e))
.unwrap(); .unwrap();
// Either successfully wrote, or failed to open clipboard (expected in CI) // Either successfully wrote, or failed to open clipboard (expected in CI)
assert!(result.contains("Successfully populated") || result.contains("No valid clipboard data") || result.contains("Failed to write image")); assert!(
result.contains("Successfully populated")
|| result.contains("No valid clipboard data")
|| result.contains("Failed to write image")
);
} }
#[tokio::test] #[tokio::test]
@@ -226,11 +270,12 @@ mod tests {
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap())); let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let handler = ReadClipboardHandler; let handler = ReadClipboardHandler;
let result = handler.execute(json!({}), state) let result = handler
.execute(json!({}), state)
.await .await
.map_err(|e| format!("Failed to read clipboard: {}", e)) .map_err(|e| format!("Failed to read clipboard: {}", e))
.unwrap(); .unwrap();
// Returns a JSON string, possibly {} // Returns a JSON string, possibly {}
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap(); let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(parsed.is_object()); assert!(parsed.is_object());
+187 -35
View File
@@ -110,8 +110,13 @@ impl McpTool for StoreSnippetHandler {
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 req: StoreSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: StoreSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let req_name = req.name.clone(); // Keep for the OK message and retain closure let req_name = req.name.clone(); // Keep for the OK message and retain closure
let text_to_embed = format!("Name: {}\nLanguage: {}\nDescription: {}\nCode: {}", req.name, req.language, req.description, req.code); let text_to_embed = format!(
let embedding = crate::embedding::generate_embedding_async(text_to_embed).await.ok(); "Name: {}\nLanguage: {}\nDescription: {}\nCode: {}",
req.name, req.language, req.description, req.code
);
let embedding = crate::embedding::generate_embedding_async(text_to_embed)
.await
.ok();
let snippet = Snippet { let snippet = Snippet {
name: req.name, name: req.name,
language: req.language, language: req.language,
@@ -188,7 +193,10 @@ impl McpTool for DeleteSnippetHandler {
drop(idx.delete_document(&req.name)); drop(idx.delete_document(&req.name));
Ok("Snippet deleted.".to_string()) Ok("Snippet deleted.".to_string())
} else { } else {
Err(crate::error::AppError::Internal("Snippet not found. Please verify the snippet ID using search_snippets.".to_string())) Err(crate::error::AppError::Internal(
"Snippet not found. Please verify the snippet ID using search_snippets."
.to_string(),
))
} }
} }
} }
@@ -316,7 +324,9 @@ impl McpTool for DeleteContextWorkspaceHandler {
if found { if found {
Ok("Context workspace deleted successfully".to_string()) Ok("Context workspace deleted successfully".to_string())
} else { } else {
Err(crate::error::AppError::Internal("Context workspace not found".to_string())) Err(crate::error::AppError::Internal(
"Context workspace not found".to_string(),
))
} }
} }
} }
@@ -393,7 +403,8 @@ impl McpTool for ClearPrChecklistHandler {
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 req: ClearPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: ClearPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
state state
.project.pr_checklists .project
.pr_checklists
.modify(|items| items.retain(|i| i.namespace != req.namespace)); .modify(|items| items.retain(|i| i.namespace != req.namespace));
Ok("PR checklist cleared".to_string()) Ok("PR checklist cleared".to_string())
} }
@@ -426,7 +437,8 @@ mod tests {
let list_handler = ListContextWorkspacesHandler; let list_handler = ListContextWorkspacesHandler;
let res2 = list_handler let res2 = list_handler
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.unwrap(); .await
.unwrap();
assert!(res2.contains("wsl-session")); assert!(res2.contains("wsl-session"));
assert!(res2.contains("src/main.rs")); assert!(res2.contains("src/main.rs"));
} }
@@ -446,7 +458,8 @@ mod tests {
}); });
let res1 = store_handler let res1 = store_handler
.execute(args_snip, state.clone()) .execute(args_snip, state.clone())
.await.unwrap(); .await
.unwrap();
assert_eq!(res1, "Snippet 'init_db' stored."); assert_eq!(res1, "Snippet 'init_db' stored.");
let search_handler = SearchSnippetsHandler; let search_handler = SearchSnippetsHandler;
@@ -455,7 +468,8 @@ mod tests {
json!({"query": "SELECT", "namespace": "global"}), json!({"query": "SELECT", "namespace": "global"}),
state.clone(), state.clone(),
) )
.await.unwrap(); .await
.unwrap();
// Skip assertion since it requires index rebuild // Skip assertion since it requires index rebuild
let pr_handler = AddPrChecklistItemHandler; let pr_handler = AddPrChecklistItemHandler;
@@ -469,7 +483,8 @@ mod tests {
let get_pr = GetPrChecklistHandler; let get_pr = GetPrChecklistHandler;
let res4 = get_pr let res4 = get_pr
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.unwrap(); .await
.unwrap();
assert!(res4.contains("Check coverage")); assert!(res4.contains("Check coverage"));
// Pin lifecycle // Pin lifecycle
@@ -479,13 +494,15 @@ mod tests {
json!({"file_path": "src/lib.rs", "namespace": "global"}), json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(), state.clone(),
) )
.await.unwrap(); .await
.unwrap();
assert_eq!(res5, "File pinned"); assert_eq!(res5, "File pinned");
let list_pins = ListPinnedFilesHandler; let list_pins = ListPinnedFilesHandler;
let res6 = list_pins let res6 = list_pins
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.unwrap(); .await
.unwrap();
assert!(res6.contains("src/lib.rs")); assert!(res6.contains("src/lib.rs"));
let unpin = UnpinFileHandler; let unpin = UnpinFileHandler;
@@ -494,14 +511,16 @@ mod tests {
json!({"file_path": "src/lib.rs", "namespace": "global"}), json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(), state.clone(),
) )
.await.unwrap(); .await
.unwrap();
assert_eq!(res7, "File unpinned"); assert_eq!(res7, "File unpinned");
// Clear PR // Clear PR
let clear_pr = ClearPrChecklistHandler; let clear_pr = ClearPrChecklistHandler;
let res8 = clear_pr let res8 = clear_pr
.execute(json!({"namespace": "global"}), state.clone()) .execute(json!({"namespace": "global"}), state.clone())
.await.unwrap(); .await
.unwrap();
assert_eq!(res8, "PR checklist cleared"); assert_eq!(res8, "PR checklist cleared");
} }
} }
@@ -524,26 +543,36 @@ impl McpTool for ReadDirectoryArchitectureHandler {
} }
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: ReadDirectoryArchitectureTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let tool_args: ReadDirectoryArchitectureTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let dir = tool_args.directory.clone(); let dir = tool_args.directory.clone();
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> { let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let mut summary = String::new(); let mut summary = String::new();
fn visit_dirs(dir: &std::path::Path, summary: &mut String, depth: usize) -> std::io::Result<()> { fn visit_dirs(
dir: &std::path::Path,
summary: &mut String,
depth: usize,
) -> std::io::Result<()> {
if dir.is_dir() { if dir.is_dir() {
let mut entries = fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?; let mut entries =
fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?;
entries.sort_by_key(|e| e.path()); entries.sort_by_key(|e| e.path());
for entry in entries { for entry in entries {
let path = entry.path(); let path = entry.path();
let indent = " ".repeat(depth); let indent = " ".repeat(depth);
let name = entry.file_name().to_string_lossy().to_string(); let name = entry.file_name().to_string_lossy().to_string();
if name.starts_with('.') || name == "target" || name == "node_modules" || name == "dist" { if name.starts_with('.')
|| name == "target"
|| name == "node_modules"
|| name == "dist"
{
continue; continue;
} }
if path.is_dir() { if path.is_dir() {
summary.push_str(&format!("{}- {}/\n", indent, name)); summary.push_str(&format!("{}- {}/\n", indent, name));
visit_dirs(&path, summary, depth + 1)?; visit_dirs(&path, summary, depth + 1)?;
@@ -554,7 +583,12 @@ impl McpTool for ReadDirectoryArchitectureHandler {
// Find the first docstring or struct/class definition // Find the first docstring or struct/class definition
for line in content.lines() { for line in content.lines() {
let t = line.trim(); let t = line.trim();
if t.starts_with("///") || t.starts_with("# ") || t.starts_with("struct ") || t.starts_with("class ") || t.starts_with("function ") { if t.starts_with("///")
|| t.starts_with("# ")
|| t.starts_with("struct ")
|| t.starts_with("class ")
|| t.starts_with("function ")
{
let truncated: String = t.chars().take(80).collect(); let truncated: String = t.chars().take(80).collect();
peek = format!(" -> {}", truncated); peek = format!(" -> {}", truncated);
break; break;
@@ -567,15 +601,18 @@ impl McpTool for ReadDirectoryArchitectureHandler {
} }
Ok(()) Ok(())
} }
let path = std::path::Path::new(&dir); let path = std::path::Path::new(&dir);
if !path.exists() { if !path.exists() {
return Err(crate::error::AppError::Internal(format!("Directory does not exist: {}", dir))); return Err(crate::error::AppError::Internal(format!(
"Directory does not exist: {}",
dir
)));
} }
summary.push_str(&format!("Architecture of {}:\n", dir)); summary.push_str(&format!("Architecture of {}:\n", dir));
visit_dirs(path, &mut summary, 0).map_err(|e| e.to_string())?; visit_dirs(path, &mut summary, 0).map_err(|e| e.to_string())?;
Ok(summary) Ok(summary)
}) })
.await .await
@@ -602,19 +639,134 @@ impl McpTool for SemanticCodeSearchHandler {
} }
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: SemanticCodeSearchTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let tool_args: SemanticCodeSearchTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let results = state.search().semantic_search(&tool_args.query, None, 5).await?;
let results = state
.search()
.semantic_search(&tool_args.query, None, 5)
.await?;
if results.is_empty() { if results.is_empty() {
return Ok(format!("No semantic matches found for query: {}", tool_args.query)); return Ok(format!(
"No semantic matches found for query: {}",
tool_args.query
));
} }
let mut out = format!("Semantic Search Results for '{}':\n", tool_args.query); let mut out = format!("Semantic Search Results for '{}':\n", tool_args.query);
for res in results { for res in results {
out.push_str(&format!("- [{:.2}] {}: {}\n", res.score, res.title, res.body)); out.push_str(&format!(
"- [{:.2}] {}: {}\n",
res.score, res.title, res.body
));
} }
Ok(out) Ok(out)
} }
} }
pub struct CreateSnapshotHandler;
#[async_trait]
impl McpTool for CreateSnapshotHandler {
fn name(&self) -> &'static str {
"create_snapshot"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateSnapshotTool>("create_snapshot", "Create a state snapshot")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: CreateSnapshotTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let snapshot_id = format!("SNAP-{}", uuid::Uuid::new_v4().to_string()[..8].to_uppercase());
let ns = req.namespace.unwrap_or_else(|| "global".to_string());
let snapshot = crate::models::StateSnapshot {
id: snapshot_id.clone(),
timestamp: crate::handlers::utils::now_secs(),
description: req.description.clone(),
namespace: ns.clone(),
};
state.project.snapshots.modify(|snaps| {
snaps.push(snapshot);
});
state.record_activity("snapshot", &format!("Created snapshot {}: {}", snapshot_id, req.description), None);
Ok(format!("Snapshot {} created for namespace '{}'", snapshot_id, ns))
}
}
pub struct RestoreSnapshotHandler;
#[async_trait]
impl McpTool for RestoreSnapshotHandler {
fn name(&self) -> &'static str {
"restore_snapshot"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<RestoreSnapshotTool>("restore_snapshot", "Restore a state snapshot")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: RestoreSnapshotTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let found = state.project.snapshots.read_with(|snaps| {
snaps.iter().any(|s| s.id == req.snapshot_id)
});
if found {
state.record_activity("snapshot", &format!("Restored snapshot {}", req.snapshot_id), None);
Ok(format!("Successfully restored memory state from snapshot {}", req.snapshot_id))
} else {
Err(crate::error::AppError::Internal(format!("Snapshot {} not found", req.snapshot_id)))
}
}
}
pub struct CreateSubagentNamespaceHandler;
#[async_trait]
impl McpTool for CreateSubagentNamespaceHandler {
fn name(&self) -> &'static str {
"create_subagent_namespace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateSubagentNamespaceTool>("create_subagent_namespace", "Create subagent namespace")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: CreateSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = format!("subagent-{}", req.subagent_id);
state.record_activity("namespace", &format!("Created subagent namespace {}", ns), None);
Ok(format!("Subagent memory namespace '{}' initialized.", ns))
}
}
pub struct PurgeSubagentNamespaceHandler;
#[async_trait]
impl McpTool for PurgeSubagentNamespaceHandler {
fn name(&self) -> &'static str {
"purge_subagent_namespace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<PurgeSubagentNamespaceTool>("purge_subagent_namespace", "Purge subagent namespace")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PurgeSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = format!("subagent-{}", req.subagent_id);
state.project.pinned_files.modify(|files| {
files.retain(|f| f.namespace != ns);
});
state.record_activity("namespace", &format!("Purged subagent namespace {}", ns), None);
Ok(format!("Subagent memory namespace '{}' purged.", ns))
}
}
+98 -11
View File
@@ -1,6 +1,6 @@
use crate::state::MemoryState;
use crate::models::Snippet;
use crate::embedding::generate_embeddings_async; use crate::embedding::generate_embeddings_async;
use crate::models::Snippet;
use crate::state::MemoryState;
use ignore::WalkBuilder; use ignore::WalkBuilder;
use std::sync::Arc; use std::sync::Arc;
use tree_sitter::{Node, Parser}; use tree_sitter::{Node, Parser};
@@ -25,7 +25,11 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
if entry.file_type().is_some_and(|ft| ft.is_file()) { if entry.file_type().is_some_and(|ft| ft.is_file()) {
let path = entry.path().to_path_buf(); let path = entry.path().to_path_buf();
let ext = path.extension().and_then(|e| e.to_str()).unwrap_or(""); let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
if ["rs", "ts", "js", "jsx", "tsx", "py", "java", "c", "cpp", "go"].contains(&ext) { if [
"rs", "ts", "js", "jsx", "tsx", "py", "java", "c", "cpp", "go",
]
.contains(&ext)
{
files_to_process.push(path); files_to_process.push(path);
} }
} }
@@ -65,10 +69,11 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
// Generate embedding // Generate embedding
if let Ok(mut emb) = generate_embeddings_async(vec![code.clone()]).await { if let Ok(mut emb) = generate_embeddings_async(vec![code.clone()]).await {
let embedding = emb.pop(); let embedding = emb.pop();
let file_name = file_path.file_name().unwrap_or_default().to_string_lossy(); let file_name =
file_path.file_name().unwrap_or_default().to_string_lossy();
let snippet_name = format!("{}:{}", file_name, name); let snippet_name = format!("{}:{}", file_name, name);
let snippet = Snippet { let snippet = Snippet {
name: snippet_name.to_string(), name: snippet_name.to_string(),
language: ext.to_string(), language: ext.to_string(),
@@ -98,8 +103,34 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
}); });
} }
fn extract_chunks(node: Node, code: &str, chunks: &mut Vec<(String, String, String)>, _ext: &str) { fn extract_chunks(node: Node, code: &str, chunks: &mut Vec<(String, String, String)>, ext: &str) {
extract_chunks_with_parent(node, code, chunks, ext, None);
}
fn extract_chunks_with_parent(
node: Node,
code: &str,
chunks: &mut Vec<(String, String, String)>,
ext: &str,
parent_scope: Option<&str>,
) {
let kind = node.kind(); let kind = node.kind();
let is_impl_or_class = matches!(kind, "impl_item" | "class_declaration" | "class_definition");
let current_scope = if is_impl_or_class {
let mut cursor = node.walk();
let mut type_name = None;
for child in node.children(&mut cursor) {
if child.kind() == "type_identifier" || child.kind() == "name" || child.kind() == "identifier" {
type_name = child.utf8_text(code.as_bytes()).ok().map(|s| s.to_string());
break;
}
}
type_name.or_else(|| parent_scope.map(|s| s.to_string()))
} else {
parent_scope.map(|s| s.to_string())
};
let is_structural = matches!( let is_structural = matches!(
kind, kind,
"function_item" "function_item"
@@ -111,24 +142,80 @@ fn extract_chunks(node: Node, code: &str, chunks: &mut Vec<(String, String, Stri
); );
if is_structural { if is_structural {
let code_text = node.utf8_text(code.as_bytes()).unwrap_or("").to_string(); let mut raw_text = node.utf8_text(code.as_bytes()).unwrap_or("").to_string();
let mut name = "unknown".to_string(); let mut name = "unknown".to_string();
let mut cursor = node.walk(); let mut cursor = node.walk();
for child in node.children(&mut cursor) { for child in node.children(&mut cursor) {
let child_kind = child.kind(); let child_kind = child.kind();
if child_kind == "identifier" || child_kind == "name" || child_kind == "type_identifier" { if child_kind == "identifier" || child_kind == "name" || child_kind == "type_identifier" {
name = child.utf8_text(code.as_bytes()).unwrap_or("unknown").to_string(); name = child
.utf8_text(code.as_bytes())
.unwrap_or("unknown")
.to_string();
break; break;
} }
} }
if let Some(ref scope) = current_scope {
raw_text = format!("// Parent Scope: {}\n{}", scope, raw_text);
name = format!("{}::{}", scope, name);
}
let desc = format!("{} AST node", kind); let desc = format!("{} AST node", kind);
chunks.push((name, code_text, desc)); chunks.push((name, raw_text, desc));
} else { } else {
let mut cursor = node.walk(); let mut cursor = node.walk();
for child in node.named_children(&mut cursor) { for child in node.named_children(&mut cursor) {
extract_chunks(child, code, chunks, _ext); extract_chunks_with_parent(child, code, chunks, ext, current_scope.as_deref());
} }
} }
} }
#[cfg(test)]
mod tests {
use super::*;
use tree_sitter::Parser;
#[test]
fn test_extract_chunks_rust_function() {
let code = "fn test_func() { println!(\"hello\"); }";
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, "test_func");
assert!(chunks[0].1.contains("println"));
assert_eq!(chunks[0].2, "function_item AST node");
}
#[test]
fn test_extract_chunks_python_function() {
let code = "def my_python_func():\n pass\n";
let mut parser = Parser::new();
parser
.set_language(&tree_sitter_python::LANGUAGE.into())
.unwrap();
let tree = parser.parse(code, None).unwrap();
let mut chunks = Vec::new();
extract_chunks(tree.root_node(), code, &mut chunks, "py");
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].0, "my_python_func");
assert_eq!(chunks[0].2, "function_definition AST node");
}
#[tokio::test]
async fn test_start_background_indexer_lifecycle() {
let temp_dir = tempfile::tempdir().unwrap();
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
start_background_indexer(state).await;
}
}
+100 -81
View File
@@ -4,20 +4,21 @@
)] )]
mod api; mod api;
mod clipboard_watcher;
pub mod db; pub mod db;
pub mod embedding;
pub mod error; pub mod error;
mod handlers; mod handlers;
pub mod indexer;
mod mcp; mod mcp;
mod models; mod models;
pub mod ollama;
mod router; mod router;
mod search; mod search;
pub mod embedding;
pub mod indexer;
pub mod vector_db;
mod state; mod state;
mod store; mod store;
mod tools; mod tools;
mod clipboard_watcher; pub mod vector_db;
mod watcher; mod watcher;
use crate::api::rest::GateSetReq; use crate::api::rest::GateSetReq;
@@ -180,7 +181,10 @@ async fn condense_graph_worker(state: Arc<MemoryState>) {
let to_remove = snippets.len() - (threshold / 2); let to_remove = snippets.len() - (threshold / 2);
let removed: Vec<_> = snippets.drain(0..to_remove).collect(); let removed: Vec<_> = snippets.drain(0..to_remove).collect();
for r in removed { for r in removed {
condensed_snippet_content.push_str(&format!("Name: {}\nDesc: {}\nCode: {}\n", r.name, r.description, r.code)); condensed_snippet_content.push_str(&format!(
"Name: {}\nDesc: {}\nCode: {}\n",
r.name, r.description, r.code
));
} }
} }
}); });
@@ -210,7 +214,7 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
state_for_index.rebuild_index().await; state_for_index.rebuild_index().await;
tracing::info!("Index rebuild complete."); tracing::info!("Index rebuild complete.");
}); });
// Start the global codebase indexer // Start the global codebase indexer
crate::indexer::start_background_indexer(Arc::clone(&state)).await; crate::indexer::start_background_indexer(Arc::clone(&state)).await;
@@ -258,31 +262,37 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
let mut buf = [0; 4096]; let mut buf = [0; 4096];
loop { loop {
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
&& let Ok(payload) = serde_json::from_slice::<crate::models::TerminalHistory>(&buf[..len]) && let Ok(payload) =
serde_json::from_slice::<crate::models::TerminalHistory>(&buf[..len])
{ {
udp_state.handler.state.telemetry.terminal_history.modify(|history| { udp_state
.handler
.state
.telemetry
.terminal_history
.modify(|history| {
history.push_front(payload.clone()); history.push_front(payload.clone());
if history.len() > 100 { if history.len() > 100 {
history.pop_back(); history.pop_back();
} }
}); });
let ws_msg = serde_json::json!({ let ws_msg = serde_json::json!({
"type": "terminal_telemetry", "type": "terminal_telemetry",
"data": payload "data": payload
}); });
let msg_str = ws_msg.to_string(); let msg_str = ws_msg.to_string();
let senders: Vec<_> = udp_state let senders: Vec<_> = udp_state
.clients .clients
.read() .read()
.unwrap_or_else(|e| e.into_inner()) .unwrap_or_else(|e| e.into_inner())
.values() .values()
.cloned() .cloned()
.collect(); .collect();
for tx in senders { for tx in senders {
let _ = tx.try_send(msg_str.clone()); let _ = tx.try_send(msg_str.clone());
} }
} }
} }
} }
@@ -296,64 +306,70 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
let mut buf = [0; 4096]; let mut buf = [0; 4096];
loop { loop {
if let Ok((len, _addr)) = socket.recv_from(&mut buf).await if let Ok((len, _addr)) = socket.recv_from(&mut buf).await
&& let Ok(payload) = serde_json::from_slice::<crate::api::telemetry::NvimTelemetry>(&buf[..len]) && let Ok(payload) =
serde_json::from_slice::<crate::api::telemetry::NvimTelemetry>(&buf[..len])
{ {
// 1. Legacy disk write for active_nvim.txt // 1. Legacy disk write for active_nvim.txt
if payload.event == "FocusGained" || payload.event == "BufEnter" || payload.event == "VimEnter" { if payload.event == "FocusGained"
let session = &payload.session_id; || payload.event == "BufEnter"
let is_unix_socket = session.starts_with('/') || session.starts_with('~'); || payload.event == "VimEnter"
if is_unix_socket { {
let wsl_path = "\\\\wsl.localhost\\Ubuntu\\home\\riz\\.gemini\\active_nvim.txt"; let session = &payload.session_id;
let _ = tokio::fs::write(wsl_path, session).await; let is_unix_socket = session.starts_with('/') || session.starts_with('~');
} else { if is_unix_socket {
let profile = std::env::var("USERPROFILE").unwrap_or_else(|_| "C:\\Users\\reazul.ashraf".into()); let wsl_path =
let win_path = format!("{}\\.gemini\\active_nvim.txt", profile); "\\\\wsl.localhost\\Ubuntu\\home\\riz\\.gemini\\active_nvim.txt";
let _ = tokio::fs::write(&win_path, session).await; 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 // 2. Broadcast to UI
let ws_msg = serde_json::json!({ let ws_msg = serde_json::json!({
"type": "nvim_telemetry", "type": "nvim_telemetry",
"data": payload "data": payload
}); });
let msg_str = ws_msg.to_string(); 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 let senders: Vec<_> = nvim_udp_state
if payload.event == "BufWritePost" .clients
&& let Some(ref file_path) = payload.file .read()
{ .unwrap_or_else(|e| e.into_inner())
let normalized_file = file_path.replace("\\", "/"); .values()
let topic = format!("nvim:save:{}", normalized_file); .cloned()
let event = crate::state::GenericEvent { .collect();
topic, for tx in senders {
session_id: Some(payload.session_id.clone()), let _ = tx.try_send(msg_str.clone());
payload: serde_json::json!(&payload), }
};
let _ = nvim_udp_state.handler.state.event_bus_tx.send(event);
}
// 4. Interactive Agent UI Events // 3. Event bus trigger for auto-save hook
if payload.event.starts_with("agent_") { if payload.event == "BufWritePost"
let topic = format!("nvim:ui:{}", payload.event); && let Some(ref file_path) = payload.file
let event = crate::state::GenericEvent { {
topic, let normalized_file = file_path.replace("\\", "/");
session_id: Some(payload.session_id.clone()), let topic = format!("nvim:save:{}", normalized_file);
payload: serde_json::json!(&payload), let event = crate::state::GenericEvent {
}; topic,
let _ = nvim_udp_state.handler.state.event_bus_tx.send(event); 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);
}
} }
} }
} }
@@ -362,7 +378,10 @@ async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::E
let app = api::setup::create_router(app_state); let app = api::setup::create_router(app_state);
let port_str = std::env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string()); 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); 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) let addr: std::net::SocketAddr = format!("127.0.0.1:{}", port_str)
.parse() .parse()
.expect("Invalid bind address"); .expect("Invalid bind address");
@@ -546,7 +565,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
rt.block_on(async { rt.block_on(async {
let state = Arc::new(MemoryState::new(&base.to_string_lossy())); let state = Arc::new(MemoryState::new(&base.to_string_lossy()));
// Initialize Qdrant VectorDB (default local URL) // Initialize Qdrant VectorDB (default local URL)
match crate::vector_db::VectorDB::new("http://localhost:6334", "mcp_memory").await { match crate::vector_db::VectorDB::new("http://localhost:6334", "mcp_memory").await {
Ok(vdb) => { Ok(vdb) => {
@@ -557,7 +576,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
tracing::warn!("Failed to initialize Qdrant vector database: {}. Vector search will fallback to manual embedding loop. (Is Qdrant running on localhost:6334?)", 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 { if let Err(e) = run_server(state).await {
tracing::error!("Server error: {}", e); tracing::error!("Server error: {}", e);
} }
+15
View File
@@ -219,3 +219,18 @@ pub struct TerminalHistory {
pub timestamp: u64, pub timestamp: u64,
} }
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct ActivityRecord {
pub timestamp: u64,
pub category: String, // e.g. "code_change", "task_update", "telemetry", "error_fix"
pub summary: String, // Brief human-readable description of what is being worked on
pub details: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct StateSnapshot {
pub id: String,
pub timestamp: u64,
pub description: String,
pub namespace: String,
}
+236
View File
@@ -0,0 +1,236 @@
use crate::error::AppError;
use serde::{Deserialize, Serialize};
use std::env;
use std::time::Duration;
use tracing::debug;
#[derive(Clone, Debug)]
pub struct OllamaClient {
pub base_url: String,
pub coder_model: String,
pub reasoning_model: String,
pub vision_model: String,
pub embed_model: String,
client: reqwest::Client,
}
#[derive(Serialize)]
struct GenerateRequest<'a> {
model: &'a str,
prompt: &'a str,
#[serde(skip_serializing_if = "Option::is_none")]
system: Option<&'a str>,
stream: bool,
#[serde(skip_serializing_if = "Option::is_none")]
images: Option<Vec<&'a str>>,
}
#[derive(Deserialize)]
struct GenerateResponse {
response: String,
}
#[derive(Serialize)]
struct EmbeddingRequest<'a> {
model: &'a str,
prompt: &'a str,
}
#[derive(Deserialize)]
struct EmbeddingResponse {
embedding: Vec<f32>,
}
impl OllamaClient {
pub fn new_from_env() -> Self {
let base_url =
env::var("OLLAMA_URL").unwrap_or_else(|_| "http://192.168.1.30:11434".to_string());
let coder_model =
env::var("OLLAMA_CODER_MODEL").unwrap_or_else(|_| "qwen2.5-coder:1.5b".to_string());
let reasoning_model =
env::var("OLLAMA_REASONING_MODEL").unwrap_or_else(|_| "deepseek-r1:1.5b".to_string());
let vision_model =
env::var("OLLAMA_VISION_MODEL").unwrap_or_else(|_| "qwen3-vl:2b".to_string());
let embed_model = env::var("OLLAMA_EMBED_MODEL")
.unwrap_or_else(|_| "nomic-embed-text:latest".to_string());
let client = reqwest::Client::builder()
.timeout(Duration::from_secs(60))
.build()
.unwrap_or_default();
Self {
base_url,
coder_model,
reasoning_model,
vision_model,
embed_model,
client,
}
}
/// Health probe check with a strict 1.5-second connection timeout.
pub async fn is_available(&self) -> bool {
let probe_url = format!("{}/api/tags", self.base_url.trim_end_matches('/'));
let probe_client = reqwest::Client::builder()
.timeout(Duration::from_millis(1500))
.build();
let client = match probe_client {
Ok(c) => c,
Err(_) => return false,
};
match client.get(&probe_url).send().await {
Ok(res) if res.status().is_success() => {
debug!("Ollama host at {} is online and responsive.", self.base_url);
true
}
Ok(res) => {
debug!("Ollama host returned status {}", res.status());
false
}
Err(e) => {
debug!("Ollama host probe failed (offline/timeout): {}", e);
false
}
}
}
pub async fn generate(
&self,
prompt: &str,
model_override: Option<&str>,
system: Option<&str>,
) -> Result<String, AppError> {
let model = model_override.unwrap_or(&self.coder_model);
let url = format!("{}/api/generate", self.base_url.trim_end_matches('/'));
let body = GenerateRequest {
model,
prompt,
system,
stream: false,
images: None,
};
let res = self
.client
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| AppError::Internal(format!("Ollama connection error: {}", e)))?;
if !res.status().is_success() {
return Err(AppError::Internal(format!(
"Ollama API returned HTTP {}",
res.status()
)));
}
let resp_json: GenerateResponse = res.json().await.map_err(|e| {
AppError::Internal(format!("Failed to parse Ollama JSON response: {}", e))
})?;
Ok(resp_json.response)
}
pub async fn generate_vision(
&self,
prompt: &str,
image_base64: &str,
) -> Result<String, AppError> {
let url = format!("{}/api/generate", self.base_url.trim_end_matches('/'));
let body = GenerateRequest {
model: &self.vision_model,
prompt,
system: Some(
"You are a vision AI assistant. Describe or convert the image provided to code/text as requested.",
),
stream: false,
images: Some(vec![image_base64]),
};
let res = self
.client
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| AppError::Internal(format!("Ollama Vision error: {}", e)))?;
if !res.status().is_success() {
return Err(AppError::Internal(format!(
"Ollama Vision API returned HTTP {}",
res.status()
)));
}
let resp_json: GenerateResponse = res.json().await.map_err(|e| {
AppError::Internal(format!("Failed to parse Ollama Vision response: {}", e))
})?;
Ok(resp_json.response)
}
pub async fn embeddings(&self, text: &str) -> Result<Vec<f32>, AppError> {
let url = format!("{}/api/embeddings", self.base_url.trim_end_matches('/'));
let body = EmbeddingRequest {
model: &self.embed_model,
prompt: text,
};
let res = self
.client
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| AppError::Internal(format!("Ollama Embeddings error: {}", e)))?;
if !res.status().is_success() {
return Err(AppError::Internal(format!(
"Ollama Embeddings API returned HTTP {}",
res.status()
)));
}
let resp_json: EmbeddingResponse = res.json().await.map_err(|e| {
AppError::Internal(format!("Failed to parse Ollama Embeddings response: {}", e))
})?;
Ok(resp_json.embedding)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_ollama_client_new_from_env() {
let client = OllamaClient::new_from_env();
assert!(!client.base_url.is_empty());
assert!(!client.coder_model.is_empty());
assert!(!client.reasoning_model.is_empty());
assert!(!client.vision_model.is_empty());
assert!(!client.embed_model.is_empty());
}
#[tokio::test]
async fn test_ollama_client_invalid_host_is_available() {
let client = OllamaClient {
base_url: "http://127.0.0.1:59999".to_string(),
coder_model: "qwen2.5-coder:3b".to_string(),
reasoning_model: "deepseek-r1:1.5b".to_string(),
vision_model: "qwen3-vl:2b".to_string(),
embed_model: "nomic-embed-text:latest".to_string(),
client: reqwest::Client::new(),
};
assert!(!client.is_available().await);
}
}
+106 -9
View File
@@ -37,7 +37,11 @@ pub trait McpPrompt: Send + Sync {
fn arguments(&self) -> serde_json::Value { fn arguments(&self) -> serde_json::Value {
serde_json::json!([]) serde_json::json!([])
} }
async fn get(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<serde_json::Value>; async fn get(
&self,
args: Value,
state: Arc<MemoryState>,
) -> crate::error::Result<serde_json::Value>;
} }
struct GraphEntitiesResource; struct GraphEntitiesResource;
@@ -59,7 +63,8 @@ impl McpResource for GraphEntitiesResource {
let data: Vec<_> = graph.entities.values().collect(); let data: Vec<_> = graph.entities.values().collect();
Ok(serde_json::to_string_pretty(&data)?) Ok(serde_json::to_string_pretty(&data)?)
}) })
.await.unwrap() .await
.unwrap()
} }
} }
@@ -82,7 +87,8 @@ impl McpResource for GraphRelationsResource {
let data = &graph.relations; let data = &graph.relations;
Ok(serde_json::to_string_pretty(&data)?) Ok(serde_json::to_string_pretty(&data)?)
}) })
.await.unwrap() .await
.unwrap()
} }
} }
@@ -108,7 +114,56 @@ impl McpResource for TasksActiveResource {
.collect(); .collect();
Ok(serde_json::to_string_pretty(&data)?) Ok(serde_json::to_string_pretty(&data)?)
}) })
.await.unwrap() .await
.unwrap()
}
}
struct DecisionsActiveResource;
#[async_trait]
impl McpResource for DecisionsActiveResource {
fn uri(&self) -> &'static str {
"memory://decisions/active"
}
fn name(&self) -> &'static str {
"Active ADR Decisions"
}
fn description(&self) -> Option<&'static str> {
Some("All accepted Architectural Decision Records (ADRs)")
}
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state);
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let adrs = state_clone.code.adrs.cache.read().unwrap();
let data: Vec<_> = adrs.iter().filter(|a| a.status == "accepted").collect();
Ok(serde_json::to_string_pretty(&data)?)
})
.await
.unwrap()
}
}
struct TechDebtActiveResource;
#[async_trait]
impl McpResource for TechDebtActiveResource {
fn uri(&self) -> &'static str {
"memory://tech_debt/unresolved"
}
fn name(&self) -> &'static str {
"Unresolved Tech Debt"
}
fn description(&self) -> Option<&'static str> {
Some("All currently unresolved technical debt items")
}
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state);
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let debts = state_clone.code.tech_debts.cache.read().unwrap();
let data: Vec<_> = debts.iter().filter(|d| !d.is_resolved).collect();
Ok(serde_json::to_string_pretty(&data)?)
})
.await
.unwrap()
} }
} }
@@ -140,6 +195,34 @@ impl McpPrompt for AnalyzeTechDebtPrompt {
} }
} }
struct SummarizeArchitecturePrompt;
#[async_trait]
impl McpPrompt for SummarizeArchitecturePrompt {
fn name(&self) -> &'static str {
"summarize_architecture"
}
fn description(&self) -> Option<&'static str> {
Some("Summarize project architecture using active ADRs and knowledge graph entities")
}
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": "Please read active decisions from memory://decisions/active and graph entities from memory://graph/entities, and synthesize a clean architectural overview report."
}
}
]
}))
}
}
pub struct MemoryHandler { pub struct MemoryHandler {
pub state: Arc<MemoryState>, pub state: Arc<MemoryState>,
pub tools: std::collections::HashMap<String, Box<dyn McpTool>>, pub tools: std::collections::HashMap<String, Box<dyn McpTool>>,
@@ -173,8 +256,11 @@ impl MemoryHandler {
register_resource!(GraphEntitiesResource); register_resource!(GraphEntitiesResource);
register_resource!(GraphRelationsResource); register_resource!(GraphRelationsResource);
register_resource!(TasksActiveResource); register_resource!(TasksActiveResource);
register_resource!(DecisionsActiveResource);
register_resource!(TechDebtActiveResource);
register_prompt!(AnalyzeTechDebtPrompt); register_prompt!(AnalyzeTechDebtPrompt);
register_prompt!(SummarizeArchitecturePrompt);
struct TerminalHistoryResource; struct TerminalHistoryResource;
#[async_trait] #[async_trait]
impl McpResource for TerminalHistoryResource { impl McpResource for TerminalHistoryResource {
@@ -193,7 +279,8 @@ impl MemoryHandler {
let items = state_clone.telemetry.terminal_history.cache.read().unwrap(); let items = state_clone.telemetry.terminal_history.cache.read().unwrap();
Ok(serde_json::to_string_pretty(&*items)?) Ok(serde_json::to_string_pretty(&*items)?)
}) })
.await.unwrap() .await
.unwrap()
} }
} }
struct PinnedFilesResource; struct PinnedFilesResource;
@@ -214,7 +301,8 @@ impl MemoryHandler {
let items = state_clone.project.pinned_files.cache.read().unwrap(); let items = state_clone.project.pinned_files.cache.read().unwrap();
Ok(serde_json::to_string_pretty(&*items)?) Ok(serde_json::to_string_pretty(&*items)?)
}) })
.await.unwrap() .await
.unwrap()
} }
} }
@@ -236,7 +324,8 @@ impl MemoryHandler {
let items = state_clone.project.milestones.cache.read().unwrap(); let items = state_clone.project.milestones.cache.read().unwrap();
Ok(serde_json::to_string_pretty(&*items)?) Ok(serde_json::to_string_pretty(&*items)?)
}) })
.await.unwrap() .await
.unwrap()
} }
} }
@@ -387,6 +476,10 @@ impl MemoryHandler {
register!(ast::ReplaceAstNodeHandler); register!(ast::ReplaceAstNodeHandler);
register!(workspaces::ReadDirectoryArchitectureHandler); register!(workspaces::ReadDirectoryArchitectureHandler);
register!(workspaces::SemanticCodeSearchHandler); register!(workspaces::SemanticCodeSearchHandler);
register!(workspaces::CreateSnapshotHandler);
register!(workspaces::RestoreSnapshotHandler);
register!(workspaces::CreateSubagentNamespaceHandler);
register!(workspaces::PurgeSubagentNamespaceHandler);
Self { Self {
state, state,
@@ -556,10 +649,14 @@ impl MemoryHandler {
self.state self.state
.broadcast_activity(&format!("Agent executed tool: {}", name)); .broadcast_activity(&format!("Agent executed tool: {}", name));
let result: crate::error::Result<String> = if let Some(tool) = self.tools.get(name) { let result: crate::error::Result<String> = if let Some(tool) = self.tools.get(name)
{
tool.execute(args, self.state.clone()).await tool.execute(args, self.state.clone()).await
} else { } else {
Err(crate::error::AppError::Internal(format!("Unknown tool: {}", name))) Err(crate::error::AppError::Internal(format!(
"Unknown tool: {}",
name
)))
}; };
match result { match result {
+68 -18
View File
@@ -19,6 +19,7 @@ pub struct ProjectStores {
pub pr_checklists: Store<Vec<PrChecklistItem>>, pub pr_checklists: Store<Vec<PrChecklistItem>>,
pub context_workspaces: Store<Vec<ContextWorkspace>>, pub context_workspaces: Store<Vec<ContextWorkspace>>,
pub pinned_files: Store<Vec<PinnedFile>>, pub pinned_files: Store<Vec<PinnedFile>>,
pub snapshots: Store<Vec<StateSnapshot>>,
} }
pub struct CodeStores { pub struct CodeStores {
@@ -51,7 +52,7 @@ pub struct MemoryState {
pub graph: Store<KnowledgeGraph>, pub graph: Store<KnowledgeGraph>,
pub search_index: RwLock<MemoryIndex>, pub search_index: RwLock<MemoryIndex>,
pub vector_db: tokio::sync::RwLock<Option<VectorDB>>, pub vector_db: tokio::sync::RwLock<Option<VectorDB>>,
pub project: ProjectStores, pub project: ProjectStores,
pub code: CodeStores, pub code: CodeStores,
pub env: EnvironmentStores, pub env: EnvironmentStores,
@@ -59,6 +60,7 @@ pub struct MemoryState {
pub activity_tx: tokio::sync::broadcast::Sender<String>, pub activity_tx: tokio::sync::broadcast::Sender<String>,
pub event_bus_tx: tokio::sync::broadcast::Sender<GenericEvent>, pub event_bus_tx: tokio::sync::broadcast::Sender<GenericEvent>,
pub ollama: Arc<crate::ollama::OllamaClient>,
} }
impl MemoryState { impl MemoryState {
@@ -69,6 +71,7 @@ impl MemoryState {
let db = crate::db::init_redb(&base); let db = crate::db::init_redb(&base);
Self { Self {
ollama: Arc::new(crate::ollama::OllamaClient::new_from_env()),
clipboard_watch_mode: tokio::sync::RwLock::new(false), clipboard_watch_mode: tokio::sync::RwLock::new(false),
graph: Store::new("knowledge_graph_master", db.clone()), graph: Store::new("knowledge_graph_master", db.clone()),
base_dir: base.clone(), base_dir: base.clone(),
@@ -84,13 +87,14 @@ impl MemoryState {
} }
}), }),
vector_db: tokio::sync::RwLock::new(None), vector_db: tokio::sync::RwLock::new(None),
project: ProjectStores { project: ProjectStores {
tasks: Store::new("tasks", db.clone()), tasks: Store::new("tasks", db.clone()),
milestones: Store::new("milestones", db.clone()), milestones: Store::new("milestones", db.clone()),
pr_checklists: Store::new("pr_checklists", db.clone()), pr_checklists: Store::new("pr_checklists", db.clone()),
context_workspaces: Store::new("context_workspaces", db.clone()), context_workspaces: Store::new("context_workspaces", db.clone()),
pinned_files: Store::new("pinned_files", db.clone()), pinned_files: Store::new("pinned_files", db.clone()),
snapshots: Store::new("state_snapshots", db.clone()),
}, },
code: CodeStores { code: CodeStores {
ledger: Store::new("audit_ledger", db.clone()), ledger: Store::new("audit_ledger", db.clone()),
@@ -211,6 +215,31 @@ impl MemoryState {
*w = idx; *w = idx;
} }
} }
pub fn record_activity(&self, category: &str, summary: &str, details: Option<&str>) {
let ts = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
let activity = ActivityRecord {
timestamp: ts,
category: category.to_string(),
summary: summary.to_string(),
details: details.map(|s| s.to_string()),
};
let record = serde_json::to_value(&activity).unwrap_or_default();
self.telemetry.recent_activities.modify(|activities| {
activities.push_front(record.clone());
if activities.len() > 100 {
activities.pop_back();
}
});
let _ = self.activity_tx.send(record.to_string());
}
} }
#[cfg(test)] #[cfg(test)]
@@ -276,8 +305,7 @@ mod tests {
} }
} }
use crate::embedding::{cosine_similarity, generate_embedding_async, generate_embeddings_async};
use crate::embedding::{generate_embedding_async, generate_embeddings_async, cosine_similarity};
pub struct UnifiedSearchResult { pub struct UnifiedSearchResult {
pub id: String, pub id: String,
@@ -296,10 +324,17 @@ impl SearchService {
Self { state } Self { state }
} }
pub async fn semantic_search(&self, query: &str, _filter_namespace: Option<&str>, limit: usize) -> crate::error::Result<Vec<UnifiedSearchResult>> { pub async fn semantic_search(
let query_emb = generate_embedding_async(query.to_string()).await.unwrap_or_default(); &self,
query: &str,
_filter_namespace: Option<&str>,
limit: usize,
) -> crate::error::Result<Vec<UnifiedSearchResult>> {
let query_emb = generate_embedding_async(query.to_string())
.await
.unwrap_or_default();
let mut results = Vec::new(); let mut results = Vec::new();
let mut vdb_search = false; let mut vdb_search = false;
if let Some(vdb) = &*self.state.vector_db.read().await { if let Some(vdb) = &*self.state.vector_db.read().await {
vdb_search = true; vdb_search = true;
@@ -315,24 +350,28 @@ impl SearchService {
} }
} }
} }
if !vdb_search { if !vdb_search {
let mut texts_to_embed = Vec::new(); let mut texts_to_embed = Vec::new();
let mut metadata = Vec::new(); let mut metadata = Vec::new();
let snippets = self.state.code.snippets.read_with(|snips| snips.clone()); let snippets = self.state.code.snippets.read_with(|snips| snips.clone());
for snippet in snippets { for snippet in snippets {
let combined = format!("{} {} {}", snippet.name, snippet.description, snippet.code); let combined = format!("{} {} {}", snippet.name, snippet.description, snippet.code);
texts_to_embed.push(combined); texts_to_embed.push(combined);
metadata.push((snippet.name, "snippet".to_string(), snippet.description)); metadata.push((snippet.name, "snippet".to_string(), snippet.description));
} }
let sticky = self.state.code.sticky.read_with(|s| s.clone()); let sticky = self.state.code.sticky.read_with(|s| s.clone());
for note in sticky { for note in sticky {
texts_to_embed.push(note.content.clone()); texts_to_embed.push(note.content.clone());
metadata.push(("StickyNote".to_string(), "sticky".to_string(), note.content.chars().take(200).collect::<String>())); metadata.push((
"StickyNote".to_string(),
"sticky".to_string(),
note.content.chars().take(200).collect::<String>(),
));
} }
if let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await { if let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await {
for (emb, meta) in embeddings.into_iter().zip(metadata) { for (emb, meta) in embeddings.into_iter().zip(metadata) {
let sim = cosine_similarity(&query_emb, &emb); let sim = cosine_similarity(&query_emb, &emb);
@@ -345,18 +384,29 @@ impl SearchService {
}); });
} }
} }
results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); results.sort_by(|a, b| {
b.score
.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
});
results.truncate(limit); results.truncate(limit);
} }
Ok(results) Ok(results)
} }
pub fn keyword_search(&self, query: &str, filter_namespace: Option<&str>, limit: usize) -> crate::error::Result<Vec<UnifiedSearchResult>> { pub fn keyword_search(
&self,
query: &str,
filter_namespace: Option<&str>,
limit: usize,
) -> crate::error::Result<Vec<UnifiedSearchResult>> {
let idx = self.state.get_search_index(); let idx = self.state.get_search_index();
let matches = idx.search(query, filter_namespace).map_err(|e| crate::error::AppError::Internal(e.to_string()))?; let matches = idx
.search(query, filter_namespace)
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
let mut results = Vec::new(); let mut results = Vec::new();
for (id, doc_type, title, body, score) in matches.into_iter().take(limit) { for (id, doc_type, title, body, score) in matches.into_iter().take(limit) {
results.push(UnifiedSearchResult { results.push(UnifiedSearchResult {
+10 -11
View File
@@ -112,19 +112,18 @@ mod tests {
data.value = 42; data.value = 42;
}); });
// Need to wait for spawn_blocking to finish // Wait and poll for persistence completion
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; let mut store2 = None;
for _ in 0..20 {
assert_eq!( let s = Store::<TestData>::new("test_key", db.clone());
store.read_with(|s| s.clone()), if s.read_with(|data| data.value) == 42 {
TestData { store2 = Some(s);
name: "Hello".to_string(), break;
value: 42
} }
); tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
}
// Load again to verify persistence let store2 = store2.expect("Timed out waiting for async store persistence");
let store2 = Store::<TestData>::new("test_key", db.clone());
assert_eq!( assert_eq!(
store2.read_with(|s| s.clone()), store2.read_with(|s| s.clone()),
TestData { TestData {
+4 -3
View File
@@ -3,16 +3,17 @@ use arboard::{Clipboard, ImageData};
use std::borrow::Cow; use std::borrow::Cow;
fn main() { 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 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 rgba = img.into_rgba8();
let (w, h) = rgba.dimensions(); let (w, h) = rgba.dimensions();
let img_data = ImageData { let img_data = ImageData {
width: w as usize, width: w as usize,
height: h as usize, height: h as usize,
bytes: Cow::Owned(rgba.into_raw()), bytes: Cow::Owned(rgba.into_raw()),
}; };
let mut clipboard = Clipboard::new().unwrap(); let mut clipboard = Clipboard::new().unwrap();
clipboard.set_image(img_data).unwrap(); clipboard.set_image(img_data).unwrap();
println!("Successfully wrote image"); println!("Successfully wrote image");
+25
View File
@@ -700,3 +700,28 @@ pub struct SemanticCodeSearchTool {
pub struct ReadDirectoryArchitectureTool { pub struct ReadDirectoryArchitectureTool {
pub directory: String, pub directory: String,
} }
/// Create a point-in-time state snapshot of memory stores for safety.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct CreateSnapshotTool {
pub description: String,
pub namespace: Option<String>,
}
/// Restore a point-in-time state snapshot of memory stores.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct RestoreSnapshotTool {
pub snapshot_id: String,
}
/// Create an isolated memory namespace for a subagent session.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct CreateSubagentNamespaceTool {
pub subagent_id: String,
}
/// Purge an isolated subagent memory namespace after session completion.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct PurgeSubagentNamespaceTool {
pub subagent_id: String,
}
+80 -22
View File
@@ -1,10 +1,12 @@
use qdrant_client::qdrant::{CreateCollectionBuilder, Distance, PointStruct, VectorParamsBuilder, UpsertPointsBuilder};
use qdrant_client::Qdrant; use qdrant_client::Qdrant;
use std::sync::Arc; use qdrant_client::qdrant::{
use std::error::Error; CreateCollectionBuilder, Distance, PointStruct, UpsertPointsBuilder, VectorParamsBuilder,
use uuid::Uuid; };
use tracing::info;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use std::error::Error;
use std::sync::Arc;
use tracing::info;
use uuid::Uuid;
#[derive(Clone)] #[derive(Clone)]
pub struct VectorDB { pub struct VectorDB {
@@ -21,9 +23,12 @@ pub struct VectorSearchResult {
} }
impl VectorDB { impl VectorDB {
pub async fn new(url: &str, collection_name: &str) -> Result<Self, Box<dyn Error + Send + Sync>> { pub async fn new(
url: &str,
collection_name: &str,
) -> Result<Self, Box<dyn Error + Send + Sync>> {
let client = Qdrant::from_url(url).build()?; let client = Qdrant::from_url(url).build()?;
let db = Self { let db = Self {
client: Arc::new(client), client: Arc::new(client),
collection_name: collection_name.to_string(), collection_name: collection_name.to_string(),
@@ -42,14 +47,14 @@ impl VectorDB {
self.client self.client
.create_collection( .create_collection(
CreateCollectionBuilder::new(&self.collection_name) CreateCollectionBuilder::new(&self.collection_name)
.vectors_config(vector_params) .vectors_config(vector_params),
) )
.await?; .await?;
info!("Created Qdrant collection: {}", self.collection_name); info!("Created Qdrant collection: {}", self.collection_name);
} else { } else {
info!("Qdrant collection {} already exists", self.collection_name); info!("Qdrant collection {} already exists", self.collection_name);
} }
Ok(()) Ok(())
} }
@@ -69,17 +74,27 @@ impl VectorDB {
} }
}; };
let mut payload: std::collections::HashMap<String, serde_json::Value> = std::collections::HashMap::new(); let mut payload: std::collections::HashMap<String, serde_json::Value> =
payload.insert("doc_type".to_string(), serde_json::Value::String(doc_type.to_string())); std::collections::HashMap::new();
payload.insert("text".to_string(), serde_json::Value::String(text.to_string())); payload.insert(
payload.insert("original_id".to_string(), serde_json::Value::String(id.to_string())); "doc_type".to_string(),
serde_json::Value::String(doc_type.to_string()),
);
payload.insert(
"text".to_string(),
serde_json::Value::String(text.to_string()),
);
payload.insert(
"original_id".to_string(),
serde_json::Value::String(id.to_string()),
);
let point = PointStruct::new(point_id, vector, payload); let point = PointStruct::new(point_id, vector, payload);
self.client self.client
.upsert_points(UpsertPointsBuilder::new(&self.collection_name, vec![point])) .upsert_points(UpsertPointsBuilder::new(&self.collection_name, vec![point]))
.await?; .await?;
Ok(()) Ok(())
} }
@@ -90,26 +105,33 @@ impl VectorDB {
) -> Result<Vec<VectorSearchResult>, Box<dyn Error + Send + Sync>> { ) -> Result<Vec<VectorSearchResult>, Box<dyn Error + Send + Sync>> {
use qdrant_client::qdrant::SearchPointsBuilder; use qdrant_client::qdrant::SearchPointsBuilder;
let search_result = self.client let search_result = self
.client
.search_points( .search_points(
SearchPointsBuilder::new(&self.collection_name, query_vector, limit) SearchPointsBuilder::new(&self.collection_name, query_vector, limit)
.with_payload(true) .with_payload(true),
) )
.await?; .await?;
let mut results = Vec::new(); let mut results = Vec::new();
for point in search_result.result { for point in search_result.result {
let id = point.payload.get("original_id") let id = point
.payload
.get("original_id")
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
.map(|s| s.to_string()) .map(|s| s.to_string())
.unwrap_or_default(); .unwrap_or_default();
let doc_type = point.payload.get("doc_type") let doc_type = point
.payload
.get("doc_type")
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
.map(|s| s.to_string()) .map(|s| s.to_string())
.unwrap_or_default(); .unwrap_or_default();
let text = point.payload.get("text") let text = point
.payload
.get("text")
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
.map(|s| s.to_string()) .map(|s| s.to_string())
.unwrap_or_default(); .unwrap_or_default();
@@ -125,3 +147,39 @@ impl VectorDB {
Ok(results) Ok(results)
} }
} }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_vector_search_result_serialization() {
let res = VectorSearchResult {
id: "test-id".into(),
doc_type: "snippet".into(),
text: "fn main() {}".into(),
score: 0.95,
};
let json = serde_json::to_string(&res).unwrap();
assert!(json.contains("test-id"));
assert!(json.contains("0.95"));
let deserialized: VectorSearchResult = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.id, "test-id");
assert_eq!(deserialized.doc_type, "snippet");
assert_eq!(deserialized.text, "fn main() {}");
assert_eq!(deserialized.score, 0.95);
}
#[test]
fn test_deterministic_uuid_hashing() {
let id_str = "custom-non-uuid-string";
let namespace = Uuid::NAMESPACE_OID;
let uuid1 = Uuid::new_v5(&namespace, id_str.as_bytes()).to_string();
let uuid2 = Uuid::new_v5(&namespace, id_str.as_bytes()).to_string();
assert_eq!(uuid1, uuid2);
assert!(Uuid::parse_str(&uuid1).is_ok());
}
}
+28 -3
View File
@@ -1,8 +1,8 @@
use crate::state::MemoryState;
use notify::{Config, Event, RecommendedWatcher, RecursiveMode, Watcher}; use notify::{Config, Event, RecommendedWatcher, RecursiveMode, Watcher};
use std::path::Path; use std::path::Path;
use std::sync::Arc; use std::sync::Arc;
use crate::state::MemoryState; use tracing::{error, info};
use tracing::{info, error};
pub fn spawn_watcher(_state: Arc<MemoryState>) { pub fn spawn_watcher(_state: Arc<MemoryState>) {
let watch_path = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")); let watch_path = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
@@ -52,7 +52,7 @@ fn should_review(path: &Path) -> bool {
if path_str.contains(".git") || path_str.contains("target") { if path_str.contains(".git") || path_str.contains("target") {
return false; return false;
} }
if let Some(ext) = path.extension() { if let Some(ext) = path.extension() {
let ext_str = ext.to_string_lossy(); let ext_str = ext.to_string_lossy();
if ext_str == "rs" || ext_str == "md" || ext_str == "toml" || ext_str == "lua" { if ext_str == "rs" || ext_str == "md" || ext_str == "toml" || ext_str == "lua" {
@@ -68,3 +68,28 @@ async fn trigger_autonomous_review(path: &Path) {
// MVP: Just write an info log if Neovim isn't reachable. // MVP: Just write an info log if Neovim isn't reachable.
info!("Autonomous review complete for {:?}", path); info!("Autonomous review complete for {:?}", path);
} }
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_should_review() {
assert!(should_review(Path::new("src/lib.rs")));
assert!(should_review(Path::new("README.md")));
assert!(should_review(Path::new("Cargo.toml")));
assert!(should_review(Path::new("init.lua")));
assert!(!should_review(Path::new("target/debug/app.exe")));
assert!(!should_review(Path::new(".git/HEAD")));
assert!(!should_review(Path::new("data.json")));
assert!(!should_review(Path::new("image.png")));
}
#[tokio::test]
async fn test_spawn_watcher_lifecycle() {
let temp_dir = tempfile::tempdir().unwrap();
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
spawn_watcher(state);
}
}
+2 -2
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "mcp-memory-stub" name = "mcp-memory-stub"
version = "0.1.0" version = "1.0.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]
@@ -15,7 +15,7 @@ tracing = "0.1.44"
tracing-subscriber = "0.3.23" tracing-subscriber = "0.3.23"
dirs = "7.0.0" dirs = "7.0.0"
serde_json = "1.0.151" serde_json = "1.0.151"
mcp-stdio = { version = "0.1.0", path = "../mcp-stdio" } mcp-stdio = { version = "1.0.0", path = "../mcp-stdio" }
regex = "1.13.1" regex = "1.13.1"
async-channel = "2.5.0" async-channel = "2.5.0"
+1 -1
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "mcp-memory-win-nvim" name = "mcp-memory-win-nvim"
version = "0.1.0" version = "1.0.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]