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

fix(mcp): propagate internal serialization errors instead of silently swallowing them to prevent LLM hallucination

Browse Source
This commit is contained in:
Riz Ashraf committed 2026-09-27 22:19:43 +01:00
1 parent f8925050db
commit 8f32a09399
13 files changed
+235 -63
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

nvim-core/src/lib.rs
+1 -1
View File
@@ -684,7 +684,7 @@ async fn execute_nvim_lua_with_args(code: &str, args: Vec<rmpv::Value>) -> Resul
if let serde_json::Value::String(s) = json_res { if let serde_json::Value::String(s) = json_res {
Ok(s) Ok(s)
} else { } else {
Ok(serde_json::to_string_pretty(&json_res).unwrap_or_default()) serde_json::to_string_pretty(&json_res).map_err(|e| e.to_string())
} }
} }
server/src/api/events.rs
+8 -6
View File
@@ -1,9 +1,9 @@
use crate::AppState; use crate::AppState;
use crate::state::GenericEvent;
use axum::extract::{Query, State}; use axum::extract::{Query, State};
use axum::response::IntoResponse; use axum::response::IntoResponse;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
use crate::state::GenericEvent;
pub async fn wait_for_event_handler( pub async fn wait_for_event_handler(
State(state): State<Arc<AppState>>, State(state): State<Arc<AppState>>,
@@ -18,7 +18,9 @@ pub async fn wait_for_event_handler(
match rx.recv().await { match rx.recv().await {
Ok(event) => { Ok(event) => {
let topic_matches = topic.as_ref().map_or(true, |t| t == &event.topic); let topic_matches = topic.as_ref().map_or(true, |t| t == &event.topic);
let session_matches = session_id.as_ref().map_or(true, |s| Some(s) == event.session_id.as_ref()); let session_matches = session_id
.as_ref()
.map_or(true, |s| Some(s) == event.session_id.as_ref());
if topic_matches && session_matches { if topic_matches && session_matches {
return axum::Json(event); return axum::Json(event);
@@ -46,14 +48,14 @@ pub async fn post_event_handler(
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::state::MemoryState;
use crate::router::MemoryHandler; use crate::router::MemoryHandler;
use tempfile::tempdir; use crate::state::MemoryState;
use axum::extract::Query; use axum::extract::Query;
use axum::extract::State; use axum::extract::State;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::atomic::AtomicUsize;
use std::sync::RwLock; use std::sync::RwLock;
use std::sync::atomic::AtomicUsize;
use tempfile::tempdir;
#[tokio::test] #[tokio::test]
async fn test_events_wait_and_post() { async fn test_events_wait_and_post() {
@@ -70,7 +72,7 @@ mod tests {
let mut params = HashMap::new(); let mut params = HashMap::new();
params.insert("topic".to_string(), "test_topic".to_string()); params.insert("topic".to_string(), "test_topic".to_string());
params.insert("session_id".to_string(), "123".to_string()); params.insert("session_id".to_string(), "123".to_string());
let wait_task = tokio::spawn(async move { let wait_task = tokio::spawn(async move {
let res = wait_for_event_handler(State(app_state_clone), Query(params)).await; let res = wait_for_event_handler(State(app_state_clone), Query(params)).await;
// axum::Json<GenericEvent> is returned, we need to extract it somehow, but just returning is enough for testing // axum::Json<GenericEvent> is returned, we need to extract it somehow, but just returning is enough for testing
server/src/api/mod.rs
+1 -1
View File
@@ -1,5 +1,5 @@
pub mod events;
pub mod rest; pub mod rest;
pub mod setup; pub mod setup;
pub mod telemetry; pub mod telemetry;
pub mod ws; pub mod ws;
pub mod events;
server/src/api/setup.rs
+80
View File
@@ -148,6 +148,86 @@ pub fn create_router(app_state: Arc<AppState>) -> Router {
} }
}), }),
) )
.route(
"/api/tech_debts",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let json = state_clone.tech_debts.read_with(|items| serde_json::to_string(items).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], json)
}
}),
)
.route(
"/api/adrs",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let json = state_clone.adrs.read_with(|items| serde_json::to_string(items).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], json)
}
}),
)
.route(
"/api/context_workspaces",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let json = state_clone.context_workspaces.read_with(|items| serde_json::to_string(items).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], json)
}
}),
)
.route(
"/api/handoff_memos",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let json = state_clone.handoff_memos.read_with(|items| serde_json::to_string(items).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], json)
}
}),
)
.route(
"/api/milestones",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let json = state_clone.milestones.read_with(|items| serde_json::to_string(items).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], json)
}
}),
)
.route(
"/api/snippets",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let json = state_clone.snippets.read_with(|items| serde_json::to_string(items).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], json)
}
}),
)
.route(
"/api/pr_checklists",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let json = state_clone.pr_checklists.read_with(|items| serde_json::to_string(items).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], json)
}
}),
)
.route(
"/api/error_fixes",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let json = state_clone.error_fixes.read_with(|items| serde_json::to_string(items).unwrap_or_else(|_| "[]".to_string()));
([(axum::http::header::CONTENT_TYPE, "application/json")], json)
}
}),
)
.route( .route(
"/api/stats", "/api/stats",
get({ get({
server/src/handlers/graph.rs
+43 -12
View File
@@ -151,8 +151,14 @@ impl McpTool for CreateRelationsHandler {
Ok(r) => r, Ok(r) => r,
Err(e) => { Err(e) => {
let err_msg = e.to_string(); let err_msg = e.to_string();
if err_msg.contains("missing field `from`") || err_msg.contains("missing field `to`") || err_msg.contains("missing field `relation_type`") { if err_msg.contains("missing field `from`")
return Err(format!("Schema error: {}. Note that the relation schema strictly uses 'from', 'to', and 'relation_type' (not 'source', 'target', or 'relationType'). Please correct your tool call arguments.", err_msg)); || err_msg.contains("missing field `to`")
|| err_msg.contains("missing field `relation_type`")
{
return Err(format!(
"Schema error: {}. Note that the relation schema strictly uses 'from', 'to', and 'relation_type' (not 'source', 'target', or 'relationType'). Please correct your tool call arguments.",
err_msg
));
} }
return Err(err_msg); return Err(err_msg);
} }
@@ -166,15 +172,22 @@ impl McpTool for CreateRelationsHandler {
if from_exists && to_exists { if from_exists && to_exists {
g.relations.push(relation); g.relations.push(relation);
} else { } else {
if !from_exists { missing_nodes.insert(relation.from); } if !from_exists {
if !to_exists { missing_nodes.insert(relation.to); } missing_nodes.insert(relation.from);
}
if !to_exists {
missing_nodes.insert(relation.to);
}
} }
} }
} }
}); });
if !missing_nodes.is_empty() { if !missing_nodes.is_empty() {
let missing: Vec<_> = missing_nodes.into_iter().collect(); let missing: Vec<_> = missing_nodes.into_iter().collect();
return Err(format!("Error: Relations dropped due to missing entities: {}", missing.join(", "))); return Err(format!(
"Error: Relations dropped due to missing entities: {}",
missing.join(", ")
));
} }
Ok("Relations created".to_string()) Ok("Relations created".to_string())
} }
@@ -205,7 +218,10 @@ impl McpTool for AddObservationsHandler {
} }
}); });
if !missing_entities.is_empty() { if !missing_entities.is_empty() {
return Err(format!("Error: Observations dropped for missing entities: {}", missing_entities.join(", "))); return Err(format!(
"Error: Observations dropped for missing entities: {}",
missing_entities.join(", ")
));
} }
Ok("Observations added".to_string()) Ok("Observations added".to_string())
} }
@@ -239,7 +255,10 @@ impl McpTool for DeleteEntitiesHandler {
}); });
if !missing.is_empty() { if !missing.is_empty() {
return Err(format!("Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.", missing.join(", "))); return Err(format!(
"Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
missing.join(", ")
));
} }
let idx = state.get_search_index(); let idx = state.get_search_index();
@@ -280,7 +299,10 @@ impl McpTool for DeleteObservationsHandler {
} }
}); });
if !missing.is_empty() { if !missing.is_empty() {
return Err(format!("Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.", missing.join(", "))); return Err(format!(
"Error: Entities not found: {}. Please use the search_nodes or read_graph tools to verify the exact entity names.",
missing.join(", ")
));
} }
Ok("Observations deleted".to_string()) Ok("Observations deleted".to_string())
} }
@@ -308,7 +330,10 @@ impl McpTool for DeleteRelationsHandler {
missing_count = to_rem.len() - (initial_len - master.relations.len()); missing_count = to_rem.len() - (initial_len - master.relations.len());
}); });
if missing_count > 0 { if missing_count > 0 {
return Err(format!("Error: {} relation(s) not found in graph. Please verify exact relation properties using read_graph.", missing_count)); return Err(format!(
"Error: {} relation(s) not found in graph. Please verify exact relation properties using read_graph.",
missing_count
));
} }
Ok("Relations deleted".to_string()) Ok("Relations deleted".to_string())
} }
@@ -536,7 +561,10 @@ impl McpTool for CondenseEntityHandler {
} }
}); });
if missing { if missing {
return Err(format!("Error: Entity '{}' not found. Please verify the exact entity name using search_nodes.", req.entity_name)); return Err(format!(
"Error: Entity '{}' not found. Please verify the exact entity name using search_nodes.",
req.entity_name
));
} }
Ok("Entity condensed".to_string()) Ok("Entity condensed".to_string())
} }
@@ -588,7 +616,10 @@ impl McpTool for MergeEntitiesHandler {
}); });
}); });
if missing { if missing {
return Err(format!("Error: Source entity '{}' not found. Please verify the exact entity name using search_nodes.", req.source_entity)); return Err(format!(
"Error: Source entity '{}' not found. Please verify the exact entity name using search_nodes.",
req.source_entity
));
} }
Ok("Entities merged".to_string()) Ok("Entities merged".to_string())
} }
@@ -691,7 +722,7 @@ mod tests {
}); });
let res = handler.execute(args, state.clone()).await.unwrap(); let res = handler.execute(args, state.clone()).await.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)
let bad_args = json!({ let bad_args = json!({
"relations": [ "relations": [
server/src/handlers/meta.rs
+27 -12
View File
@@ -98,11 +98,15 @@ impl McpTool for DeleteDecisionHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<crate::tools::DeleteDecisionTool>("delete_decision", "Delete an architectural decision record") crate::mcp::tool_def::<crate::tools::DeleteDecisionTool>(
"delete_decision",
"Delete an architectural decision record",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
let req: crate::tools::DeleteDecisionTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: crate::tools::DeleteDecisionTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut found = false; let mut found = false;
state.adrs.modify(|adrs| { state.adrs.modify(|adrs| {
if let Some(pos) = adrs.iter().position(|a| a.id == req.id) { if let Some(pos) = adrs.iter().position(|a| a.id == req.id) {
@@ -110,7 +114,7 @@ impl McpTool for DeleteDecisionHandler {
found = true; found = true;
} }
}); });
if found { if found {
state.rebuild_index().await; state.rebuild_index().await;
Ok("Decision deleted successfully".to_string()) Ok("Decision deleted successfully".to_string())
@@ -337,7 +341,10 @@ impl McpTool for ResolveTechDebtHandler {
if found { if found {
Ok("Tech debt resolved".to_string()) Ok("Tech debt resolved".to_string())
} else { } else {
Err("Tech debt not found. Please verify the tech debt ID using list_tech_debt.".to_string()) Err(
"Tech debt not found. Please verify the tech debt ID using list_tech_debt."
.to_string(),
)
} }
} }
} }
@@ -385,7 +392,10 @@ 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 = match state.get_search_index().search(&req.query, req.namespace.as_deref()) { let matches = match state
.get_search_index()
.search(&req.query, req.namespace.as_deref())
{
Ok(m) => m, Ok(m) => m,
Err(e) => { Err(e) => {
return Err(format!( return Err(format!(
@@ -403,7 +413,9 @@ impl McpTool for OmniSearchHandler {
if doc_type == "entity" if doc_type == "entity"
&& let Some(e) = full.entities.get(id) && let Some(e) = full.entities.get(id)
{ {
if count >= limit { continue; } if count >= limit {
continue;
}
count += 1; count += 1;
if !include_body { if !include_body {
let mut summary = e.clone(); let mut summary = e.clone();
@@ -792,13 +804,13 @@ mod tests {
dependencies: vec![], dependencies: vec![],
acceptance_criteria: vec![], acceptance_criteria: vec![],
}; };
{ {
state.tasks.modify(|t| { state.tasks.modify(|t| {
t.push(task.clone()); t.push(task.clone());
}); });
} }
state.rebuild_index().await; state.rebuild_index().await;
state.get_search_index().reader.reload().unwrap(); state.get_search_index().reader.reload().unwrap();
@@ -808,21 +820,24 @@ mod tests {
.await .await
.unwrap(); .unwrap();
println!("OMNI RES: {}", omni_res); println!("OMNI RES: {}", omni_res);
assert!(omni_res.contains("omni-1"), "omni search should return results containing the task id"); assert!(
omni_res.contains("omni-1"),
"omni search should return results containing the task id"
);
} }
#[tokio::test] #[tokio::test]
async fn test_omni_search_malformed_query() { async fn test_omni_search_malformed_query() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::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 omni = OmniSearchHandler; let omni = OmniSearchHandler;
// Pass a malformed Lucene query (unclosed parenthesis) // Pass a malformed Lucene query (unclosed parenthesis)
let omni_res = omni let omni_res = omni
.execute(json!({"query": "title: (unclosed"}), state.clone()) .execute(json!({"query": "title: (unclosed"}), state.clone())
.await; .await;
assert!(omni_res.is_err()); assert!(omni_res.is_err());
let err_msg = omni_res.unwrap_err(); let err_msg = omni_res.unwrap_err();
assert!(err_msg.contains("malformed Lucene syntax")); assert!(err_msg.contains("malformed Lucene syntax"));
server/src/handlers/tasks.rs
+4 -1
View File
@@ -444,7 +444,10 @@ impl McpTool for UpdateMilestoneHandler {
if found { if found {
Ok("Milestone updated".to_string()) Ok("Milestone updated".to_string())
} else { } else {
Err("Milestone not found. Please verify the milestone ID using list_milestones.".to_string()) Err(
"Milestone not found. Please verify the milestone ID using list_milestones."
.to_string(),
)
} }
} }
} }
server/src/handlers/workspaces.rs
+13 -4
View File
@@ -185,7 +185,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("Snippet not found. Please verify the snippet ID using search_snippets.".to_string()) Err(
"Snippet not found. Please verify the snippet ID using search_snippets."
.to_string(),
)
} }
} }
} }
@@ -270,7 +273,10 @@ impl McpTool for ListContextWorkspacesHandler {
let req: ListContextWorkspacesTool = let req: ListContextWorkspacesTool =
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.context_workspaces.read_with(|ws| { let data = state.context_workspaces.read_with(|ws| {
let filtered: Vec<_> = ws.iter().filter(|w| req.namespace.as_ref().map_or(true, |ns| &w.namespace == ns)).collect(); let filtered: Vec<_> = ws
.iter()
.filter(|w| req.namespace.as_ref().map_or(true, |ns| &w.namespace == ns))
.collect();
serde_json::to_string(&filtered).map_err(|e| e.to_string()) serde_json::to_string(&filtered).map_err(|e| e.to_string())
})?; })?;
Ok(data) Ok(data)
@@ -295,10 +301,13 @@ impl McpTool for DeleteContextWorkspaceHandler {
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
let req: crate::tools::DeleteContextWorkspaceTool = let req: crate::tools::DeleteContextWorkspaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?; serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut found = false; let mut found = false;
state.context_workspaces.modify(|ws| { state.context_workspaces.modify(|ws| {
if let Some(pos) = ws.iter().position(|w| w.namespace == req.namespace && w.name == req.name) { if let Some(pos) = ws
.iter()
.position(|w| w.namespace == req.namespace && w.name == req.name)
{
ws.remove(pos); ws.remove(pos);
found = true; found = true;
} }
server/src/mcp.rs
+3 -2
View File
@@ -23,8 +23,9 @@ pub fn error(id: serde_json::Value, code: i32, message: &str) -> serde_json::Val
pub fn tool_def<T: JsonSchema>(name: &str, fallback_desc: &str) -> serde_json::Value { pub fn tool_def<T: JsonSchema>(name: &str, fallback_desc: &str) -> serde_json::Value {
let schema = schemars::schema_for!(T); let schema = schemars::schema_for!(T);
let schema_val = serde_json::to_value(&schema).unwrap_or_default(); let schema_val = serde_json::to_value(&schema).unwrap_or_default();
let desc = schema_val.get("description") let desc = schema_val
.get("description")
.and_then(|d| d.as_str()) .and_then(|d| d.as_str())
.unwrap_or(fallback_desc) .unwrap_or(fallback_desc)
.to_string(); .to_string();
server/src/router.rs
+5 -5
View File
@@ -57,7 +57,7 @@ impl McpResource for GraphEntitiesResource {
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let graph = state_clone.graph.cache.read().unwrap(); let graph = state_clone.graph.cache.read().unwrap();
let data: Vec<_> = graph.entities.values().collect(); let data: Vec<_> = graph.entities.values().collect();
Ok(serde_json::to_string_pretty(&data).unwrap_or_default()) serde_json::to_string_pretty(&data).map_err(|e| e.to_string())
}) })
.await .await
.unwrap() .unwrap()
@@ -81,7 +81,7 @@ impl McpResource for GraphRelationsResource {
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let graph = state_clone.graph.cache.read().unwrap(); let graph = state_clone.graph.cache.read().unwrap();
let data = &graph.relations; let data = &graph.relations;
Ok(serde_json::to_string_pretty(&data).unwrap_or_default()) serde_json::to_string_pretty(&data).map_err(|e| e.to_string())
}) })
.await .await
.unwrap() .unwrap()
@@ -108,7 +108,7 @@ impl McpResource for TasksActiveResource {
.iter() .iter()
.filter(|t| t.status != "completed" && t.status != "done") .filter(|t| t.status != "completed" && t.status != "done")
.collect(); .collect();
Ok(serde_json::to_string_pretty(&data).unwrap_or_default()) serde_json::to_string_pretty(&data).map_err(|e| e.to_string())
}) })
.await .await
.unwrap() .unwrap()
@@ -194,7 +194,7 @@ impl MemoryHandler {
let state_clone = Arc::clone(&state); let state_clone = Arc::clone(&state);
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let items = state_clone.pinned_files.cache.read().unwrap(); let items = state_clone.pinned_files.cache.read().unwrap();
Ok(serde_json::to_string_pretty(&*items).unwrap_or_default()) serde_json::to_string_pretty(&*items).map_err(|e| e.to_string())
}) })
.await .await
.unwrap() .unwrap()
@@ -217,7 +217,7 @@ impl MemoryHandler {
let state_clone = Arc::clone(&state); let state_clone = Arc::clone(&state);
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let items = state_clone.milestones.cache.read().unwrap(); let items = state_clone.milestones.cache.read().unwrap();
Ok(serde_json::to_string_pretty(&*items).unwrap_or_default()) serde_json::to_string_pretty(&*items).map_err(|e| e.to_string())
}) })
.await .await
.unwrap() .unwrap()
server/src/search.rs
+10 -5
View File
@@ -244,14 +244,16 @@ impl MemoryIndex {
self.type_field => "entity", self.type_field => "entity",
self.namespace_field => e.namespace.as_str() self.namespace_field => e.namespace.as_str()
)); ));
self.needs_commit.store(true, std::sync::atomic::Ordering::SeqCst); self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
} }
} }
pub fn delete_all(&self) { pub fn delete_all(&self) {
if let Ok(writer) = self.writer.lock() { if let Ok(writer) = self.writer.lock() {
let _ = writer.delete_all_documents(); let _ = writer.delete_all_documents();
self.needs_commit.store(true, std::sync::atomic::Ordering::SeqCst); self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
} }
} }
@@ -266,7 +268,8 @@ impl MemoryIndex {
self.namespace_field => "global" self.namespace_field => "global"
)); ));
println!("Writer add_document returned id/result"); println!("Writer add_document returned id/result");
self.needs_commit.store(true, std::sync::atomic::Ordering::SeqCst); self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
println!("Needs_commit set to true in add_task_sync"); println!("Needs_commit set to true in add_task_sync");
} else { } else {
println!("Failed to acquire writer lock in add_task_sync"); println!("Failed to acquire writer lock in add_task_sync");
@@ -282,7 +285,8 @@ impl MemoryIndex {
self.type_field => "snippet", self.type_field => "snippet",
self.namespace_field => "global" self.namespace_field => "global"
)); ));
self.needs_commit.store(true, std::sync::atomic::Ordering::SeqCst); self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
} }
} }
@@ -295,7 +299,8 @@ impl MemoryIndex {
self.type_field => "adr", self.type_field => "adr",
self.namespace_field => "global" self.namespace_field => "global"
)); ));
self.needs_commit.store(true, std::sync::atomic::Ordering::SeqCst); self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst);
} }
} }
} }
server/src/state.rs
+21 -7
View File
@@ -84,7 +84,6 @@ impl MemoryState {
} }
} }
pub fn broadcast_activity(&self, message: &str) { pub fn broadcast_activity(&self, message: &str) {
let time = std::time::SystemTime::now() let time = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH) .duration_since(std::time::UNIX_EPOCH)
@@ -124,19 +123,28 @@ impl MemoryState {
} }
pub fn get_search_index(&self) -> MemoryIndex { pub fn get_search_index(&self) -> MemoryIndex {
self.search_index.read().unwrap_or_else(|e| e.into_inner()).clone() self.search_index
.read()
.unwrap_or_else(|e| e.into_inner())
.clone()
} }
pub async fn rebuild_index(self: &Arc<Self>) { pub async fn rebuild_index(self: &Arc<Self>) {
let idx = self.search_index.read().unwrap().clone(); let idx = self.search_index.read().unwrap().clone();
idx.delete_all(); idx.delete_all();
let entities: Vec<_> = self.graph.read_with(|g| g.entities.values().cloned().collect()); let entities: Vec<_> = self
.graph
.read_with(|g| g.entities.values().cloned().collect());
let tasks = self.tasks.read_with(|t| t.clone()); let tasks = self.tasks.read_with(|t| t.clone());
let snippets = self.snippets.read_with(|s| s.clone()); let snippets = self.snippets.read_with(|s| s.clone());
let adrs = self.adrs.read_with(|a| a.clone()); let adrs = self.adrs.read_with(|a| a.clone());
println!("rebuild_index: found {} entities, {} tasks", entities.len(), tasks.len()); println!(
"rebuild_index: found {} entities, {} tasks",
entities.len(),
tasks.len()
);
let idx_clone = idx.clone(); let idx_clone = idx.clone();
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
@@ -209,15 +217,21 @@ mod tests {
let idx = arc_state.search_index.read().unwrap(); let idx = arc_state.search_index.read().unwrap();
// Force reload reader to ensure it sees the commit made by rebuild_index // Force reload reader to ensure it sees the commit made by rebuild_index
idx.reader.reload().unwrap(); idx.reader.reload().unwrap();
println!("Index reader doc count: {}", idx.reader.searcher().num_docs()); println!(
"Index reader doc count: {}",
idx.reader.searcher().num_docs()
);
let all_docs = idx.search("Test", None).expect("Search failed"); let all_docs = idx.search("Test", None).expect("Search failed");
println!("All docs for 'Test': {:?}", all_docs); println!("All docs for 'Test': {:?}", all_docs);
// Verify the task added synchronously is actually searchable // Verify the task added synchronously is actually searchable
let results = idx.search("Test", None).expect("Search failed"); let results = idx.search("Test", None).expect("Search failed");
assert_eq!(results.len(), 1, "Expected exactly 1 search result"); assert_eq!(results.len(), 1, "Expected exactly 1 search result");
assert_eq!(results[0].0, "123", "Expected the result to be the task we just added"); assert_eq!(
results[0].0, "123",
"Expected the result to be the task we just added"
);
assert_eq!(results[0].1, "task", "Expected document type to be task"); assert_eq!(results[0].1, "task", "Expected document type to be task");
} }
} }
server/src/tools.rs
+19 -7
View File
@@ -612,16 +612,28 @@ mod tests {
fn test_schema_extraction_includes_descriptions() { fn test_schema_extraction_includes_descriptions() {
let schema = schema_for!(SetAcceptanceCriteriaTool); let schema = schema_for!(SetAcceptanceCriteriaTool);
let schema_json = serde_json::to_value(&schema).unwrap(); let schema_json = serde_json::to_value(&schema).unwrap();
let desc = schema_json.get("description").and_then(|d| d.as_str()).unwrap_or(""); let desc = schema_json
assert!(desc.contains("Define a strict checklist of acceptance criteria"), "Schema should include struct docstring as description"); .get("description")
.and_then(|d| d.as_str())
.unwrap_or("");
assert!(
desc.contains("Define a strict checklist of acceptance criteria"),
"Schema should include struct docstring as description"
);
let schema2 = schema_for!(LogCodeChangeTool); let schema2 = schema_for!(LogCodeChangeTool);
let schema2_json = serde_json::to_value(&schema2).unwrap(); let schema2_json = serde_json::to_value(&schema2).unwrap();
let props = schema2_json.get("properties").expect("Missing properties"); let props = schema2_json.get("properties").expect("Missing properties");
let file_path_prop = props.get("file_path").expect("Missing file_path property"); let file_path_prop = props.get("file_path").expect("Missing file_path property");
let field_desc = file_path_prop.get("description").and_then(|d| d.as_str()).unwrap_or(""); let field_desc = file_path_prop
assert!(field_desc.contains("The path of the file that was changed"), "Schema should include field docstring as description"); .get("description")
.and_then(|d| d.as_str())
.unwrap_or("");
assert!(
field_desc.contains("The path of the file that was changed"),
"Schema should include field docstring as description"
);
} }
} }