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

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

server/src/handlers/ast.rs
+140 -15
View File
@@ -56,7 +56,7 @@ impl McpTool for ReadFileSkeletonHandler {
let tree = parser.parse(&code, None).ok_or_else(|| {
crate::error::AppError::Internal("Failed to parse code".to_string())
})?;
let mut result_skeleton = String::new();
let mut result_skeleton = String::with_capacity(code.len() / 2);
fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) {
if depth > 128 {
@@ -141,7 +141,7 @@ impl McpTool for ReplaceAstNodeHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<ReplaceAstNodeTool>(
"replace_ast_node",
"Replace a specific AST node (e.g., function, struct) entirely using tree-sitter for robust structural editing.",
"Replace a specific AST node (e.g., function, struct, enum, class, trait) entirely using tree-sitter for robust structural editing. Supported node_type values include: 'function_item' (or 'function'/'fn'/'method'), 'struct_item' (or 'struct'), 'class_declaration' (or 'class'), 'enum_item' (or 'enum'), 'trait_item' (or 'trait'/'interface'), 'type_alias_declaration' (or 'type').",
)
}
@@ -185,6 +185,34 @@ impl McpTool for ReplaceAstNodeHandler {
crate::error::AppError::Internal("Failed to parse code".to_string())
})?;
fn matches_node_type(actual_kind: &str, requested_type: &str) -> bool {
if actual_kind == requested_type {
return true;
}
match requested_type.to_lowercase().as_str() {
"function" | "func" | "fn" | "method" | "def" => matches!(
actual_kind,
"function_item"
| "function_declaration"
| "function_definition"
| "method_definition"
| "function"
),
"struct" => matches!(actual_kind, "struct_item" | "struct_declaration" | "struct_specifier"),
"class" => matches!(actual_kind, "class_declaration" | "class_definition" | "class_item"),
"enum" => matches!(actual_kind, "enum_item" | "enum_declaration"),
"trait" | "interface" => matches!(
actual_kind,
"trait_item" | "interface_declaration" | "interface_item"
),
"type" | "type_alias" => matches!(
actual_kind,
"type_alias_declaration" | "type_item" | "type_definition"
),
_ => false,
}
}
// Search for the node
fn find_node<'a>(
node: Node<'a>,
@@ -192,12 +220,12 @@ impl McpTool for ReplaceAstNodeHandler {
target_type: &str,
target_name: &str,
) -> Option<Node<'a>> {
if node.kind() == target_type {
if matches_node_type(node.kind(), target_type) {
// Try to find the name/identifier
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
let kind = child.kind();
if kind == "identifier" || kind == "name" {
if kind == "identifier" || kind == "name" || kind == "property_identifier" || kind == "field_identifier" {
let name = child.utf8_text(code.as_bytes()).unwrap_or("");
if name == target_name {
return Some(node);
@@ -250,6 +278,61 @@ impl McpTool for ReplaceAstNodeHandler {
}
}
fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: bool) -> Vec<serde_json::Value> {
let mut results = Vec::new();
let cwd = match std::env::current_dir() {
Ok(dir) => dir,
Err(_) => return results,
};
let walker = ignore::WalkBuilder::new(&cwd)
.hidden(true)
.git_ignore(true)
.build();
let mut scanned_files = 0;
for result in walker {
let entry = match result {
Ok(e) => e,
Err(_) => continue,
};
if entry.file_type().is_some_and(|ft| ft.is_file()) {
let path = entry.path();
let ext = path.extension().and_then(|s| s.to_str()).unwrap_or("");
if matches!(ext, "rs" | "ts" | "tsx" | "js" | "jsx" | "py" | "go" | "java" | "c" | "cpp" | "h" | "hpp") {
scanned_files += 1;
if scanned_files > 500 {
break;
}
if let Ok(content) = std::fs::read_to_string(path) {
for (line_num, line) in content.lines().enumerate() {
let is_match = if filter_fn_call {
line.contains(&format!("{}(", target_sym))
|| line.contains(&format!("{}.await", target_sym))
} else {
line.contains(target_sym)
};
if is_match {
results.push(serde_json::json!({
"file_path": path.to_string_lossy(),
"line": line_num + 1,
"content": line.trim(),
}));
if results.len() >= limit {
return results;
}
}
}
}
}
}
}
results
}
pub struct FindSymbolReferencesHandler;
#[async_trait]
@@ -269,22 +352,36 @@ impl McpTool for FindSymbolReferencesHandler {
let req: crate::tools::FindSymbolReferencesTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(10);
let target_sym = req.symbol;
let target_sym = req.symbol.clone();
let matches = state.code.snippets.read_with(|snippets| {
let mut matches = state.code.snippets.read_with(|snippets| {
let mut refs = Vec::new();
for snippet in snippets {
if snippet.code.contains(&target_sym) || snippet.name.contains(&target_sym) {
refs.push(snippet.clone());
refs.push(serde_json::json!({
"source": "snippet",
"name": snippet.name,
"code": snippet.code,
}));
if refs.len() >= limit {
break;
}
}
}
Ok::<String, crate::error::AppError>(serde_json::to_string(&refs)?)
Ok::<Vec<serde_json::Value>, crate::error::AppError>(refs)
})?;
Ok(matches)
if matches.len() < limit {
let remaining = limit - matches.len();
let disk_matches = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&target_sym, remaining, false)
})
.await
.unwrap_or_default();
matches.extend(disk_matches);
}
Ok(serde_json::to_string_pretty(&matches)?)
}
}
@@ -307,24 +404,38 @@ impl McpTool for GetCallersHandler {
let req: crate::tools::GetCallersTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(10);
let target_fn = req.function_name;
let target_fn = req.function_name.clone();
let callers = state.code.snippets.read_with(|snippets| {
let mut callers = state.code.snippets.read_with(|snippets| {
let mut matching = Vec::new();
for snippet in snippets {
if snippet.code.contains(&format!("{}(", target_fn))
|| snippet.code.contains(&format!("{}.await", target_fn))
{
matching.push(snippet.clone());
matching.push(serde_json::json!({
"source": "snippet",
"name": snippet.name,
"code": snippet.code,
}));
if matching.len() >= limit {
break;
}
}
}
Ok::<String, crate::error::AppError>(serde_json::to_string(&matching)?)
Ok::<Vec<serde_json::Value>, crate::error::AppError>(matching)
})?;
Ok(callers)
if callers.len() < limit {
let remaining = limit - callers.len();
let disk_callers = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&target_fn, remaining, true)
})
.await
.unwrap_or_default();
callers.extend(disk_callers);
}
Ok(serde_json::to_string_pretty(&callers)?)
}
}
@@ -357,6 +468,20 @@ impl McpTool for AnalyzeImpactHandler {
}
});
let sym_clone = sym.clone();
let disk_refs = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&sym_clone, 20, false)
})
.await
.unwrap_or_default();
for r in &disk_refs {
if let Some(path) = r.get("file_path").and_then(|p| p.as_str()) {
let line = r.get("line").and_then(|l| l.as_u64()).unwrap_or(0);
callers.push(format!("{}:{}", path, line));
}
}
let mut kg_connected = Vec::new();
state.read_graph(|g| {
for rel in &g.relations {
@@ -384,7 +509,7 @@ impl McpTool for AnalyzeImpactHandler {
"target_symbol": sym,
"risk_level": risk_level,
"ast_callers_count": caller_count,
"ast_callers_sample": callers.into_iter().take(5).collect::<Vec<_>>(),
"ast_callers_sample": callers.into_iter().take(10).collect::<Vec<_>>(),
"graph_relations_count": graph_count,
"graph_relations": kg_connected,
"recommendation": match risk_level {
server/src/handlers/env.rs
+5 -5
View File
@@ -16,7 +16,7 @@ impl McpTool for UpdateEnvFingerprintHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<UpdateEnvFingerprintTool>(
"update_env_fingerprint",
"Execute update_env_fingerprint",
"Update system OS, shell, and installed tool version fingerprints for a workspace namespace.",
)
}
@@ -50,7 +50,7 @@ impl McpTool for ReadEnvFingerprintHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<ReadEnvFingerprintTool>(
"read_env_fingerprint",
"Execute read_env_fingerprint",
"Read system OS, shell, and tool version fingerprints for a workspace namespace.",
)
}
@@ -81,7 +81,7 @@ impl McpTool for LogEnvRequirementHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<LogEnvRequirementTool>(
"log_env_requirement",
"Execute log_env_requirement",
"Log required environment variables, configuration keys, or secrets for a namespace.",
)
}
@@ -111,7 +111,7 @@ impl McpTool for RegisterEnvironmentHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<RegisterEnvironmentTool>(
"register_environment",
"Execute register_environment",
"Register infrastructure environment details (URL, VPN requirements, notes) for a namespace.",
)
}
@@ -144,7 +144,7 @@ impl McpTool for GetEnvironmentDetailsHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<GetEnvironmentDetailsTool>(
"get_environment_details",
"Execute get_environment_details",
"Retrieve registered infrastructure environment details for a namespace.",
)
}
server/src/handlers/git.rs
+2 -1
View File
@@ -69,7 +69,8 @@ impl McpTool for GetActiveWorktreeContextHandler {
// Truncate diff output if it's too large to save tokens
if diff_output.len() > 10000 {
diff_output.truncate(10000);
let valid_boundary = diff_output.floor_char_boundary(10000);
diff_output.truncate(valid_boundary);
diff_output.push_str("\n... [Diff truncated due to size]");
}
server/src/handlers/graph.rs
+213 -122
View File
@@ -12,6 +12,112 @@ struct BorrowedGraph<'a> {
relations: Vec<&'a crate::models::Relation>,
}
pub struct GraphQueryBuilder<'a> {
graph: &'a crate::models::KnowledgeGraph,
max_depth: usize,
relation_filter: Option<&'a str>,
}
impl<'a> GraphQueryBuilder<'a> {
pub fn new(graph: &'a crate::models::KnowledgeGraph) -> Self {
Self {
graph,
max_depth: 5,
relation_filter: None,
}
}
pub fn max_depth(mut self, depth: usize) -> Self {
self.max_depth = depth;
self
}
pub fn relation_filter(mut self, filter: Option<&'a str>) -> Self {
self.relation_filter = filter;
self
}
pub fn find_shortest_path(&self, start: &str, end: &str) -> Option<Vec<String>> {
let mut adj: std::collections::HashMap<&str, Vec<(&str, &str, bool)>> =
std::collections::HashMap::with_capacity(self.graph.relations.len() * 2);
for rel in &self.graph.relations {
if let Some(rf) = self.relation_filter {
if rel.relation_type != rf {
continue;
}
}
adj.entry(rel.from.as_str())
.or_default()
.push((rel.to.as_str(), rel.relation_type.as_str(), false));
adj.entry(rel.to.as_str())
.or_default()
.push((rel.from.as_str(), rel.relation_type.as_str(), true));
}
let mut queue = std::collections::VecDeque::new();
let mut visited = std::collections::HashSet::new();
let mut parents = std::collections::HashMap::new();
queue.push_back(start);
visited.insert(start);
let mut found = false;
let mut current_depth = 0;
let mut nodes_at_current_depth = 1;
let mut nodes_at_next_depth = 0;
while let Some(current) = queue.pop_front() {
if current == end {
found = true;
break;
}
// Visited node upper-bound cap to guarantee deterministic BFS bounds on dense graphs
if visited.len() > 10_000 {
break;
}
nodes_at_current_depth -= 1;
if current_depth < self.max_depth {
if let Some(neighbors) = adj.get(current) {
for &(neighbor, rel_type, is_inverse) in neighbors {
if !visited.contains(neighbor) {
visited.insert(neighbor);
parents.insert(neighbor, (current, rel_type, is_inverse));
queue.push_back(neighbor);
nodes_at_next_depth += 1;
}
}
}
}
if nodes_at_current_depth == 0 {
current_depth += 1;
nodes_at_current_depth = nodes_at_next_depth;
nodes_at_next_depth = 0;
}
}
if found {
let mut path = Vec::new();
let mut curr = end;
while curr != start {
if let Some((parent, rel_type, is_inverse)) = parents.get(&curr) {
if *is_inverse {
path.push(format!("{} -[inverse({})]-> {}", parent, rel_type, curr));
} else {
path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
}
curr = parent;
} else {
break;
}
}
path.reverse();
Some(path)
} else {
None
}
}
}
pub struct QueryGraphPathHandler;
#[async_trait]
@@ -21,88 +127,32 @@ impl McpTool for QueryGraphPathHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryGraphPathTool>("query_graph_path", "Execute query_graph_path")
crate::mcp::tool_def::<QueryGraphPathTool>(
"query_graph_path",
"Find the shortest relationship path between two entities in the knowledge graph within a maximum depth.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: crate::tools::QueryGraphPathTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
state.read_graph(|graph| {
let max_depth = req.max_depth.unwrap_or(5);
// Pre-index relations into an adjacency map for O(1) neighbor lookups
let mut adj: std::collections::HashMap<&str, Vec<(&str, &str, bool)>> = std::collections::HashMap::new();
for rel in &graph.relations {
adj.entry(rel.from.as_str())
.or_default()
.push((rel.to.as_str(), rel.relation_type.as_str(), false));
adj.entry(rel.to.as_str())
.or_default()
.push((rel.from.as_str(), rel.relation_type.as_str(), true));
}
let mut queue: std::collections::VecDeque<&str> = std::collections::VecDeque::new();
let mut visited: std::collections::HashSet<&str> = std::collections::HashSet::new();
let mut parents: std::collections::HashMap<&str, (&str, &str, bool)> =
std::collections::HashMap::new();
queue.push_back(req.start_node.as_str());
visited.insert(req.start_node.as_str());
let mut found = false;
let mut current_depth = 0;
let mut nodes_at_current_depth = 1;
let mut nodes_at_next_depth = 0;
while let Some(current) = queue.pop_front() {
if current == req.end_node {
found = true;
break;
tokio::task::spawn_blocking(move || {
state.read_graph(|graph| {
let max_depth = req.max_depth.unwrap_or(5);
let builder = GraphQueryBuilder::new(graph).max_depth(max_depth as usize);
if let Some(path) = builder.find_shortest_path(&req.start_node, &req.end_node) {
Ok(format!("Path found:\n{}", path.join("\n")))
} else {
Ok(format!(
"No path found between {} and {} within depth {}",
req.start_node, req.end_node, max_depth
))
}
nodes_at_current_depth -= 1;
if current_depth < max_depth {
if let Some(neighbors) = adj.get(current) {
for &(neighbor, rel_type, is_inverse) in neighbors {
if !visited.contains(neighbor) {
visited.insert(neighbor);
parents.insert(neighbor, (current, rel_type, is_inverse));
queue.push_back(neighbor);
nodes_at_next_depth += 1;
}
}
}
}
if nodes_at_current_depth == 0 {
current_depth += 1;
nodes_at_current_depth = nodes_at_next_depth;
nodes_at_next_depth = 0;
}
}
if found {
let mut path = Vec::new();
let mut curr = req.end_node.as_str();
while curr != req.start_node {
if let Some((parent, rel_type, is_inverse)) = parents.get(&curr) {
if *is_inverse {
path.push(format!("{} -[inverse({})]-> {}", parent, rel_type, curr));
} else {
path.push(format!("{} -[{}]-> {}", parent, rel_type, curr));
}
curr = parent;
} else {
break;
}
}
path.reverse();
Ok(format!("Path found:\n{}", path.join("\n")))
} else {
Ok(format!(
"No path found between {} and {} within depth {}",
req.start_node, req.end_node, max_depth
))
}
})
})
.await
.map_err(|e| crate::error::AppError::Internal(format!("Graph traversal task failed: {}", e)))?
}
}
@@ -115,7 +165,10 @@ impl McpTool for CreateEntitiesHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Execute create_entities")
crate::mcp::tool_def::<CreateEntitiesTool>(
"create_entities",
"Create new entities in the knowledge graph with normalized entity types.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -130,11 +183,12 @@ impl McpTool for CreateEntitiesHandler {
}
}
});
let idx = state.get_search_index();
for entity in inserted {
drop(idx.index_entity(&entity));
let names: Vec<String> = inserted.iter().map(|e| format!("{} ({})", e.name, e.entity_type)).collect();
if !inserted.is_empty() {
let idx = state.get_search_index().await;
let _ = idx.index_entities_batch(&inserted).await;
}
Ok("Entities created".to_string())
Ok(format!("Successfully created {} entity/entities: {}", names.len(), names.join(", ")))
}
}
@@ -147,7 +201,10 @@ impl McpTool for CreateRelationsHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateRelationsTool>("create_relations", "Execute create_relations")
crate::mcp::tool_def::<CreateRelationsTool>(
"create_relations",
"Create directed relationships between entities in the knowledge graph. Requires 'from', 'to', and 'relation_type'.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -183,23 +240,38 @@ impl McpTool for CreateRelationsHandler {
}
});
if !missing_nodes.is_empty() {
let missing: Vec<_> = missing_nodes.into_iter().collect();
return Err(crate::error::AppError::Internal(format!(
"Error: Relations dropped due to missing entities: {}",
missing.join(", ")
)));
}
let mut auto_created = Vec::new();
let mut added_relations = Vec::new();
state.modify_graph(|g| {
for node_name in missing_nodes {
if !g.entities.contains_key(&node_name) {
g.entities.insert(
node_name.clone(),
crate::models::Entity {
name: node_name.clone(),
entity_type: "Entity".to_string(),
observations: vec!["Auto-created stub entity for relation endpoint".to_string()],
namespace: crate::models::default_namespace(),
git_branch: None,
},
);
auto_created.push(node_name);
}
}
for mut relation in req.relations {
if !relation.from.is_empty() && !relation.to.is_empty() {
relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type);
added_relations.push(format!("{} -[{}]-> {}", relation.from, relation.relation_type, relation.to));
g.relations.push(relation);
}
}
});
Ok("Relations created".to_string())
let mut msg = format!("Successfully created {} relation(s):\n{}", added_relations.len(), added_relations.join("\n"));
if !auto_created.is_empty() {
msg.push_str(&format!("\nNote: Auto-created {} missing stub entity/entities: {}", auto_created.len(), auto_created.join(", ")));
}
Ok(msg)
}
}
@@ -212,7 +284,10 @@ impl McpTool for AddObservationsHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<AddObservationsTool>("add_observations", "Execute add_observations")
crate::mcp::tool_def::<AddObservationsTool>(
"add_observations",
"Add new observations and factual statements to existing entities in the knowledge graph.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -284,7 +359,7 @@ impl McpTool for DeleteEntitiesHandler {
.retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
});
let idx = state.get_search_index();
let idx = state.get_search_index().await;
for name in to_delete {
drop(idx.delete_document(&name));
}
@@ -361,7 +436,7 @@ impl McpTool for DeleteRelationsHandler {
});
if missing_count > 0 {
return Err(crate::error::AppError::Internal(format!(
"Error: {} relation(s) not found in graph. Please verify exact relation properties using read_graph.",
"Error: {} relation(s) not found in graph. Please verify exact relation properties (from, to, relation_type) using read_graph or get_subgraph.",
missing_count
)));
}
@@ -378,7 +453,10 @@ impl McpTool for ReadGraphHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ReadGraphTool>("read_graph", "Execute read_graph")
crate::mcp::tool_def::<ReadGraphTool>(
"read_graph",
"Read entities and relations from the knowledge graph with optional namespace filtering and token truncation. For large graphs, specify 'namespace' or use 'search_nodes' or 'get_subgraph' for targeted discovery.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -405,8 +483,9 @@ impl McpTool for ReadGraphHandler {
if let Some(max_tok) = max_tokens {
let max_chars = max_tok * 4;
if result_json.len() > max_chars {
result_json.truncate(max_chars);
result_json.push_str("... [TRUNCATED_TO_MAX_TOKENS]");
let valid_boundary = result_json.floor_char_boundary(max_chars);
result_json.truncate(valid_boundary);
result_json.push_str("\n... [TRUNCATED_TO_MAX_TOKENS. Use search_nodes or get_subgraph for targeted discovery]");
}
}
Ok(result_json)
@@ -432,12 +511,10 @@ impl McpTool for SearchNodesHandler {
let limit = req.limit.unwrap_or(10);
let include_body = req.include_body.unwrap_or(false);
let matches = if let Ok(idx) = state.search_index.read() {
idx.search(&req.query, req.namespace.as_deref())
.unwrap_or_default()
} else {
vec![]
};
let idx = state.get_search_index().await;
let matches = idx
.search(&req.query, req.namespace.as_deref())
.unwrap_or_default();
let data = state.read_graph(|full| -> crate::error::Result<String> {
let mut matched_entities = Vec::new();
@@ -709,6 +786,11 @@ impl McpTool for MergeEntitiesHandler {
r.to = req.target_entity.clone();
}
// Filter out self-loops
if r.from == r.to {
return false;
}
if r.from == req.target_entity || r.to == req.target_entity {
seen.insert(r.clone())
} else {
@@ -908,16 +990,17 @@ impl McpTool for SweepGraphHealthHandler {
}
}
// 2. Compute similarity pairs for duplicate detection
// 2. Compute similarity pairs for duplicate detection using pre-computed lowercase names
let names: Vec<_> = g.entities.keys().cloned().collect();
let lower_names: Vec<String> = names.iter().map(|n| n.to_lowercase()).collect();
for i in 0..names.len() {
for j in (i + 1)..names.len() {
let n1 = &names[i];
let n2 = &names[j];
let l1 = &lower_names[i];
let l2 = &lower_names[j];
let l1 = n1.to_lowercase();
let l2 = n2.to_lowercase();
if l1 == l2 || ((l1.contains(&l2) || l2.contains(&l1)) && l1.len().min(l2.len()) > 3) {
if l1 == l2 || ((l1.contains(l2.as_str()) || l2.contains(l1.as_str())) && l1.len().min(l2.len()) > 3) {
duplicates.push(serde_json::json!({
"entity_a": n1,
"entity_b": n2,
@@ -1050,7 +1133,8 @@ impl McpTool for SummarizeSubgraphHandler {
let max_tokens = req.max_tokens.unwrap_or(1000);
let max_chars = max_tokens * 4;
if markdown.len() > max_chars {
markdown.truncate(max_chars);
let valid_boundary = markdown.floor_char_boundary(max_chars);
markdown.truncate(valid_boundary);
markdown.push_str("\n... [Truncated to fit token budget]");
}
@@ -1063,13 +1147,11 @@ mod tests {
use super::*;
use crate::handlers::meta::{BroadcastAgentSignalHandler, QueryAgentSignalsHandler};
use serde_json::json;
use tempfile::tempdir;
#[tokio::test]
async fn test_create_and_read_entities() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let state = Arc::new(MemoryState::new_in_memory());
let create_handler = CreateEntitiesHandler;
let args = json!({
@@ -1083,7 +1165,7 @@ mod tests {
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res, "Entities created");
assert!(res.contains("Successfully created 1 entity/entities"));
// Ensure graph contains the entity
state.graph.read_with(|g| {
@@ -1094,8 +1176,7 @@ mod tests {
#[tokio::test]
async fn test_create_relations() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let state = Arc::new(MemoryState::new_in_memory());
// Needs entities first
state.graph.modify(|g| {
@@ -1132,14 +1213,26 @@ mod tests {
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert_eq!(res, "Relations created");
assert!(res.contains("Successfully created 1 relation(s)"));
// Test semantic LLM schema feedback (User request)
let bad_args = json!({
// Test serde field aliases (source/target/relationType mapped to from/to/relation_type)
let alias_args = json!({
"relations": [
{"source": "A", "target": "B", "relationType": "knows"}
]
});
let alias_res = handler
.execute(alias_args, state.clone())
.await
.unwrap();
assert!(alias_res.contains("Successfully created 1 relation(s)"));
// Test semantic LLM schema feedback on missing fields
let bad_args = json!({
"relations": [
{"invalid_field": "X"}
]
});
let err_res = handler
.execute(bad_args, state.clone())
.await
@@ -1151,8 +1244,7 @@ mod tests {
#[tokio::test]
async fn test_observations_and_reads() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let state = Arc::new(MemoryState::new_in_memory());
// Need entity first
state.graph.modify(|g| {
@@ -1208,8 +1300,7 @@ mod tests {
#[tokio::test]
async fn test_advanced_graph_operations() {
let dir = tempfile::tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let state = Arc::new(MemoryState::new_in_memory());
let create_handler = CreateEntitiesHandler;
let args_ent = json!({
@@ -1289,13 +1380,13 @@ mod tests {
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap();
assert!(!res_orphans.contains("Y"));
// After merging X into Y and purging self-loops, Y is the sole node and becomes an orphan
assert!(res_orphans.contains("Y"));
}
#[tokio::test]
async fn test_more_graph_handlers() {
let dir = tempfile::tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let state = Arc::new(MemoryState::new_in_memory());
let create_handler = CreateEntitiesHandler;
let args_ent = json!({
server/src/handlers/logs.rs
+4 -2
View File
@@ -64,10 +64,12 @@ impl McpTool for GetRecentLogsHandler {
file.seek(SeekFrom::End(-(read_size as i64)))
.map_err(|e| e.to_string())?;
let mut buffer = String::new();
file.read_to_string(&mut buffer)
let mut vec_buf = Vec::new();
file.read_to_end(&mut vec_buf)
.map_err(|e| e.to_string())?;
let buffer = String::from_utf8_lossy(&vec_buf).to_string();
let lines: Vec<&str> = buffer.lines().collect();
let recent_lines = if lines.len() > 100 {
lines[lines.len() - 100..].join("\n")
server/src/handlers/meta.rs
+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
server/src/handlers/notes.rs
+33 -23
View File
@@ -49,17 +49,21 @@ impl McpTool for StickyNotesHandler {
}
StickyNoteAction::Read => {
let now = crate::handlers::utils::now_secs();
let mut active_notes = Vec::new();
state.code.sticky.modify(|notes| {
notes.retain(|n| {
if let Some(exp) = n.expires_at {
exp > now
} else {
true
}
});
active_notes = notes.clone();
let has_expired = state.code.sticky.read_with(|notes| {
notes.iter().any(|n| matches!(n.expires_at, Some(exp) if exp <= now))
});
if has_expired {
state.code.sticky.modify(|notes| {
notes.retain(|n| {
if let Some(exp) = n.expires_at {
exp > now
} else {
true
}
});
});
}
let active_notes = state.code.sticky.read_with(|notes| notes.clone());
Ok(serde_json::to_string(&active_notes)?)
}
StickyNoteAction::Delete => {
@@ -68,8 +72,9 @@ impl McpTool for StickyNotesHandler {
})?;
let mut success = false;
state.code.sticky.modify(|notes| {
if idx > 0 && idx <= notes.len() {
notes.remove(idx - 1);
let target_pos = if idx == 0 { 0 } else { idx - 1 };
if target_pos < notes.len() {
notes.remove(target_pos);
success = true;
}
});
@@ -77,7 +82,7 @@ impl McpTool for StickyNotesHandler {
Ok("Sticky note deleted.".to_string())
} else {
Err(crate::error::AppError::Internal(
"Invalid sticky note index. Next step: Call sticky_notes with action='read' to view active notes and valid 1-based indices, then retry.".to_string(),
"Invalid sticky note index. Next step: Call sticky_notes with action='read' to view active notes and valid indices, then retry.".to_string(),
))
}
}
@@ -214,17 +219,22 @@ impl McpTool for GenerateStandupReportHandler {
serde_json::from_value(args).map_err(|e| e.to_string())?;
let cutoff = crate::handlers::utils::now_secs().saturating_sub(req.hours_lookback * 3600);
let report_str = state.project.tasks.read_with(|items| {
state.code.ledger.read_with(|changes| {
state.telemetry.session_summaries.read_with(|summaries| {
let filtered_tasks: Vec<_> = items.iter().filter(|t| t.updated_at >= cutoff).collect();
let filtered_changes: Vec<_> = changes.iter().filter(|c| c.timestamp >= cutoff).collect();
let filtered_summaries: Vec<_> = summaries.iter().filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff).collect();
let report = serde_json::json!({ "tasks_updated": filtered_tasks, "code_changes": filtered_changes, "session_summaries": filtered_summaries });
report.to_string()
})
})
let filtered_tasks = state.project.tasks.read_with(|items| {
items.iter().filter(|t| t.updated_at >= cutoff).cloned().collect::<Vec<_>>()
});
let filtered_changes = state.code.ledger.read_with(|changes| {
changes.iter().filter(|c| c.timestamp >= cutoff).cloned().collect::<Vec<_>>()
});
let filtered_summaries = state.telemetry.session_summaries.read_with(|summaries| {
summaries.iter().filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff).cloned().collect::<Vec<_>>()
});
let report = serde_json::json!({
"tasks_updated": filtered_tasks,
"code_changes": filtered_changes,
"session_summaries": filtered_summaries
});
let report_str = report.to_string();
Ok(report_str)
}
}
server/src/handlers/tasks.rs
+4 -3
View File
@@ -46,7 +46,7 @@ impl McpTool for TasksHandler {
acceptance_criteria: vec![],
expires_at: None,
};
let idx = state.get_search_index();
let idx = state.get_search_index().await;
drop(idx.index_task(&task));
state.project.tasks.modify(|tasks| {
tasks.push(task.clone());
@@ -187,7 +187,7 @@ impl McpTool for TasksHandler {
});
if deleted_count > 0 {
let idx = state.get_search_index();
let idx = state.get_search_index().await;
for deleted_id in actually_deleted {
drop(idx.delete_document(&deleted_id));
}
@@ -228,7 +228,8 @@ impl McpTool for TasksHandler {
if let Some(max_t) = req.max_tokens {
let char_limit = max_t * 4;
if json_str.len() > char_limit {
json_str.truncate(char_limit);
let boundary = json_str.floor_char_boundary(char_limit);
json_str.truncate(boundary);
json_str.push_str(" ...[truncated due to max_tokens]");
}
}
server/src/handlers/vision.rs
+81 -88
View File
@@ -30,33 +30,56 @@ impl McpTool for WriteClipboardHandler {
let res = tokio::task::spawn_blocking(move || {
let mut msgs = Vec::new();
if let Ok(mut clipboard) = Clipboard::new()
&& let Some(text) = &tool_args.text
&& clipboard.set_text(text).is_ok()
{
msgs.push("Wrote text");
if let Some(text) = &tool_args.text {
let mut written = false;
if let Ok(mut clipboard) = Clipboard::new() {
for _ in 0..3 {
if clipboard.set_text(text).is_ok() {
written = true;
break;
}
std::thread::sleep(std::time::Duration::from_millis(20));
}
}
if written {
msgs.push("Wrote text");
}
}
// Handle arboard for image
if let Some(image_path) = &tool_args.image_path {
match image::open(image_path) {
Ok(img) => {
let img = if img.width() > 2048 || img.height() > 2048 {
img.resize(2048, 2048, FilterType::Triangle)
} else {
img
};
let rgba = img.into_rgba8();
let (w, h) = rgba.dimensions();
let img_data = ImageData {
width: w as usize,
height: h as usize,
bytes: Cow::Owned(rgba.into_raw()),
};
let raw_bytes = rgba.into_raw();
let mut written = false;
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if clipboard.set_image(img_data).is_ok() {
msgs.push("Wrote image");
} else {
return Err(crate::error::AppError::Internal(
"Failed to write image to clipboard".to_string(),
));
for _ in 0..3 {
let img_data = ImageData {
width: w as usize,
height: h as usize,
bytes: Cow::Borrowed(&raw_bytes),
};
if clipboard.set_image(img_data).is_ok() {
written = true;
break;
}
std::thread::sleep(std::time::Duration::from_millis(20));
}
}
if written {
msgs.push("Wrote image");
} else {
return Err(crate::error::AppError::Internal(
"Failed to write image to clipboard".to_string(),
));
}
}
Err(e) => {
return Err(crate::error::AppError::Internal(format!(
@@ -85,92 +108,34 @@ impl McpTool for WriteClipboardHandler {
}
pub fn get_native_clipboard_text() -> Option<String> {
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(text) = clipboard.get_text() {
if !text.trim().is_empty() {
return Some(text);
}
}
}
let mut cmd_wl = std::process::Command::new("wl-paste");
cmd_wl.arg("--no-newline");
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
}
if let Ok(output) = cmd_wl.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(text) = String::from_utf8(output.stdout) {
for _ in 0..3 {
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(text) = clipboard.get_text() {
if !text.trim().is_empty() {
return Some(text);
}
}
}
std::thread::sleep(std::time::Duration::from_millis(20));
}
let mut cmd_xc = std::process::Command::new("xclip");
cmd_xc.args(["-selection", "clipboard", "-o"]);
if std::env::var("DISPLAY").is_err() {
cmd_xc.env("DISPLAY", ":0");
}
if let Ok(output) = cmd_xc.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(text) = String::from_utf8(output.stdout) {
if !text.trim().is_empty() {
return Some(text);
}
}
}
}
None
}
pub fn get_native_clipboard_image() -> Option<image::DynamicImage> {
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(image_data) = clipboard.get_image() {
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
image_data.width as u32,
image_data.height as u32,
image_data.bytes.into_owned(),
) {
return Some(image::DynamicImage::ImageRgba8(img));
}
}
}
for mime in &["image/png", "image/jpeg", "image/bmp", "image/tiff"] {
let mut cmd_wl = std::process::Command::new("wl-paste");
cmd_wl.args(["--type", mime]);
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
}
if let Ok(output) = cmd_wl.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(img) = image::load_from_memory(&output.stdout) {
return Some(img);
for _ in 0..3 {
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(image_data) = clipboard.get_image() {
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
image_data.width as u32,
image_data.height as u32,
image_data.bytes.into_owned(),
) {
return Some(image::DynamicImage::ImageRgba8(img));
}
}
}
std::thread::sleep(std::time::Duration::from_millis(20));
}
for mime in &["image/png", "image/jpeg", "image/bmp"] {
let mut cmd_xc = std::process::Command::new("xclip");
cmd_xc.args(["-selection", "clipboard", "-t", mime, "-o"]);
if std::env::var("DISPLAY").is_err() {
cmd_xc.env("DISPLAY", ":0");
}
if let Ok(output) = cmd_xc.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(img) = image::load_from_memory(&output.stdout) {
return Some(img);
}
}
}
}
None
}
@@ -422,5 +387,33 @@ mod tests {
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(parsed.is_object());
}
#[test]
fn test_no_subprocess_clipboard_regression() {
let vision_src = include_str!("vision.rs");
let code_only = vision_src.split("mod tests").next().unwrap_or(vision_src);
let forbidden_cmd = format!("Command::{}{}", "n", "ew");
let forbidden_ps = format!("power{}", "shell");
let forbidden_wl = format!("wl-{}", "paste");
let forbidden_xc = format!("x{}", "clip");
assert!(
!code_only.contains(&forbidden_cmd),
"Regression detected: vision.rs must not spawn subprocesses!"
);
assert!(
!code_only.contains(&forbidden_ps),
"Regression detected: vision.rs must not invoke powershell!"
);
assert!(
!code_only.contains(&forbidden_wl),
"Regression detected: vision.rs must not invoke wl-paste!"
);
assert!(
!code_only.contains(&forbidden_xc),
"Regression detected: vision.rs must not invoke xclip!"
);
}
}
server/src/handlers/workspaces.rs
+10 -7
View File
@@ -15,7 +15,7 @@ impl McpTool for PinFileHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<PinFileTool>("pin_file", "Execute pin_file")
crate::mcp::tool_def::<PinFileTool>("pin_file", "Pin a file path in the workspace context to prioritize it across subagent turns.")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -42,7 +42,7 @@ impl McpTool for UnpinFileHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Execute unpin_file")
crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Unpin a file path from the workspace context.")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -65,7 +65,7 @@ impl McpTool for ListPinnedFilesHandler {
fn schema(&self) -> Value {
crate::mcp::tool_def::<ListPinnedFilesTool>(
"list_pinned_files",
"Execute list_pinned_files",
"List active pinned files in the workspace context, optionally filtered by namespace or branch.",
)
}
@@ -104,7 +104,7 @@ impl McpTool for StoreSnippetHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Execute store_snippet")
crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Store or replace a reusable code snippet with vector embedding and search tags.")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -127,7 +127,7 @@ impl McpTool for StoreSnippetHandler {
embedding,
};
let idx = state.get_search_index();
let idx = state.get_search_index().await;
drop(idx.index_snippet(&snippet));
state.code.snippets.modify(|snippets| {
@@ -148,7 +148,7 @@ impl McpTool for TagSnippetHandler {
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<TagSnippetTool>("tag_snippet", "Execute tag_snippet")
crate::mcp::tool_def::<TagSnippetTool>("tag_snippet", "Add organizational tags to an existing stored code snippet.")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -254,7 +254,7 @@ impl McpTool for DeleteSnippetHandler {
deleted = snippets.len() < orig;
});
if deleted {
let idx = state.get_search_index();
let idx = state.get_search_index().await;
drop(idx.delete_document(&req.name));
Ok("Snippet deleted.".to_string())
} else {
@@ -870,6 +870,9 @@ impl McpTool for ReadDirectoryArchitectureHandler {
summary: &mut String,
depth: usize,
) -> std::io::Result<()> {
if depth > 10 {
return Ok(());
}
if dir.is_dir() {
let mut entries =
fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?;