agentroot_mcp/
server.rs

1//! MCP server implementation
2
3use crate::protocol::*;
4use crate::tools;
5use agentroot_core::Database;
6use anyhow::Result;
7use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
8
9pub struct McpServer<'a> {
10    db: &'a Database,
11}
12
13impl<'a> McpServer<'a> {
14    pub fn new(db: &'a Database) -> Self {
15        Self { db }
16    }
17
18    pub async fn run(&self) -> Result<()> {
19        let stdin = tokio::io::stdin();
20        let stdout = tokio::io::stdout();
21
22        let mut reader = BufReader::new(stdin);
23        let mut writer = BufWriter::new(stdout);
24        let mut line = String::new();
25
26        loop {
27            line.clear();
28            let bytes_read = reader.read_line(&mut line).await?;
29
30            if bytes_read == 0 {
31                break;
32            }
33
34            let trimmed = line.trim();
35            if trimmed.is_empty() {
36                continue;
37            }
38
39            let request: JsonRpcRequest = match serde_json::from_str(trimmed) {
40                Ok(r) => r,
41                Err(e) => {
42                    let response =
43                        JsonRpcResponse::error(None, -32700, &format!("Parse error: {}", e));
44                    self.write_response(&mut writer, &response).await?;
45                    continue;
46                }
47            };
48
49            let response = self.handle_request(&request).await;
50            self.write_response(&mut writer, &response).await?;
51        }
52
53        Ok(())
54    }
55
56    async fn write_response<W: AsyncWriteExt + Unpin>(
57        &self,
58        writer: &mut W,
59        response: &JsonRpcResponse,
60    ) -> Result<()> {
61        let json = serde_json::to_string(response)?;
62        writer.write_all(json.as_bytes()).await?;
63        writer.write_all(b"\n").await?;
64        writer.flush().await?;
65        Ok(())
66    }
67
68    async fn handle_request(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
69        match request.method.as_str() {
70            "initialize" => self.handle_initialize(request),
71            "tools/list" => self.handle_tools_list(request),
72            "tools/call" => self.handle_tools_call(request).await,
73            "resources/list" => self.handle_resources_list(request),
74            "prompts/list" => self.handle_prompts_list(request),
75            _ => JsonRpcResponse::error(
76                request.id.clone(),
77                -32601,
78                &format!("Method not found: {}", request.method),
79            ),
80        }
81    }
82
83    fn handle_initialize(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
84        let result = serde_json::json!({
85            "protocolVersion": "2024-11-05",
86            "capabilities": {
87                "tools": {},
88                "resources": { "subscribe": false },
89                "prompts": {}
90            },
91            "serverInfo": {
92                "name": "agentroot",
93                "version": env!("CARGO_PKG_VERSION")
94            }
95        });
96        JsonRpcResponse::success(request.id.clone(), result)
97    }
98
99    fn handle_tools_list(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
100        let tools = vec![
101            tools::search_tool_definition(),
102            tools::vsearch_tool_definition(),
103            tools::query_tool_definition(),
104            tools::smart_search_tool_definition(),
105            tools::get_tool_definition(),
106            tools::multi_get_tool_definition(),
107            tools::status_tool_definition(),
108            tools::collection_add_tool_definition(),
109            tools::collection_remove_tool_definition(),
110            tools::collection_update_tool_definition(),
111            tools::metadata_add_tool_definition(),
112            tools::metadata_get_tool_definition(),
113            tools::metadata_query_tool_definition(),
114        ];
115
116        JsonRpcResponse::success(request.id.clone(), serde_json::json!({ "tools": tools }))
117    }
118
119    async fn handle_tools_call(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
120        let name = request
121            .params
122            .get("name")
123            .and_then(|v| v.as_str())
124            .unwrap_or("");
125
126        let arguments = request
127            .params
128            .get("arguments")
129            .cloned()
130            .unwrap_or(serde_json::json!({}));
131
132        let result = match name {
133            "search" => tools::handle_search(self.db, arguments).await,
134            "vsearch" => tools::handle_vsearch(self.db, arguments).await,
135            "query" => tools::handle_query(self.db, arguments).await,
136            "smart_search" => tools::handle_smart_search(self.db, arguments).await,
137            "get" => tools::handle_get(self.db, arguments).await,
138            "multi_get" => tools::handle_multi_get(self.db, arguments).await,
139            "status" => tools::handle_status(self.db).await,
140            "collection_add" => tools::handle_collection_add(self.db, arguments).await,
141            "collection_remove" => tools::handle_collection_remove(self.db, arguments).await,
142            "collection_update" => tools::handle_collection_update(self.db, arguments).await,
143            "metadata_add" => tools::handle_metadata_add(self.db, arguments).await,
144            "metadata_get" => tools::handle_metadata_get(self.db, arguments).await,
145            "metadata_query" => tools::handle_metadata_query(self.db, arguments).await,
146            _ => Err(anyhow::anyhow!("Unknown tool: {}", name)),
147        };
148
149        match result {
150            Ok(tool_result) => JsonRpcResponse::success(
151                request.id.clone(),
152                serde_json::to_value(tool_result).unwrap(),
153            ),
154            Err(e) => {
155                let error_result = ToolResult {
156                    content: vec![Content::Text {
157                        text: format!("Error: {}", e),
158                    }],
159                    structured_content: None,
160                    is_error: Some(true),
161                };
162                JsonRpcResponse::success(
163                    request.id.clone(),
164                    serde_json::to_value(error_result).unwrap(),
165                )
166            }
167        }
168    }
169
170    fn handle_resources_list(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
171        JsonRpcResponse::success(request.id.clone(), serde_json::json!({ "resources": [] }))
172    }
173
174    fn handle_prompts_list(&self, request: &JsonRpcRequest) -> JsonRpcResponse {
175        let prompts = vec![serde_json::json!({
176            "name": "query",
177            "title": "Agentroot Query Guide",
178            "description": "How to effectively search your knowledge base"
179        })];
180        JsonRpcResponse::success(
181            request.id.clone(),
182            serde_json::json!({ "prompts": prompts }),
183        )
184    }
185}
186
187pub async fn start_server(db: &Database) -> Result<()> {
188    let server = McpServer::new(db);
189    server.run().await
190}
agentroot_mcp/server.rs

agentroot_mcp/
server.rs
