Files
mcp-memory/server/src/handlers.rs
T

3056 lines
131 KiB
Rust

use crate::models::*;
use crate::state::MemoryState;
use crate::tools::*;
macro_rules! parse_tool {
($args:expr, $id:expr, $type:ty) => {
match parse_args::<$type>($args) {
Ok(r) => r,
Err(e) => {
let response = Some(crate::mcp::success(
$id.clone(),
serde_json::json!({"isError": true, "content": [{"type": "text", "text": format!("Invalid args: {}", e)}] }),
));
tracing::trace!("Returning response from handle_request: {:?}", response);
return response;
}
}
};
}
use serde::de::DeserializeOwned;
use std::collections::HashSet;
use std::sync::Arc;
use std::time::{SystemTime, UNIX_EPOCH};
fn parse_args<T: DeserializeOwned>(args: serde_json::Value) -> Result<T, String> {
serde_json::from_value(args).map_err(|e| format!("Invalid args: {}", e))
}
pub struct MemoryHandler {
pub state: Arc<MemoryState>,
}
impl MemoryHandler {
pub async fn handle_request(&self, req: serde_json::Value) -> Option<serde_json::Value> {
let start_time = std::time::Instant::now();
let id = req.get("id").cloned().unwrap_or(serde_json::Value::Null);
let id_clone = id.clone();
let method = req.get("method").and_then(|m| m.as_str()).unwrap_or("");
let tool_name = if method == "tools/call" {
req.get("params").and_then(|p| p.get("name")).and_then(|n| n.as_str()).unwrap_or("unknown")
} else {
""
};
if method == "tools/call" {
tracing::info!(">>> [Server] Handling MCP tool call: {} (id: {})", tool_name, id);
} else {
tracing::debug!(">>> [Server] Handling MCP request method: {}", method);
}
tracing::trace!(">>> [Server] Full MCP Request payload: {}", req.to_string());
let response = match method {
"server/discover" => {
let payload = serde_json::json!({
"resultType": "complete",
"ttlMs": 0,
"cacheScope": "public",
"supportedVersions": ["2026-07-28", "2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05"],
"capabilities": {
"tools": serde_json::json!({})
},
"_meta": {
"io.modelcontextprotocol/serverInfo": {
"name": "gemini-mcp-memory",
"version": "3.0.0"
}
}
});
tracing::debug!(
"<<< [Server] Replying to server/discover with payload: {}",
payload.to_string()
);
Some(crate::mcp::success(id, payload))
}
"initialize" => {
let init = rmcp::model::InitializeResult::new(
rmcp::model::ServerCapabilities::builder()
.enable_tools()
.build(),
)
.with_server_info(rmcp::model::Implementation::new(
"gemini-mcp-memory",
"3.0.0",
));
tracing::debug!("<<< [Server] Replying to initialize with rmcp payload");
Some(crate::mcp::success(
id,
serde_json::to_value(&init).unwrap(),
))
}
"notifications/initialized" => None,
"tools/list" => {
let tools = vec![
crate::mcp::tool_def::<crate::tools::QueryGraphPathTool>(
"query_graph_path",
"Traverse the knowledge graph to find a path between two entities.",
),
crate::mcp::tool_def::<CreateEntitiesTool>(
"create_entities",
"Create new entities in the knowledge graph.",
),
crate::mcp::tool_def::<CreateRelationsTool>(
"create_relations",
"Create new relations between entities in the knowledge graph.",
),
crate::mcp::tool_def::<AddObservationsTool>(
"add_observations",
"Add new observations to existing entities in the knowledge graph.",
),
crate::mcp::tool_def::<DeleteEntitiesTool>(
"delete_entities",
"Delete entities from the knowledge graph.",
),
crate::mcp::tool_def::<DeleteObservationsTool>(
"delete_observations",
"Delete observations from existing entities.",
),
crate::mcp::tool_def::<DeleteRelationsTool>(
"delete_relations",
"Delete relations between entities.",
),
crate::mcp::tool_def::<ReadGraphTool>(
"read_graph",
"Read the entire knowledge graph.",
),
crate::mcp::tool_def::<SearchNodesTool>(
"search_nodes",
"Search for entities in the knowledge graph by name or type.",
),
crate::mcp::tool_def::<OpenNodesTool>(
"open_nodes",
"Open and retrieve full details of specific nodes in the knowledge graph.",
),
crate::mcp::tool_def::<LogCodeChangeTool>(
"log_code_change",
"Log a significant code change or refactor in the memory system.",
),
crate::mcp::tool_def::<QueryRecentChangesTool>(
"query_recent_changes",
"Query recently logged code changes.",
),
crate::mcp::tool_def::<VisualizeGraphTool>(
"visualize_graph",
"Generate a visual representation of the knowledge graph.",
),
crate::mcp::tool_def::<AddStickyNoteTool>(
"add_sticky_note",
"Add a sticky note for unstructured thoughts or reminders.",
),
crate::mcp::tool_def::<ReadStickyNotesTool>(
"read_sticky_notes",
"Read all active sticky notes.",
),
crate::mcp::tool_def::<DeleteStickyNoteTool>(
"delete_sticky_note",
"Delete a specific sticky note by its 1-indexed position.",
),
crate::mcp::tool_def::<ClearStickyNotesTool>(
"clear_sticky_notes",
"Clear all active sticky notes.",
),
crate::mcp::tool_def::<CondenseEntityTool>(
"condense_entity",
"Condense or summarize an entity's observations to reduce size.",
),
crate::mcp::tool_def::<AddTaskTool>(
"add_task",
"Add a new task to the task tracker.",
),
crate::mcp::tool_def::<UpdateTaskStatusTool>(
"update_task_status",
"Update the status of an existing task.",
),
crate::mcp::tool_def::<DeleteTaskTool>(
"delete_task",
"Delete a task and all its children.",
),
crate::mcp::tool_def::<ListActiveTasksTool>(
"list_active_tasks",
"List all currently active tasks.",
),
crate::mcp::tool_def::<SetAcceptanceCriteriaTool>(
"set_acceptance_criteria",
"Define a strict checklist of acceptance criteria for a given task.",
),
crate::mcp::tool_def::<VerifyAcceptanceCriteriaTool>(
"verify_acceptance_criteria",
"Mark a previously defined acceptance criteria as met.",
),
crate::mcp::tool_def::<StoreSnippetTool>(
"store_snippet",
"Store a reusable code snippet.",
),
crate::mcp::tool_def::<SearchSnippetsTool>(
"search_snippets",
"Search through stored code snippets.",
),
crate::mcp::tool_def::<DeleteSnippetTool>(
"delete_snippet",
"Delete a stored code snippet.",
),
crate::mcp::tool_def::<LogDecisionTool>(
"log_decision",
"Log an architectural decision record (ADR).",
),
crate::mcp::tool_def::<QueryDecisionsTool>(
"query_decisions",
"Query architectural decision records.",
),
crate::mcp::tool_def::<MergeEntitiesTool>(
"merge_entities",
"Merge two entities in the knowledge graph into one.",
),
crate::mcp::tool_def::<FindOrphansTool>(
"find_orphans",
"Find orphaned entities (entities without any relations) in the graph.",
),
crate::mcp::tool_def::<LearnPreferenceTool>(
"learn_preference",
"Record a user preference or behavior to adapt future interactions.",
),
crate::mcp::tool_def::<ReadPreferencesTool>(
"read_preferences",
"Read all learned user preferences.",
),
crate::mcp::tool_def::<LogErrorFixTool>(
"log_error_fix",
"Log a complex error and its fix for future reference.",
),
crate::mcp::tool_def::<SearchErrorFixesTool>(
"search_error_fixes",
"Search through previously logged error fixes.",
),
crate::mcp::tool_def::<PinFileTool>(
"pin_file",
"Pin a file to keep it explicitly in the context workspace.",
),
crate::mcp::tool_def::<UnpinFileTool>(
"unpin_file",
"Unpin a file from the context workspace.",
),
crate::mcp::tool_def::<ListPinnedFilesTool>(
"list_pinned_files",
"List all currently pinned files.",
),
crate::mcp::tool_def::<AddSessionSummaryTool>(
"add_session_summary",
"Add a summary of the current session.",
),
crate::mcp::tool_def::<GetProjectTimelineTool>(
"get_project_timeline",
"Get a timeline of major project events.",
),
crate::mcp::tool_def::<LeaveHandoffMemoTool>(
"leave_handoff_memo",
"Leave a memo for the next session or agent.",
),
crate::mcp::tool_def::<ReadHandoffMemosTool>(
"read_handoff_memos",
"Read pending handoff memos.",
),
crate::mcp::tool_def::<ClearHandoffMemosTool>(
"clear_handoff_memos",
"Clear handoff memos after reading.",
),
crate::mcp::tool_def::<UpdateEnvFingerprintTool>(
"update_env_fingerprint",
"Update the environment fingerprint (e.g., OS, tool versions).",
),
crate::mcp::tool_def::<ReadEnvFingerprintTool>(
"read_env_fingerprint",
"Read the current environment fingerprint.",
),
crate::mcp::tool_def::<LogEnvRequirementTool>(
"log_env_requirement",
"Log a required tool or package for the environment.",
),
crate::mcp::tool_def::<AddMilestoneTool>(
"add_milestone",
"Add a new project milestone.",
),
crate::mcp::tool_def::<UpdateMilestoneTool>(
"update_milestone",
"Update the status of a project milestone.",
),
crate::mcp::tool_def::<ListMilestonesTool>(
"list_milestones",
"List all project milestones.",
),
crate::mcp::tool_def::<GenerateStandupReportTool>(
"generate_standup_report",
"Generate a standup report summarizing recent work, blockers, and next steps.",
),
crate::mcp::tool_def::<RegisterEnvironmentTool>(
"register_environment",
"Register details about a specific deployment environment.",
),
crate::mcp::tool_def::<GetEnvironmentDetailsTool>(
"get_environment_details",
"Get detailed information about a specific deployment environment.",
),
crate::mcp::tool_def::<AddPrChecklistItemTool>(
"add_pr_checklist_item",
"Add an item to the PR checklist.",
),
crate::mcp::tool_def::<GetPrChecklistTool>(
"get_pr_checklist",
"Get the current PR checklist.",
),
crate::mcp::tool_def::<ClearPrChecklistTool>(
"clear_pr_checklist",
"Clear the PR checklist.",
),
crate::mcp::tool_def::<LogTechDebtTool>(
"log_tech_debt",
"Log identified technical debt.",
),
crate::mcp::tool_def::<ResolveTechDebtTool>(
"resolve_tech_debt",
"Mark a logged technical debt as resolved.",
),
crate::mcp::tool_def::<ListTechDebtTool>(
"list_tech_debt",
"List all unresolved technical debt.",
),
crate::mcp::tool_def::<SaveContextWorkspaceTool>(
"save_context_workspace",
"Save the current set of pinned files and context.",
),
crate::mcp::tool_def::<LoadContextWorkspaceTool>(
"load_context_workspace",
"Load a previously saved context workspace.",
),
crate::mcp::tool_def::<ListContextWorkspacesTool>(
"list_context_workspaces",
"List all saved context workspaces.",
),
crate::mcp::tool_def::<OmniSearchTool>(
"omni_search",
"Search across all memory sources (graph, tasks, snippets, ADRs, etc.) at once.",
),
crate::mcp::tool_def::<GetProjectHealthTool>(
"get_project_health",
"Get a synthesized health report of the project based on memory data.",
),
];
Some(crate::mcp::success(
id,
serde_json::json!({ "tools": tools }),
))
}
"tools/call" => {
let params = req.get("params").unwrap_or(&serde_json::Value::Null);
let name = params.get("name").and_then(|n| n.as_str()).unwrap_or("");
let args = params
.get("arguments")
.cloned()
.unwrap_or(serde_json::Value::Object(Default::default()));
self.state
.broadcast_activity(&format!("Agent executed tool: {}", name));
let result: Result<String, String> = match name {
"query_graph_path" => {
let req = parse_tool!(args.clone(), id, crate::tools::QueryGraphPathTool);
let graph = self.state.get_full_graph();
let max_depth = req.max_depth.unwrap_or(5);
let mut queue = std::collections::VecDeque::new();
let mut visited = std::collections::HashSet::new();
let mut parents: std::collections::HashMap<String, (String, String)> =
std::collections::HashMap::new();
queue.push_back(req.start_node.clone());
visited.insert(req.start_node.clone());
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;
}
nodes_at_current_depth -= 1;
if current_depth < max_depth {
for rel in &graph.relations {
if rel.from == current && !visited.contains(&rel.to) {
visited.insert(rel.to.clone());
parents.insert(
rel.to.clone(),
(current.clone(), rel.relation_type.clone()),
);
queue.push_back(rel.to.clone());
nodes_at_next_depth += 1;
} else if rel.to == current && !visited.contains(&rel.from) {
visited.insert(rel.from.clone());
parents.insert(
rel.from.clone(),
(
current.clone(),
format!("inverse({})", rel.relation_type),
),
);
queue.push_back(rel.from.clone());
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.clone();
while curr != req.start_node {
let (parent, rel) = parents.get(&curr).unwrap();
path.push(format!("({}) --[{}]--> ({})", parent, rel, curr));
curr = parent.clone();
}
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
))
}
}
"create_entities" => {
let req = parse_tool!(args.clone(), id, CreateEntitiesTool);
self.state.modify_graph(|g| {
for entity in req.entities {
if !entity.name.is_empty() {
if let Ok(idx) = self.state.search_index.read() {
drop(idx.index_entity(&entity));
}
g.entities.insert(entity.name.clone(), entity);
}
}
});
Ok("Entities created".to_string())
}
"create_relations" => {
let req = parse_tool!(args.clone(), id, CreateRelationsTool);
self.state.modify_graph(|g| {
for relation in req.relations {
if !relation.from.is_empty() && !relation.to.is_empty() {
g.relations.push(relation);
}
}
});
Ok("Relations created".to_string())
}
"add_observations" => {
let req = parse_tool!(args.clone(), id, AddObservationsTool);
let full = self.state.get_full_graph();
self.state.modify_graph(|g| {
for o in req.observations {
if let Some(full_e) = full.entities.get(&o.entity_name) {
let mut e =
g.entities.get(&o.entity_name).cloned().unwrap_or_else(
|| Entity {
name: o.entity_name.clone(),
entity_type: full_e.entity_type.clone(),
observations: vec![],
namespace: full_e.namespace.clone(),
git_branch: None,
},
);
e.observations.extend(o.contents);
g.entities.insert(o.entity_name, e);
}
}
});
Ok("Observations added".to_string())
}
"delete_entities" => {
let req = parse_tool!(args.clone(), id, DeleteEntitiesTool);
let to_delete: HashSet<_> = req.entity_names.into_iter().collect();
self.state.modify_graph(|master| {
for name in &to_delete {
master.entities.remove(name);
}
master.relations.retain(|r| {
!to_delete.contains(&r.from) && !to_delete.contains(&r.to)
});
});
Ok("Entities deleted".to_string())
}
"delete_observations" => {
let req = parse_tool!(args.clone(), id, DeleteObservationsTool);
self.state.modify_graph(|master| {
for d in req.deletions {
if let Some(e) = master.entities.get_mut(&d.entity_name) {
let to_rem: HashSet<_> = d.observations.into_iter().collect();
e.observations.retain(|o| !to_rem.contains(o));
}
}
});
Ok("Observations deleted".to_string())
}
"delete_relations" => {
let req = parse_tool!(args.clone(), id, DeleteRelationsTool);
self.state.modify_graph(|master| {
let to_rem: HashSet<_> = req.relations.into_iter().collect();
master.relations.retain(|r| !to_rem.contains(r));
});
Ok("Relations deleted".to_string())
}
"read_graph" => {
let req = parse_tool!(args.clone(), id, ReadGraphTool);
let mut full = self.state.get_full_graph();
if let Some(ns) = req.namespace {
full.entities.retain(|_, e| e.namespace == ns);
full.relations.retain(|r| r.namespace == ns);
}
let data = serde_json::to_string(&full).unwrap_or_default();
Ok(data.to_string())
}
"search_nodes" => {
let req = parse_tool!(args.clone(), id, SearchNodesTool);
let matches = if let Ok(idx) = self.state.search_index.read() {
idx.search(&req.query, req.namespace.as_deref())
.unwrap_or_default()
} else {
vec![]
};
let mut result = KnowledgeGraph::default();
let full = self.state.get_full_graph();
for (id, doc_type, _, _, _) in matches {
if doc_type == "entity"
&& let Some(e) = full.entities.get(&id)
{
result.entities.insert(id, e.clone());
}
}
let data = serde_json::to_string(&result).unwrap_or_default();
Ok(data.to_string())
}
"open_nodes" => {
let req = parse_tool!(args.clone(), id, OpenNodesTool);
let targets: HashSet<_> = req.names.into_iter().collect();
let full = self.state.get_full_graph();
let mut result = KnowledgeGraph::default();
let mut connected = HashSet::new();
for r in &full.relations {
if targets.contains(&r.from) {
connected.insert(r.to.clone());
result.relations.push(r.clone());
} else if targets.contains(&r.to) {
connected.insert(r.from.clone());
result.relations.push(r.clone());
}
}
for (name, e) in full.entities {
if targets.contains(&name) || connected.contains(&name) {
result.entities.insert(name, e);
}
}
let data = serde_json::to_string(&result).unwrap_or_default();
Ok(data.to_string())
}
"log_code_change" => {
let req = parse_tool!(args.clone(), id, LogCodeChangeTool);
self.state.ledger.modify(|ledger| {
ledger.push(CodeChange {
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
file_path: req.file_path,
description: req.description,
git_commit: req.git_commit,
git_branch: req.git_branch,
});
});
Ok("Code change logged".to_string())
}
"query_recent_changes" => {
let data = serde_json::to_string(&self.state.ledger.read())
.unwrap_or_else(|_| "[]".to_string());
Ok(data.to_string())
}
"visualize_graph" => {
let req = parse_tool!(args.clone(), id, VisualizeGraphTool);
let query = req.query.unwrap_or_default().to_lowercase();
let full = self.state.get_full_graph();
let mut included = HashSet::new();
for (name, e) in &full.entities {
if let Some(ns) = &req.namespace
&& e.namespace != *ns
{
continue;
}
if query.is_empty()
|| name.to_lowercase().contains(&query)
|| e.entity_type.to_lowercase().contains(&query)
{
included.insert(name.clone());
}
}
let mut to_draw = Vec::new();
for r in &full.relations {
if let Some(ns) = &req.namespace
&& r.namespace != *ns
{
continue;
}
if query.is_empty()
|| included.contains(&r.from)
|| included.contains(&r.to)
{
included.insert(r.from.clone());
included.insert(r.to.clone());
to_draw.push(r);
}
}
let mut output = String::from("graph TD;\n");
for name in &included {
let safe = name.replace("\"", "").replace("(", "").replace(")", "");
output.push_str(&format!(
" id_{}[\"{}\"];\n",
safe.replace(" ", "_").replace("-", "_").replace(".", "_"),
safe
));
}
for r in to_draw {
let safe_f = r
.from
.replace(" ", "_")
.replace("-", "_")
.replace(".", "_")
.replace("\"", "")
.replace("(", "")
.replace(")", "");
let safe_t =
r.to.replace(" ", "_")
.replace("-", "_")
.replace(".", "_")
.replace("\"", "")
.replace("(", "")
.replace(")", "");
let safe_r = r.relation_type.replace("\"", "");
output.push_str(&format!(
" id_{}-->|\"{}\"|id_{};\n",
safe_f, safe_r, safe_t
));
}
if output == "graph TD;\n" {
output = "No nodes found to visualize.".to_string();
}
Ok(output.to_string())
}
"add_sticky_note" => {
let req = parse_tool!(args.clone(), id, AddStickyNoteTool);
self.state.sticky.modify(|notes| {
notes.push(StickyNote {
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
content: req.content,
});
});
Ok("Sticky note added.".to_string())
}
"read_sticky_notes" => {
let data = serde_json::to_string(&self.state.sticky.read())
.unwrap_or_else(|_| "[]".to_string());
Ok(data.to_string())
}
"delete_sticky_note" => {
let req = parse_tool!(args.clone(), id, DeleteStickyNoteTool);
let mut success = false;
self.state.sticky.modify(|notes| {
if req.index > 0 && req.index <= notes.len() {
notes.remove(req.index - 1);
success = true;
}
});
if success {
Ok("Sticky note deleted.".to_string())
} else {
Err("Invalid sticky note index.".to_string())
}
}
"clear_sticky_notes" => {
self.state.sticky.modify(|notes| {
notes.clear();
});
Ok("All sticky notes cleared.".to_string())
}
"condense_entity" => {
let req = parse_tool!(args.clone(), id, CondenseEntityTool);
self.state.modify_graph(|master| {
if let Some(e) = master.entities.get_mut(&req.entity_name) {
e.observations = req.summarized_observations;
}
});
Ok("Entity condensed".to_string())
}
"add_task" => {
let req = parse_tool!(args.clone(), id, AddTaskTool);
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
let task_id = uuid::Uuid::new_v4().to_string();
let parent_id = req.parent_id.clone();
let deps = req.dependencies.clone().unwrap_or_default();
let task = Task {
id: task_id.clone(),
title: req.title,
status: "pending".to_string(),
description: req.description,
created_at: now,
updated_at: now,
git_branch: req.git_branch,
parent_id,
dependencies: deps,
acceptance_criteria: vec![],
};
if let Ok(idx) = self.state.search_index.read() {
drop(idx.index_task(&task));
}
self.state.tasks.modify(|tasks| {
tasks.push(task);
});
Ok(format!("Task added with ID: {}", task_id).to_string())
}
"delete_task" => {
let req = parse_tool!(args.clone(), id, DeleteTaskTool);
let mut deleted_count = 0;
self.state.tasks.modify(|tasks| {
let initial_len = tasks.len();
// Collect IDs of tasks to delete (this task + all its recursive children)
let mut to_delete = std::collections::HashSet::new();
to_delete.insert(req.id.clone());
let mut added_new = true;
while added_new {
added_new = false;
for t in tasks.iter() {
if let Some(pid) = &t.parent_id
&& to_delete.contains(pid) && !to_delete.contains(&t.id) {
to_delete.insert(t.id.clone());
added_new = true;
}
}
}
tasks.retain(|t| !to_delete.contains(&t.id));
deleted_count = initial_len - tasks.len();
});
if deleted_count > 0 {
Ok(vec![
format!("Deleted task and its children ({} total).", deleted_count)
.to_string(),
][0]
.clone())
} else {
Ok("Task not found.".to_string())
}
}
"update_task_status" => {
let req = parse_tool!(args.clone(), id, UpdateTaskStatusTool);
let mut found = false;
let mut blocked = false;
let mut blocker_details = String::new();
let target_status = req.status.to_lowercase();
self.state.tasks.modify(|tasks| {
// Find target task
let mut target_id = String::new();
if let Some(t) =
tasks.iter().find(|t| t.id == req.id || t.title == req.id)
{
target_id = t.id.clone();
}
if target_id.is_empty() {
return;
}
found = true;
if target_status == "done" || target_status == "completed" {
// 1. Check Acceptance Criteria
if let Some(t) = tasks.iter().find(|t| t.id == target_id)
&& t.acceptance_criteria.iter().any(|c| !c.is_met) {
blocked = true;
blocker_details =
"Unmet acceptance criteria exist.".to_string();
}
// 2. Check dependencies
if !blocked {
let mut uncompleted_deps = Vec::new();
if let Some(t) = tasks.iter().find(|t| t.id == target_id) {
for dep_id in &t.dependencies {
if let Some(dep_task) =
tasks.iter().find(|dt| dt.id == *dep_id)
&& dep_task.status != "completed"
&& dep_task.status != "done"
{
uncompleted_deps.push(dep_task.title.clone());
}
}
}
if !uncompleted_deps.is_empty() {
blocked = true;
blocker_details = format!(
"Blocked by dependencies: {}",
uncompleted_deps.join(", ")
);
}
}
// 3. Check child tasks
if !blocked {
let mut uncompleted_children = Vec::new();
for child in tasks
.iter()
.filter(|t| t.parent_id.as_ref() == Some(&target_id))
{
if child.status != "completed" && child.status != "done" {
uncompleted_children.push(child.title.clone());
}
}
if !uncompleted_children.is_empty() {
blocked = true;
blocker_details = format!(
"Blocked by child tasks: {}",
uncompleted_children.join(", ")
);
}
}
}
if !blocked {
// Apply update
if let Some(t) = tasks.iter_mut().find(|t| t.id == target_id) {
t.status = target_status.clone();
t.updated_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
}
// Cascade cancellation to children
if target_status == "cancelled" || target_status == "abandoned" {
let mut to_cancel = vec![target_id.clone()];
let mut i = 0;
while i < to_cancel.len() {
let current_pid = to_cancel[i].clone();
for t in tasks.iter_mut() {
if t.parent_id.as_ref() == Some(&current_pid)
&& t.status != "completed"
{
t.status = target_status.clone();
to_cancel.push(t.id.clone());
}
}
i += 1;
}
}
}
});
if blocked {
Ok(vec![format!(
"Error: Cannot transition task. {}",
blocker_details
)]
.into_iter()
.next()
.unwrap())
} else if found {
Ok("Task status updated.".to_string())
} else {
Ok("Task not found.".to_string())
}
}
"list_active_tasks" => {
let req = parse_tool!(args.clone(), id, ListActiveTasksTool);
let mut tasks = self.state.tasks.read();
tasks.retain(|t| t.status != "done");
if let Some(branch) = req.git_branch {
tasks.retain(|t| {
t.git_branch.is_none()
|| t.git_branch.as_deref() == Some(branch.as_str())
});
}
let data = serde_json::to_string(&tasks).unwrap_or_default();
Ok(data.to_string())
}
"set_acceptance_criteria" => {
let req = parse_tool!(args.clone(), id, SetAcceptanceCriteriaTool);
let mut success = false;
self.state.tasks.modify(|tasks| {
if let Some(task) =
tasks.iter_mut().rev().find(|t| t.title == req.task_title)
{
task.acceptance_criteria = req
.criteria
.into_iter()
.map(|desc| crate::models::AcceptanceCriteria {
id: uuid::Uuid::new_v4().to_string(),
description: desc,
is_met: false,
})
.collect();
task.updated_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
success = true;
}
});
if success {
Ok(
vec!["Acceptance criteria set successfully.".to_string()][0]
.clone(),
)
} else {
Ok("Task not found.".to_string())
}
}
"verify_acceptance_criteria" => {
let req = parse_tool!(args.clone(), id, VerifyAcceptanceCriteriaTool);
let mut success = false;
let mut already_met = false;
self.state.tasks.modify(|tasks| {
if let Some(task) = tasks.iter_mut().find(|t| t.id == req.task_id)
&& let Some(ac) = task
.acceptance_criteria
.iter_mut()
.find(|c| c.id == req.criteria || c.description == req.criteria)
{
if ac.is_met {
already_met = true;
} else {
ac.is_met = true;
success = true;
task.updated_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
}
}
});
if success {
Ok(vec![format!(
"Acceptance criteria verified with proof: {}",
req.proof
)][0]
.clone())
} else if already_met {
Ok("Acceptance criteria was already met.".to_string())
} else {
Ok(
vec!["Acceptance criteria or task not found.".to_string()][0]
.clone(),
)
}
}
"store_snippet" => {
let req = parse_tool!(args.clone(), id, StoreSnippetTool);
let snippet = Snippet {
name: req.name.clone(),
language: req.language,
code: req.code,
description: req.description,
updated_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
};
let s_clone = snippet.clone();
self.state.snippets.modify(|snippets| {
snippets.retain(|s| s.name != req.name);
snippets.push(s_clone);
});
if let Ok(idx) = self.state.search_index.read() {
drop(idx.index_snippet(&snippet));
}
Ok(format!("Snippet '{}' stored.", req.name).to_string())
}
"search_snippets" => {
let req = parse_tool!(args.clone(), id, SearchSnippetsTool);
let query = req.query.to_lowercase();
let snippets = self.state.snippets.read();
let mut results = Vec::new();
for s in snippets {
if s.name.to_lowercase().contains(&query)
|| s.description.to_lowercase().contains(&query)
|| s.language.to_lowercase().contains(&query)
{
results.push(s);
}
}
let data = serde_json::to_string(&results).unwrap_or_default();
Ok(data.to_string())
}
"delete_snippet" => {
let req = parse_tool!(args.clone(), id, DeleteSnippetTool);
let mut deleted = false;
self.state.snippets.modify(|snippets| {
let orig = snippets.len();
snippets.retain(|s| s.name != req.name);
deleted = snippets.len() < orig;
});
if deleted {
Ok("Snippet deleted.".to_string())
} else {
Ok("Snippet not found.".to_string())
}
}
"log_decision" => {
let req = parse_tool!(args.clone(), id, LogDecisionTool);
let mut adr_id = String::new();
let mut new_adr = None;
self.state.adrs.modify(|adrs| {
adr_id = format!("ADR-{:04}", adrs.len() + 1);
let a = Adr {
id: adr_id.clone(),
title: req.title,
context: req.context,
decision: req.decision,
consequence: req.consequence,
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
};
new_adr = Some(a.clone());
adrs.push(a);
});
if let Some(adr) = new_adr {
if let Ok(idx) = self.state.search_index.read() {
drop(idx.index_adr(&adr));
}
}
Ok(format!("Decision logged as {}", adr_id).to_string())
}
"query_decisions" => {
let req = parse_tool!(args.clone(), id, QueryDecisionsTool);
let mut adrs = self.state.adrs.read();
if let Some(q) = req.query {
let q = q.to_lowercase();
adrs.retain(|a| {
a.title.to_lowercase().contains(&q)
|| a.context.to_lowercase().contains(&q)
|| a.decision.to_lowercase().contains(&q)
});
}
let data = serde_json::to_string(&adrs).unwrap_or_default();
Ok(data.to_string())
}
"merge_entities" => {
let req = parse_tool!(args.clone(), id, MergeEntitiesTool);
self.state.modify_graph(|master| {
if let Some(src) = master.entities.remove(&req.source_entity) {
if let Some(tgt) = master.entities.get_mut(&req.target_entity) {
tgt.observations.extend(src.observations);
tgt.observations =
MemoryState::unique_items(tgt.observations.clone());
} else {
let mut new_tgt = src.clone();
new_tgt.name = req.target_entity.clone();
master.entities.insert(req.target_entity.clone(), new_tgt);
}
}
for r in &mut master.relations {
if r.from == req.source_entity {
r.from = req.target_entity.clone();
}
if r.to == req.source_entity {
r.to = req.target_entity.clone();
}
}
master.relations = MemoryState::unique_items(master.relations.clone());
});
Ok("Entities merged".to_string())
}
"find_orphans" => {
let full = self.state.get_full_graph();
let mut connected = std::collections::HashSet::new();
for r in &full.relations {
connected.insert(r.from.clone());
connected.insert(r.to.clone());
}
let orphans: Vec<String> = full
.entities
.keys()
.filter(|k| !connected.contains(*k))
.cloned()
.collect();
let data = serde_json::to_string(&orphans).unwrap_or_default();
Ok(data.to_string())
}
"learn_preference" => {
let req = parse_tool!(args.clone(), id, LearnPreferenceTool);
self.state.prefs.modify(|prefs| {
prefs.insert(
req.key.clone(),
crate::models::Preference {
key: req.key.clone(),
value: req.value,
updated_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
},
);
});
Ok("Preference learned".to_string())
}
"read_preferences" => {
let prefs = self.state.prefs.read();
let data = serde_json::to_string(&prefs).unwrap_or_default();
Ok(data.to_string())
}
"log_error_fix" => {
let req = parse_tool!(args.clone(), id, LogErrorFixTool);
self.state.error_fixes.modify(|fixes| {
fixes.push(crate::models::ErrorFix {
signature: req.signature,
solution: req.solution,
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
git_commit: req.git_commit,
git_branch: req.git_branch,
})
});
Ok("Error fix logged".to_string())
}
"search_error_fixes" => {
let req = parse_tool!(args.clone(), id, SearchErrorFixesTool);
let q = req.query.to_lowercase();
let mut fixes = self.state.error_fixes.read();
fixes.retain(|f| {
f.signature.to_lowercase().contains(&q)
|| f.solution.to_lowercase().contains(&q)
});
let data = serde_json::to_string(&fixes).unwrap_or_default();
Ok(data.to_string())
}
"pin_file" => {
let req = parse_tool!(args.clone(), id, PinFileTool);
self.state.pinned_files.modify(|pinned| {
pinned.retain(|p| {
!(p.namespace == req.namespace && p.file_path == req.file_path)
});
pinned.push(crate::models::PinnedFile {
namespace: req.namespace,
file_path: req.file_path,
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
git_branch: req.git_branch,
});
});
Ok("File pinned".to_string())
}
"unpin_file" => {
let req = parse_tool!(args.clone(), id, UnpinFileTool);
self.state.pinned_files.modify(|pinned| {
pinned.retain(|p| {
!(p.namespace == req.namespace && p.file_path == req.file_path)
})
});
Ok("File unpinned".to_string())
}
"list_pinned_files" => {
let req = parse_tool!(args.clone(), id, ListPinnedFilesTool);
let mut pinned = self.state.pinned_files.read();
if let Some(ns) = req.namespace {
pinned.retain(|p| p.namespace == ns);
}
if let Some(branch) = req.git_branch {
pinned.retain(|p| {
p.git_branch.is_none()
|| p.git_branch.as_deref() == Some(branch.as_str())
});
}
let data = serde_json::to_string(&pinned).unwrap_or_default();
Ok(data.to_string())
}
"add_session_summary" => {
let req = parse_tool!(args.clone(), id, AddSessionSummaryTool);
self.state.session_summaries.modify(|summaries| {
summaries.push(crate::models::SessionSummary {
summary: req.summary,
namespace: req.namespace,
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
})
});
Ok("Session summary added".to_string())
}
"get_project_timeline" => {
let req = parse_tool!(args.clone(), id, GetProjectTimelineTool);
let mut summaries = self.state.session_summaries.read();
if let Some(ns) = req.namespace {
summaries.retain(|s| s.namespace == ns);
}
summaries.sort_by_key(|s| s.timestamp);
let data = serde_json::to_string(&summaries).unwrap_or_default();
Ok(data.to_string())
}
"leave_handoff_memo" => {
let req = parse_tool!(args.clone(), id, LeaveHandoffMemoTool);
self.state.handoff_memos.modify(|memos| {
memos.push(crate::models::HandoffMemo {
id: uuid::Uuid::new_v4().to_string(),
author: "agy".to_string(),
content: req.content,
namespace: req.namespace,
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
})
});
Ok("Handoff memo left".to_string())
}
"read_handoff_memos" => {
let req = parse_tool!(args.clone(), id, ReadHandoffMemosTool);
let mut memos = self.state.handoff_memos.read();
if let Some(ns) = req.namespace {
memos.retain(|m| m.namespace == ns);
}
let data = serde_json::to_string(&memos).unwrap_or_default();
Ok(data.to_string())
}
"clear_handoff_memos" => {
let req = parse_tool!(args.clone(), id, ClearHandoffMemosTool);
let ids: HashSet<_> = req.ids.into_iter().collect();
self.state
.handoff_memos
.modify(|memos| memos.retain(|m| !ids.contains(&m.id)));
Ok("Handoff memos cleared".to_string())
}
"update_env_fingerprint" => {
let req = parse_tool!(args.clone(), id, UpdateEnvFingerprintTool);
self.state.env_fingerprints.modify(|fps| {
fps.insert(
req.namespace.clone(),
crate::models::EnvFingerprint {
namespace: req.namespace.clone(),
os: std::env::consts::OS.to_string(),
shell: std::env::var("SHELL")
.unwrap_or_else(|_| "unknown".to_string()),
tool_versions: req.tool_versions,
updated_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
},
);
});
Ok("Env fingerprint updated".to_string())
}
"read_env_fingerprint" => {
let req = parse_tool!(args.clone(), id, ReadEnvFingerprintTool);
let fps = self.state.env_fingerprints.read();
if let Some(fp) = fps.get(&req.namespace) {
let data = serde_json::to_string(fp).unwrap_or_default();
Ok(data.to_string())
} else {
Ok("{}".to_string())
}
}
"log_env_requirement" => {
let req = parse_tool!(args.clone(), id, LogEnvRequirementTool);
self.state.env_requirements.modify(|reqs| {
reqs.retain(|r| !(r.namespace == req.namespace && r.key == req.key));
reqs.push(crate::models::EnvRequirement {
namespace: req.namespace,
key: req.key,
description: req.description,
is_secret: req.is_secret,
});
});
Ok("Env requirement logged".to_string())
}
"add_milestone" => {
let req = parse_tool!(args.clone(), id, AddMilestoneTool);
self.state.milestones.modify(|ms| {
ms.push(crate::models::Milestone {
id: uuid::Uuid::new_v4().to_string(),
title: req.title,
status: "pending".to_string(),
namespace: req.namespace,
target_date: None,
})
});
Ok("Milestone added".to_string())
}
"update_milestone" => {
let req = parse_tool!(args.clone(), id, UpdateMilestoneTool);
let mut found = false;
self.state.milestones.modify(|ms| {
for m in ms.iter_mut() {
if m.id == req.id {
m.status = req.status.clone();
found = true;
break;
}
}
});
if found {
Ok("Milestone updated".to_string())
} else {
Ok("Milestone not found".to_string())
}
}
"list_milestones" => {
let req = parse_tool!(args.clone(), id, ListMilestonesTool);
let mut ms = self.state.milestones.read();
if let Some(ns) = req.namespace {
ms.retain(|m| m.namespace == ns);
}
let data = serde_json::to_string(&ms).unwrap_or_default();
Ok(data.to_string())
}
"generate_standup_report" => {
let req = parse_tool!(args.clone(), id, GenerateStandupReportTool);
let cutoff = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
.saturating_sub(req.hours_lookback * 3600);
let tasks = self
.state
.tasks
.read()
.into_iter()
.filter(|t| t.updated_at >= cutoff)
.collect::<Vec<_>>();
let changes = self
.state
.ledger
.read()
.into_iter()
.filter(|c| c.timestamp >= cutoff)
.collect::<Vec<_>>();
let summaries = self
.state
.session_summaries
.read()
.into_iter()
.filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff)
.collect::<Vec<_>>();
let report = serde_json::json!({ "tasks_updated": tasks, "code_changes": changes, "session_summaries": summaries });
Ok(report.to_string())
}
"register_environment" => {
let req = parse_tool!(args.clone(), id, RegisterEnvironmentTool);
self.state.environments.modify(|envs| {
envs.retain(|e| !(e.namespace == req.namespace && e.name == req.name));
envs.push(crate::models::EnvironmentDetail {
namespace: req.namespace,
name: req.name,
url: req.url,
description: req.description,
requires_vpn: req.requires_vpn,
updated_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
});
});
Ok("Environment registered".to_string())
}
"get_environment_details" => {
let req = parse_tool!(args.clone(), id, GetEnvironmentDetailsTool);
let mut envs = self.state.environments.read();
envs.retain(|e| e.namespace == req.namespace);
let data = serde_json::to_string(&envs).unwrap_or_default();
Ok(data.to_string())
}
"add_pr_checklist_item" => {
let req = parse_tool!(args.clone(), id, AddPrChecklistItemTool);
self.state.pr_checklists.modify(|items| {
items.push(crate::models::PrChecklistItem {
namespace: req.namespace,
id: uuid::Uuid::new_v4().to_string(),
description: req.description,
})
});
Ok("PR checklist item added".to_string())
}
"get_pr_checklist" => {
let req = parse_tool!(args.clone(), id, GetPrChecklistTool);
let mut items = self.state.pr_checklists.read();
items.retain(|i| i.namespace == req.namespace);
let data = serde_json::to_string(&items).unwrap_or_default();
Ok(data.to_string())
}
"clear_pr_checklist" => {
let req = parse_tool!(args.clone(), id, ClearPrChecklistTool);
self.state
.pr_checklists
.modify(|items| items.retain(|i| i.namespace != req.namespace));
Ok("PR checklist cleared".to_string())
}
"log_tech_debt" => {
let req = parse_tool!(args.clone(), id, LogTechDebtTool);
self.state.tech_debts.modify(|debts| {
debts.push(crate::models::TechDebt {
id: uuid::Uuid::new_v4().to_string(),
namespace: req.namespace,
description: req.description,
ideal_solution: req.ideal_solution,
is_resolved: false,
created_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
git_commit: req.git_commit,
git_branch: req.git_branch,
})
});
Ok("Tech debt logged".to_string())
}
"resolve_tech_debt" => {
let req = parse_tool!(args.clone(), id, ResolveTechDebtTool);
let mut found = false;
self.state.tech_debts.modify(|debts| {
for d in debts.iter_mut() {
if d.id == req.id {
d.is_resolved = true;
found = true;
break;
}
}
});
if found {
Ok("Tech debt resolved".to_string())
} else {
Ok("Tech debt not found".to_string())
}
}
"list_tech_debt" => {
let req = parse_tool!(args.clone(), id, ListTechDebtTool);
let mut debts = self.state.tech_debts.read();
debts.retain(|d| {
d.namespace == req.namespace && (req.include_resolved || !d.is_resolved)
});
let data = serde_json::to_string(&debts).unwrap_or_default();
Ok(data.to_string())
}
"save_context_workspace" => {
let req = parse_tool!(args.clone(), id, SaveContextWorkspaceTool);
self.state.context_workspaces.modify(|ws| {
ws.retain(|w| !(w.namespace == req.namespace && w.name == req.name));
ws.push(crate::models::ContextWorkspace {
namespace: req.namespace,
name: req.name,
pinned_files: req.pinned_files,
active_task_ids: req.active_task_ids,
saved_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
});
});
Ok("Context workspace saved".to_string())
}
"load_context_workspace" => {
let req = parse_tool!(args.clone(), id, LoadContextWorkspaceTool);
let mut ws = self.state.context_workspaces.read();
ws.retain(|w| w.namespace == req.namespace && w.name == req.name);
let data = serde_json::to_string(&ws.first()).unwrap_or_default();
Ok(data.to_string())
}
"list_context_workspaces" => {
let req = parse_tool!(args.clone(), id, ListContextWorkspacesTool);
let mut ws = self.state.context_workspaces.read();
ws.retain(|w| w.namespace == req.namespace);
let data = serde_json::to_string(&ws).unwrap_or_default();
Ok(data.to_string())
}
"omni_search" => {
let req = parse_tool!(args.clone(), id, OmniSearchTool);
let matches = if let Ok(idx) = self.state.search_index.read() {
idx.search(&req.query, req.namespace.as_deref())
.unwrap_or_default()
} else {
vec![]
};
let mut kg = KnowledgeGraph::default();
let mut tasks = Vec::new();
let mut snippets = Vec::new();
let mut adrs = Vec::new();
let full = self.state.get_full_graph();
for (id, doc_type, _, _, _) in &matches {
if doc_type == "entity"
&& let Some(e) = full.entities.get(id)
{
kg.entities.insert(id.clone(), e.clone());
}
}
for t in self.state.tasks.read() {
if matches.iter().any(|(id, typ, _, _, _)| id == &t.id && typ == "task") {
tasks.push(t);
}
}
for s in self.state.snippets.read() {
if matches
.iter()
.any(|(id, typ, _, _, _)| id == &s.name && typ == "snippet")
{
snippets.push(s);
}
}
for a in self.state.adrs.read() {
if matches.iter().any(|(id, typ, _, _, _)| id == &a.id && typ == "adr") {
adrs.push(a);
}
}
let q = req.query.to_lowercase();
let tech_debts: Vec<_> = self
.state
.tech_debts
.read()
.into_iter()
.filter(|d| {
(req.namespace.is_none()
|| d.namespace == *req.namespace.as_ref().unwrap())
&& (d.description.to_lowercase().contains(&q)
|| d.ideal_solution.to_lowercase().contains(&q))
})
.collect();
let memos: Vec<_> = self
.state
.handoff_memos
.read()
.into_iter()
.filter(|m| {
(req.namespace.is_none()
|| m.namespace == *req.namespace.as_ref().unwrap())
&& m.content.to_lowercase().contains(&q)
})
.collect();
let error_fixes: Vec<_> = self
.state
.error_fixes
.read()
.into_iter()
.filter(|f| {
f.signature.to_lowercase().contains(&q)
|| f.solution.to_lowercase().contains(&q)
})
.collect();
let report = serde_json::json!({
"knowledge_graph": kg.entities,
"tasks": tasks,
"snippets": snippets,
"adrs": adrs,
"tech_debts": tech_debts,
"handoff_memos": memos,
"error_fixes": error_fixes
});
Ok(report.to_string())
}
"get_project_health" => {
let req = parse_tool!(args.clone(), id, GetProjectHealthTool);
let active_tasks = self
.state
.tasks
.read()
.into_iter()
.filter(|t| t.status != "done")
.count();
let unresolved_debt = self
.state
.tech_debts
.read()
.into_iter()
.filter(|d| d.namespace == req.namespace && !d.is_resolved)
.count();
let unread_memos = self
.state
.handoff_memos
.read()
.into_iter()
.filter(|m| m.namespace == req.namespace)
.count();
let active_milestones = self
.state
.milestones
.read()
.into_iter()
.filter(|m| m.namespace == req.namespace && m.status != "done")
.count();
let remaining_checklists = self
.state
.pr_checklists
.read()
.into_iter()
.filter(|c| c.namespace == req.namespace)
.count();
let report = serde_json::json!({
"active_tasks": active_tasks,
"unresolved_tech_debt": unresolved_debt,
"unread_handoff_memos": unread_memos,
"active_milestones": active_milestones,
"remaining_pr_checklist_items": remaining_checklists
});
Ok(report.to_string())
}
_ => Err(format!("Unknown tool: {}", name)),
};
match result {
Ok(text) => Some(crate::mcp::success(
id,
serde_json::json!({
"content": [{ "type": "text", "text": text }]
}),
)),
Err(e) => Some(crate::mcp::success(
id,
serde_json::json!({
"isError": true,
"content": [{ "type": "text", "text": e }]
}),
)),
}
}
_ => {
if id != serde_json::Value::Null {
Some(crate::mcp::error(id, -32601, "Method not found"))
} else {
None
}
}
};
let elapsed = start_time.elapsed();
if method == "tools/call" {
let is_error = response.as_ref().is_some_and(|r| r.get("error").is_some() || r.get("result").and_then(|res| res.get("isError")).and_then(|e| e.as_bool()).unwrap_or(false));
tracing::info!("<<< [Server] MCP tool call {} (id: {}) completed in {:?} [Error: {}]", tool_name, id_clone, elapsed, is_error);
// Broadcast completion latency to the UI Activity Feed
let status_msg = if is_error { "with error" } else { "successfully" };
self.state.broadcast_activity(&format!("Tool {} completed {} in {:?}", tool_name, status_msg, elapsed));
} else {
tracing::debug!("<<< [Server] MCP request method {} completed in {:?}", method, elapsed);
}
tracing::trace!("Returning response from handle_request: {:?}", response);
response
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::state::MemoryState;
use serde_json::json;
use std::sync::Arc;
#[tokio::test]
async fn test_handle_initialize() {
let store_dir = std::env::temp_dir().join(format!(
"mcp_test_handlers_{}",
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs()
));
std::fs::create_dir_all(&store_dir).unwrap();
let redb_path = store_dir.join("mcp_store.redb");
let db = Arc::new(redb::Database::create(&redb_path).unwrap());
{
let write_txn = db.begin_write().unwrap();
let _ = write_txn.open_table(crate::store::STORE_TABLE);
write_txn.commit().unwrap();
}
let state = Arc::new(MemoryState {
base_dir: store_dir.clone(),
graph: crate::store::Store::new("knowledge_graph_master", db.clone()),
search_index: std::sync::RwLock::new(
crate::search::MemoryIndex::new(&store_dir).unwrap(),
),
ledger: crate::store::Store::new("audit_ledger", db.clone()),
sticky: crate::store::Store::new("sticky_notes", db.clone()),
tasks: crate::store::Store::new("tasks", db.clone()),
snippets: crate::store::Store::new("snippets", db.clone()),
adrs: crate::store::Store::new("adrs", db.clone()),
prefs: crate::store::Store::new("preferences", db.clone()),
error_fixes: crate::store::Store::new("error_fixes", db.clone()),
pinned_files: crate::store::Store::new("pinned_files", db.clone()),
session_summaries: crate::store::Store::new("session_summaries", db.clone()),
handoff_memos: crate::store::Store::new("handoff_memos", db.clone()),
env_fingerprints: crate::store::Store::new("env_fingerprints", db.clone()),
env_requirements: crate::store::Store::new("env_requirements", db.clone()),
milestones: crate::store::Store::new("milestones", db.clone()),
environments: crate::store::Store::new("environments", db.clone()),
pr_checklists: crate::store::Store::new("pr_checklists", db.clone()),
tech_debts: crate::store::Store::new("tech_debts", db.clone()),
gates: crate::store::Store::new("gates", db.clone()),
context_workspaces: crate::store::Store::new("context_workspaces", db.clone()),
activity_tx: tokio::sync::broadcast::channel(100).0,
});
let handler = MemoryHandler { state };
let req = json!({
"jsonrpc": "2.0",
"id": 1,
"method": "initialize",
"params": {
"protocolVersion": "2024-11-05",
"capabilities": {},
"clientInfo": {
"name": "test-client",
"version": "1.0.0"
}
}
});
let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 1);
assert!(response.get("result").is_some());
let result = &response["result"];
// assert_eq!(result["protocolVersion"], "2024-11-05");
// CRITICAL BUG FIX CHECK: capabilities MUST contain an empty tools object
// Note: Currently it is set to `{}` which may cause proxy dropping tools. Let's verify it matches the actual behavior.
assert_eq!(result["capabilities"], serde_json::json!({"tools": {}}));
assert_eq!(result["serverInfo"]["name"], "gemini-mcp-memory");
}
fn setup_test_handler(test_name: &str) -> MemoryHandler {
let store_dir = std::env::temp_dir().join(format!(
"mcp_test_handlers_{}_{}",
test_name,
uuid::Uuid::new_v4()
));
std::fs::create_dir_all(&store_dir).unwrap();
let redb_path = store_dir.join("mcp_store.redb");
let db = Arc::new(redb::Database::create(&redb_path).unwrap());
{
let write_txn = db.begin_write().unwrap();
let _ = write_txn.open_table(crate::store::STORE_TABLE);
write_txn.commit().unwrap();
}
let state = Arc::new(MemoryState {
base_dir: store_dir.clone(),
graph: crate::store::Store::new("knowledge_graph_master", db.clone()),
search_index: std::sync::RwLock::new(
crate::search::MemoryIndex::new(&store_dir).unwrap(),
),
ledger: crate::store::Store::new("audit_ledger", db.clone()),
sticky: crate::store::Store::new("sticky_notes", db.clone()),
tasks: crate::store::Store::new("tasks", db.clone()),
snippets: crate::store::Store::new("snippets", db.clone()),
adrs: crate::store::Store::new("adrs", db.clone()),
prefs: crate::store::Store::new("preferences", db.clone()),
error_fixes: crate::store::Store::new("error_fixes", db.clone()),
pinned_files: crate::store::Store::new("pinned_files", db.clone()),
session_summaries: crate::store::Store::new("session_summaries", db.clone()),
handoff_memos: crate::store::Store::new("handoff_memos", db.clone()),
env_fingerprints: crate::store::Store::new("env_fingerprints", db.clone()),
env_requirements: crate::store::Store::new("env_requirements", db.clone()),
milestones: crate::store::Store::new("milestones", db.clone()),
environments: crate::store::Store::new("environments", db.clone()),
pr_checklists: crate::store::Store::new("pr_checklists", db.clone()),
tech_debts: crate::store::Store::new("tech_debts", db.clone()),
gates: crate::store::Store::new("gates", db.clone()),
context_workspaces: crate::store::Store::new("context_workspaces", db.clone()),
activity_tx: tokio::sync::broadcast::channel(100).0,
});
MemoryHandler { state }
}
#[tokio::test]
async fn test_handle_tools_list() {
let handler = setup_test_handler("tools_list");
let req = json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/list",
"params": {}
});
let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 2);
let tools = response["result"]["tools"]
.as_array()
.expect("Tools must be an array");
assert!(!tools.is_empty());
// Verify a specific tool is registered
let add_task_tool = tools
.iter()
.find(|t| t["name"] == "add_task")
.expect("add_task tool missing");
assert_eq!(
add_task_tool["description"],
"Add a new task to the task tracker."
);
}
#[tokio::test]
async fn test_handle_add_task() {
let handler = setup_test_handler("add_task");
let req = json!({
"jsonrpc": "2.0",
"id": 3,
"method": "tools/call",
"params": {
"name": "add_task",
"arguments": {
"title": "Fix bug in handlers",
"description": "The proxy drops capabilities.",
"git_branch": "master"
}
}
});
let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 3);
let content = &response["result"]["content"][0];
assert_eq!(content["type"], "text");
assert!(
content["text"]
.as_str()
.unwrap()
.starts_with("Task added with ID: ")
);
// Verify task was actually added to store
let tasks = handler.state.tasks.read();
assert_eq!(tasks.len(), 1);
assert_eq!(tasks[0].title, "Fix bug in handlers");
assert_eq!(tasks[0].status, "pending");
}
#[tokio::test]
async fn test_handle_create_entities() {
let handler = setup_test_handler("create_entities");
let req = json!({
"jsonrpc": "2.0",
"id": 4,
"method": "tools/call",
"params": {
"name": "create_entities",
"arguments": {
"entities": [
{
"name": "MemoryHandler",
"entityType": "struct",
"observations": ["Handles MCP requests natively"],
"namespace": "core"
}
]
}
}
});
let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 4);
let content = &response["result"]["content"][0];
assert_eq!(content["text"], "Entities created");
// Verify entity was actually added to state
let session_graph = handler.state.graph.read();
let entity = session_graph
.entities
.get("MemoryHandler")
.expect("Entity should be in session graph");
assert_eq!(entity.entity_type, "struct");
assert_eq!(entity.observations, vec!["Handles MCP requests natively"]);
assert_eq!(entity.namespace, "core".to_string());
}
#[tokio::test]
async fn test_handle_store_snippet() {
let handler = setup_test_handler("store_snippet");
let req = json!({
"jsonrpc": "2.0",
"id": 5,
"method": "tools/call",
"params": {
"name": "store_snippet",
"arguments": {
"name": "Test Snippet",
"description": "A snippet used for testing",
"language": "rust",
"code": "fn main() { println!(\"Hello, World!\"); }"
}
}
});
let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 5);
let snippets = handler.state.snippets.read();
assert_eq!(snippets.len(), 1);
assert_eq!(snippets[0].name, "Test Snippet");
assert_eq!(snippets[0].language, "rust");
}
#[tokio::test]
async fn test_handle_add_sticky_note() {
let handler = setup_test_handler("add_sticky_note");
let req = json!({
"jsonrpc": "2.0",
"id": 6,
"method": "tools/call",
"params": {
"name": "add_sticky_note",
"arguments": {
"content": "Don't forget to check coverage!"
}
}
});
let response = handler
.handle_request(req)
.await
.expect("Expected a response");
assert_eq!(response["id"], 6);
let notes = handler.state.sticky.read();
assert_eq!(notes.len(), 1);
assert_eq!(notes[0].content, "Don't forget to check coverage!");
}
#[tokio::test]
async fn test_handle_create_relations() {
let handler = setup_test_handler("create_relations");
let req = json!({
"jsonrpc": "2.0",
"id": 7,
"method": "tools/call",
"params": {
"name": "create_relations",
"arguments": {
"relations": [
{
"from": "NodeA",
"to": "NodeB",
"relationType": "depends_on",
"namespace": "core"
}
]
}
}
});
let response = handler.handle_request(req).await.unwrap();
assert_eq!(response["id"], 7);
let session = handler.state.graph.read();
assert_eq!(session.relations.len(), 1);
assert_eq!(session.relations[0].from, "NodeA");
assert_eq!(session.relations[0].to, "NodeB");
}
#[tokio::test]
async fn test_handle_add_observations() {
let handler = setup_test_handler("add_observations");
// Pre-populate entity
handler.state.graph.modify(|session| {
session.entities.insert(
"NodeA".to_string(),
crate::models::Entity {
name: "NodeA".to_string(),
entity_type: "class".to_string(),
observations: vec!["Initial".to_string()],
namespace: "".to_string(),
git_branch: None,
},
);
});
let req = json!({
"jsonrpc": "2.0",
"id": 8,
"method": "tools/call",
"params": {
"name": "add_observations",
"arguments": {
"observations": [
{
"entityName": "NodeA",
"contents": ["New observation"]
}
]
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let session = handler.state.graph.read();
let entity = session.entities.get("NodeA").unwrap();
assert_eq!(entity.observations, vec!["Initial", "New observation"]);
}
#[tokio::test]
async fn test_handle_delete_entities() {
let handler = setup_test_handler("delete_entities");
handler.state.graph.modify(|session| {
session.entities.insert(
"ToDelete".to_string(),
crate::models::Entity {
name: "ToDelete".to_string(),
entity_type: "var".to_string(),
observations: vec![],
namespace: "".to_string(),
git_branch: None,
},
);
});
// Force flush session to master
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",
"id": 9,
"method": "tools/call",
"params": {
"name": "delete_entities",
"arguments": {
"entityNames": ["ToDelete"]
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let full_graph = handler.state.get_full_graph();
assert!(full_graph.entities.get("ToDelete").is_none());
}
#[tokio::test]
async fn test_handle_delete_observations() {
let handler = setup_test_handler("delete_observations");
handler.state.graph.modify(|session| {
session.entities.insert(
"NodeA".to_string(),
crate::models::Entity {
name: "NodeA".to_string(),
entity_type: "class".to_string(),
observations: vec!["Keep".to_string(), "Drop".to_string()],
namespace: "".to_string(),
git_branch: None,
},
);
});
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",
"id": 10,
"method": "tools/call",
"params": {
"name": "delete_observations",
"arguments": {
"deletions": [
{
"entityName": "NodeA",
"observations": ["Drop"]
}
]
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let full = handler.state.get_full_graph();
let entity = full.entities.get("NodeA").unwrap();
assert_eq!(entity.observations, vec!["Keep"]);
}
#[tokio::test]
async fn test_handle_log_code_change() {
let handler = setup_test_handler("log_code_change");
let req = json!({
"jsonrpc": "2.0",
"id": 11,
"method": "tools/call",
"params": {
"name": "log_code_change",
"arguments": {
"filePath": "server/src/handlers.rs",
"description": "Added some unit tests",
"git_commit": "1234567"
}
}
});
let response = handler.handle_request(req).await.unwrap();
assert_eq!(response["id"], 11);
let ledger = handler.state.ledger.read();
assert_eq!(ledger.len(), 1);
assert_eq!(ledger[0].file_path, "server/src/handlers.rs");
assert_eq!(ledger[0].git_commit.as_deref(), Some("1234567"));
}
#[tokio::test]
async fn test_handle_list_active_tasks() {
let handler = setup_test_handler("list_active_tasks");
handler.state.tasks.modify(|tasks| {
tasks.push(crate::models::Task {
id: "1".to_string(),
title: "Active Task".to_string(),
status: "pending".to_string(),
description: "".to_string(),
created_at: 0,
updated_at: 0,
git_branch: None,
acceptance_criteria: vec![],
dependencies: vec![],
parent_id: None,
});
tasks.push(crate::models::Task {
id: "2".to_string(),
title: "Completed Task".to_string(),
status: "done".to_string(),
description: "".to_string(),
created_at: 0,
updated_at: 0,
git_branch: None,
acceptance_criteria: vec![],
dependencies: vec![],
parent_id: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 12,
"method": "tools/call",
"params": {
"name": "list_active_tasks",
"arguments": {}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("Active Task"));
assert!(!content.contains("Completed Task"));
}
#[tokio::test]
async fn test_handle_search_snippets() {
let handler = setup_test_handler("search_snippets");
handler.state.snippets.modify(|snippets| {
snippets.push(crate::models::Snippet {
name: "React hook".to_string(),
language: "typescript".to_string(),
code: "useMemo(() => {}, [])".to_string(),
description: "React memoization".to_string(),
updated_at: 0,
});
snippets.push(crate::models::Snippet {
name: "Rust struct".to_string(),
language: "rust".to_string(),
code: "struct A {}".to_string(),
description: "Rust code".to_string(),
updated_at: 0,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 13,
"method": "tools/call",
"params": {
"name": "search_snippets",
"arguments": {
"query": "React"
}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("React hook"));
assert!(!content.contains("Rust struct"));
}
#[tokio::test]
async fn test_handle_read_sticky_notes() {
let handler = setup_test_handler("read_sticky_notes");
handler.state.sticky.modify(|sticky| {
sticky.push(crate::models::StickyNote {
content: "Remember to commit".to_string(),
timestamp: 0,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 14,
"method": "tools/call",
"params": {
"name": "read_sticky_notes",
"arguments": {}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("Remember to commit"));
}
#[tokio::test]
async fn test_handle_delete_relations() {
let handler = setup_test_handler("delete_relations");
handler.state.graph.modify(|session| {
session.relations.push(crate::models::Relation {
from: "A".to_string(),
to: "B".to_string(),
relation_type: "calls".to_string(),
namespace: "".to_string(),
});
});
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",
"id": 15,
"method": "tools/call",
"params": {
"name": "delete_relations",
"arguments": {
"relations": [
{
"from": "A",
"to": "B",
"relationType": "calls",
"namespace": ""
}
]
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let full = handler.state.get_full_graph();
assert!(full.relations.is_empty());
}
#[tokio::test]
async fn test_handle_read_graph() {
let handler = setup_test_handler("read_graph");
handler.state.graph.modify(|session| {
session.entities.insert(
"NodeA".to_string(),
crate::models::Entity {
name: "NodeA".to_string(),
entity_type: "var".to_string(),
observations: vec![],
namespace: "".to_string(),
git_branch: None,
},
);
});
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",
"id": 16,
"method": "tools/call",
"params": {
"name": "read_graph",
"arguments": {}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("NodeA"));
}
#[tokio::test]
async fn test_handle_open_nodes() {
let handler = setup_test_handler("open_nodes");
let entity = crate::models::Entity {
name: "UserRepository".to_string(),
entity_type: "class".to_string(),
observations: vec!["Handles user data".to_string()],
namespace: "".to_string(),
git_branch: None,
};
handler.state.graph.modify(|session| {
session
.entities
.insert("UserRepository".to_string(), entity);
});
handler.state.modify_graph(|_| {});
let req = json!({
"jsonrpc": "2.0",
"id": 17,
"method": "tools/call",
"params": {
"name": "open_nodes",
"arguments": {
"names": ["UserRepository"]
}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("UserRepository"));
assert!(content.contains("Handles user data"));
}
#[tokio::test]
async fn test_handle_log_decision() {
let handler = setup_test_handler("log_decision");
let req = json!({
"jsonrpc": "2.0",
"id": 20,
"method": "tools/call",
"params": {
"name": "log_decision",
"arguments": {
"title": "Use async I/O",
"context": "Need better throughput",
"decision": "Use tokio",
"consequence": "Requires async all the way down"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let adrs = handler.state.adrs.read();
assert_eq!(adrs.len(), 1);
assert_eq!(adrs[0].title, "Use async I/O");
}
#[tokio::test]
async fn test_handle_query_decisions() {
let handler = setup_test_handler("query_decisions");
handler.state.adrs.modify(|adrs| {
adrs.push(crate::models::Adr {
id: "adr-1".to_string(),
title: "Use PostgreSQL".to_string(),
context: "Need relational data".to_string(),
decision: "Use pg".to_string(),
consequence: "Maintenance overhead".to_string(),
timestamp: 0,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 21,
"method": "tools/call",
"params": {
"name": "query_decisions",
"arguments": {
"query": "Postgre"
}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("PostgreSQL"));
}
#[tokio::test]
async fn test_handle_log_error_fix() {
let handler = setup_test_handler("log_error_fix");
let req = json!({
"jsonrpc": "2.0",
"id": 22,
"method": "tools/call",
"params": {
"name": "log_error_fix",
"arguments": {
"signature": "IndexOutOfBounds",
"solution": "Check array length"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let fixes = handler.state.error_fixes.read();
assert_eq!(fixes.len(), 1);
assert_eq!(fixes[0].signature, "IndexOutOfBounds");
}
#[tokio::test]
async fn test_handle_search_error_fixes() {
let handler = setup_test_handler("search_error_fixes");
handler.state.error_fixes.modify(|fixes| {
fixes.push(crate::models::ErrorFix {
signature: "NullPointerException".to_string(),
solution: "Initialize the pointer".to_string(),
timestamp: 0,
git_branch: None,
git_commit: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 23,
"method": "tools/call",
"params": {
"name": "search_error_fixes",
"arguments": {
"query": "NullPointer"
}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("Initialize the pointer"));
}
#[tokio::test]
async fn test_handle_list_pinned_files() {
let handler = setup_test_handler("list_pinned_files");
handler.state.pinned_files.modify(|files| {
files.push(crate::models::PinnedFile {
file_path: "src/important.rs".to_string(),
timestamp: 0,
namespace: "".to_string(),
git_branch: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 24,
"method": "tools/call",
"params": {
"name": "list_pinned_files",
"arguments": {}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("important.rs"));
}
#[tokio::test]
async fn test_handle_add_session_summary() {
let handler = setup_test_handler("add_session_summary");
let req = json!({
"jsonrpc": "2.0",
"id": 25,
"method": "tools/call",
"params": {
"name": "add_session_summary",
"arguments": {
"namespace": "",
"summary": "Finished writing tests"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let summaries = handler.state.session_summaries.read();
assert_eq!(summaries.len(), 1);
assert_eq!(summaries[0].summary, "Finished writing tests");
}
#[tokio::test]
async fn test_handle_get_project_timeline() {
let handler = setup_test_handler("get_project_timeline");
handler.state.session_summaries.modify(|summaries| {
summaries.push(crate::models::SessionSummary {
summary: "Day 1: Setup project".to_string(),
namespace: "".to_string(),
timestamp: 0,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 26,
"method": "tools/call",
"params": {
"name": "get_project_timeline",
"arguments": {
"namespace": ""
}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("Day 1: Setup project"));
}
#[tokio::test]
async fn test_handle_log_tech_debt() {
let handler = setup_test_handler("log_tech_debt");
let req = json!({
"jsonrpc": "2.0",
"id": 27,
"method": "tools/call",
"params": {
"name": "log_tech_debt",
"arguments": {
"namespace": "",
"description": "Hardcoded values",
"ideal_solution": "Remove magic numbers"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let debt = handler.state.tech_debts.read();
assert_eq!(debt.len(), 1);
assert_eq!(debt[0].description, "Hardcoded values");
}
#[tokio::test]
async fn test_handle_list_tech_debt() {
let handler = setup_test_handler("list_tech_debt");
handler.state.tech_debts.modify(|debts| {
debts.push(crate::models::TechDebt {
id: "debt-1".to_string(),
description: "Bad naming".to_string(),
ideal_solution: "Rename x to num_elements".to_string(),
namespace: "".to_string(),
is_resolved: false,
created_at: 0,
git_branch: None,
git_commit: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 28,
"method": "tools/call",
"params": {
"name": "list_tech_debt",
"arguments": {
"namespace": "",
"include_resolved": false
}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("Bad naming"));
}
#[tokio::test]
async fn test_handle_get_project_health() {
let handler = setup_test_handler("get_project_health");
let req = json!({
"jsonrpc": "2.0",
"id": 29,
"method": "tools/call",
"params": {
"name": "get_project_health",
"arguments": {
"namespace": ""
}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("\"active_tasks\""));
assert!(content.contains("\"unresolved_tech_debt\""));
}
#[tokio::test]
async fn test_handle_resolve_tech_debt() {
let handler = setup_test_handler("resolve_tech_debt");
handler.state.tech_debts.modify(|debts| {
debts.push(crate::models::TechDebt {
id: "debt-2".to_string(),
description: "Old api".to_string(),
ideal_solution: "Use new api".to_string(),
namespace: "".to_string(),
is_resolved: false,
created_at: 0,
git_branch: None,
git_commit: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 30,
"method": "tools/call",
"params": {
"name": "resolve_tech_debt",
"arguments": {
"id": "debt-2"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let debts = handler.state.tech_debts.read();
assert!(debts[0].is_resolved);
}
#[tokio::test]
async fn test_handle_leave_handoff_memo() {
let handler = setup_test_handler("leave_handoff_memo");
let req = json!({
"jsonrpc": "2.0",
"id": 31,
"method": "tools/call",
"params": {
"name": "leave_handoff_memo",
"arguments": {
"namespace": "",
"content": "Make sure to check the logs.",
"author": "Riz"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let memos = handler.state.handoff_memos.read();
assert_eq!(memos.len(), 1);
assert_eq!(memos[0].content, "Make sure to check the logs.");
}
#[tokio::test]
async fn test_handle_query_recent_changes() {
let handler = setup_test_handler("query_recent_changes");
handler.state.ledger.modify(|ledger| {
ledger.push(crate::models::CodeChange {
timestamp: 0,
file_path: "src/main.rs".to_string(),
description: "Fix bug".to_string(),
git_commit: None,
git_branch: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 18,
"method": "tools/call",
"params": {
"name": "query_recent_changes",
"arguments": {}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("Fix bug"));
assert!(content.contains("src/main.rs"));
}
#[tokio::test]
async fn test_handle_update_task_status() {
let handler = setup_test_handler("update_task_status");
handler.state.tasks.modify(|tasks| {
tasks.push(crate::models::Task {
id: "test-task-123".to_string(),
title: "In progress task".to_string(),
status: "pending".to_string(),
description: "".to_string(),
created_at: 0,
updated_at: 0,
git_branch: None,
acceptance_criteria: vec![],
dependencies: vec![],
parent_id: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 19,
"method": "tools/call",
"params": {
"name": "update_task_status",
"arguments": {
"id": "test-task-123",
"status": "in_progress"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let tasks = handler.state.tasks.read();
assert_eq!(tasks[0].status, "in_progress");
}
#[tokio::test]
async fn test_handle_pin_file() {
let handler = setup_test_handler("pin_file");
let req = json!({
"jsonrpc": "2.0",
"id": 100,
"method": "tools/call",
"params": {
"name": "pin_file",
"arguments": {
"file_path": "/path/to/pinned.rs",
"namespace": "global"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let pinned = handler.state.pinned_files.read();
assert!(pinned.iter().any(|p| p.file_path == "/path/to/pinned.rs"));
}
#[tokio::test]
async fn test_handle_unpin_file() {
let handler = setup_test_handler("unpin_file");
handler.state.pinned_files.modify(|files| {
files.push(crate::models::PinnedFile {
namespace: "global".to_string(),
file_path: "/path/to/unpin.rs".to_string(),
timestamp: 0,
git_branch: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 101,
"method": "tools/call",
"params": {
"name": "unpin_file",
"arguments": {
"file_path": "/path/to/unpin.rs",
"namespace": "global"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let pinned = handler.state.pinned_files.read();
assert!(!pinned.iter().any(|p| p.file_path == "/path/to/unpin.rs"));
}
#[tokio::test]
async fn test_handle_learn_preference() {
let handler = setup_test_handler("learn_preference");
let req = json!({
"jsonrpc": "2.0",
"id": 102,
"method": "tools/call",
"params": {
"name": "learn_preference",
"arguments": {
"key": "formatting",
"value": "use spaces"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let prefs = handler.state.prefs.read();
assert!(prefs.values().any(|p| p.key == "formatting" && p.value == "use spaces"));
}
#[tokio::test]
async fn test_handle_read_preferences() {
let handler = setup_test_handler("read_preferences");
handler.state.prefs.modify(|prefs| {
prefs.insert("theme".to_string(), crate::models::Preference {
key: "theme".to_string(),
value: "dark".to_string(),
updated_at: 0,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 103,
"method": "tools/call",
"params": {
"name": "read_preferences",
"arguments": {}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("theme"));
assert!(content.contains("dark"));
}
#[tokio::test]
async fn test_handle_delete_task() {
let handler = setup_test_handler("delete_task");
handler.state.tasks.modify(|tasks| {
tasks.push(crate::models::Task {
id: "task-to-delete".to_string(),
title: "".to_string(),
status: "".to_string(),
description: "".to_string(),
created_at: 0,
updated_at: 0,
git_branch: None,
acceptance_criteria: vec![],
dependencies: vec![],
parent_id: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 104,
"method": "tools/call",
"params": {
"name": "delete_task",
"arguments": {
"id": "task-to-delete"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let tasks = handler.state.tasks.read();
assert!(tasks.is_empty());
}
#[tokio::test]
async fn test_handle_add_milestone() {
let handler = setup_test_handler("add_milestone");
let req = json!({
"jsonrpc": "2.0",
"id": 105,
"method": "tools/call",
"params": {
"name": "add_milestone",
"arguments": {
"title": "v1.0",
"namespace": "global"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let milestones = handler.state.milestones.read();
assert_eq!(milestones.len(), 1);
assert_eq!(milestones[0].title, "v1.0");
}
#[tokio::test]
async fn test_handle_update_milestone() {
let handler = setup_test_handler("update_milestone");
handler.state.milestones.modify(|ms| {
ms.push(crate::models::Milestone {
id: "v1.0".to_string(),
title: "v1.0".to_string(),
status: "pending".to_string(),
namespace: "".to_string(),
target_date: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 106,
"method": "tools/call",
"params": {
"name": "update_milestone",
"arguments": {
"id": "v1.0",
"status": "completed"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let ms = handler.state.milestones.read();
assert_eq!(ms[0].status, "completed");
}
#[tokio::test]
async fn test_handle_list_milestones() {
let handler = setup_test_handler("list_milestones");
handler.state.milestones.modify(|ms| {
ms.push(crate::models::Milestone {
id: "v2.0".to_string(),
title: "v2.0".to_string(),
status: "pending".to_string(),
namespace: "".to_string(),
target_date: None,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 107,
"method": "tools/call",
"params": {
"name": "list_milestones",
"arguments": {}
}
});
let response = handler.handle_request(req).await.unwrap();
let content = response["result"]["content"][0]["text"].as_str().unwrap();
assert!(content.contains("v2.0"));
}
#[tokio::test]
async fn test_handle_add_pr_checklist_item() {
let handler = setup_test_handler("add_pr_checklist_item");
let req = json!({
"jsonrpc": "2.0",
"id": 108,
"method": "tools/call",
"params": {
"name": "add_pr_checklist_item",
"arguments": {
"description": "Check tests",
"namespace": "global"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let cl = handler.state.pr_checklists.read();
assert_eq!(cl.len(), 1);
assert_eq!(cl[0].description, "Check tests");
}
#[tokio::test]
async fn test_handle_clear_pr_checklist() {
let handler = setup_test_handler("clear_pr_checklist");
handler.state.pr_checklists.modify(|cl| {
cl.push(crate::models::PrChecklistItem {
namespace: "global".to_string(),
id: "item".to_string(),
description: "".to_string(),
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 109,
"method": "tools/call",
"params": {
"name": "clear_pr_checklist",
"arguments": {
"namespace": "global"
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let cl = handler.state.pr_checklists.read();
assert!(cl.is_empty());
}
#[tokio::test]
async fn test_handle_clear_handoff_memos() {
let handler = setup_test_handler("clear_handoff_memos");
handler.state.handoff_memos.modify(|ms| {
ms.push(crate::models::HandoffMemo {
id: "memo".to_string(),
namespace: "".to_string(),
content: "".to_string(),
author: "".to_string(),
timestamp: 0,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 110,
"method": "tools/call",
"params": {
"name": "clear_handoff_memos",
"arguments": {
"ids": ["memo"]
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let ms = handler.state.handoff_memos.read();
assert!(ms.is_empty());
}
#[tokio::test]
async fn test_handle_clear_sticky_notes() {
let handler = setup_test_handler("clear_sticky_notes");
handler.state.sticky.modify(|ns| {
ns.push(crate::models::StickyNote {
content: "".to_string(),
timestamp: 0,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 111,
"method": "tools/call",
"params": {
"name": "clear_sticky_notes",
"arguments": {}
}
});
let _ = handler.handle_request(req).await.unwrap();
let ns = handler.state.sticky.read();
assert!(ns.is_empty());
}
#[tokio::test]
async fn test_handle_delete_sticky_note() {
let handler = setup_test_handler("delete_sticky_note");
handler.state.sticky.modify(|ns| {
ns.push(crate::models::StickyNote {
content: "note1".to_string(),
timestamp: 0,
});
ns.push(crate::models::StickyNote {
content: "note2".to_string(),
timestamp: 0,
});
});
let req = json!({
"jsonrpc": "2.0",
"id": 112,
"method": "tools/call",
"params": {
"name": "delete_sticky_note",
"arguments": {
"index": 1
}
}
});
let _ = handler.handle_request(req).await.unwrap();
let ns = handler.state.sticky.read();
assert_eq!(ns.len(), 1);
assert_eq!(ns[0].content, "note2");
}
}