ddb_core/mcp/
server.rs

1//! DDB MCP Server implementation using JSON-RPC over stdio
2
3use crate::config::{Config, TableCatalog};
4use anyhow::Result;
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use std::io::{self, BufRead, Write};
8use std::sync::Arc;
9use tokio::sync::RwLock;
10use tracing::{info, debug};
11
12use super::tools::ToolHandler;
13use super::resources::ResourceHandler;
14use super::prompts::PromptHandler;
15
16#[derive(Debug, Serialize, Deserialize)]
17struct JsonRpcRequest {
18    jsonrpc: String,
19    id: Option<Value>,
20    method: String,
21    params: Option<Value>,
22}
23
24#[derive(Debug, Serialize, Deserialize)]
25struct JsonRpcResponse {
26    jsonrpc: String,
27    id: Option<Value>,
28    #[serde(skip_serializing_if = "Option::is_none")]
29    result: Option<Value>,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    error: Option<JsonRpcError>,
32}
33
34#[derive(Debug, Serialize, Deserialize)]
35struct JsonRpcError {
36    code: i32,
37    message: String,
38    #[serde(skip_serializing_if = "Option::is_none")]
39    data: Option<Value>,
40}
41
42/// DDB MCP Server
43pub struct DdbMcpServer {
44    _config: Arc<RwLock<Config>>,
45    _catalog: Arc<RwLock<TableCatalog>>,
46    tool_handler: ToolHandler,
47    resource_handler: ResourceHandler,
48    prompt_handler: PromptHandler,
49}
50
51impl DdbMcpServer {
52    /// Create a new DDB MCP server
53    pub fn new() -> Result<Self> {
54        // Load configuration
55        let config = Config::load().unwrap_or_else(|e| {
56            eprintln!("Warning: Failed to load config: {}", e);
57            eprintln!("Using default configuration");
58            Config::default()
59        });
60
61        // Load table catalog
62        let catalog = TableCatalog::load_from_config(&config).unwrap_or_else(|e| {
63            eprintln!("Warning: Failed to load table catalog: {}", e);
64            TableCatalog::new()
65        });
66
67        let config = Arc::new(RwLock::new(config));
68        let catalog = Arc::new(RwLock::new(catalog));
69
70        Ok(Self {
71            _config: config.clone(),
72            _catalog: catalog.clone(),
73            tool_handler: ToolHandler::new(config.clone(), catalog.clone()),
74            resource_handler: ResourceHandler::new(config.clone(), catalog.clone()),
75            prompt_handler: PromptHandler::new(),
76        })
77    }
78
79    /// Run the MCP server (processes stdio)
80    pub async fn run(self) -> Result<()> {
81        info!("Starting DDB MCP Server v{}", crate::VERSION);
82        info!("Reading from stdin, writing to stdout");
83
84        let stdin = io::stdin();
85        let mut stdout = io::stdout();
86
87        for line in stdin.lock().lines() {
88            let line = line?;
89            if line.trim().is_empty() {
90                continue;
91            }
92
93            debug!("Received: {}", line);
94
95            // Parse JSON-RPC request
96            let request: Result<JsonRpcRequest, _> = serde_json::from_str(&line);
97
98            let response = match request {
99                Ok(req) => self.handle_request(req).await,
100                Err(e) => JsonRpcResponse {
101                    jsonrpc: "2.0".to_string(),
102                    id: None,
103                    result: None,
104                    error: Some(JsonRpcError {
105                        code: -32700,
106                        message: format!("Parse error: {}", e),
107                        data: None,
108                    }),
109                },
110            };
111
112            // Send response
113            let response_str = serde_json::to_string(&response)?;
114            debug!("Sending: {}", response_str);
115            writeln!(stdout, "{}", response_str)?;
116            stdout.flush()?;
117        }
118
119        Ok(())
120    }
121
122    async fn handle_request(&self, request: JsonRpcRequest) -> JsonRpcResponse {
123        let result = match request.method.as_str() {
124            "initialize" => self.handle_initialize(request.params).await,
125            "tools/list" => self.handle_tools_list().await,
126            "tools/call" => self.handle_tools_call(request.params).await,
127            "resources/list" => self.handle_resources_list().await,
128            "resources/read" => self.handle_resources_read(request.params).await,
129            "prompts/list" => self.handle_prompts_list().await,
130            "prompts/get" => self.handle_prompts_get(request.params).await,
131            _ => Err(format!("Method not found: {}", request.method)),
132        };
133
134        match result {
135            Ok(value) => JsonRpcResponse {
136                jsonrpc: "2.0".to_string(),
137                id: request.id,
138                result: Some(value),
139                error: None,
140            },
141            Err(msg) => JsonRpcResponse {
142                jsonrpc: "2.0".to_string(),
143                id: request.id,
144                result: None,
145                error: Some(JsonRpcError {
146                    code: -32603,
147                    message: msg,
148                    data: None,
149                }),
150            },
151        }
152    }
153
154    async fn handle_initialize(&self, _params: Option<Value>) -> Result<Value, String> {
155        Ok(serde_json::json!({
156            "protocolVersion": "2024-11-05",
157            "capabilities": {
158                "tools": {},
159                "resources": {},
160                "prompts": {}
161            },
162            "serverInfo": {
163                "name": "ddb-mcp-server",
164                "version": crate::VERSION
165            }
166        }))
167    }
168
169    async fn handle_tools_list(&self) -> Result<Value, String> {
170        let tools = self.tool_handler.list_tools();
171        Ok(serde_json::json!({
172            "tools": tools
173        }))
174    }
175
176    async fn handle_tools_call(&self, params: Option<Value>) -> Result<Value, String> {
177        let params = params.ok_or("Missing parameters")?;
178        let name = params["name"].as_str().ok_or("Missing tool name")?;
179        let arguments = params.get("arguments").cloned();
180
181        self.tool_handler
182            .call_tool(name, arguments)
183            .await
184            .map_err(|e| e.to_string())
185    }
186
187    async fn handle_resources_list(&self) -> Result<Value, String> {
188        let resources = self.resource_handler
189            .list_resources()
190            .await
191            .map_err(|e| e.to_string())?;
192
193        Ok(serde_json::json!({
194            "resources": resources
195        }))
196    }
197
198    async fn handle_resources_read(&self, params: Option<Value>) -> Result<Value, String> {
199        let params = params.ok_or("Missing parameters")?;
200        let uri = params["uri"].as_str().ok_or("Missing resource URI")?;
201
202        let content = self.resource_handler
203            .read_resource(uri)
204            .await
205            .map_err(|e| e.to_string())?;
206
207        Ok(serde_json::json!({
208            "contents": [
209                {
210                    "uri": uri,
211                    "mimeType": "application/json",
212                    "text": serde_json::to_string_pretty(&content).unwrap()
213                }
214            ]
215        }))
216    }
217
218    async fn handle_prompts_list(&self) -> Result<Value, String> {
219        let prompts = self.prompt_handler.list_prompts();
220        Ok(serde_json::json!({
221            "prompts": prompts
222        }))
223    }
224
225    async fn handle_prompts_get(&self, params: Option<Value>) -> Result<Value, String> {
226        let params = params.ok_or("Missing parameters")?;
227        let name = params["name"].as_str().ok_or("Missing prompt name")?;
228        let arguments = params.get("arguments").cloned();
229
230        self.prompt_handler
231            .get_prompt(name, arguments)
232            .await
233            .map_err(|e| e.to_string())
234    }
235}
236
237impl Default for DdbMcpServer {
238    fn default() -> Self {
239        Self::new().expect("Failed to create DDB MCP Server")
240    }
241}
ddb_core/mcp/server.rs

ddb_core/mcp/
server.rs
