codex_memory/mcp_server/
mod.rs

1//! Minimal MCP server implementation
2pub mod handlers;
3pub mod tools;
4pub mod transport;
5
6// Re-export for tests
7pub use handlers::MCPHandlers;
8
9use crate::config::Config;
10use crate::error::Result;
11use crate::storage::Storage;
12use std::sync::Arc;
13use tokio::io::{AsyncReadExt, AsyncWriteExt};
14use tracing::{error, info};
15
16/// Simple MCP server
17pub struct MCPServer {
18    _config: Config,
19    handlers: Arc<MCPHandlers>,
20}
21
22impl MCPServer {
23    /// Create a new MCP server
24    pub fn new(config: Config, storage: Arc<Storage>) -> Self {
25        let handlers = Arc::new(MCPHandlers::new(storage));
26        Self {
27            _config: config,
28            handlers,
29        }
30    }
31
32    /// Run in stdio mode for Claude Desktop
33    pub async fn run_stdio(&self) -> Result<()> {
34        info!("MCP server running in stdio mode");
35
36        let stdin = tokio::io::stdin();
37        let stdout = tokio::io::stdout();
38        let mut stdin = stdin;
39        let mut stdout = stdout;
40
41        let mut buffer = String::new();
42        let mut temp_buf = [0u8; 8192]; // 8KB buffer for reading chunks
43
44        loop {
45            // Read chunk from stdin
46            match stdin.read(&mut temp_buf).await {
47                Ok(0) => {
48                    info!("Received EOF, shutting down MCP server");
49                    break; // EOF
50                }
51                Ok(n) => {
52                    // Convert bytes to string and append to buffer
53                    let chunk = String::from_utf8_lossy(&temp_buf[..n]);
54                    info!(
55                        "Read {} bytes: {:?}",
56                        n,
57                        &chunk[..std::cmp::min(100, chunk.len())]
58                    );
59                    buffer.push_str(&chunk);
60
61                    // Process complete JSON objects from buffer
62                    while let Some(json_end) = self.find_complete_json(&buffer) {
63                        let json_str = buffer[..json_end].trim().to_string();
64                        buffer.drain(..json_end);
65
66                        if !json_str.is_empty() {
67                            info!("Processing JSON request ({} chars)", json_str.len());
68                            let response = self.handle_request(&json_str).await;
69                            stdout.write_all(response.as_bytes()).await?;
70                            stdout.write_all(b"\n").await?;
71                            stdout.flush().await?;
72                        }
73                    }
74                }
75                Err(e) => {
76                    error!("Error reading input: {}", e);
77                    break;
78                }
79            }
80        }
81
82        info!("MCP server shutting down gracefully");
83        // The pool will be dropped automatically when Storage is dropped
84        Ok(())
85    }
86
87    /// Find the end of a complete JSON object in the buffer
88    fn find_complete_json(&self, buffer: &str) -> Option<usize> {
89        let mut brace_count = 0;
90        let mut in_string = false;
91        let mut escape_next = false;
92        let mut start_found = false;
93
94        for (i, ch) in buffer.char_indices() {
95            if escape_next {
96                escape_next = false;
97                continue;
98            }
99
100            match ch {
101                '\\' if in_string => escape_next = true,
102                '"' => in_string = !in_string,
103                '{' if !in_string => {
104                    brace_count += 1;
105                    start_found = true;
106                }
107                '}' if !in_string => {
108                    brace_count -= 1;
109                    if start_found && brace_count == 0 {
110                        return Some(i + 1);
111                    }
112                }
113                _ => {}
114            }
115        }
116
117        None
118    }
119
120    async fn handle_request(&self, request: &str) -> String {
121        let request: serde_json::Value = match serde_json::from_str(request) {
122            Ok(v) => v,
123            Err(e) => {
124                return format!(
125                    r#"{{"jsonrpc":"2.0","error":{{"code":-32700,"message":"Parse error: {}"}}}}"#,
126                    e
127                );
128            }
129        };
130
131        let method = request["method"].as_str().unwrap_or("");
132        let params = request.get("params").cloned().unwrap_or_default();
133        let id = request.get("id").cloned();
134
135        let result = match method {
136            "initialize" => Ok(serde_json::json!({
137                "protocolVersion": "2024-11-05",
138                "capabilities": {
139                    "tools": {
140                        "list": tools::MCPTools::get_tools_list()
141                    }
142                },
143                "serverInfo": {
144                    "name": "codex-store",
145                    "version": env!("CARGO_PKG_VERSION")
146                }
147            })),
148            "tools/list" => Ok(serde_json::json!({
149                "tools": tools::MCPTools::get_tools_list()
150            })),
151            "tools/call" => {
152                let tool_name = params["name"].as_str().unwrap_or("");
153                let tool_params = params.get("arguments").cloned().unwrap_or_default();
154                self.handlers.handle_tool_call(tool_name, tool_params).await
155            }
156            _ => Err(crate::error::Error::Other(format!(
157                "Unknown method: {}",
158                method
159            ))),
160        };
161
162        match result {
163            Ok(value) => {
164                if let Some(id) = id {
165                    format!(r#"{{"jsonrpc":"2.0","id":{},"result":{}}}"#, id, value)
166                } else {
167                    format!(r#"{{"jsonrpc":"2.0","result":{}}}"#, value)
168                }
169            }
170            Err(e) => {
171                if let Some(id) = id {
172                    format!(
173                        r#"{{"jsonrpc":"2.0","id":{},"error":{{"code":-32000,"message":"{}"}}}}"#,
174                        id, e
175                    )
176                } else {
177                    format!(
178                        r#"{{"jsonrpc":"2.0","error":{{"code":-32000,"message":"{}"}}}}"#,
179                        e
180                    )
181                }
182            }
183        }
184    }
185}
codex_memory/mcp_server/mod.rs

codex_memory/mcp_server/
mod.rs
