1use serde_json::{Value, json};
2
3use crate::{
4 tool::{Tool, ToolInfo},
5 types::{JsonRpcError, JsonRpcRequest, JsonRpcResponse, McpError, ToolError},
6};
7
8pub struct McpServer<T: Tool> {
9 tool: T,
10}
11
12impl<T: Tool> McpServer<T> {
13 pub fn new(tool: T) -> Self {
14 Self { tool }
15 }
16
17 pub fn handle_request(&self, request: JsonRpcRequest) -> JsonRpcResponse {
18 let request_id = request.id.clone().unwrap_or(serde_json::Value::Null);
19 match self.process_request(request) {
20 Ok(result) => JsonRpcResponse {
21 jsonrpc: "2.0".to_string(),
22 result: Some(result),
23 error: None,
24 id: request_id,
25 },
26 Err(e) => JsonRpcResponse {
27 jsonrpc: "2.0".to_string(),
28 result: None,
29 error: Some(self.error_to_jsonrpc(e)),
30 id: request_id,
31 },
32 }
33 }
34
35 fn process_request(&self, request: JsonRpcRequest) -> Result<Value, McpError> {
36 match request.method.as_str() {
37 "initialize" => self.handle_initialize(request.params),
38 "tools/list" => self.handle_tools_list(),
39 "tools/call" => self.handle_tools_call(request.params),
40 method => Err(McpError::MethodNotFound(method.to_string())),
41 }
42 }
43
44 fn handle_initialize(&self, _params: Option<Value>) -> Result<Value, McpError> {
45 Ok(json!({
46 "protocolVersion": "2025-03-26",
47 "serverInfo": {
48 "name": self.tool.server_name(),
49 "version": self.tool.server_version()
50 },
51 "capabilities": self.tool.capabilities()
52 }))
53 }
54
55 fn handle_tools_list(&self) -> Result<Value, McpError> {
56 let tool_info = ToolInfo::from(&self.tool);
57 Ok(json!({
58 "tools": [{
59 "name": tool_info.name,
60 "description": tool_info.description,
61 "inputSchema": tool_info.input_schema
62 }]
63 }))
64 }
65
66 fn handle_tools_call(&self, params: Option<Value>) -> Result<Value, McpError> {
67 let params = params
68 .ok_or_else(|| McpError::InvalidRequest("Missing params for tools/call".to_string()))?;
69
70 let tool_name = params["name"]
71 .as_str()
72 .ok_or_else(|| McpError::InvalidRequest("Missing tool name".to_string()))?;
73
74 if tool_name != self.tool.name() {
75 return Err(McpError::ToolError(ToolError::InvalidArguments(format!(
76 "Unknown tool: {tool_name}"
77 ))));
78 }
79
80 let arguments = params.get("arguments").cloned().unwrap_or(json!({}));
81
82 match self.tool.call(&arguments) {
83 Ok(result) => Ok(json!(result)),
84 Err(e) => Err(McpError::ToolError(e)),
85 }
86 }
87
88 fn error_to_jsonrpc(&self, error: McpError) -> JsonRpcError {
89 match error {
90 McpError::InvalidRequest(msg) => JsonRpcError::invalid_request(msg),
91 McpError::MethodNotFound(method) => JsonRpcError::method_not_found(method),
92 McpError::ToolError(e) => JsonRpcError::internal_error(e.to_string()),
93 McpError::JsonError(e) => JsonRpcError::internal_error(e.to_string()),
94 }
95 }
96}
97
98#[cfg(test)]
99mod tests {
100 use serde_json::json;
101
102 use super::*;
103 use crate::{
104 tool::Tool,
105 types::{ToolError, ToolResult},
106 };
107
108 #[derive(Clone)]
109 struct TestTool;
110
111 impl Tool for TestTool {
112 fn name(&self) -> &'static str {
113 "test_tool"
114 }
115
116 fn description(&self) -> &'static str {
117 "A test tool"
118 }
119
120 fn input_schema(&self) -> Value {
121 json!({
122 "type": "object",
123 "properties": {
124 "input": {"type": "string"}
125 },
126 "required": ["input"]
127 })
128 }
129
130 fn call(&self, args: &Value) -> Result<ToolResult, ToolError> {
131 if let Some(input) = args.get("input").and_then(|v| v.as_str()) {
132 Ok(ToolResult::text(format!("Processed: {input}")))
133 } else {
134 Err(ToolError::InvalidArguments("Missing input".to_string()))
135 }
136 }
137 }
138
139 #[test]
140 fn test_initialize() {
141 let server = McpServer::new(TestTool);
142 let request = JsonRpcRequest {
143 jsonrpc: "2.0".to_string(),
144 method: "initialize".to_string(),
145 params: None,
146 id: Some(json!(1)),
147 };
148
149 let response = server.handle_request(request);
150 assert!(response.error.is_none());
151
152 let result = response.result.unwrap();
153 assert_eq!(result["protocolVersion"], "2025-03-26");
154 assert_eq!(result["serverInfo"]["name"], "ftl-test_tool");
155 }
156
157 #[test]
158 fn test_tools_list() {
159 let server = McpServer::new(TestTool);
160 let request = JsonRpcRequest {
161 jsonrpc: "2.0".to_string(),
162 method: "tools/list".to_string(),
163 params: None,
164 id: Some(json!(1)),
165 };
166
167 let response = server.handle_request(request);
168 assert!(response.error.is_none());
169
170 let result = response.result.unwrap();
171 let tools = result["tools"].as_array().unwrap();
172 assert_eq!(tools.len(), 1);
173 assert_eq!(tools[0]["name"], "test_tool");
174 assert_eq!(tools[0]["description"], "A test tool");
175 }
176
177 #[test]
178 fn test_tools_call_success() {
179 let server = McpServer::new(TestTool);
180 let request = JsonRpcRequest {
181 jsonrpc: "2.0".to_string(),
182 method: "tools/call".to_string(),
183 params: Some(json!({
184 "name": "test_tool",
185 "arguments": {
186 "input": "hello"
187 }
188 })),
189 id: Some(json!(1)),
190 };
191
192 let response = server.handle_request(request);
193 assert!(response.error.is_none());
194
195 let result = response.result.unwrap();
196 assert_eq!(result["content"][0]["text"], "Processed: hello");
197 }
198
199 #[test]
200 fn test_tools_call_missing_argument() {
201 let server = McpServer::new(TestTool);
202 let request = JsonRpcRequest {
203 jsonrpc: "2.0".to_string(),
204 method: "tools/call".to_string(),
205 params: Some(json!({
206 "name": "test_tool",
207 "arguments": {}
208 })),
209 id: Some(json!(1)),
210 };
211
212 let response = server.handle_request(request);
213 assert!(response.error.is_some());
214 assert_eq!(response.error.unwrap().code, -32603);
215 }
216
217 #[test]
218 fn test_unknown_method() {
219 let server = McpServer::new(TestTool);
220 let request = JsonRpcRequest {
221 jsonrpc: "2.0".to_string(),
222 method: "unknown/method".to_string(),
223 params: None,
224 id: Some(json!(1)),
225 };
226
227 let response = server.handle_request(request);
228 assert!(response.error.is_some());
229 assert_eq!(response.error.unwrap().code, -32601);
230 }
231
232 #[test]
233 fn test_missing_params() {
234 let server = McpServer::new(TestTool);
235 let request = JsonRpcRequest {
236 jsonrpc: "2.0".to_string(),
237 method: "tools/call".to_string(),
238 params: None,
239 id: Some(json!(1)),
240 };
241
242 let response = server.handle_request(request);
243 assert!(response.error.is_some());
244 assert_eq!(response.error.unwrap().code, -32600);
245 }
246
247 #[test]
248 fn test_wrong_tool_name() {
249 let server = McpServer::new(TestTool);
250 let request = JsonRpcRequest {
251 jsonrpc: "2.0".to_string(),
252 method: "tools/call".to_string(),
253 params: Some(json!({
254 "name": "wrong_tool",
255 "arguments": {"input": "test"}
256 })),
257 id: Some(json!(1)),
258 };
259
260 let response = server.handle_request(request);
261 assert!(response.error.is_some());
262 }
263}