diff --git a/crates/trapfall-mcp/src/lib.rs b/crates/trapfall-mcp/src/lib.rs index 11c2cff..8a99eec 100644 --- a/crates/trapfall-mcp/src/lib.rs +++ b/crates/trapfall-mcp/src/lib.rs @@ -3,60 +3,143 @@ //! Implements JSON-RPC 2.0 over stdin/stdout for Model Context Protocol. //! Spawned by AI agents as subprocess. No auth needed (process isolation). -use std::io::{BufRead, Write}; - use anyhow::Result; use serde_json::{Value, json}; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use trapfall_core::Store; use trapfall_proto::IssueStatus; -/// Process a single JSON-RPC message and return the response string. -/// Used by both the stdin server loop and integration tests. -pub async fn handle_message(input: &str, store: &Store) -> String { +/// MCP protocol versions this server can speak, newest first. +const SUPPORTED_PROTOCOL_VERSIONS: &[&str] = &["2025-06-18", "2025-03-26", "2024-11-05"]; + +// JSON-RPC 2.0 error codes. +const PARSE_ERROR: i64 = -32700; +const INVALID_REQUEST: i64 = -32600; +const METHOD_NOT_FOUND: i64 = -32601; +const INVALID_PARAMS: i64 = -32602; + +/// A JSON-RPC protocol-level error. +#[derive(Debug)] +struct RpcError { + code: i64, + message: String, +} + +impl RpcError { + fn new(code: i64, message: impl Into) -> Self { + Self { code, message: message.into() } + } + + fn invalid_params(message: impl Into) -> Self { + Self::new(INVALID_PARAMS, message) + } +} + +fn error_response(id: Value, err: RpcError) -> Value { + json!({ "jsonrpc": "2.0", "error": { "code": err.code, "message": err.message }, "id": id }) +} + +/// Process a JSON-RPC message (single request, notification or batch). +/// +/// Returns `None` when no response must be sent (notifications, or a batch +/// consisting solely of notifications). Used by both the stdin server loop +/// and integration tests. +pub async fn handle_message(input: &str, store: &Store) -> Option { let msg: Value = match serde_json::from_str(input.trim()) { Ok(v) => v, Err(e) => { - return serde_json::to_string(&json!({ - "jsonrpc": "2.0", - "error": {"code": -32700, "message": format!("Parse error: {e}")}, - "id": null - })) - .unwrap(); + return Some( + error_response(Value::Null, RpcError::new(PARSE_ERROR, format!("Parse error: {e}"))).to_string(), + ); + } + }; + + match msg { + Value::Array(items) => { + if items.is_empty() { + return Some( + error_response(Value::Null, RpcError::new(INVALID_REQUEST, "Invalid Request: empty batch")) + .to_string(), + ); + } + let mut responses = Vec::new(); + for item in items { + if let Some(resp) = handle_single(item, store).await { + responses.push(resp); + } + } + if responses.is_empty() { None } else { Some(Value::Array(responses).to_string()) } } + other => handle_single(other, store).await.map(|v| v.to_string()), + } +} + +/// Handle one JSON-RPC message object. `None` means "send nothing". +async fn handle_single(msg: Value, store: &Store) -> Option { + let Some(obj) = msg.as_object() else { + return Some(error_response(Value::Null, RpcError::new(INVALID_REQUEST, "Invalid Request: not an object"))); + }; + + // `id` must be a string, number or null; its absence marks a notification. + let (id, is_notification) = match obj.get("id") { + None => (Value::Null, true), + Some(v @ (Value::String(_) | Value::Number(_) | Value::Null)) => (v.clone(), false), + Some(_) => { + return Some(error_response( + Value::Null, + RpcError::new(INVALID_REQUEST, "Invalid Request: id must be a string, number or null"), + )); + } + }; + + // A response from the peer (no method, has result/error) needs no reply. + if !obj.contains_key("method") && (obj.contains_key("result") || obj.contains_key("error")) { + return None; + } + + if obj.get("jsonrpc").and_then(Value::as_str) != Some("2.0") { + return Some(error_response(id, RpcError::new(INVALID_REQUEST, "Invalid Request: jsonrpc must be \"2.0\""))); + } + + let Some(method) = obj.get("method").and_then(Value::as_str) else { + return Some(error_response(id, RpcError::new(INVALID_REQUEST, "Invalid Request: method must be a string"))); }; - let id = msg.get("id").cloned(); - let method = msg.get("method").and_then(|m| m.as_str()).unwrap_or(""); - let params = msg.get("params").cloned().unwrap_or(json!({})); + let params = match obj.get("params") { + None => json!({}), + Some(p @ (Value::Object(_) | Value::Array(_))) => p.clone(), + Some(_) => { + return (!is_notification) + .then(|| error_response(id, RpcError::invalid_params("params must be an object or array"))); + } + }; let result = handle_request(method, params, store).await; - match result { - Ok(val) => json!({ "jsonrpc": "2.0", "result": val, "id": id }), - Err(code_msg) => json!({ "jsonrpc": "2.0", "error": {"code": -32603, "message": code_msg}, "id": id }), + if is_notification { + return None; } - .to_string() + Some(match result { + Ok(val) => json!({ "jsonrpc": "2.0", "result": val, "id": id }), + Err(err) => error_response(id, err), + }) } -/// Run the MCP server, reading JSON-RPC from stdin and writing to stdout. +/// Run the MCP server, reading newline-delimited JSON-RPC from stdin and +/// writing responses to stdout. pub async fn run_server(store: Store) -> Result<()> { - let stdin = std::io::stdin(); - let stdout = std::io::stdout(); - let mut stdout = stdout.lock(); - - let mut line = String::new(); - let stdin = stdin.lock(); - - for result in stdin.lines() { - line.clear(); - match result { - Ok(l) => line.push_str(&l), - Err(_) => break, - } + let mut lines = BufReader::new(tokio::io::stdin()).lines(); + let mut stdout = tokio::io::stdout(); - let response = handle_message(&line, &store).await; - writeln!(stdout, "{}", response)?; - stdout.flush()?; + while let Some(line) = lines.next_line().await? { + if line.trim().is_empty() { + continue; + } + if let Some(response) = handle_message(&line, &store).await { + stdout.write_all(response.as_bytes()).await?; + stdout.write_all(b"\n").await?; + stdout.flush().await?; + } } Ok(()) @@ -64,22 +147,45 @@ pub async fn run_server(store: Store) -> Result<()> { // ── JSON-RPC Dispatcher ──────────────────────────────────────────────── -async fn handle_request(method: &str, params: Value, store: &Store) -> Result { +async fn handle_request(method: &str, params: Value, store: &Store) -> Result { match method { - "initialize" => Ok(json!({ - "protocolVersion": "2024-11-05", - "capabilities": { "tools": {} }, - "serverInfo": { "name": "trapfall", "version": env!("CARGO_PKG_VERSION") } - })), - "notifications/initialized" => Ok(Value::Null), + "initialize" => { + let requested = params + .get("protocolVersion") + .and_then(Value::as_str) + .ok_or_else(|| RpcError::invalid_params("missing protocolVersion"))?; + let version = SUPPORTED_PROTOCOL_VERSIONS + .iter() + .find(|v| **v == requested) + .unwrap_or(&SUPPORTED_PROTOCOL_VERSIONS[0]); + Ok(json!({ + "protocolVersion": version, + "capabilities": { "tools": {} }, + "serverInfo": { "name": "trapfall", "version": env!("CARGO_PKG_VERSION") } + })) + } "tools/list" => Ok(json!({ "tools": tools_list() })), "tools/call" => { - let tool_name = params.get("name").and_then(|n| n.as_str()).ok_or("missing tool name")?; - let args = params.get("arguments").cloned().unwrap_or(json!({})); - call_tool(tool_name, args, store).await + let tool_name = params + .get("name") + .and_then(Value::as_str) + .ok_or_else(|| RpcError::invalid_params("missing tool name"))?; + let args = match params.get("arguments") { + None | Some(Value::Null) => json!({}), + Some(a @ Value::Object(_)) => a.clone(), + Some(_) => return Err(RpcError::invalid_params("arguments must be an object")), + }; + if !tools_list().iter().any(|t| t["name"] == tool_name) { + return Err(RpcError::invalid_params(format!("Unknown tool: {tool_name}"))); + } + // Tool execution failures are reported in-band (`isError`), not as protocol errors. + Ok(match call_tool(tool_name, args, store).await { + Ok(v) => v, + Err(msg) => json!({ "content": [{ "type": "text", "text": msg }], "isError": true }), + }) } "ping" => Ok(json!({})), - _ => Err(format!("Unknown method: {method}")), + _ => Err(RpcError::new(METHOD_NOT_FOUND, format!("Unknown method: {method}"))), } } @@ -399,7 +505,7 @@ async fn tool_rotate_dsn(args: Value, store: &Store) -> Result { let slug = args.get("project_slug").and_then(|v| v.as_str()).ok_or("missing project_slug")?; let project = resolve_slug(slug, store).await?; let new_key = store.rotate_dsn(&project.id).await.map_err(|e| e.to_string())?; - Ok(text_msg(format!("DSN rotated for {}. New key: {}...{}", slug, &new_key[..8], &new_key[new_key.len() - 4..]))) + Ok(text_msg(format!("DSN rotated for {}. New key: {}", slug, abbreviate_key(&new_key)))) } async fn tool_healthcheck(store: &Store) -> Result { @@ -415,10 +521,45 @@ fn mask_dsn(dsn: &str) -> String { Some(p) => p, None => return "***masked***".to_string(), }; - let key = &dsn[..at_pos].trim_start_matches("https://"); - if key.len() <= 12 { - format!("https://***@{}", &dsn[at_pos + 1..]) + let key = dsn[..at_pos].trim_start_matches("https://"); + let host = &dsn[at_pos + 1..]; + if key.chars().count() <= 12 { + format!("https://***@{host}") } else { - format!("https://{}...{}@{}", &key[..8], &key[key.len() - 4..], &dsn[at_pos + 1..]) + format!("https://{}@{host}", abbreviate_key(key)) + } +} + +/// Abbreviate a key to its first 8 and last 4 characters (char-boundary safe). +/// Keys too short to abbreviate meaningfully are fully masked. +fn abbreviate_key(key: &str) -> String { + let count = key.chars().count(); + if count <= 12 { + return "***".to_string(); + } + let head: String = key.chars().take(8).collect(); + let tail: String = key.chars().skip(count - 4).collect(); + format!("{head}...{tail}") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn abbreviate_key_handles_short_and_multibyte() { + assert_eq!(abbreviate_key(""), "***"); + assert_eq!(abbreviate_key("abc"), "***"); + assert_eq!(abbreviate_key("0123456789abcdef"), "01234567...cdef"); + // Multibyte chars must not panic on slicing. + assert_eq!(abbreviate_key("ééééééééééééééé"), "éééééééé...éééé"); + } + + #[test] + fn mask_dsn_handles_edge_cases() { + assert_eq!(mask_dsn("no-at-sign"), "***masked***"); + assert_eq!(mask_dsn("https://short@host/1"), "https://***@host/1"); + assert_eq!(mask_dsn("https://0123456789abcdef@host/1"), "https://01234567...cdef@host/1"); + assert_eq!(mask_dsn("https://éééééééééééééééé@host/1"), "https://éééééééé...éééé@host/1"); } } diff --git a/crates/trapfall-mcp/tests/mcp_test.rs b/crates/trapfall-mcp/tests/mcp_test.rs index ee66699..49b6ece 100644 --- a/crates/trapfall-mcp/tests/mcp_test.rs +++ b/crates/trapfall-mcp/tests/mcp_test.rs @@ -12,8 +12,13 @@ fn rpc_request(method: &str, params: Value) -> Value { } /// Helper: parse JSON-RPC response. -fn parse_response(output: &str) -> Value { - serde_json::from_str(output).expect("valid JSON response") +fn parse_response(output: &Option) -> Value { + serde_json::from_str(output.as_deref().expect("a response was expected")).expect("valid JSON response") +} + +/// Send a raw JSON value through the handler. +async fn send(store: &Store, msg: Value) -> Option { + trapfall_mcp::handle_message(&msg.to_string(), store).await } /// Set up a store with migrated SQLite in-memory database. @@ -92,7 +97,7 @@ async fn test_unknown_method() { let result = trapfall_mcp::handle_message(&input, &store).await; let resp = parse_response(&result); - assert!(resp["error"]["code"].is_number()); + assert_eq!(resp["error"]["code"], -32601); assert!(resp["error"]["message"].as_str().unwrap().contains("Unknown")); } @@ -127,3 +132,150 @@ async fn test_list_projects_tool() { // Empty project list is valid — we just verify the tool doesn't crash assert_eq!(projects.len(), 0); } + +// ── JSON-RPC 2.0 / MCP compliance (#336) ─────────────────────────────── + +#[tokio::test] +async fn test_notification_gets_no_response() { + let store = setup_store().await; + let out = send(&store, json!({"jsonrpc": "2.0", "method": "notifications/initialized"})).await; + assert!(out.is_none(), "notifications must not get a response, got {out:?}"); + + // Unknown-method notifications are also silent. + let out = send(&store, json!({"jsonrpc": "2.0", "method": "nope"})).await; + assert!(out.is_none()); +} + +#[tokio::test] +async fn test_invalid_params_is_32602() { + let store = setup_store().await; + let resp = parse_response(&send(&store, rpc_request("tools/call", json!({}))).await); + assert_eq!(resp["error"]["code"], -32602); + assert_eq!(resp["id"], 1); +} + +#[tokio::test] +async fn test_unknown_tool_is_32602() { + let store = setup_store().await; + let resp = parse_response(&send(&store, rpc_request("tools/call", json!({"name": "nope"}))).await); + assert_eq!(resp["error"]["code"], -32602); +} + +#[tokio::test] +async fn test_tool_failure_is_is_error_result() { + let store = setup_store().await; + let resp = parse_response( + &send(&store, rpc_request("tools/call", json!({"name": "get_issue", "arguments": {"issue_id": "missing"}}))) + .await, + ); + assert!(resp.get("error").is_none(), "tool failures are not protocol errors: {resp}"); + assert_eq!(resp["result"]["isError"], true); + assert_eq!(resp["result"]["content"][0]["type"], "text"); + + // Missing argument is also a tool-level error. + let resp = parse_response(&send(&store, rpc_request("tools/call", json!({"name": "get_issue"}))).await); + assert_eq!(resp["result"]["isError"], true); +} + +#[tokio::test] +async fn test_invalid_request_is_32600() { + let store = setup_store().await; + // wrong jsonrpc version + let resp = parse_response(&send(&store, json!({"jsonrpc": "1.0", "id": 1, "method": "ping"})).await); + assert_eq!(resp["error"]["code"], -32600); + // missing jsonrpc + let resp = parse_response(&send(&store, json!({"id": 1, "method": "ping"})).await); + assert_eq!(resp["error"]["code"], -32600); + // missing method + let resp = parse_response(&send(&store, json!({"jsonrpc": "2.0", "id": 7})).await); + assert_eq!(resp["error"]["code"], -32600); + assert_eq!(resp["id"], 7); + // non-object + let resp = parse_response(&send(&store, json!(42)).await); + assert_eq!(resp["error"]["code"], -32600); + assert_eq!(resp["id"], Value::Null); + // invalid id type + let resp = parse_response(&send(&store, json!({"jsonrpc": "2.0", "id": {"a": 1}, "method": "ping"})).await); + assert_eq!(resp["error"]["code"], -32600); + assert_eq!(resp["id"], Value::Null); + // non-string method + let resp = parse_response(&send(&store, json!({"jsonrpc": "2.0", "id": 1, "method": 5})).await); + assert_eq!(resp["error"]["code"], -32600); +} + +#[tokio::test] +async fn test_string_id_echoed() { + let store = setup_store().await; + let resp = parse_response(&send(&store, json!({"jsonrpc": "2.0", "id": "abc", "method": "ping"})).await); + assert_eq!(resp["id"], "abc"); + assert_eq!(resp["result"], json!({})); +} + +#[tokio::test] +async fn test_batch() { + let store = setup_store().await; + let out = send( + &store, + json!([ + {"jsonrpc": "2.0", "id": 1, "method": "ping"}, + {"jsonrpc": "2.0", "method": "notifications/initialized"}, + {"jsonrpc": "2.0", "id": 2, "method": "nope"}, + 5 + ]), + ) + .await; + let resp = parse_response(&out); + let arr = resp.as_array().expect("batch response is an array"); + assert_eq!(arr.len(), 3, "notification produces no entry: {resp}"); + assert_eq!(arr[0]["id"], 1); + assert_eq!(arr[0]["result"], json!({})); + assert_eq!(arr[1]["error"]["code"], -32601); + assert_eq!(arr[2]["error"]["code"], -32600); +} + +#[tokio::test] +async fn test_batch_edge_cases() { + let store = setup_store().await; + // empty batch + let resp = parse_response(&send(&store, json!([])).await); + assert_eq!(resp["error"]["code"], -32600); + // all notifications -> nothing + let out = send(&store, json!([{"jsonrpc": "2.0", "method": "notifications/initialized"}])).await; + assert!(out.is_none()); +} + +#[tokio::test] +async fn test_initialize_negotiates_protocol_version() { + let store = setup_store().await; + let init = |v: Value| rpc_request("initialize", json!({"protocolVersion": v, "capabilities": {}})); + + let resp = parse_response(&send(&store, init(json!("2025-03-26"))).await); + assert_eq!(resp["result"]["protocolVersion"], "2025-03-26"); + + let resp = parse_response(&send(&store, init(json!("2024-11-05"))).await); + assert_eq!(resp["result"]["protocolVersion"], "2024-11-05"); + + // unsupported -> server's latest supported version + let resp = parse_response(&send(&store, init(json!("1999-01-01"))).await); + assert_eq!(resp["result"]["protocolVersion"], "2025-06-18"); + + // missing protocolVersion is invalid params + let resp = parse_response(&send(&store, rpc_request("initialize", json!({}))).await); + assert_eq!(resp["error"]["code"], -32602); +} + +#[tokio::test] +async fn test_rotate_dsn_succeeds_and_masks_key() { + let store = setup_store().await; + let project = store.create_project("demo", "Demo").await.unwrap(); + let resp = parse_response( + &send( + &store, + rpc_request("tools/call", json!({"name": "rotate_dsn", "arguments": {"project_slug": project.slug}})), + ) + .await, + ); + assert_ne!(resp["result"]["isError"], true, "{resp}"); + let text = resp["result"]["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("New key:"), "{text}"); +}