From 47794e992a02cc25b7f78e4ccb2308fa68f144a3 Mon Sep 17 00:00:00 2001 From: Kayshen-X Date: Thu, 14 May 2026 22:43:14 +0800 Subject: [PATCH] fix(desktop/mcp): implement initialize + tools/list handshake MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Codex stop-time review on 23b0bfba flagged that real MCP clients (Claude Code, Codex, etc.) won't dispatch tools/call cold — they open with \`initialize\` and \`tools/list\` to discover the server. The previous --mcp implementation responded only to tools/call, so real clients hung at handshake. Add in-binary handlers for the MCP control-plane methods: - \`initialize\` → protocolVersion "2024-11-05" + capabilities (tools.listChanged=false) + serverInfo - \`notifications/initialized\` (and bare \`initialized\`) → absorbed silently (spec: notifications have no response) - \`tools/list\` → response with all 14 tools + their JSON inputSchemas (name / description / required args / enums for kind + drop_children) - \`ping\` → empty result - everything else falls through to shell-core's run_stdio_with_applier so the parser hardening + apply discipline still apply Top-level method/id sniffing uses the same key-walker pattern as shell-core's arguments_field, so a nested key called "method" or a string value of "initialize" can't shadow the real top-level method. Four sniff/response unit tests cover the discipline. End-to-end roundtrip verified with real stdin/stdout: initialize → notifications/initialized → tools/list → tools/call insert_node → tools/call list_pages → all five frames return correct JSON-RPC + the inserted rect persists to disk. --- crates/openpencil-desktop/src/mcp_serve.rs | 417 ++++++++++++++++++--- 1 file changed, 369 insertions(+), 48 deletions(-) diff --git a/crates/openpencil-desktop/src/mcp_serve.rs b/crates/openpencil-desktop/src/mcp_serve.rs index 74384d3ff..902a35448 100644 --- a/crates/openpencil-desktop/src/mcp_serve.rs +++ b/crates/openpencil-desktop/src/mcp_serve.rs @@ -6,17 +6,20 @@ //! Gemini / Copilot) can spawn the binary in this mode to drive //! the Rust editor exactly the way they drive TS pen-mcp today. //! -//! Lifecycle per dispatched call: -//! 1. Read a JSON-RPC line from stdin. -//! 2. Re-build the `ToolRegistry` against the current document -//! so read-tool snapshots reflect the latest state. -//! 3. Dispatch. -//! 4. If the tool emitted an `McpCommand`, the host applier -//! calls `Document::apply_mcp_command(&cmd)` and saves the -//! file on success. -//! 5. Write the JSON-RPC response. +//! Per-line dispatch: +//! - `initialize` → respond with protocol version + server +//! capabilities so the client completes its handshake. +//! - `tools/list` → respond with the 14-tool catalog + JSON +//! schemas for each input. +//! - `notifications/initialized` / `ping` → handled inline +//! (no body / pong). +//! - `tools/call` / legacy direct dispatch → routed through +//! shell-core's `run_stdio_with_applier` so the parser-level +//! hardening (structured-arg rejection / top-level walker / +//! no-hang error) protects this binary too. The applier +//! mutates the live document + saves on success. -use std::io::{BufReader, BufWriter, Write}; +use std::io::{BufRead, BufReader, BufWriter, Write}; use std::path::PathBuf; use openpencil_shell_core::document::Document; @@ -32,58 +35,71 @@ use openpencil_shell_core::mcp::{ use crate::persistence::{load_from_path, save_to_path}; /// Run the stdio MCP server against `path`. Returns Ok(()) on EOF, -/// Err on unrecoverable IO. The function blocks the calling -/// thread for the lifetime of the stdio connection. +/// Err on unrecoverable IO. Blocks the calling thread for the +/// lifetime of the stdio connection. pub fn run(path: PathBuf) -> Result<(), String> { let mut doc = Document::empty(); load_from_path(&mut doc, &path).map_err(|e| format!("load {}: {e}", path.display()))?; - // Re-built per call so read-tool snapshots reflect the latest - // state — write commands mutate the doc between dispatches. - // The registry holds Box, so re-registration is - // the simplest path until snapshots become incremental. - let rebuild_registry = |doc: &Document| { - let mut r = ToolRegistry::default(); - r.register(Box::new(document_info_snapshot(doc))); - r.register(Box::new(selection_snapshot(doc))); - r.register(Box::new(get_node_snapshot(doc))); - r.register(Box::new(list_pages_snapshot(doc))); - r.register(Box::new(list_variables_snapshot(doc))); - r.register(Box::new(get_active_theme_snapshot(doc))); - r.register(Box::new(set_variable_color_snapshot(doc))); - r.register(Box::new(set_active_axis_value_snapshot(doc))); - r.register(Box::new(insert_node_snapshot())); - r.register(Box::new(update_node_snapshot())); - r.register(Box::new(delete_node_snapshot())); - r.register(Box::new(move_node_snapshot())); - r.register(Box::new(copy_node_snapshot())); - r.register(Box::new(replace_node_snapshot())); - r - }; - let stdin = std::io::stdin(); let stdout = std::io::stdout(); let mut reader = BufReader::new(stdin.lock()); let mut writer = BufWriter::new(stdout.lock()); - // `run_stdio_with_applier` re-uses the same registry for every - // dispatched line, but each write-command apply mutates `doc`, - // making the snapshots stale. To keep the contract simple we - // wrap the loop ourselves: dispatch one call at a time so the - // registry can be rebuilt against the latest document between - // calls. let mut line = String::new(); loop { line.clear(); - let n = std::io::BufRead::read_line(&mut reader, &mut line) + let n = reader + .read_line(&mut line) .map_err(|e| format!("stdin read: {e}"))?; if n == 0 { return Ok(()); // EOF } + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + + // MCP handshake / discovery methods short-circuit the + // tool dispatcher. Detect them via a cheap method-field + // sniff so we don't need to drag JSON parsing in here — + // shell-core's parser stays the single source for the + // `tools/call` envelope. + let method = sniff_method(trimmed); + match method.as_deref() { + Some("initialize") => { + if let Some(id_raw) = sniff_id_raw(trimmed) { + writeln!(writer, "{}", initialize_response(&id_raw)) + .map_err(|e| format!("stdout write: {e}"))?; + writer.flush().map_err(|e| format!("stdout flush: {e}"))?; + } + continue; + } + Some("tools/list") => { + if let Some(id_raw) = sniff_id_raw(trimmed) { + writeln!(writer, "{}", tools_list_response(&id_raw)) + .map_err(|e| format!("stdout write: {e}"))?; + writer.flush().map_err(|e| format!("stdout flush: {e}"))?; + } + continue; + } + Some("notifications/initialized") | Some("initialized") => { + // Notification — no response required by spec. + continue; + } + Some("ping") => { + if let Some(id_raw) = sniff_id_raw(trimmed) { + writeln!(writer, "{}", ping_response(&id_raw)) + .map_err(|e| format!("stdout write: {e}"))?; + writer.flush().map_err(|e| format!("stdout flush: {e}"))?; + } + continue; + } + _ => {} + } + + // Fall through: tools/call or legacy direct dispatch. let registry = rebuild_registry(&doc); - // Reuse the shell-core dispatcher for parser + apply - // discipline; feed it a one-line cursor + a closure that - // mutates the live document and saves on success. let mut applier_failed: Option = None; let path_for_save = path.clone(); { @@ -91,9 +107,6 @@ pub fn run(path: PathBuf) -> Result<(), String> { let path_ref = &path_for_save; let applier_failed_ref = &mut applier_failed; let mut input = std::io::Cursor::new(line.as_bytes()); - // The shell-core loop reads to EOF; the Cursor over a - // single line terminates after one iteration so we - // dispatch + apply + respond exactly once per turn. run_stdio_with_applier(®istry, &mut input, &mut writer, |cmd| { if !doc_ref.apply_mcp_command(cmd) { return false; @@ -112,3 +125,311 @@ pub fn run(path: PathBuf) -> Result<(), String> { } } } + +/// Re-build the registry against the latest document so read-tool +/// snapshots reflect every prior write command's mutations. +fn rebuild_registry(doc: &Document) -> ToolRegistry { + let mut r = ToolRegistry::default(); + r.register(Box::new(document_info_snapshot(doc))); + r.register(Box::new(selection_snapshot(doc))); + r.register(Box::new(get_node_snapshot(doc))); + r.register(Box::new(list_pages_snapshot(doc))); + r.register(Box::new(list_variables_snapshot(doc))); + r.register(Box::new(get_active_theme_snapshot(doc))); + r.register(Box::new(set_variable_color_snapshot(doc))); + r.register(Box::new(set_active_axis_value_snapshot(doc))); + r.register(Box::new(insert_node_snapshot())); + r.register(Box::new(update_node_snapshot())); + r.register(Box::new(delete_node_snapshot())); + r.register(Box::new(move_node_snapshot())); + r.register(Box::new(copy_node_snapshot())); + r.register(Box::new(replace_node_snapshot())); + r +} + +/// Cheap top-level "method" field extractor. Returns the unquoted +/// string value; None if the field is missing or unparseable. +/// Walks the line key by key so a nested or string-valued +/// "method" in another field can't shadow the real top-level +/// method (mirrors `arguments_field`'s discipline in shell-core). +fn sniff_method(line: &str) -> Option { + let bytes = line.as_bytes(); + // Skip past the leading `{` if present. + let mut i = 0usize; + while i < bytes.len() && bytes[i].is_ascii_whitespace() { + i += 1; + } + if i >= bytes.len() || bytes[i] != b'{' { + return None; + } + i += 1; + walk_top_level_for_string_value(bytes, &mut i, "method") +} + +/// Return the JSON token (verbatim — with quotes if string) that +/// follows `"id":` at the top level. Preserves the original +/// representation so the response carries the same id type the +/// client sent. +fn sniff_id_raw(line: &str) -> Option { + let bytes = line.as_bytes(); + let mut i = 0usize; + while i < bytes.len() && bytes[i].is_ascii_whitespace() { + i += 1; + } + if i >= bytes.len() || bytes[i] != b'{' { + return None; + } + i += 1; + walk_top_level_for_raw_value(bytes, &mut i, "id") +} + +/// Generic top-level key walker — returns the value for `target` +/// when seen at depth 0 of the object body starting at `*i`. +/// `string_only` extracts the inner contents (without quotes); +/// the verbatim variant returns the full literal. +fn walk_top_level_for_string_value( + bytes: &[u8], + i: &mut usize, + target: &str, +) -> Option { + walk_top_level(bytes, i, target, /*string_only=*/ true) +} + +fn walk_top_level_for_raw_value(bytes: &[u8], i: &mut usize, target: &str) -> Option { + walk_top_level(bytes, i, target, /*string_only=*/ false) +} + +fn walk_top_level( + bytes: &[u8], + i: &mut usize, + target: &str, + string_only: bool, +) -> Option { + loop { + // Skip whitespace + commas. + while *i < bytes.len() && (bytes[*i].is_ascii_whitespace() || bytes[*i] == b',') { + *i += 1; + } + if *i >= bytes.len() || bytes[*i] == b'}' { + return None; + } + if bytes[*i] != b'"' { + return None; + } + *i += 1; + let key_start = *i; + while *i < bytes.len() && bytes[*i] != b'"' { + if bytes[*i] == b'\\' { + *i = i.saturating_add(2); + } else { + *i += 1; + } + } + if *i >= bytes.len() { + return None; + } + let key = std::str::from_utf8(&bytes[key_start..*i]).ok()?; + *i += 1; + while *i < bytes.len() && bytes[*i].is_ascii_whitespace() { + *i += 1; + } + if *i >= bytes.len() || bytes[*i] != b':' { + return None; + } + *i += 1; + while *i < bytes.len() && bytes[*i].is_ascii_whitespace() { + *i += 1; + } + if *i >= bytes.len() { + return None; + } + let val_start = *i; + match bytes[*i] { + b'"' => { + *i += 1; + let inner_start = *i; + while *i < bytes.len() && bytes[*i] != b'"' { + if bytes[*i] == b'\\' { + *i = i.saturating_add(2); + } else { + *i += 1; + } + } + if *i >= bytes.len() { + return None; + } + let inner_end = *i; + *i += 1; + if key == target { + if string_only { + return std::str::from_utf8(&bytes[inner_start..inner_end]) + .ok() + .map(|s| s.to_string()); + } else { + return std::str::from_utf8(&bytes[val_start..*i]) + .ok() + .map(|s| s.to_string()); + } + } + } + b'{' | b'[' => { + // Walk past structured value, depth-tracked. + let open = bytes[*i]; + let close = if open == b'{' { b'}' } else { b']' }; + let mut depth = 1i32; + *i += 1; + let mut in_str = false; + let mut escape = false; + while *i < bytes.len() && depth > 0 { + let c = bytes[*i]; + if in_str { + if escape { + escape = false; + } else if c == b'\\' { + escape = true; + } else if c == b'"' { + in_str = false; + } + } else if c == b'"' { + in_str = true; + } else if c == open { + depth += 1; + } else if c == close { + depth -= 1; + } + *i += 1; + } + if key == target { + // Structured value where caller asked for a + // scalar. Treat as absent — caller may fall + // back to a default response shape. + return None; + } + } + _ => { + while *i < bytes.len() + && !matches!(bytes[*i], b',' | b'}' | b' ' | b'\t' | b'\n' | b'\r') + { + *i += 1; + } + if key == target { + return std::str::from_utf8(&bytes[val_start..*i]) + .ok() + .map(|s| s.to_string()); + } + } + } + } +} + +fn initialize_response(id_raw: &str) -> String { + // Spec: `initialize` returns protocolVersion + capabilities + + // serverInfo. We declare only `tools` capabilities — no + // resources / prompts / completion are exposed yet. + format!( + r#"{{"jsonrpc":"2.0","id":{id_raw},"result":{{"protocolVersion":"2024-11-05","capabilities":{{"tools":{{"listChanged":false}}}},"serverInfo":{{"name":"openpencil-mcp","version":"0.1.0"}}}}}}"# + ) +} + +fn ping_response(id_raw: &str) -> String { + format!(r#"{{"jsonrpc":"2.0","id":{id_raw},"result":{{}}}}"#) +} + +fn tools_list_response(id_raw: &str) -> String { + // The tool catalog must match what `rebuild_registry` + // installs. Schemas are minimal but sufficient for an MCP + // client to render a tool picker + validate calls. + format!( + r#"{{"jsonrpc":"2.0","id":{id_raw},"result":{{"tools":[{}]}}}}"#, + TOOL_SCHEMAS.join(",") + ) +} + +/// Per-tool JSON schemas — kept as string literals to avoid +/// pulling serde into this binary's MCP path. Each entry is a +/// complete JSON object: name + description + inputSchema. +const TOOL_SCHEMAS: &[&str] = &[ + // --- read tools --- + r#"{"name":"get_document_info","description":"Summarize the open document (page count, active page, etc).","inputSchema":{"type":"object","properties":{},"additionalProperties":false}}"#, + r#"{"name":"get_selection","description":"Return the current selection state (ids, count).","inputSchema":{"type":"object","properties":{},"additionalProperties":false}}"#, + r#"{"name":"get_node","description":"Read a node by id with depth-limited descendants.","inputSchema":{"type":"object","properties":{"node_id":{"type":"string","description":"u64 node id"}},"required":["node_id"]}}"#, + r#"{"name":"list_pages","description":"List page ids + names.","inputSchema":{"type":"object","properties":{},"additionalProperties":false}}"#, + r#"{"name":"list_variables","description":"List design variables with kinds.","inputSchema":{"type":"object","properties":{},"additionalProperties":false}}"#, + r#"{"name":"get_active_theme","description":"Return the active theme axis pinning per axis.","inputSchema":{"type":"object","properties":{},"additionalProperties":false}}"#, + // --- write tools --- + r##"{"name":"set_variable_color","description":"Set a Color-kind variable's value.","inputSchema":{"type":"object","properties":{"name":{"type":"string"},"hex":{"type":"string","description":"#rgb / #rrggbb / #rrggbbaa"}},"required":["name","hex"]}}"##, + r#"{"name":"set_active_axis_value","description":"Pin a theme axis to one of its allowed values.","inputSchema":{"type":"object","properties":{"axis":{"type":"string"},"value":{"type":"string"}},"required":["axis","value"]}}"#, + r#"{"name":"insert_node","description":"Create a new leaf node on the active page.","inputSchema":{"type":"object","properties":{"kind":{"type":"string","enum":["frame","group","rect","ellipse","polygon","line","text","path"]},"name":{"type":"string"},"x":{"type":"string"},"y":{"type":"string"},"width":{"type":"string"},"height":{"type":"string"},"fill_hex":{"type":"string"}},"required":["kind","name","x","y","width","height"]}}"#, + r#"{"name":"update_node","description":"Patch fields on an existing node. Pass any subset of x/y/width/height/name/fill_hex.","inputSchema":{"type":"object","properties":{"node_id":{"type":"string"},"x":{"type":"string"},"y":{"type":"string"},"width":{"type":"string"},"height":{"type":"string"},"name":{"type":"string"},"fill_hex":{"type":"string"}},"required":["node_id"]}}"#, + r#"{"name":"delete_node","description":"Remove a node + descendants from its parent.","inputSchema":{"type":"object","properties":{"node_id":{"type":"string"}},"required":["node_id"]}}"#, + r#"{"name":"move_node","description":"Reparent a node. target_parent_id=0 puts it at the active page root.","inputSchema":{"type":"object","properties":{"node_id":{"type":"string"},"target_parent_id":{"type":"string"}},"required":["node_id","target_parent_id"]}}"#, + r#"{"name":"copy_node","description":"Deep-clone a subtree with fresh ids under a new parent.","inputSchema":{"type":"object","properties":{"node_id":{"type":"string"},"target_parent_id":{"type":"string"}},"required":["node_id","target_parent_id"]}}"#, + r#"{"name":"replace_node","description":"Swap an existing node at the same parent slot with a freshly-built leaf. Set drop_children=true to discard a container's subtree.","inputSchema":{"type":"object","properties":{"node_id":{"type":"string"},"kind":{"type":"string","enum":["frame","group","rect","ellipse","polygon","line","text","path"]},"name":{"type":"string"},"x":{"type":"string"},"y":{"type":"string"},"width":{"type":"string"},"height":{"type":"string"},"fill_hex":{"type":"string"},"drop_children":{"type":"string","enum":["true","false"]}},"required":["node_id","kind","name","x","y","width","height"]}}"#, +]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sniff_method_walks_top_level() { + assert_eq!( + sniff_method(r#"{"id":1,"method":"initialize","params":{}}"#), + Some("initialize".into()) + ); + assert_eq!( + sniff_method(r#"{"id":1,"method":"tools/call","params":{"name":"x"}}"#), + Some("tools/call".into()) + ); + // Nested `method` keys must not shadow the real one. + assert_eq!( + sniff_method(r#"{"id":1,"method":"tools/list","params":{"method":"fake"}}"#), + Some("tools/list".into()) + ); + assert_eq!(sniff_method("not even json"), None); + } + + #[test] + fn sniff_id_raw_preserves_type() { + assert_eq!( + sniff_id_raw(r#"{"id":42,"method":"x"}"#), + Some("42".into()) + ); + assert_eq!( + sniff_id_raw(r#"{"id":"abc","method":"x"}"#), + Some(r#""abc""#.into()) + ); + } + + #[test] + fn initialize_response_includes_protocol_and_capabilities() { + let r = initialize_response("7"); + assert!(r.contains(r#""id":7"#)); + assert!(r.contains(r#""protocolVersion""#)); + assert!(r.contains(r#""tools""#)); + assert!(r.contains(r#""serverInfo""#)); + } + + #[test] + fn tools_list_response_includes_all_fourteen_tools() { + let r = tools_list_response("3"); + for name in [ + "get_document_info", + "get_selection", + "get_node", + "list_pages", + "list_variables", + "get_active_theme", + "set_variable_color", + "set_active_axis_value", + "insert_node", + "update_node", + "delete_node", + "move_node", + "copy_node", + "replace_node", + ] { + assert!(r.contains(name), "tools/list must include {name}: {r}"); + } + } +}