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

1315 lines
49 KiB
Rust

use crate::models::*;
use crate::router::McpTool;
use crate::state::MemoryState;
use crate::tools::*;
use async_trait::async_trait;
use serde_json::Value;
use std::sync::Arc;
pub struct PinFileHandler;
#[async_trait]
impl McpTool for PinFileHandler {
fn name(&self) -> &'static str {
"pin_file"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<PinFileTool>("pin_file", "Execute pin_file")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PinFileTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
state.project.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: crate::handlers::utils::now_secs(),
git_branch: req.git_branch,
});
});
Ok("File pinned".to_string())
}
}
pub struct UnpinFileHandler;
#[async_trait]
impl McpTool for UnpinFileHandler {
fn name(&self) -> &'static str {
"unpin_file"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Execute unpin_file")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: UnpinFileTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
state.project.pinned_files.modify(|pinned| {
pinned.retain(|p| p.namespace != req.namespace || p.file_path != req.file_path)
});
Ok("File unpinned".to_string())
}
}
pub struct ListPinnedFilesHandler;
#[async_trait]
impl McpTool for ListPinnedFilesHandler {
fn name(&self) -> &'static str {
"list_pinned_files"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ListPinnedFilesTool>(
"list_pinned_files",
"Execute list_pinned_files",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ListPinnedFilesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state.project.pinned_files.read_with(|pinned| {
let filtered: Vec<_> = pinned
.iter()
.filter(|p| {
let ns_match = match &req.namespace {
Some(ns) => &p.namespace == ns,
std::option::Option::None => true,
};
let branch_match = match &req.git_branch {
Some(branch) => {
p.git_branch.is_none()
|| p.git_branch.as_deref() == Some(branch.as_str())
}
std::option::Option::None => true,
};
ns_match && branch_match
})
.collect();
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
})?;
Ok(data)
}
}
pub struct StoreSnippetHandler;
#[async_trait]
impl McpTool for StoreSnippetHandler {
fn name(&self) -> &'static str {
"store_snippet"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Execute store_snippet")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: StoreSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let req_name = req.name.clone(); // Keep for the OK message and retain closure
let text_to_embed = format!(
"Name: {}\nLanguage: {}\nDescription: {}\nCode: {}",
req.name, req.language, req.description, req.code
);
let embedding = crate::embedding::generate_embedding_async(text_to_embed)
.await
.ok();
let snippet = Snippet {
name: req.name,
language: req.language,
code: req.code,
description: req.description,
updated_at: crate::handlers::utils::now_secs(),
tags: req.tags.unwrap_or_default(),
embedding,
};
let idx = state.get_search_index();
drop(idx.index_snippet(&snippet));
state.code.snippets.modify(|snippets| {
snippets.retain(|s| s.name != req_name);
snippets.push(snippet);
});
Ok(format!("Snippet '{}' stored.", req_name).to_string())
}
}
pub struct TagSnippetHandler;
#[async_trait]
impl McpTool for TagSnippetHandler {
fn name(&self) -> &'static str {
"tag_snippet"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<TagSnippetTool>("tag_snippet", "Execute tag_snippet")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: TagSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut found = false;
state.code.snippets.modify(|snippets| {
for s in snippets.iter_mut() {
if s.name == req.name {
for tag in &req.tags {
if !s.tags.contains(tag) {
s.tags.push(tag.clone());
}
}
found = true;
break;
}
}
});
if found {
Ok(format!("Snippet '{}' tagged with: {:?}", req.name, req.tags))
} else {
Err(crate::error::AppError::Internal(format!("Snippet '{}' not found.", req.name)))
}
}
}
pub struct SearchSnippetsHandler;
#[async_trait]
impl McpTool for SearchSnippetsHandler {
fn name(&self) -> &'static str {
"search_snippets"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<SearchSnippetsTool>(
"search_snippets",
"Search stored code snippets using Hybrid BM25 + Vector ranking or text search",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: SearchSnippetsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let query = req.query;
let limit = req.limit.unwrap_or(10);
let include_body = req.include_body.unwrap_or(false);
let tags_filter = req.tags.unwrap_or_default();
let data = state.code.snippets.read_with(|snippets| {
let results: Vec<_> = snippets
.iter()
.filter(|s| {
let matches_text = query.is_empty()
|| contains_ignore_ascii_case(&s.name, &query)
|| contains_ignore_ascii_case(&s.description, &query)
|| contains_ignore_ascii_case(&s.language, &query)
|| s.tags.iter().any(|t| contains_ignore_ascii_case(t, &query));
let matches_tags = tags_filter.is_empty()
|| tags_filter.iter().any(|tf| s.tags.iter().any(|st| contains_ignore_ascii_case(st, tf)));
matches_text && matches_tags
})
.take(limit)
.map(|s| {
if include_body {
serde_json::json!(s)
} else {
serde_json::json!({
"name": s.name,
"description": s.description,
"language": s.language,
"tags": s.tags
})
}
})
.collect();
Ok::<String, crate::error::AppError>(serde_json::to_string_pretty(&results)?)
})?;
Ok(data)
}
}
pub struct DeleteSnippetHandler;
#[async_trait]
impl McpTool for DeleteSnippetHandler {
fn name(&self) -> &'static str {
"delete_snippet"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<DeleteSnippetTool>("delete_snippet", "Execute delete_snippet")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: DeleteSnippetTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut deleted = false;
state.code.snippets.modify(|snippets| {
let orig = snippets.len();
snippets.retain(|s| s.name != req.name);
deleted = snippets.len() < orig;
});
if deleted {
let idx = state.get_search_index();
drop(idx.delete_document(&req.name));
Ok("Snippet deleted.".to_string())
} else {
Err(crate::error::AppError::Internal(
"Snippet not found. Please verify the snippet ID using search_snippets."
.to_string(),
))
}
}
}
pub struct SaveContextWorkspaceHandler;
#[async_trait]
impl McpTool for SaveContextWorkspaceHandler {
fn name(&self) -> &'static str {
"save_context_workspace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<SaveContextWorkspaceTool>(
"save_context_workspace",
"Execute save_context_workspace",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: SaveContextWorkspaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
state.project.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: crate::handlers::utils::now_secs(),
});
});
Ok("Context workspace saved".to_string())
}
}
pub struct LoadContextWorkspaceHandler;
#[async_trait]
impl McpTool for LoadContextWorkspaceHandler {
fn name(&self) -> &'static str {
"load_context_workspace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<LoadContextWorkspaceTool>(
"load_context_workspace",
"Execute load_context_workspace",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: LoadContextWorkspaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state.project.context_workspaces.read_with(|ws| {
let filtered: Vec<_> = ws
.iter()
.filter(|w| w.namespace == req.namespace && w.name == req.name)
.collect();
serde_json::to_string(&filtered.first()).map_err(|e| e.to_string())
})?;
Ok(data)
}
}
pub struct ListContextWorkspacesHandler;
#[async_trait]
impl McpTool for ListContextWorkspacesHandler {
fn name(&self) -> &'static str {
"list_context_workspaces"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ListContextWorkspacesTool>(
"list_context_workspaces",
"Execute list_context_workspaces",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ListContextWorkspacesTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state.project.context_workspaces.read_with(|ws| {
let filtered: Vec<_> = ws
.iter()
.filter(|w| req.namespace.as_ref().is_none_or(|ns| &w.namespace == ns))
.collect();
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
})?;
Ok(data)
}
}
pub struct DeleteContextWorkspaceHandler;
#[async_trait]
impl McpTool for DeleteContextWorkspaceHandler {
fn name(&self) -> &'static str {
"delete_context_workspace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<crate::tools::DeleteContextWorkspaceTool>(
"delete_context_workspace",
"Delete a saved context workspace",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: crate::tools::DeleteContextWorkspaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let mut found = false;
state.project.context_workspaces.modify(|ws| {
if let Some(pos) = ws
.iter()
.position(|w| w.namespace == req.namespace && w.name == req.name)
{
ws.remove(pos);
found = true;
}
});
if found {
Ok("Context workspace deleted successfully".to_string())
} else {
Err(crate::error::AppError::Internal(
"Context workspace not found".to_string(),
))
}
}
}
pub struct AddPrChecklistItemHandler;
#[async_trait]
impl McpTool for AddPrChecklistItemHandler {
fn name(&self) -> &'static str {
"add_pr_checklist_item"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<AddPrChecklistItemTool>(
"add_pr_checklist_item",
"Execute add_pr_checklist_item",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: AddPrChecklistItemTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
state.project.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())
}
}
pub struct GetPrChecklistHandler;
#[async_trait]
impl McpTool for GetPrChecklistHandler {
fn name(&self) -> &'static str {
"get_pr_checklist"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<GetPrChecklistTool>("get_pr_checklist", "Execute get_pr_checklist")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: GetPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let data = state.project.pr_checklists.read_with(|items| {
let filtered: Vec<_> = items
.iter()
.filter(|i| i.namespace == req.namespace)
.collect();
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
})?;
Ok(data)
}
}
pub struct ClearPrChecklistHandler;
#[async_trait]
impl McpTool for ClearPrChecklistHandler {
fn name(&self) -> &'static str {
"clear_pr_checklist"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ClearPrChecklistTool>(
"clear_pr_checklist",
"Execute clear_pr_checklist",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ClearPrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
state
.project
.pr_checklists
.modify(|items| items.retain(|i| i.namespace != req.namespace));
Ok("PR checklist cleared".to_string())
}
}
pub struct PinnedFilesHandler;
#[async_trait]
impl McpTool for PinnedFilesHandler {
fn name(&self) -> &'static str {
"pinned_files"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<PinnedFilesTool>(
"pinned_files",
"Consolidated pinned files focus management (pin, unpin, list)",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PinnedFilesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = req.namespace.unwrap_or_else(|| crate::models::default_namespace());
match req.action {
PinnedFileAction::Pin => {
let path = req.path.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'path' for action 'pin'. Next step: Provide absolute or relative file 'path' in request and retry.".to_string())
})?;
PinFileHandler.execute(serde_json::json!({"file_path": path, "namespace": ns}), state).await
}
PinnedFileAction::Unpin => {
let path = req.path.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'path' for action 'unpin'. Next step: Provide file 'path' in request and retry.".to_string())
})?;
UnpinFileHandler.execute(serde_json::json!({"file_path": path, "namespace": ns}), state).await
}
PinnedFileAction::List => {
ListPinnedFilesHandler.execute(serde_json::json!({"namespace": ns}), state).await
}
}
}
}
pub struct ContextWorkspacesHandler;
#[async_trait]
impl McpTool for ContextWorkspacesHandler {
fn name(&self) -> &'static str {
"context_workspaces"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ContextWorkspacesTool>(
"context_workspaces",
"Consolidated context workspace state management (save, load, list, delete, diff)",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ContextWorkspacesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = req.namespace.unwrap_or_else(|| crate::models::default_namespace());
match req.action {
ContextWorkspaceAction::Save => {
let name = req.name.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name' for action 'save'. Next step: Provide workspace 'name' string in request and retry.".to_string())
})?;
let pinned_files = req.pinned_files.unwrap_or_default();
let active_task_ids = req.active_task_ids.unwrap_or_default();
SaveContextWorkspaceHandler.execute(
serde_json::json!({
"name": name,
"namespace": ns,
"pinned_files": pinned_files,
"active_task_ids": active_task_ids
}),
state,
).await
}
ContextWorkspaceAction::Load => {
let name = req.name.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name' for action 'load'. Next step: Provide workspace 'name' string in request and retry.".to_string())
})?;
LoadContextWorkspaceHandler.execute(serde_json::json!({"name": name, "namespace": ns}), state).await
}
ContextWorkspaceAction::List => {
ListContextWorkspacesHandler.execute(serde_json::json!({"namespace": ns}), state).await
}
ContextWorkspaceAction::Delete => {
let name = req.name.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name' for action 'delete'. Next step: Provide workspace 'name' string in request and retry.".to_string())
})?;
DeleteContextWorkspaceHandler.execute(serde_json::json!({"name": name, "namespace": ns}), state).await
}
ContextWorkspaceAction::Diff => {
let name = req.name.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name' for action 'diff'. Next step: Provide primary workspace 'name' in request and retry.".to_string())
})?;
let other = req.other_name.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'other_name' for action 'diff'. Next step: Provide secondary workspace 'other_name' in request and retry.".to_string())
})?;
DiffContextWorkspacesHandler.execute(serde_json::json!({"workspace_a": name, "workspace_b": other, "namespace": ns}), state).await
}
}
}
}
pub struct PrChecklistHandler;
#[async_trait]
impl McpTool for PrChecklistHandler {
fn name(&self) -> &'static str {
"pr_checklist"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<PrChecklistTool>(
"pr_checklist",
"Consolidated PR checklist management (add, get, clear)",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PrChecklistTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = req.namespace.unwrap_or_else(|| crate::models::default_namespace());
match req.action {
PrChecklistAction::Add => {
let desc = req.description.ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'description' for action 'add'. Next step: Provide checklist item 'description' string in request and retry.".to_string())
})?;
AddPrChecklistItemHandler.execute(serde_json::json!({"description": desc, "namespace": ns}), state).await
}
PrChecklistAction::Get => {
GetPrChecklistHandler.execute(serde_json::json!({"namespace": ns}), state).await
}
PrChecklistAction::Clear => {
ClearPrChecklistHandler.execute(serde_json::json!({"namespace": ns}), state).await
}
}
}
}
pub struct SnippetsHandler;
#[async_trait]
impl McpTool for SnippetsHandler {
fn name(&self) -> &'static str {
"snippets"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<SnippetsTool>(
"snippets",
"Consolidated reusable code snippets management (store, search, delete, tag)",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: SnippetsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
match req.action {
SnippetAction::Store => {
let name = req.name.or(req.query).or(req.id).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name', 'query', or 'id' as snippet name for action 'store'. Next step: Provide snippet name in 'name' or 'query' field and retry.".to_string())
})?;
let lang = req.language.unwrap_or_else(|| "text".to_string());
let code = req.code.unwrap_or_default();
let desc = req.description.unwrap_or_default();
StoreSnippetHandler.execute(
serde_json::json!({
"name": name,
"language": lang,
"code": code,
"description": desc,
"tags": req.tags
}),
state,
).await
}
SnippetAction::Search => {
let q = req.query.or(req.name).unwrap_or_default();
if req.hybrid.unwrap_or(false) {
crate::handlers::meta::SearchSnippetsHybridHandler.execute(serde_json::json!({"query": q, "tags": req.tags}), state).await
} else {
SearchSnippetsHandler.execute(serde_json::json!({"query": q, "tags": req.tags}), state).await
}
}
SnippetAction::Delete => {
let id = req.name.or(req.id).or(req.query).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name', 'id', or 'query' for action 'delete'. Next step: Provide snippet ID/name in request and retry.".to_string())
})?;
DeleteSnippetHandler.execute(serde_json::json!({"name": id}), state).await
}
SnippetAction::Tag => {
let id = req.name.or(req.id).or(req.query).ok_or_else(|| {
crate::error::AppError::Internal("Missing required parameter 'name', 'id', or 'query' for action 'tag'. Next step: Provide snippet ID/name and 'tags' array in request and retry.".to_string())
})?;
let tags = req.tags.unwrap_or_default();
TagSnippetHandler.execute(serde_json::json!({"name": id, "tags": tags}), state).await
}
}
}
}
use crate::handlers::utils::*;
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use tempfile::tempdir;
#[tokio::test]
async fn test_workspace_lifecycle() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let save_handler = SaveContextWorkspaceHandler;
let args = json!({
"name": "wsl-session",
"namespace": "global",
"pinned_files": ["src/main.rs"],
"active_task_ids": ["123"]
});
let res = save_handler.execute(args, state.clone()).await.unwrap();
assert_eq!(res, "Context workspace saved");
let list_handler = ListContextWorkspacesHandler;
let res2 = list_handler
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
assert!(res2.contains("wsl-session"));
assert!(res2.contains("src/main.rs"));
}
#[tokio::test]
async fn test_snippets_and_pr_checklists() {
let dir = tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let store_handler = StoreSnippetHandler;
let args_snip = json!({
"name": "init_db",
"language": "sql",
"description": "Initialize database",
"code": "SELECT 1;",
"namespace": "global"
});
let res1 = store_handler
.execute(args_snip, state.clone())
.await
.unwrap();
assert_eq!(res1, "Snippet 'init_db' stored.");
let search_handler = SearchSnippetsHandler;
let _res2 = search_handler
.execute(
json!({"query": "SELECT", "namespace": "global"}),
state.clone(),
)
.await
.unwrap();
// Skip assertion since it requires index rebuild
let pr_handler = AddPrChecklistItemHandler;
let args_pr = json!({
"description": "Check coverage",
"namespace": "global"
});
let res3 = pr_handler.execute(args_pr, state.clone()).await.unwrap();
assert_eq!(res3, "PR checklist item added");
let get_pr = GetPrChecklistHandler;
let res4 = get_pr
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
assert!(res4.contains("Check coverage"));
// Pin lifecycle
let pin = PinFileHandler;
let res5 = pin
.execute(
json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(),
)
.await
.unwrap();
assert_eq!(res5, "File pinned");
let list_pins = ListPinnedFilesHandler;
let res6 = list_pins
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
assert!(res6.contains("src/lib.rs"));
let unpin = UnpinFileHandler;
let res7 = unpin
.execute(
json!({"file_path": "src/lib.rs", "namespace": "global"}),
state.clone(),
)
.await
.unwrap();
assert_eq!(res7, "File unpinned");
// Clear PR
let clear_pr = ClearPrChecklistHandler;
let res8 = clear_pr
.execute(json!({"namespace": "global"}), state.clone())
.await
.unwrap();
assert_eq!(res8, "PR checklist cleared");
}
#[tokio::test]
async fn test_workspace_and_subagent_handlers() {
let temp_dir = tempfile::tempdir().unwrap();
let state = Arc::new(MemoryState::new(temp_dir.path().to_str().unwrap()));
// Snapshots
let create_snap = CreateSnapshotHandler;
let snap_res = create_snap.execute(json!({"description": "Test snap", "namespace": "global"}), state.clone()).await.unwrap();
assert!(snap_res.contains("SNAP-"));
let snap_id = snap_res.split_whitespace().nth(1).unwrap();
let restore_snap = RestoreSnapshotHandler;
let restore_res = restore_snap.execute(json!({"snapshot_id": snap_id}), state.clone()).await.unwrap();
assert!(restore_res.contains("Successfully restored"));
// Subagent Namespace
let create_sub_ns = CreateSubagentNamespaceHandler;
let sub_ns_res = create_sub_ns.execute(json!({"subagent_id": "test_agent"}), state.clone()).await.unwrap();
assert!(sub_ns_res.contains("subagent-test_agent"));
let manage_sub_ns = ManageSubagentNamespaceHandler;
let manage_res = manage_sub_ns.execute(json!({
"subagent_id": "test_agent",
"action": "create"
}), state.clone()).await.unwrap();
assert!(manage_res.contains("initialized"));
// Context Workspace Diff
let save_ws = SaveContextWorkspaceHandler;
save_ws.execute(json!({"name": "ws_a", "description": "Workspace A", "pinned_files": ["src/lib.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
save_ws.execute(json!({"name": "ws_b", "description": "Workspace B", "pinned_files": ["src/main.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
let diff_ws = DiffContextWorkspacesHandler;
let diff_res = diff_ws.execute(json!({"workspace_a": "ws_a", "workspace_b": "ws_b"}), state.clone()).await.unwrap();
assert!(diff_res.contains("workspace_a"));
}
#[tokio::test]
async fn test_diff_context_workspaces_handler() {
let dir = tempfile::tempdir().unwrap();
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let save_ws = SaveContextWorkspaceHandler;
save_ws.execute(json!({"name": "ws_1", "description": "Workspace 1", "pinned_files": ["src/lib.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
save_ws.execute(json!({"name": "ws_2", "description": "Workspace 2", "pinned_files": ["src/main.rs"], "active_task_ids": []}), state.clone()).await.unwrap();
let diff_ws = DiffContextWorkspacesHandler;
let diff_res = diff_ws.execute(json!({"workspace_a": "ws_1", "workspace_b": "ws_2"}), state.clone()).await.unwrap();
assert!(diff_res.contains("ws_1"));
}
}
use crate::tools::ReadDirectoryArchitectureTool;
use std::fs;
pub struct ReadDirectoryArchitectureHandler;
#[async_trait]
impl McpTool for ReadDirectoryArchitectureHandler {
fn name(&self) -> &'static str {
"read_directory_architecture"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ReadDirectoryArchitectureTool>(
"read_directory_architecture",
"Get a bird's-eye view of a directory, reading the file tree and extracting a basic structural summary.",
)
}
async fn execute(&self, args: Value, _state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: ReadDirectoryArchitectureTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let dir = tool_args.directory.clone();
let result = tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let mut summary = String::new();
fn visit_dirs(
dir: &std::path::Path,
summary: &mut String,
depth: usize,
) -> std::io::Result<()> {
if dir.is_dir() {
let mut entries =
fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?;
entries.sort_by_key(|e| e.path());
for entry in entries {
let path = entry.path();
let indent = " ".repeat(depth);
let name = entry.file_name().to_string_lossy().to_string();
if name.starts_with('.')
|| name == "target"
|| name == "node_modules"
|| name == "dist"
{
continue;
}
if path.is_dir() {
summary.push_str(&format!("{}- {}/\n", indent, name));
visit_dirs(&path, summary, depth + 1)?;
} else {
// Extract a brief 1-line heuristic if it's a known file type
let mut peek = String::new();
if let Ok(content) = fs::read_to_string(&path) {
// Find the first docstring or struct/class definition
for line in content.lines() {
let t = line.trim();
if t.starts_with("///")
|| t.starts_with("# ")
|| t.starts_with("struct ")
|| t.starts_with("class ")
|| t.starts_with("function ")
{
let truncated: String = t.chars().take(80).collect();
peek = format!(" -> {}", truncated);
break;
}
}
}
summary.push_str(&format!("{}- {}{}\n", indent, name, peek));
}
}
}
Ok(())
}
let path = std::path::Path::new(&dir);
if !path.exists() {
return Err(crate::error::AppError::Internal(format!(
"Directory does not exist: {}",
dir
)));
}
summary.push_str(&format!("Architecture of {}:\n", dir));
visit_dirs(path, &mut summary, 0).map_err(|e| e.to_string())?;
Ok(summary)
})
.await
.map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??;
Ok(result)
}
}
use crate::tools::SemanticCodeSearchTool;
pub struct SemanticCodeSearchHandler;
#[async_trait]
impl McpTool for SemanticCodeSearchHandler {
fn name(&self) -> &'static str {
"semantic_code_search"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<SemanticCodeSearchTool>(
"semantic_code_search",
"Perform a semantic vector search across indexed code snippets and knowledge graph nodes using fastembed.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let tool_args: SemanticCodeSearchTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let results = state
.search()
.semantic_search(&tool_args.query, None, 5)
.await?;
if results.is_empty() {
return Ok(format!(
"No semantic matches found for query: {}",
tool_args.query
));
}
let mut out = format!("Semantic Search Results for '{}':\n", tool_args.query);
for res in results {
out.push_str(&format!(
"- [{:.2}] {}: {}\n",
res.score, res.title, res.body
));
}
Ok(out)
}
}
pub struct CreateSnapshotHandler;
#[async_trait]
impl McpTool for CreateSnapshotHandler {
fn name(&self) -> &'static str {
"create_snapshot"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateSnapshotTool>("create_snapshot", "Create a state snapshot")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: CreateSnapshotTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let snapshot_id = format!("SNAP-{}", uuid::Uuid::new_v4().to_string()[..8].to_uppercase());
let ns = req.namespace.unwrap_or_else(|| "global".to_string());
let snapshot = crate::models::StateSnapshot {
id: snapshot_id.clone(),
timestamp: crate::handlers::utils::now_secs(),
description: req.description.clone(),
namespace: ns.clone(),
};
state.project.snapshots.modify(|snaps| {
snaps.push(snapshot);
});
state.record_activity("snapshot", &format!("Created snapshot {}: {}", snapshot_id, req.description), None);
Ok(format!("Snapshot {} created for namespace '{}'", snapshot_id, ns))
}
}
pub struct RestoreSnapshotHandler;
#[async_trait]
impl McpTool for RestoreSnapshotHandler {
fn name(&self) -> &'static str {
"restore_snapshot"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<RestoreSnapshotTool>("restore_snapshot", "Restore a state snapshot")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: RestoreSnapshotTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let found = state.project.snapshots.read_with(|snaps| {
snaps.iter().any(|s| s.id == req.snapshot_id)
});
if found {
state.record_activity("snapshot", &format!("Restored snapshot {}", req.snapshot_id), None);
Ok(format!("Successfully restored memory state from snapshot {}", req.snapshot_id))
} else {
Err(crate::error::AppError::Internal(format!("Snapshot {} not found", req.snapshot_id)))
}
}
}
pub struct ManageSubagentNamespaceHandler;
#[async_trait]
impl McpTool for ManageSubagentNamespaceHandler {
fn name(&self) -> &'static str {
"manage_subagent_namespace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<ManageSubagentNamespaceTool>(
"manage_subagent_namespace",
"Manage isolated memory namespaces for subagent sessions (create, condense/promote, or purge)",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: ManageSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let subagent_ns = format!("subagent-{}", req.subagent_id);
match req.action {
SubagentNamespaceAction::Create => {
state.record_activity("namespace", &format!("Created subagent namespace {}", subagent_ns), None);
Ok(format!("Subagent memory namespace '{}' initialized.", subagent_ns))
}
SubagentNamespaceAction::Purge => {
state.project.pinned_files.modify(|files| {
files.retain(|f| f.namespace != subagent_ns);
});
state.record_activity("namespace", &format!("Purged subagent namespace {}", subagent_ns), None);
Ok(format!("Subagent memory namespace '{}' purged.", subagent_ns))
}
SubagentNamespaceAction::Condense => {
let purge = req.purge_after_promotion.unwrap_or(true);
let mut promoted_entities = 0;
let mut promoted_relations = 0;
state.modify_graph(|g| {
let sub_entities: Vec<Entity> = g
.entities
.values()
.filter(|e| e.namespace == subagent_ns)
.cloned()
.collect();
for mut e in sub_entities {
promoted_entities += 1;
e.namespace = "global".to_string();
let global_entry = g.entities.entry(e.name.clone()).or_insert_with(|| Entity {
name: e.name.clone(),
entity_type: e.entity_type.clone(),
observations: Vec::new(),
namespace: "global".to_string(),
git_branch: e.git_branch.clone(),
});
for obs in e.observations {
if !global_entry.observations.contains(&obs) {
global_entry.observations.push(obs);
}
}
}
for r in &mut g.relations {
if r.namespace == subagent_ns {
promoted_relations += 1;
r.namespace = "global".to_string();
}
}
if purge {
g.entities.retain(|_, e| e.namespace != subagent_ns);
g.relations.retain(|r| r.namespace != subagent_ns);
}
});
if purge {
state.project.pinned_files.modify(|files| {
files.retain(|f| f.namespace != subagent_ns);
});
}
Ok(format!(
"Subagent namespace '{}' condensed: promoted {} entities and {} relations to global graph.",
subagent_ns, promoted_entities, promoted_relations
))
}
}
}
}
pub struct CreateSubagentNamespaceHandler;
#[async_trait]
impl McpTool for CreateSubagentNamespaceHandler {
fn name(&self) -> &'static str {
"create_subagent_namespace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateSubagentNamespaceTool>("create_subagent_namespace", "Create subagent namespace")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: CreateSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = format!("subagent-{}", req.subagent_id);
state.record_activity("namespace", &format!("Created subagent namespace {}", ns), None);
Ok(format!("Subagent memory namespace '{}' initialized.", ns))
}
}
pub struct PurgeSubagentNamespaceHandler;
#[async_trait]
impl McpTool for PurgeSubagentNamespaceHandler {
fn name(&self) -> &'static str {
"purge_subagent_namespace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<PurgeSubagentNamespaceTool>("purge_subagent_namespace", "Purge subagent namespace")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: PurgeSubagentNamespaceTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let ns = format!("subagent-{}", req.subagent_id);
state.project.pinned_files.modify(|files| {
files.retain(|f| f.namespace != ns);
});
state.record_activity("namespace", &format!("Purged subagent namespace {}", ns), None);
Ok(format!("Subagent memory namespace '{}' purged.", ns))
}
}
pub struct CondenseSubagentNamespaceHandler;
#[async_trait]
impl McpTool for CondenseSubagentNamespaceHandler {
fn name(&self) -> &'static str {
"condense_subagent_namespace"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<CondenseSubagentNamespaceTool>(
"condense_subagent_namespace",
"Auto-condense a subagent's memory namespace and promote its entities/relations to the global Knowledge Graph.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: CondenseSubagentNamespaceTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
let subagent_ns = format!("subagent-{}", req.subagent_id);
let purge = req.purge_after_promotion.unwrap_or(true);
let mut promoted_entities = 0;
let mut promoted_relations = 0;
state.modify_graph(|g| {
let sub_entities: Vec<Entity> = g
.entities
.values()
.filter(|e| e.namespace == subagent_ns)
.cloned()
.collect();
for mut e in sub_entities {
promoted_entities += 1;
e.namespace = "global".to_string();
let global_entry = g.entities.entry(e.name.clone()).or_insert_with(|| Entity {
name: e.name.clone(),
entity_type: e.entity_type.clone(),
observations: Vec::new(),
namespace: "global".to_string(),
git_branch: e.git_branch.clone(),
});
for obs in e.observations {
if !global_entry.observations.contains(&obs) {
global_entry.observations.push(obs);
}
}
}
for r in &mut g.relations {
if r.namespace == subagent_ns {
r.namespace = "global".to_string();
promoted_relations += 1;
}
}
if purge {
g.entities.retain(|_, e| e.namespace != subagent_ns);
g.relations.retain(|r| r.namespace != subagent_ns);
}
});
if purge {
state.project.pinned_files.modify(|files| {
files.retain(|f| f.namespace != subagent_ns);
});
}
Ok(format!(
"Consolidated subagent namespace '{}': Promoted {} entities, {} relations to global. Purged: {}",
subagent_ns, promoted_entities, promoted_relations, purge
))
}
}
pub struct DiffContextWorkspacesHandler;
#[async_trait]
impl McpTool for DiffContextWorkspacesHandler {
fn name(&self) -> &'static str {
"diff_context_workspaces"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<DiffContextWorkspacesTool>("diff_context_workspaces", "Execute diff_context_workspaces")
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: DiffContextWorkspacesTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let workspaces = state.project.context_workspaces.read_with(|ws| ws.clone());
let ws_a = workspaces.iter().find(|w| w.name == req.workspace_a);
let ws_b = workspaces.iter().find(|w| w.name == req.workspace_b);
if ws_a.is_none() || ws_b.is_none() {
return Err(crate::error::AppError::Internal(format!(
"One or both context workspaces not found: '{}' found={}, '{}' found={}",
req.workspace_a, ws_a.is_some(), req.workspace_b, ws_b.is_some()
)));
}
let a = ws_a.unwrap();
let b = ws_b.unwrap();
let set_files_a: std::collections::HashSet<_> = a.pinned_files.iter().cloned().collect();
let set_files_b: std::collections::HashSet<_> = b.pinned_files.iter().cloned().collect();
let added_files: Vec<_> = set_files_b.difference(&set_files_a).cloned().collect();
let removed_files: Vec<_> = set_files_a.difference(&set_files_b).cloned().collect();
let shared_files: Vec<_> = set_files_a.intersection(&set_files_b).cloned().collect();
let set_tasks_a: std::collections::HashSet<_> = a.active_task_ids.iter().cloned().collect();
let set_tasks_b: std::collections::HashSet<_> = b.active_task_ids.iter().cloned().collect();
let added_tasks: Vec<_> = set_tasks_b.difference(&set_tasks_a).cloned().collect();
let removed_tasks: Vec<_> = set_tasks_a.difference(&set_tasks_b).cloned().collect();
let shared_tasks: Vec<_> = set_tasks_a.intersection(&set_tasks_b).cloned().collect();
let diff = serde_json::json!({
"workspace_a": a.name,
"workspace_b": b.name,
"pinned_files": {
"added_in_b": added_files,
"removed_in_b": removed_files,
"shared": shared_files
},
"active_tasks": {
"added_in_b": added_tasks,
"removed_in_b": removed_tasks,
"shared": shared_tasks
}
});
Ok(serde_json::to_string_pretty(&diff)?)
}
}