1248 lines
44 KiB
Rust
1248 lines
44 KiB
Rust
use crate::state::MemoryState;
|
|
use async_trait::async_trait;
|
|
use serde_json::Value;
|
|
use std::sync::Arc;
|
|
|
|
#[async_trait]
|
|
pub trait McpTool: Send + Sync {
|
|
/// The unique name of the tool
|
|
fn name(&self) -> &'static str;
|
|
|
|
/// The JSON schema for the tool
|
|
fn schema(&self) -> Value;
|
|
|
|
/// Execute the tool with the given arguments
|
|
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String>;
|
|
}
|
|
|
|
#[async_trait]
|
|
pub trait McpResource: Send + Sync {
|
|
fn uri(&self) -> &'static str;
|
|
fn name(&self) -> &'static str;
|
|
fn description(&self) -> Option<&'static str> {
|
|
None
|
|
}
|
|
fn mime_type(&self) -> Option<&'static str> {
|
|
Some("application/json")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String>;
|
|
}
|
|
|
|
#[async_trait]
|
|
pub trait McpPrompt: Send + Sync {
|
|
fn name(&self) -> &'static str;
|
|
fn description(&self) -> Option<&'static str> {
|
|
None
|
|
}
|
|
fn arguments(&self) -> serde_json::Value {
|
|
serde_json::json!([])
|
|
}
|
|
async fn get(
|
|
&self,
|
|
args: Value,
|
|
state: Arc<MemoryState>,
|
|
) -> crate::error::Result<serde_json::Value>;
|
|
}
|
|
|
|
struct GraphEntitiesResource;
|
|
#[async_trait]
|
|
impl McpResource for GraphEntitiesResource {
|
|
fn uri(&self) -> &'static str {
|
|
"memory://graph/entities"
|
|
}
|
|
fn name(&self) -> &'static str {
|
|
"Graph Entities"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("All nodes and entities currently stored in the knowledge graph")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let graph = state.graph.cache.read().unwrap_or_else(|e| e.into_inner());
|
|
let data: Vec<_> = graph.entities.values().collect();
|
|
if data.len() > 100 {
|
|
let truncated: Vec<_> = data.iter().take(100).cloned().collect();
|
|
let res = serde_json::json!({
|
|
"_meta": format!("Showing 100 of {} total entities. Use omni_search or get_subgraph for targeted entity discovery.", data.len()),
|
|
"entities": truncated
|
|
});
|
|
Ok(serde_json::to_string_pretty(&res)?)
|
|
} else {
|
|
Ok(serde_json::to_string_pretty(&data)?)
|
|
}
|
|
}
|
|
}
|
|
|
|
struct GraphRelationsResource;
|
|
#[async_trait]
|
|
impl McpResource for GraphRelationsResource {
|
|
fn uri(&self) -> &'static str {
|
|
"memory://graph/relations"
|
|
}
|
|
fn name(&self) -> &'static str {
|
|
"Graph Relations"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("All relationships between entities currently stored in the knowledge graph")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let graph = state.graph.cache.read().unwrap_or_else(|e| e.into_inner());
|
|
let data = &graph.relations;
|
|
if data.len() > 200 {
|
|
let truncated: Vec<_> = data.iter().take(200).cloned().collect();
|
|
let res = serde_json::json!({
|
|
"_meta": format!("Showing 200 of {} total relations. Use query_graph_path or get_subgraph for targeted relation discovery.", data.len()),
|
|
"relations": truncated
|
|
});
|
|
Ok(serde_json::to_string_pretty(&res)?)
|
|
} else {
|
|
Ok(serde_json::to_string_pretty(&data)?)
|
|
}
|
|
}
|
|
}
|
|
|
|
struct TasksActiveResource;
|
|
#[async_trait]
|
|
impl McpResource for TasksActiveResource {
|
|
fn uri(&self) -> &'static str {
|
|
"memory://tasks/active"
|
|
}
|
|
fn name(&self) -> &'static str {
|
|
"Active Tasks"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("List of currently active tasks")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let tasks = state
|
|
.project
|
|
.tasks
|
|
.cache
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
let data: Vec<_> = tasks.iter().filter(|t| t.is_active()).collect();
|
|
Ok(serde_json::to_string_pretty(&data)?)
|
|
}
|
|
}
|
|
|
|
struct DecisionsActiveResource;
|
|
#[async_trait]
|
|
impl McpResource for DecisionsActiveResource {
|
|
fn uri(&self) -> &'static str {
|
|
"memory://decisions/active"
|
|
}
|
|
fn name(&self) -> &'static str {
|
|
"Active ADR Decisions"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("All accepted Architectural Decision Records (ADRs)")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let adrs = state
|
|
.code
|
|
.adrs
|
|
.cache
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
let data: Vec<_> = adrs.iter().filter(|a| a.status == "accepted").collect();
|
|
Ok(serde_json::to_string_pretty(&data)?)
|
|
}
|
|
}
|
|
|
|
struct TechDebtActiveResource;
|
|
#[async_trait]
|
|
impl McpResource for TechDebtActiveResource {
|
|
fn uri(&self) -> &'static str {
|
|
"memory://tech_debt/unresolved"
|
|
}
|
|
fn name(&self) -> &'static str {
|
|
"Unresolved Tech Debt"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("All currently unresolved technical debt items")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let debts = state
|
|
.code
|
|
.tech_debts
|
|
.cache
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
let data: Vec<_> = debts.iter().filter(|d| !d.is_resolved).collect();
|
|
Ok(serde_json::to_string_pretty(&data)?)
|
|
}
|
|
}
|
|
|
|
struct AnalyzeTechDebtPrompt;
|
|
#[async_trait]
|
|
impl McpPrompt for AnalyzeTechDebtPrompt {
|
|
fn name(&self) -> &'static str {
|
|
"analyze_tech_debt"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("Analyze the project's current technical debt")
|
|
}
|
|
async fn get(
|
|
&self,
|
|
_args: Value,
|
|
_state: Arc<MemoryState>,
|
|
) -> crate::error::Result<serde_json::Value> {
|
|
Ok(serde_json::json!({
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"type": "text",
|
|
"text": "Please read the current technical debt from the knowledge graph (or use the list_tech_debt tool) and provide a prioritization plan for addressing it."
|
|
}
|
|
}
|
|
]
|
|
}))
|
|
}
|
|
}
|
|
|
|
struct SummarizeArchitecturePrompt;
|
|
#[async_trait]
|
|
impl McpPrompt for SummarizeArchitecturePrompt {
|
|
fn name(&self) -> &'static str {
|
|
"summarize_architecture"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("Summarize project architecture using active ADRs and knowledge graph entities")
|
|
}
|
|
async fn get(
|
|
&self,
|
|
_args: Value,
|
|
_state: Arc<MemoryState>,
|
|
) -> crate::error::Result<serde_json::Value> {
|
|
Ok(serde_json::json!({
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"type": "text",
|
|
"text": "Please read active decisions from memory://decisions/active and graph entities from memory://graph/entities, and synthesize a clean architectural overview report."
|
|
}
|
|
}
|
|
]
|
|
}))
|
|
}
|
|
}
|
|
|
|
struct SessionDeltaResource;
|
|
#[async_trait]
|
|
impl McpResource for SessionDeltaResource {
|
|
fn uri(&self) -> &'static str {
|
|
"memory://session/delta"
|
|
}
|
|
fn name(&self) -> &'static str {
|
|
"Session Delta"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("Recent changes, active tasks, code edits, and notes created in the current session")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let now = crate::handlers::utils::now_secs();
|
|
let lookback = 7200; // last 2 hours
|
|
let cutoff = now.saturating_sub(lookback);
|
|
|
|
let changes = state
|
|
.code
|
|
.ledger
|
|
.cache
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
let recent_changes: Vec<_> = changes.iter().filter(|c| c.timestamp >= cutoff).collect();
|
|
|
|
let tasks = state
|
|
.project
|
|
.tasks
|
|
.cache
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
let active_tasks: Vec<_> = tasks.iter().filter(|t| t.is_active()).collect();
|
|
|
|
let delta = serde_json::json!({
|
|
"recent_changes": recent_changes,
|
|
"active_tasks": active_tasks,
|
|
"timestamp": now,
|
|
});
|
|
|
|
Ok(serde_json::to_string_pretty(&delta)?)
|
|
}
|
|
}
|
|
|
|
pub struct ContextWarmupPrompt;
|
|
#[async_trait]
|
|
impl McpPrompt for ContextWarmupPrompt {
|
|
fn name(&self) -> &'static str {
|
|
"context_warmup"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some(
|
|
"Consolidated context warmup prompt providing active tasks, active decisions, and unresolved tech debt",
|
|
)
|
|
}
|
|
async fn get(
|
|
&self,
|
|
_args: Value,
|
|
_state: Arc<MemoryState>,
|
|
) -> crate::error::Result<serde_json::Value> {
|
|
Ok(serde_json::json!({
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"type": "text",
|
|
"text": "Perform a complete context warmup. Please read memory://session/delta, memory://tasks/active, memory://decisions/active, and memory://tech_debt/unresolved to synthesize current project focus and active workspace constraints."
|
|
}
|
|
}
|
|
]
|
|
}))
|
|
}
|
|
}
|
|
|
|
pub struct MemoryHandler {
|
|
pub state: Arc<MemoryState>,
|
|
pub tools: std::collections::HashMap<String, Box<dyn McpTool>>,
|
|
pub resources: std::collections::HashMap<String, Box<dyn McpResource>>,
|
|
pub prompts: std::collections::HashMap<String, Box<dyn McpPrompt>>,
|
|
}
|
|
|
|
impl MemoryHandler {
|
|
pub fn new(state: Arc<MemoryState>) -> Self {
|
|
let mut tools: std::collections::HashMap<String, Box<dyn McpTool>> =
|
|
std::collections::HashMap::new();
|
|
let mut resources: std::collections::HashMap<String, Box<dyn McpResource>> =
|
|
std::collections::HashMap::new();
|
|
let mut prompts: std::collections::HashMap<String, Box<dyn McpPrompt>> =
|
|
std::collections::HashMap::new();
|
|
|
|
macro_rules! register_resource {
|
|
($handler:ident) => {
|
|
let h = $handler;
|
|
resources.insert(h.uri().to_string(), Box::new(h));
|
|
};
|
|
}
|
|
|
|
macro_rules! register_prompt {
|
|
($handler:ident) => {
|
|
let h = $handler;
|
|
prompts.insert(h.name().to_string(), Box::new(h));
|
|
};
|
|
}
|
|
|
|
register_resource!(GraphEntitiesResource);
|
|
register_resource!(GraphRelationsResource);
|
|
register_resource!(TasksActiveResource);
|
|
register_resource!(DecisionsActiveResource);
|
|
register_resource!(TechDebtActiveResource);
|
|
register_resource!(SessionDeltaResource);
|
|
|
|
register_prompt!(AnalyzeTechDebtPrompt);
|
|
register_prompt!(SummarizeArchitecturePrompt);
|
|
register_prompt!(ContextWarmupPrompt);
|
|
struct TerminalHistoryResource;
|
|
#[async_trait]
|
|
impl McpResource for TerminalHistoryResource {
|
|
fn uri(&self) -> &'static str {
|
|
"memory://terminal/recent"
|
|
}
|
|
fn name(&self) -> &'static str {
|
|
"Terminal History"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("Recent terminal execution history and exit codes")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let items = state
|
|
.telemetry
|
|
.terminal_history
|
|
.cache
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
Ok(serde_json::to_string_pretty(&*items)?)
|
|
}
|
|
}
|
|
struct RecentActivitiesResource;
|
|
#[async_trait]
|
|
impl McpResource for RecentActivitiesResource {
|
|
fn uri(&self) -> &'static str {
|
|
"memory://activity/recent"
|
|
}
|
|
fn name(&self) -> &'static str {
|
|
"Recent Activity"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("Recent IDE and developer activity logs")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let items = state
|
|
.telemetry
|
|
.recent_activities
|
|
.cache
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
Ok(serde_json::to_string_pretty(&*items)?)
|
|
}
|
|
}
|
|
|
|
struct MilestonesResource;
|
|
#[async_trait]
|
|
impl McpResource for MilestonesResource {
|
|
fn uri(&self) -> &'static str {
|
|
"memory://milestones"
|
|
}
|
|
fn name(&self) -> &'static str {
|
|
"Milestones"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("Project milestones and their status")
|
|
}
|
|
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
|
|
let items = state
|
|
.project
|
|
.milestones
|
|
.cache
|
|
.read()
|
|
.unwrap_or_else(|e| e.into_inner());
|
|
Ok(serde_json::to_string_pretty(&*items)?)
|
|
}
|
|
}
|
|
|
|
struct HandoffRoutinePrompt;
|
|
#[async_trait]
|
|
impl McpPrompt for HandoffRoutinePrompt {
|
|
fn name(&self) -> &'static str {
|
|
"handoff_routine"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("Initiate the end-of-session handoff and standup report generation")
|
|
}
|
|
async fn get(
|
|
&self,
|
|
_args: Value,
|
|
_state: Arc<MemoryState>,
|
|
) -> crate::error::Result<serde_json::Value> {
|
|
Ok(serde_json::json!({
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"type": "text",
|
|
"text": "I am logging off. Please invoke the DevOpsSRE subagent to generate a standup report and leave a handoff memo for the next session. Make sure to check active tasks and recent code changes."
|
|
}
|
|
}
|
|
]
|
|
}))
|
|
}
|
|
}
|
|
|
|
register_resource!(TerminalHistoryResource);
|
|
register_resource!(RecentActivitiesResource);
|
|
register_resource!(MilestonesResource);
|
|
struct ArchiveRoutinePrompt;
|
|
#[async_trait]
|
|
impl McpPrompt for ArchiveRoutinePrompt {
|
|
fn name(&self) -> &'static str {
|
|
"archive_routine"
|
|
}
|
|
fn description(&self) -> Option<&'static str> {
|
|
Some("Compress old session summaries into a milestone retrospective")
|
|
}
|
|
async fn get(
|
|
&self,
|
|
_args: Value,
|
|
_state: Arc<MemoryState>,
|
|
) -> crate::error::Result<serde_json::Value> {
|
|
Ok(serde_json::json!({"
|
|
messages": [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"type": "text",
|
|
"text": "Please read the old session summaries, synthesize them into a dense Milestone Retrospective entity, and then use the appropriate tools to delete the old session summaries."
|
|
}
|
|
}
|
|
]
|
|
}))
|
|
}
|
|
}
|
|
register_prompt!(HandoffRoutinePrompt);
|
|
register_prompt!(ArchiveRoutinePrompt);
|
|
|
|
macro_rules! register {
|
|
($module:ident::$handler:ident) => {
|
|
let h = crate::handlers::$module::$handler;
|
|
tools.insert(h.name().to_string(), Box::new(h));
|
|
};
|
|
}
|
|
|
|
register!(graph::QueryGraphPathHandler);
|
|
register!(graph::CreateEntitiesHandler);
|
|
register!(graph::CreateRelationsHandler);
|
|
register!(graph::AddObservationsHandler);
|
|
register!(graph::DeleteEntitiesHandler);
|
|
register!(graph::DeleteObservationsHandler);
|
|
register!(graph::DeleteRelationsHandler);
|
|
register!(graph::ReadGraphHandler);
|
|
register!(graph::SearchNodesHandler);
|
|
register!(graph::OpenNodesHandler);
|
|
register!(graph::VisualizeGraphHandler);
|
|
register!(graph::CondenseEntityHandler);
|
|
register!(graph::MergeEntitiesHandler);
|
|
register!(graph::FindOrphansHandler);
|
|
|
|
register!(tasks::TasksHandler);
|
|
register!(tasks::MilestonesHandler);
|
|
register!(notes::HandoffMemosHandler);
|
|
register!(notes::AddSessionSummaryHandler);
|
|
register!(notes::GenerateStandupReportHandler);
|
|
register!(notes::PromoteToEntityHandler);
|
|
|
|
register!(meta::DecisionsHandler);
|
|
register!(meta::TechDebtHandler);
|
|
register!(meta::LogErrorFixHandler);
|
|
register!(meta::SearchErrorFixesHandler);
|
|
register!(meta::LogCodeChangeHandler);
|
|
register!(meta::QueryRecentChangesHandler);
|
|
register!(meta::OmniSearchHandler);
|
|
register!(meta::GetProjectHealthHandler);
|
|
|
|
register!(env::EnvironmentHandler);
|
|
|
|
register!(workspaces::SnippetsHandler);
|
|
register!(vision::ClipboardHandler);
|
|
|
|
register!(git::GetActiveWorktreeContextHandler);
|
|
register!(git::QueryGitDiffsHandler);
|
|
register!(logs::ProcessLogsHandler);
|
|
register!(ast::ReadFileSkeletonHandler);
|
|
register!(ast::ReplaceAstNodeHandler);
|
|
register!(ast::FindSymbolReferencesHandler);
|
|
register!(ast::GetCallersHandler);
|
|
register!(ast::AnalyzeImpactHandler);
|
|
register!(workspaces::ReadDirectoryArchitectureHandler);
|
|
register!(workspaces::SemanticCodeSearchHandler);
|
|
register!(workspaces::ManageSubagentNamespaceHandler);
|
|
register!(graph::GetSubgraphHandler);
|
|
register!(meta::ManageCheckpointHandler);
|
|
register!(graph::SweepGraphHealthHandler);
|
|
register!(meta::QueryLineageHandler);
|
|
register!(meta::GetNextActionableTasksHandler);
|
|
register!(meta::HypothesesHandler);
|
|
register!(meta::GetPreflightContextHandler);
|
|
register!(graph::ResolveStaleSymbolsHandler);
|
|
register!(graph::SummarizeSubgraphHandler);
|
|
register!(meta::AgentSignalsHandler);
|
|
register!(meta::AutoSessionCheckpointHandler);
|
|
|
|
Self {
|
|
state,
|
|
tools,
|
|
resources,
|
|
prompts,
|
|
}
|
|
}
|
|
|
|
pub async fn handle_request(&self, req: serde_json::Value) -> Option<serde_json::Value> {
|
|
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("");
|
|
|
|
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!({}),
|
|
"resources": serde_json::json!({}),
|
|
"prompts": serde_json::json!({})
|
|
},
|
|
"_meta": {
|
|
"io.modelcontextprotocol/serverInfo": {
|
|
"name": "gemini-mcp-memory",
|
|
"version": "3.0.0"
|
|
}
|
|
}
|
|
});
|
|
Some(crate::mcp::success(id, payload))
|
|
}
|
|
"initialize" => {
|
|
let init = rmcp::model::InitializeResult::new(
|
|
rmcp::model::ServerCapabilities::builder()
|
|
.enable_tools()
|
|
.enable_resources()
|
|
.enable_prompts()
|
|
.build(),
|
|
)
|
|
.with_server_info(rmcp::model::Implementation::new(
|
|
"gemini-mcp-memory",
|
|
"3.0.0",
|
|
))
|
|
.with_instructions(include_str!("instructions.md"));
|
|
Some(crate::mcp::success(
|
|
id,
|
|
serde_json::to_value(&init).unwrap_or_default(),
|
|
))
|
|
}
|
|
"notifications/initialized" => None,
|
|
"tools/list" => {
|
|
let mut tools: Vec<serde_json::Value> =
|
|
self.tools.values().map(|t| t.schema()).collect();
|
|
tools.sort_by_key(|t| {
|
|
t.get("name")
|
|
.and_then(|n| n.as_str())
|
|
.unwrap_or("")
|
|
.to_string()
|
|
});
|
|
Some(crate::mcp::success(
|
|
id,
|
|
serde_json::json!({ "tools": tools }),
|
|
))
|
|
}
|
|
"resources/list" => {
|
|
let resources: Vec<_> = self
|
|
.resources
|
|
.values()
|
|
.map(|r| {
|
|
let mut obj = serde_json::json!({
|
|
"uri": r.uri(),
|
|
"name": r.name(),
|
|
});
|
|
if let Some(desc) = r.description() {
|
|
obj["description"] = serde_json::json!(desc);
|
|
}
|
|
if let Some(mime) = r.mime_type() {
|
|
obj["mimeType"] = serde_json::json!(mime);
|
|
}
|
|
obj
|
|
})
|
|
.collect();
|
|
|
|
let payload = serde_json::json!({ "resources": resources });
|
|
Some(crate::mcp::success(id, payload))
|
|
}
|
|
|
|
"resources/templates/list" => {
|
|
let payload = serde_json::json!({
|
|
"resourceTemplates": []
|
|
});
|
|
Some(crate::mcp::success(id, payload))
|
|
}
|
|
"resources/read" => {
|
|
let params = req.get("params").unwrap_or(&serde_json::Value::Null);
|
|
let uri = params.get("uri").and_then(|u| u.as_str()).unwrap_or("");
|
|
|
|
if let Some(resource) = self.resources.get(uri) {
|
|
match resource.read(Arc::clone(&self.state)).await {
|
|
Ok(text) => {
|
|
let payload = serde_json::json!({
|
|
"contents": [{
|
|
"uri": uri,
|
|
"mimeType": resource.mime_type().unwrap_or("application/json"),
|
|
"text": text
|
|
}]
|
|
});
|
|
Some(crate::mcp::success(id, payload))
|
|
}
|
|
Err(e) => Some(crate::mcp::error(id, -32603, &e.to_string())),
|
|
}
|
|
} else {
|
|
Some(crate::mcp::error(id, -32602, "Resource not found"))
|
|
}
|
|
}
|
|
|
|
"prompts/list" => {
|
|
let prompts: Vec<_> = self
|
|
.prompts
|
|
.values()
|
|
.map(|p| {
|
|
let mut obj = serde_json::json!({
|
|
"name": p.name(),
|
|
"arguments": p.arguments(),
|
|
});
|
|
if let Some(desc) = p.description() {
|
|
obj["description"] = serde_json::json!(desc);
|
|
}
|
|
obj
|
|
})
|
|
.collect();
|
|
|
|
let payload = serde_json::json!({ "prompts": prompts });
|
|
Some(crate::mcp::success(id, payload))
|
|
}
|
|
|
|
"prompts/get" => {
|
|
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_else(|| serde_json::json!({}));
|
|
|
|
if let Some(prompt) = self.prompts.get(name) {
|
|
match prompt.get(args, Arc::clone(&self.state)).await {
|
|
Ok(messages) => Some(crate::mcp::success(id, messages)),
|
|
Err(e) => Some(crate::mcp::error(id, -32603, &e.to_string())),
|
|
}
|
|
} else {
|
|
Some(crate::mcp::error(id, -32602, "Prompt not found"))
|
|
}
|
|
}
|
|
|
|
"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()));
|
|
|
|
let category = match name {
|
|
"clipboard" => "CLIPBOARD",
|
|
"create_entities"
|
|
| "create_relations"
|
|
| "add_observations"
|
|
| "delete_entities"
|
|
| "delete_observations"
|
|
| "delete_relations"
|
|
| "read_graph"
|
|
| "get_subgraph"
|
|
| "search_graph"
|
|
| "get_schema"
|
|
| "search_nodes"
|
|
| "open_nodes"
|
|
| "visualize_graph"
|
|
| "condense_entity"
|
|
| "merge_entities"
|
|
| "find_orphans"
|
|
| "sweep_graph_health"
|
|
| "summarize_subgraph"
|
|
| "query_graph_path"
|
|
| "resolve_stale_symbols" => "GRAPH",
|
|
"decisions" => "DECISION",
|
|
"log_code_change" => "CODE",
|
|
"log_error_fix" => "ERROR_FIX",
|
|
"tech_debt" => "TECH_DEBT",
|
|
"tasks" | "milestones" => "TASK",
|
|
"handoff_memos" => "HANDOFF_MEMO",
|
|
"manage_checkpoint" => "CHECKPOINT",
|
|
"manage_subagent_namespace" => "SUBAGENT",
|
|
"snippets" => "SNIPPET",
|
|
"search_web" => "WEB_SEARCH",
|
|
"omni_search" => "OMNI_SEARCH",
|
|
"environment" => "ENVIRONMENT",
|
|
"hypotheses" => "HYPOTHESIS",
|
|
"agent_signals" => "AGENT_SIGNAL",
|
|
"process_logs" => "PROCESS_LOG",
|
|
_ => "TOOL",
|
|
};
|
|
|
|
let activity_msg = format_tool_activity_description(name, &args);
|
|
self.state.broadcast_activity(category, &activity_msg);
|
|
|
|
let result: crate::error::Result<String> = if let Some(tool) = self.tools.get(name)
|
|
{
|
|
execute_tool_span(name, tool.as_ref(), args, self.state.clone()).await
|
|
} else {
|
|
Err(crate::error::AppError::Internal(format!(
|
|
"Unknown tool: {}",
|
|
name
|
|
)))
|
|
};
|
|
|
|
match result {
|
|
Ok(text) => {
|
|
let payload = serde_json::json!({
|
|
"content": [{"type": "text", "text": text}],
|
|
"isError": false
|
|
});
|
|
Some(crate::mcp::success(id_clone, payload))
|
|
}
|
|
Err(e) => {
|
|
tracing::error!("Tool {} failed: {}", name, e);
|
|
let payload = serde_json::json!({
|
|
"content": [{"type": "text", "text": e.to_string()}],
|
|
"isError": true
|
|
});
|
|
Some(crate::mcp::success(id_clone, payload))
|
|
}
|
|
}
|
|
}
|
|
m if m.starts_with("notifications/") => None,
|
|
"ping" => Some(crate::mcp::success(id, serde_json::json!({}))),
|
|
_ => {
|
|
if id.is_null() {
|
|
None
|
|
} else {
|
|
Some(crate::mcp::error(
|
|
id,
|
|
-32601,
|
|
&format!("Method {} not found", method),
|
|
))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) -> String {
|
|
let (action, detail) = match name {
|
|
"tasks" => {
|
|
let act = args
|
|
.get("action")
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("manage");
|
|
let title = args
|
|
.get("title")
|
|
.or_else(|| args.get("id"))
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("");
|
|
(
|
|
"Tasks",
|
|
format!("{}: {}", act, title)
|
|
.trim_end_matches(": ")
|
|
.to_string(),
|
|
)
|
|
}
|
|
"decisions" => {
|
|
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("log");
|
|
let title = args
|
|
.get("title")
|
|
.or_else(|| args.get("query"))
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("");
|
|
(
|
|
"Decisions",
|
|
format!("{}: {}", act, title)
|
|
.trim_end_matches(": ")
|
|
.to_string(),
|
|
)
|
|
}
|
|
"tech_debt" => {
|
|
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("log");
|
|
let desc = args
|
|
.get("description")
|
|
.or_else(|| args.get("id"))
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("");
|
|
(
|
|
"Tech Debt",
|
|
format!("{}: {}", act, desc)
|
|
.trim_end_matches(": ")
|
|
.to_string(),
|
|
)
|
|
}
|
|
"create_entities" => {
|
|
let names = args
|
|
.get("entities")
|
|
.and_then(|v| v.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|e| e.get("name").and_then(|n| n.as_str()))
|
|
.collect::<Vec<_>>()
|
|
.join(", ")
|
|
})
|
|
.unwrap_or_default();
|
|
("Create Entities", names)
|
|
}
|
|
"hypotheses" => {
|
|
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("manage");
|
|
let hyp = args
|
|
.get("hypothesis")
|
|
.or_else(|| args.get("query"))
|
|
.and_then(|v| v.as_str())
|
|
.unwrap_or("");
|
|
(
|
|
"Hypotheses",
|
|
format!("{}: {}", act, hyp).trim_end_matches(": ").to_string(),
|
|
)
|
|
}
|
|
"agent_signals" => {
|
|
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("signal");
|
|
let sender = args.get("sender").and_then(|v| v.as_str()).unwrap_or("");
|
|
let st = args.get("signal_type").and_then(|v| v.as_str()).unwrap_or("");
|
|
(
|
|
"Agent Signals",
|
|
format!("{}: {} [{}]", act, sender, st).trim_end_matches(": ").to_string(),
|
|
)
|
|
}
|
|
"process_logs" => {
|
|
let act = args.get("action").and_then(|v| v.as_str()).unwrap_or("logs");
|
|
let file = args.get("file_path").and_then(|v| v.as_str()).unwrap_or("");
|
|
(
|
|
"Process Logs",
|
|
format!("{}: {}", act, file).trim_end_matches(": ").to_string(),
|
|
)
|
|
}
|
|
_ => return name.to_string(),
|
|
};
|
|
if detail.is_empty() {
|
|
action.to_string()
|
|
} else {
|
|
format!("{}: {}", action, detail)
|
|
}
|
|
}
|
|
|
|
#[tracing::instrument(name = "mcp_tool_execution", skip(tool, args, state), fields(tool_name = %name))]
|
|
async fn execute_tool_span(
|
|
name: &str,
|
|
tool: &dyn McpTool,
|
|
args: Value,
|
|
state: Arc<MemoryState>,
|
|
) -> crate::error::Result<String> {
|
|
tool.execute(args, state).await
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use tempfile::tempdir;
|
|
use serde_json::json;
|
|
|
|
#[test]
|
|
fn test_format_tool_activity_description() {
|
|
let msg = format_tool_activity_description(
|
|
"tasks",
|
|
&serde_json::json!({ "action": "add", "title": "Refactor router" }),
|
|
);
|
|
assert_eq!(msg, "Tasks: add: Refactor router");
|
|
|
|
let msg2 = format_tool_activity_description(
|
|
"create_entities",
|
|
&serde_json::json!({ "entities": [{ "name": "EntityA" }, { "name": "EntityB" }] }),
|
|
);
|
|
assert_eq!(msg2, "Create Entities: EntityA, EntityB");
|
|
|
|
let msg3 = format_tool_activity_description("read_graph", &json!({}));
|
|
assert_eq!(msg3, "read_graph");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_memory_handler_tools_registration() {
|
|
let dir = tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
let handler = MemoryHandler::new(state);
|
|
|
|
// Assert some known tools are registered
|
|
assert!(handler.tools.contains_key("create_entities"));
|
|
assert!(handler.tools.contains_key("tasks"));
|
|
|
|
// Ensure we can fetch list of tools
|
|
let list_tools_req = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 1,
|
|
"method": "tools/list",
|
|
"params": {}
|
|
});
|
|
|
|
let res_list = handler.handle_request(list_tools_req).await.unwrap();
|
|
assert_eq!(res_list["jsonrpc"], "2.0");
|
|
assert_eq!(res_list["id"], 1);
|
|
assert!(res_list["result"]["tools"].as_array().unwrap().len() > 10);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_resources_and_prompts_endpoints() {
|
|
let dir = tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
let handler = MemoryHandler::new(state);
|
|
|
|
// Test resources/list
|
|
let req_list_res = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 10,
|
|
"method": "resources/list",
|
|
"params": {}
|
|
});
|
|
let res_list = handler.handle_request(req_list_res).await.unwrap();
|
|
let resources_arr = res_list["result"]["resources"].as_array().unwrap();
|
|
assert!(
|
|
resources_arr
|
|
.iter()
|
|
.any(|r| r["uri"] == "memory://tasks/active")
|
|
);
|
|
assert!(
|
|
resources_arr
|
|
.iter()
|
|
.any(|r| r["uri"] == "memory://decisions/active")
|
|
);
|
|
assert!(
|
|
resources_arr
|
|
.iter()
|
|
.any(|r| r["uri"] == "memory://activity/recent")
|
|
);
|
|
|
|
// Test resources/read
|
|
let req_read_res = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 11,
|
|
"method": "resources/read",
|
|
"params": {
|
|
"uri": "memory://tasks/active"
|
|
}
|
|
});
|
|
let res_read = handler.handle_request(req_read_res).await.unwrap();
|
|
assert_eq!(
|
|
res_read["result"]["contents"][0]["uri"],
|
|
"memory://tasks/active"
|
|
);
|
|
assert!(
|
|
res_read["result"]["contents"][0]["text"]
|
|
.as_str()
|
|
.unwrap()
|
|
.contains("[]")
|
|
); // Empty tasks
|
|
|
|
// Test prompts/list
|
|
let req_list_prompts = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 12,
|
|
"method": "prompts/list",
|
|
"params": {}
|
|
});
|
|
let res_prompts = handler.handle_request(req_list_prompts).await.unwrap();
|
|
let prompts_arr = res_prompts["result"]["prompts"].as_array().unwrap();
|
|
assert!(prompts_arr.iter().any(|p| p["name"] == "handoff_routine"));
|
|
|
|
// Test prompts/get
|
|
let req_get_prompt = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 13,
|
|
"method": "prompts/get",
|
|
"params": {
|
|
"name": "handoff_routine",
|
|
"arguments": {}
|
|
}
|
|
});
|
|
let res_get = handler.handle_request(req_get_prompt).await.unwrap();
|
|
let messages = res_get["result"]["messages"].as_array().unwrap();
|
|
assert_eq!(messages[0]["role"], "user");
|
|
assert!(
|
|
messages[0]["content"]["text"]
|
|
.as_str()
|
|
.unwrap()
|
|
.contains("standup report")
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tool_call_success_and_error_responses() {
|
|
let dir = tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
let handler = MemoryHandler::new(state);
|
|
|
|
// 1. Test successful tool call (e.g. read_graph)
|
|
let req_success = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 2,
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": "read_graph",
|
|
"arguments": {}
|
|
}
|
|
});
|
|
let res_success = handler.handle_request(req_success).await.unwrap();
|
|
assert_eq!(res_success["jsonrpc"], "2.0");
|
|
assert_eq!(res_success["id"], 2);
|
|
// A successful tool call should return a result with isError: false
|
|
assert_eq!(res_success["result"]["isError"], false);
|
|
assert!(res_success["result"]["content"].as_array().is_some());
|
|
|
|
// 2. Test tool call semantic failure (e.g. updating a task that does not exist)
|
|
let req_fail = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 3,
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": "tasks",
|
|
"arguments": {
|
|
"action": "update",
|
|
"id": "nonexistent_task_123",
|
|
"status": "completed"
|
|
}
|
|
}
|
|
});
|
|
let res_fail = handler.handle_request(req_fail).await.unwrap();
|
|
assert_eq!(res_fail["jsonrpc"], "2.0");
|
|
assert_eq!(res_fail["id"], 3);
|
|
// Semantic failures must explicitly return isError: true inside the result to halt the LLM
|
|
assert_eq!(res_fail["result"]["isError"], true);
|
|
assert!(
|
|
res_fail["result"]["content"][0]["text"]
|
|
.as_str()
|
|
.unwrap()
|
|
.contains("not found")
|
|
);
|
|
|
|
// 3. Test unknown JSON-RPC method returns JSON-RPC protocol error
|
|
let req_unknown = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 4,
|
|
"method": "unknown_method_xyz"
|
|
});
|
|
let res_unknown = handler.handle_request(req_unknown).await.unwrap();
|
|
assert!(res_unknown.get("error").is_some());
|
|
assert_eq!(res_unknown["error"]["code"], -32601);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_handle_request() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
let handler = MemoryHandler::new(state);
|
|
|
|
let req1 = serde_json::json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 5,
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": "create_entities",
|
|
"arguments": {
|
|
"entities": [{ "name": "EntityA" }, { "name": "EntityB" }]
|
|
}
|
|
}
|
|
});
|
|
let res1 = handler.handle_request(req1).await.unwrap();
|
|
assert_eq!(res1["jsonrpc"], "2.0");
|
|
assert_eq!(res1["id"], 5);
|
|
|
|
let req2 = serde_json::json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 6,
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": "read_graph",
|
|
"arguments": {}
|
|
}
|
|
});
|
|
let res2 = handler.handle_request(req2).await.unwrap();
|
|
assert_eq!(res2["jsonrpc"], "2.0");
|
|
assert_eq!(res2["id"], 6);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tools_list_schema_and_descriptions() {
|
|
let dir = tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
let handler = MemoryHandler::new(state);
|
|
|
|
let list_req = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 100,
|
|
"method": "tools/list",
|
|
"params": {}
|
|
});
|
|
|
|
let res = handler.handle_request(list_req).await.unwrap();
|
|
assert_eq!(res["jsonrpc"], "2.0");
|
|
assert_eq!(res["id"], 100);
|
|
|
|
let tools = res["result"]["tools"]
|
|
.as_array()
|
|
.expect("tools array expected");
|
|
assert!(
|
|
tools.len() >= 30,
|
|
"Expected at least 30 registered MCP tools, found {}",
|
|
tools.len()
|
|
);
|
|
|
|
for tool in tools {
|
|
let name = tool["name"].as_str().expect("tool name must be string");
|
|
assert!(!name.is_empty(), "Tool name cannot be empty");
|
|
|
|
let desc = tool["description"]
|
|
.as_str()
|
|
.expect("tool description must be string");
|
|
assert!(!desc.is_empty(), "Tool '{}' missing description", name);
|
|
|
|
let schema = tool["inputSchema"]
|
|
.as_object()
|
|
.expect("inputSchema must be object");
|
|
assert!(
|
|
schema.contains_key("properties") || schema.contains_key("type"),
|
|
"Tool '{}' schema must contain type/properties metadata",
|
|
name
|
|
);
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_all_tools_llm_request_ok_and_invalid() {
|
|
let dir = tempdir().unwrap();
|
|
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
|
|
let handler = MemoryHandler::new(state);
|
|
|
|
// 1. Verify invalid LLM requests fail gracefully with isError: true or structured response
|
|
for (tool_name, _tool) in &handler.tools {
|
|
let invalid_req = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 200,
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": tool_name,
|
|
"arguments": {
|
|
"invalid_field_unlikely_to_match": "bogus_value",
|
|
"action": "non_existent_action_xyz_999"
|
|
}
|
|
}
|
|
});
|
|
|
|
let res = handler
|
|
.handle_request(invalid_req)
|
|
.await
|
|
.expect("handler must respond");
|
|
assert_eq!(res["jsonrpc"], "2.0");
|
|
assert!(res["result"]["content"].is_array());
|
|
}
|
|
|
|
// 2. Verify valid LLM requests return isError: false
|
|
let sample_valid_calls = vec![
|
|
("read_graph", json!({})),
|
|
("search_nodes", json!({"query": "test"})),
|
|
("tasks", json!({"action": "list"})),
|
|
("decisions", json!({"action": "query"})),
|
|
("tech_debt", json!({"action": "list"})),
|
|
("snippets", json!({"action": "search", "query": "test"})),
|
|
("clipboard", json!({"action": "read"})),
|
|
("environment", json!({"action": "read_fingerprint"})),
|
|
("omni_search", json!({"query": "test"})),
|
|
("get_project_health", json!({})),
|
|
("find_orphans", json!({})),
|
|
("sweep_graph_health", json!({})),
|
|
("get_next_actionable_tasks", json!({})),
|
|
("get_preflight_context", json!({})),
|
|
(
|
|
"auto_session_checkpoint",
|
|
json!({"session_summary": "Test summary"}),
|
|
),
|
|
];
|
|
|
|
for (name, args) in sample_valid_calls {
|
|
let valid_req = json!({
|
|
"jsonrpc": "2.0",
|
|
"id": 201,
|
|
"method": "tools/call",
|
|
"params": {
|
|
"name": name,
|
|
"arguments": args
|
|
}
|
|
});
|
|
|
|
let res = handler
|
|
.handle_request(valid_req)
|
|
.await
|
|
.expect("handler must respond");
|
|
assert_eq!(res["jsonrpc"], "2.0");
|
|
assert_eq!(
|
|
res["result"]["isError"], false,
|
|
"Tool '{}' failed valid execution. Result: {:?}",
|
|
name, res
|
|
);
|
|
}
|
|
}
|
|
}
|