Skip to main content

somatize_mcp/
server.rs

1//! MCP stdio server: reads JSON-RPC from stdin, writes to stdout.
2
3use crate::context::SomaContext;
4use crate::protocol::*;
5use crate::tools;
6use serde_json::json;
7use std::io::{self, BufRead, Write};
8
9/// Run the MCP server on stdin/stdout.
10pub fn run_stdio(mut ctx: SomaContext) {
11    let stdin = io::stdin();
12    let mut stdout = io::stdout();
13
14    for line in stdin.lock().lines() {
15        let line = match line {
16            Ok(l) => l,
17            Err(_) => break,
18        };
19
20        if line.trim().is_empty() {
21            continue;
22        }
23
24        let request: JsonRpcRequest = match serde_json::from_str(&line) {
25            Ok(r) => r,
26            Err(e) => {
27                let resp = JsonRpcResponse::error(
28                    serde_json::Value::Null,
29                    -32700,
30                    format!("Parse error: {e}"),
31                );
32                write_response(&mut stdout, &resp);
33                continue;
34            }
35        };
36
37        let response = handle_request(&mut ctx, &request);
38        write_response(&mut stdout, &response);
39    }
40}
41
42fn handle_request(ctx: &mut SomaContext, req: &JsonRpcRequest) -> JsonRpcResponse {
43    match req.method.as_str() {
44        "initialize" => {
45            let result = InitializeResult {
46                protocol_version: "2024-11-05".to_string(),
47                capabilities: ServerCapabilities {
48                    tools: ToolsCapability {
49                        list_changed: false,
50                    },
51                },
52                server_info: ServerInfo {
53                    name: "soma-mcp".to_string(),
54                    version: env!("CARGO_PKG_VERSION").to_string(),
55                },
56            };
57            JsonRpcResponse::success(req.id.clone(), serde_json::to_value(result).unwrap())
58        }
59
60        "notifications/initialized" => {
61            // Client acknowledgment, no response needed for notifications
62            // but we return success anyway for non-notification calls
63            JsonRpcResponse::success(req.id.clone(), json!({}))
64        }
65
66        "tools/list" => {
67            let tool_list = tools::all_tools();
68            JsonRpcResponse::success(req.id.clone(), json!({ "tools": tool_list }))
69        }
70
71        "tools/call" => {
72            let tool_name = req
73                .params
74                .get("name")
75                .and_then(|v| v.as_str())
76                .unwrap_or("");
77            let arguments = req.params.get("arguments").cloned().unwrap_or(json!({}));
78
79            let result = tools::dispatch(ctx, tool_name, &arguments);
80            JsonRpcResponse::success(req.id.clone(), serde_json::to_value(result).unwrap())
81        }
82
83        _ => JsonRpcResponse::error(
84            req.id.clone(),
85            METHOD_NOT_FOUND,
86            format!("Method not found: {}", req.method),
87        ),
88    }
89}
90
91fn write_response(out: &mut impl Write, resp: &JsonRpcResponse) {
92    if let Ok(json) = serde_json::to_string(resp) {
93        let _ = writeln!(out, "{json}");
94        let _ = out.flush();
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101    use std::env;
102
103    fn test_ctx() -> SomaContext {
104        SomaContext::new(env::temp_dir())
105    }
106
107    fn call(ctx: &mut SomaContext, method: &str, params: serde_json::Value) -> JsonRpcResponse {
108        let req = JsonRpcRequest {
109            jsonrpc: "2.0".into(),
110            id: json!(1),
111            method: method.into(),
112            params,
113        };
114        handle_request(ctx, &req)
115    }
116
117    #[test]
118    fn initialize() {
119        let mut ctx = test_ctx();
120        let resp = call(&mut ctx, "initialize", json!({}));
121        assert!(resp.result.is_some());
122        let result = resp.result.unwrap();
123        assert_eq!(result["serverInfo"]["name"], "soma-mcp");
124    }
125
126    #[test]
127    fn tools_list() {
128        let mut ctx = test_ctx();
129        let resp = call(&mut ctx, "tools/list", json!({}));
130        let result = resp.result.unwrap();
131        let tools = result["tools"].as_array().unwrap();
132        assert!(tools.len() >= 10);
133
134        let names: Vec<&str> = tools.iter().filter_map(|t| t["name"].as_str()).collect();
135        assert!(names.contains(&"record_experiment"));
136        assert!(names.contains(&"query_knowledge_base"));
137        assert!(names.contains(&"generate_report"));
138        assert!(names.contains(&"list_filters"));
139    }
140
141    #[test]
142    fn record_and_query_experiment() {
143        let mut ctx = test_ctx();
144
145        // Record
146        let resp = call(
147            &mut ctx,
148            "tools/call",
149            json!({
150                "name": "record_experiment",
151                "arguments": {
152                    "id": "exp_001",
153                    "name": "SVM test",
154                    "hypothesis": "SVM works well on iris",
155                    "research_line": "svm_exploration",
156                    "metrics": { "f1": 0.85 },
157                    "tags": ["svm", "classification"]
158                }
159            }),
160        );
161        assert!(resp.error.is_none());
162
163        // Query
164        let resp = call(
165            &mut ctx,
166            "tools/call",
167            json!({
168                "name": "query_knowledge_base",
169                "arguments": { "query": "SVM", "max_results": 5 }
170            }),
171        );
172        let result = resp.result.unwrap();
173        let text = result["content"][0]["text"].as_str().unwrap();
174        assert!(text.contains("exp_001"));
175    }
176
177    #[test]
178    fn research_line_workflow() {
179        let mut ctx = test_ctx();
180
181        // Record experiments (line is created implicitly)
182        for (id, f1) in [("e1", 0.7), ("e2", 0.8), ("e3", 0.9)] {
183            call(
184                &mut ctx,
185                "tools/call",
186                json!({
187                    "name": "record_experiment",
188                    "arguments": {
189                        "id": id,
190                        "name": format!("Experiment {id}"),
191                        "research_line": "norm_study",
192                        "metrics": { "f1": f1 }
193                    }
194                }),
195            );
196        }
197
198        // Get trajectory
199        let resp = call(
200            &mut ctx,
201            "tools/call",
202            json!({
203                "name": "get_trajectory",
204                "arguments": { "research_line": "norm_study", "metric": "f1" }
205            }),
206        );
207        let text = resp.result.unwrap()["content"][0]["text"]
208            .as_str()
209            .unwrap()
210            .to_string();
211        assert!(text.contains("0.7"));
212        assert!(text.contains("0.9"));
213
214        // List lines
215        let resp = call(
216            &mut ctx,
217            "tools/call",
218            json!({"name": "list_research_lines", "arguments": {}}),
219        );
220        let text = resp.result.unwrap()["content"][0]["text"]
221            .as_str()
222            .unwrap()
223            .to_string();
224        assert!(text.contains("norm_study"));
225
226        // Generate report
227        let resp = call(
228            &mut ctx,
229            "tools/call",
230            json!({"name": "generate_report", "arguments": {"research_line": "norm_study"}}),
231        );
232        let text = resp.result.unwrap()["content"][0]["text"]
233            .as_str()
234            .unwrap()
235            .to_string();
236        assert!(text.contains("# Research Report: norm_study"));
237        assert!(text.contains("Trajectory"));
238    }
239
240    #[test]
241    fn unknown_method() {
242        let mut ctx = test_ctx();
243        let resp = call(&mut ctx, "unknown/method", json!({}));
244        assert!(resp.error.is_some());
245    }
246
247    #[test]
248    fn unknown_tool() {
249        let mut ctx = test_ctx();
250        let resp = call(
251            &mut ctx,
252            "tools/call",
253            json!({"name": "nonexistent_tool", "arguments": {}}),
254        );
255        let result = resp.result.unwrap();
256        assert_eq!(result["isError"], true);
257    }
258}