refactor: apply 5-pass audit optimizations across mcp-memory codebase

This commit is contained in:
Riz Ashraf committed 2026-10-06 06:05:38 +01:00
1 parent 924b6d09fa
commit 5bd8b1587a
43 files changed
+1866 -1658

No files matched your search

+99 -28
View File
@@ -15,13 +15,16 @@ impl McpTool for LogDecisionHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<LogDecisionTool>("log_decision", "Execute log_decision")
crate::mcp::tool_def::<LogDecisionTool>(
"log_decision",
"Log an architectural decision record (ADR) with title, context, and consequences",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: LogDecisionTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let idx = state.get_search_index();
let idx = state.get_search_index().await;
let mut final_id = String::new();
let mut adrs_to_index = Vec::new();
@@ -71,11 +74,16 @@ impl McpTool for QueryDecisionsHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryDecisionsTool>("query_decisions", "Execute query_decisions")
crate::mcp::tool_def::<QueryDecisionsTool>(
"query_decisions",
"Query or search existing architectural decision records (ADRs)",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: QueryDecisionsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(20);
let include_body = req.include_body.unwrap_or(true);
let data = state.code.adrs.read_with(|adrs| {
let filtered: Vec<_> = adrs
.iter()
@@ -88,8 +96,20 @@ impl McpTool for QueryDecisionsHandler {
true
}
})
.take(limit)
.collect();
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
if include_body {
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
} else {
let compact: Vec<_> = filtered.iter().map(|a| serde_json::json!({
"id": a.id,
"title": a.title,
"status": a.status,
"timestamp": a.timestamp,
})).collect();
Ok::<String, crate::error::AppError>(serde_json::to_string(&compact)?)
}
})?;
Ok(data)
}
@@ -122,7 +142,8 @@ impl McpTool for DeleteDecisionHandler {
});
if found {
state.rebuild_index().await;
let idx = state.get_search_index().await;
let _ = idx.delete_document(&req.id).await;
Ok("Decision deleted successfully".to_string())
} else {
Err(crate::error::AppError::Internal(
@@ -141,7 +162,10 @@ impl McpTool for LogErrorFixHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<LogErrorFixTool>("log_error_fix", "Execute log_error_fix")
crate::mcp::tool_def::<LogErrorFixTool>(
"log_error_fix",
"Log an error signature and its verified solution/fix for future diagnostic retrieval",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -288,7 +312,10 @@ impl McpTool for LogCodeChangeHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<LogCodeChangeTool>("log_code_change", "Execute log_code_change")
crate::mcp::tool_def::<LogCodeChangeTool>(
"log_code_change",
"Log a significant code change or refactor with file path and description",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -325,6 +352,9 @@ impl McpTool for LogCodeChangeHandler {
git_commit: req.git_commit,
git_branch: req.git_branch,
});
if ledger.len() > 1000 {
ledger.remove(0);
}
});
state.record_activity("code_change", &format!("Modified {}", req.file_path), Some(&description));
Ok(format!("Logged code change for {}: {}", req.file_path, description))
@@ -342,7 +372,7 @@ impl McpTool for QueryRecentChangesHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryRecentChangesTool>(
"query_recent_changes",
"Execute query_recent_changes",
"Query recent code changes and refactoring audit logs",
)
}
@@ -364,7 +394,10 @@ impl McpTool for LearnPreferenceHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<LearnPreferenceTool>("learn_preference", "Execute learn_preference")
crate::mcp::tool_def::<LearnPreferenceTool>(
"learn_preference",
"Learn or update a key-value user preference",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -392,7 +425,10 @@ impl McpTool for ReadPreferencesHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ReadPreferencesTool>("read_preferences", "Execute read_preferences")
crate::mcp::tool_def::<ReadPreferencesTool>(
"read_preferences",
"Read all stored user and environment preferences",
)
}
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -412,7 +448,10 @@ impl McpTool for LogTechDebtHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<LogTechDebtTool>("log_tech_debt", "Execute log_tech_debt")
crate::mcp::tool_def::<LogTechDebtTool>(
"log_tech_debt",
"Log a new technical debt item with description and ideal solution",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -452,7 +491,7 @@ impl McpTool for ResolveTechDebtHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<ResolveTechDebtTool>(
"resolve_tech_debt",
"Execute resolve_tech_debt",
"Mark an existing technical debt item as resolved by ID",
)
}
@@ -488,7 +527,10 @@ impl McpTool for ListTechDebtHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ListTechDebtTool>("list_tech_debt", "Execute list_tech_debt")
crate::mcp::tool_def::<ListTechDebtTool>(
"list_tech_debt",
"List technical debt items filtered by namespace or resolution status",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -641,7 +683,10 @@ impl McpTool for OmniSearchHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<OmniSearchTool>("omni_search", "Execute omni_search")
crate::mcp::tool_def::<OmniSearchTool>(
"omni_search",
"Unified search across entities, subgraphs, tasks, code snippets, ADRs, and technical debt",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -649,7 +694,7 @@ impl McpTool for OmniSearchHandler {
let limit = req.limit.unwrap_or(5);
let include_body = req.include_body.unwrap_or(false);
let idx = state.get_search_index();
let idx = state.get_search_index().await;
let keyword_matches = idx
.search(&req.query, req.namespace.as_deref())
.unwrap_or_default();
@@ -717,7 +762,7 @@ impl McpTool for OmniSearchHandler {
}
let mut ranked_items: Vec<_> = rrf_scores.into_values().collect();
ranked_items.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
ranked_items.sort_by(|a, b| b.0.total_cmp(&a.0));
let matches: Vec<MatchItem> =
ranked_items.into_iter().map(|(_, item)| item).collect();
@@ -982,7 +1027,7 @@ impl McpTool for GetProjectHealthHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<GetProjectHealthTool>(
"get_project_health",
"Execute get_project_health",
"Retrieve project health metrics including active tasks, technical debt, and PR checklist progress",
)
}
@@ -1037,7 +1082,10 @@ impl McpTool for SuggestErrorFixHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<SuggestErrorFixTool>("suggest_error_fix", "Execute suggest_error_fix")
crate::mcp::tool_def::<SuggestErrorFixTool>(
"suggest_error_fix",
"Suggest historical error fixes based on stack trace vector similarity",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1066,7 +1114,7 @@ impl McpTool for SuggestErrorFixHandler {
.filter(|(_, score)| *score > 0.2)
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
scored.sort_by(|a, b| b.1.total_cmp(&a.1));
let suggestions: Vec<_> = scored.into_iter().take(limit).map(|(f, score)| {
serde_json::json!({
@@ -1214,7 +1262,10 @@ impl McpTool for CheckpointStateHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CheckpointStateTool>("checkpoint_state", "Execute checkpoint_state")
crate::mcp::tool_def::<CheckpointStateTool>(
"checkpoint_state",
"Save point-in-time memory state snapshot checkpoint",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1248,7 +1299,10 @@ impl McpTool for RestoreStateHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<RestoreStateTool>("restore_state", "Execute restore_state")
crate::mcp::tool_def::<RestoreStateTool>(
"restore_state",
"Restore memory state from a point-in-time snapshot checkpoint",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1290,7 +1344,10 @@ impl McpTool for QueryLineageHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryLineageTool>("query_lineage", "Execute query_lineage")
crate::mcp::tool_def::<QueryLineageTool>(
"query_lineage",
"Query historical lineage and timeline of tasks, ADRs, and code changes",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1375,7 +1432,10 @@ impl McpTool for GetNextActionableTasksHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<GetNextActionableTasksTool>("get_next_actionable_tasks", "Execute get_next_actionable_tasks")
crate::mcp::tool_def::<GetNextActionableTasksTool>(
"get_next_actionable_tasks",
"Get unblocked pending tasks ready for execution",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1431,7 +1491,10 @@ impl McpTool for LogHypothesisHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<LogHypothesisTool>("log_hypothesis", "Execute log_hypothesis")
crate::mcp::tool_def::<LogHypothesisTool>(
"log_hypothesis",
"Log a diagnostic hypothesis and associated evidence for a task",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1463,7 +1526,10 @@ impl McpTool for QueryHypothesesHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryHypothesesTool>("query_hypotheses", "Execute query_hypotheses")
crate::mcp::tool_def::<QueryHypothesesTool>(
"query_hypotheses",
"Query active diagnostic hypotheses and evidence by task ID or keyword",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1701,7 +1767,12 @@ impl McpTool for AutoSessionCheckpointHandler {
timestamp,
};
state.telemetry.handoff_memos.modify(|m| m.push(memo));
state.telemetry.handoff_memos.modify(|m| {
m.push(memo);
if m.len() > 100 {
m.remove(0);
}
});
state.record_activity("checkpoint", &format!("Created auto session checkpoint {}", memo_id), None);
Ok(format!("Session checkpoint created with memo ID '{}'.", memo_id))
@@ -1759,7 +1830,7 @@ impl McpTool for SearchSnippetsHybridHandler {
.collect();
scored.sort_by(|a, b| {
b["score"].as_f64().unwrap_or(0.0).partial_cmp(&a["score"].as_f64().unwrap_or(0.0)).unwrap_or(std::cmp::Ordering::Equal)
b["score"].as_f64().unwrap_or(0.0).total_cmp(&a["score"].as_f64().unwrap_or(0.0))
});
let limit = req.limit.unwrap_or(10);
@@ -1974,7 +2045,7 @@ mod tests {
}
state.rebuild_index().await;
state.get_search_index().reader.reload().unwrap();
state.get_search_index().await.reader.reload().unwrap();
let omni = OmniSearchHandler;
let omni_res = omni