perf: fix memory leaks, optimize tracing logs, fix networking eof fragmentation bug in nvim-core and stub, and deduplicate json-rpc handlers
This commit is contained in:
Riz Ashraf committed
1 parent
ce88bef45e
commit
f76d8ccce8
23 files changed
+1373
-803
No files matched your search
@@ -14,3 +14,197 @@ pub trait McpTool: Send + Sync {
|
||||
/// Execute the tool with the given arguments
|
||||
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String>;
|
||||
}
|
||||
|
||||
pub struct MemoryHandler {
|
||||
pub state: Arc<MemoryState>,
|
||||
pub tools: std::collections::HashMap<String, Box<dyn McpTool>>,
|
||||
}
|
||||
|
||||
impl MemoryHandler {
|
||||
pub fn new(state: Arc<MemoryState>) -> Self {
|
||||
let mut tools: std::collections::HashMap<String, Box<dyn McpTool>> =
|
||||
std::collections::HashMap::new();
|
||||
|
||||
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::AddTaskHandler);
|
||||
register!(tasks::DeleteTaskHandler);
|
||||
register!(tasks::UpdateTaskStatusHandler);
|
||||
register!(tasks::ListActiveTasksHandler);
|
||||
register!(tasks::SetAcceptanceCriteriaHandler);
|
||||
register!(tasks::VerifyAcceptanceCriteriaHandler);
|
||||
register!(tasks::AddMilestoneHandler);
|
||||
register!(tasks::UpdateMilestoneHandler);
|
||||
register!(tasks::ListMilestonesHandler);
|
||||
|
||||
register!(notes::AddStickyNoteHandler);
|
||||
register!(notes::ReadStickyNotesHandler);
|
||||
register!(notes::DeleteStickyNoteHandler);
|
||||
register!(notes::ClearStickyNotesHandler);
|
||||
register!(notes::LeaveHandoffMemoHandler);
|
||||
register!(notes::ReadHandoffMemosHandler);
|
||||
register!(notes::ClearHandoffMemosHandler);
|
||||
register!(notes::AddSessionSummaryHandler);
|
||||
register!(notes::GenerateStandupReportHandler);
|
||||
|
||||
register!(meta::LogDecisionHandler);
|
||||
register!(meta::QueryDecisionsHandler);
|
||||
register!(meta::LogErrorFixHandler);
|
||||
register!(meta::SearchErrorFixesHandler);
|
||||
register!(meta::LogCodeChangeHandler);
|
||||
register!(meta::QueryRecentChangesHandler);
|
||||
register!(meta::LearnPreferenceHandler);
|
||||
register!(meta::ReadPreferencesHandler);
|
||||
register!(meta::LogTechDebtHandler);
|
||||
register!(meta::ResolveTechDebtHandler);
|
||||
register!(meta::ListTechDebtHandler);
|
||||
register!(meta::OmniSearchHandler);
|
||||
register!(meta::GetProjectHealthHandler);
|
||||
|
||||
register!(env::UpdateEnvFingerprintHandler);
|
||||
register!(env::ReadEnvFingerprintHandler);
|
||||
register!(env::LogEnvRequirementHandler);
|
||||
register!(env::RegisterEnvironmentHandler);
|
||||
register!(env::GetEnvironmentDetailsHandler);
|
||||
|
||||
register!(workspaces::PinFileHandler);
|
||||
register!(workspaces::UnpinFileHandler);
|
||||
register!(workspaces::ListPinnedFilesHandler);
|
||||
register!(workspaces::StoreSnippetHandler);
|
||||
register!(workspaces::SearchSnippetsHandler);
|
||||
register!(workspaces::DeleteSnippetHandler);
|
||||
register!(workspaces::SaveContextWorkspaceHandler);
|
||||
register!(workspaces::LoadContextWorkspaceHandler);
|
||||
register!(workspaces::ListContextWorkspacesHandler);
|
||||
register!(workspaces::AddPrChecklistItemHandler);
|
||||
register!(workspaces::GetPrChecklistHandler);
|
||||
register!(workspaces::ClearPrChecklistHandler);
|
||||
|
||||
Self { state, tools }
|
||||
}
|
||||
|
||||
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!({})
|
||||
},
|
||||
"_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()
|
||||
.build(),
|
||||
)
|
||||
.with_server_info(rmcp::model::Implementation::new(
|
||||
"gemini-mcp-memory",
|
||||
"3.0.0",
|
||||
));
|
||||
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 }),
|
||||
))
|
||||
}
|
||||
"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> = if let Some(tool) = self.tools.get(name) {
|
||||
tool.execute(args, self.state.clone()).await
|
||||
} else {
|
||||
Err(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}],
|
||||
"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),
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user