| 1 | //! A test MCP server over stdio (newline-delimited JSON-RPC). Stands in for third-party servers |
| 2 | //! and for built-ins such as `demucs --mcp`: it ignores its arguments. A product exposes it as a |
| 3 | //! binary whose `main` calls [`run`] and points `ed_acp_testkit::Target::mcp_stub` at it. Behind the |
| 4 | //! `test-server` feature. |
| 5 | //! |
| 6 | //! Set `STUB_PID_FILE` to have it write its process id, so a test can check it was stopped. |
| 7 | //! |
| 8 | //! Tools: `echo` (read-only), `slow_progress` (three progress notifications), |
| 9 | //! `wait_for_cancel` (answers only after `notifications/cancelled`), and the demucs-shaped |
| 10 | //! `list_models` / `separate_stems`. |
| 11 | |
| 12 | use std::collections::HashSet; |
| 13 | use std::sync::{Arc, Mutex}; |
| 14 | use std::time::Duration; |
| 15 | |
| 16 | use serde_json::{Value, json}; |
| 17 | use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; |
| 18 | |
| 19 | fn tools() -> Value { |
| 20 | json!([ |
| 21 | { |
| 22 | "name": "echo", |
| 23 | "description": "Echo the text back.", |
| 24 | "inputSchema": {"type": "object", "properties": {"text": {"type": "string"}}, "required": ["text"]}, |
| 25 | "annotations": {"readOnlyHint": true} |
| 26 | }, |
| 27 | { |
| 28 | "name": "slow_progress", |
| 29 | "description": "Report progress three times, then finish.", |
| 30 | "inputSchema": {"type": "object", "properties": {}} |
| 31 | }, |
| 32 | { |
| 33 | "name": "wait_for_cancel", |
| 34 | "description": "Blocks until the call is cancelled.", |
| 35 | "inputSchema": {"type": "object", "properties": {}} |
| 36 | }, |
| 37 | { |
| 38 | "name": "list_models", |
| 39 | "description": "List stem separation models.", |
| 40 | "inputSchema": {"type": "object", "properties": {}}, |
| 41 | "annotations": {"readOnlyHint": true} |
| 42 | }, |
| 43 | { |
| 44 | "name": "separate_stems", |
| 45 | "description": "Pretend to separate stems.", |
| 46 | "inputSchema": {"type": "object", "properties": {"input": {"type": "string"}, "model": {"type": "string"}}, "required": ["input"]} |
| 47 | } |
| 48 | ]) |
| 49 | } |
| 50 | |
| 51 | /// Serve on stdin/stdout until stdin closes. |
| 52 | pub async fn run() { |
| 53 | // Lets tests find out whether this process is still alive after the agent closes it. |
| 54 | if let Some(path) = std::env::var_os("STUB_PID_FILE") { |
| 55 | let _ = std::fs::write(path, std::process::id().to_string()); |
| 56 | } |
| 57 | let out = Arc::new(tokio::sync::Mutex::new(tokio::io::stdout())); |
| 58 | let cancelled: Arc<Mutex<HashSet<String>>> = Arc::default(); |
| 59 | let mut lines = BufReader::new(tokio::io::stdin()).lines(); |
| 60 | |
| 61 | let send = { |
| 62 | let out = out.clone(); |
| 63 | move |msg: Value| { |
| 64 | let out = out.clone(); |
| 65 | async move { |
| 66 | let mut o = out.lock().await; |
| 67 | let _ = o.write_all(format!("{msg}\n").as_bytes()).await; |
| 68 | let _ = o.flush().await; |
| 69 | } |
| 70 | } |
| 71 | }; |
| 72 | |
| 73 | while let Ok(Some(line)) = lines.next_line().await { |
| 74 | let Ok(msg) = serde_json::from_str::<Value>(&line) else { |
| 75 | continue; |
| 76 | }; |
| 77 | let method = msg |
| 78 | .get("method") |
| 79 | .and_then(Value::as_str) |
| 80 | .unwrap_or("") |
| 81 | .to_string(); |
| 82 | let id = msg.get("id").cloned(); |
| 83 | let params = msg.get("params").cloned().unwrap_or(Value::Null); |
| 84 | let Some(id) = id else { |
| 85 | if method == "notifications/cancelled" |
| 86 | && let Some(rid) = params.get("requestId") |
| 87 | { |
| 88 | cancelled.lock().unwrap().insert(rid.to_string()); |
| 89 | } |
| 90 | continue; |
| 91 | }; |
| 92 | match method.as_str() { |
| 93 | "initialize" => { |
| 94 | let version = params |
| 95 | .get("protocolVersion") |
| 96 | .cloned() |
| 97 | .unwrap_or(json!("2025-06-18")); |
| 98 | send(json!({"jsonrpc": "2.0", "id": id, "result": { |
| 99 | "protocolVersion": version, |
| 100 | "capabilities": {"tools": {}}, |
| 101 | "serverInfo": {"name": "stub", "version": "0.0.0"} |
| 102 | }})) |
| 103 | .await; |
| 104 | } |
| 105 | "ping" => send(json!({"jsonrpc": "2.0", "id": id, "result": {}})).await, |
| 106 | "tools/list" => { |
| 107 | send(json!({"jsonrpc": "2.0", "id": id, "result": {"tools": tools()}})).await |
| 108 | } |
| 109 | "tools/call" => { |
| 110 | let send = send.clone(); |
| 111 | let cancelled = cancelled.clone(); |
| 112 | tokio::spawn(async move { |
| 113 | let name = params |
| 114 | .get("name") |
| 115 | .and_then(Value::as_str) |
| 116 | .unwrap_or("") |
| 117 | .to_string(); |
| 118 | let args = params.get("arguments").cloned().unwrap_or(json!({})); |
| 119 | let token = params.pointer("/_meta/progressToken").cloned(); |
| 120 | let progress = |n: u32, msg: &str| { |
| 121 | token.as_ref().map(|t| { |
| 122 | json!({"jsonrpc": "2.0", "method": "notifications/progress", |
| 123 | "params": {"progressToken": t, "progress": n, "total": 100, "message": msg}}) |
| 124 | }) |
| 125 | }; |
| 126 | let result = match name.as_str() { |
| 127 | "echo" => { |
| 128 | let text = args.get("text").and_then(Value::as_str).unwrap_or(""); |
| 129 | json!({"content": [{"type": "text", "text": format!("echo: {text}")}], "isError": false}) |
| 130 | } |
| 131 | "slow_progress" | "separate_stems" => { |
| 132 | for (n, m) in [(10, "starting"), (50, "halfway"), (90, "almost")] { |
| 133 | if let Some(p) = progress(n, m) { |
| 134 | send(p).await; |
| 135 | } |
| 136 | tokio::time::sleep(Duration::from_millis(40)).await; |
| 137 | } |
| 138 | if name == "slow_progress" { |
| 139 | json!({"content": [{"type": "text", "text": "done"}], "isError": false}) |
| 140 | } else { |
| 141 | let input = args |
| 142 | .get("input") |
| 143 | .and_then(Value::as_str) |
| 144 | .unwrap_or("/tmp/x.wav"); |
| 145 | let dir = format!("{}_stems", input.trim_end_matches(".wav")); |
| 146 | let stems: Vec<Value> = ["drums", "bass", "vocals", "other"] |
| 147 | .iter() |
| 148 | .map(|s| json!({"id": s, "path": format!("{dir}/{s}.wav")})) |
| 149 | .collect(); |
| 150 | json!({ |
| 151 | "content": [{"type": "text", "text": format!("Separated into {} stems", stems.len())}], |
| 152 | "structuredContent": {"stems": stems}, |
| 153 | "isError": false |
| 154 | }) |
| 155 | } |
| 156 | } |
| 157 | "list_models" => json!({ |
| 158 | "content": [{"type": "text", "text": "htdemucs 80 MB"}], |
| 159 | "structuredContent": {"models": [ |
| 160 | {"id": "htdemucs", "size_mb": 80, "cached": false}, |
| 161 | {"id": "htdemucs_6s", "size_mb": 55, "cached": true} |
| 162 | ]}, |
| 163 | "isError": false |
| 164 | }), |
| 165 | "wait_for_cancel" => { |
| 166 | let key = id.to_string(); |
| 167 | while !cancelled.lock().unwrap().contains(&key) { |
| 168 | tokio::time::sleep(Duration::from_millis(20)).await; |
| 169 | } |
| 170 | json!({"content": [{"type": "text", "text": "cancelled"}], "isError": true}) |
| 171 | } |
| 172 | other => { |
| 173 | send(json!({"jsonrpc": "2.0", "id": id, "error": {"code": -32602, "message": format!("unknown tool {other}")}})).await; |
| 174 | return; |
| 175 | } |
| 176 | }; |
| 177 | send(json!({"jsonrpc": "2.0", "id": id, "result": result})).await; |
| 178 | }); |
| 179 | } |
| 180 | other => { |
| 181 | send(json!({"jsonrpc": "2.0", "id": id, "error": {"code": -32601, "message": format!("method not found: {other}")}})).await; |
| 182 | } |
| 183 | } |
| 184 | } |
| 185 | } |