agentmux_srv/backend/
mcp_probe.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! MCP server health/prerequisite probe
5//! (docs/specs/SPEC_MCP_INTEGRATION_PARITY_ABLETON_PILOT_2026_07_08.md §4.4).
6//!
7//! Opens a short-lived MCP client connection to a configured server — spawns
8//! the stdio process (or POSTs to the http url), performs the `initialize`
9//! handshake, then `tools/list` — and reports whether it's reachable at all.
10//!
11//! This is a protocol-level probe only. For a server like ableton-mcp, a
12//! successful probe means the MCP process itself started and speaks the
13//! protocol — NOT that the underlying app (Ableton Live) is running, since
14//! that dependency is only exercised when a tool is actually called. That
15//! gap is why the catalog's `prereqNote` (static remediation text, §4.6)
16//! exists alongside this dynamic check, not instead of it.
17//!
18//! Deliberately hand-rolled rather than pulling in an MCP client SDK crate
19//! (e.g. `rmcp`): `crates/mcp/src/main.rs` already hand-rolls the *server*
20//! side of this exact newline-delimited JSON-RPC 2.0 wire format, so a
21//! probe-only client follows the same house style instead of adding a new
22//! dependency for a handshake this small.
23
24#[cfg(windows)]
25use agentmux_common::win32::NoWindow;
26use std::process::Stdio;
27use std::time::Duration;
28
29use serde::Serialize;
30use serde_json::{json, Value};
31use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader};
32use tokio::process::Command;
33use tokio::time::timeout;
34
35const PROBE_TIMEOUT: Duration = Duration::from_secs(8);
36const PROTOCOL_VERSION: &str = "2024-11-05";
37
38#[derive(Debug, Clone, Copy, Serialize, ts_rs::TS)]
39#[ts(export, export_to = "../../../frontend/types/rpc/")]
40#[serde(rename_all = "snake_case")]
41pub enum ProbeStatus {
42    Connected,
43    Unreachable,
44    HandshakeFailed,
45    InvalidConfig,
46}
47
48#[derive(Debug, Clone, Serialize, ts_rs::TS)]
49#[ts(export, export_to = "../../../frontend/types/rpc/")]
50pub struct ProbeResult {
51    pub status: ProbeStatus,
52    #[ts(type = "number | null")]
53    pub tool_count: Option<usize>,
54    pub server_name: Option<String>,
55    pub server_version: Option<String>,
56    pub error: Option<String>,
57}
58
59impl ProbeResult {
60    fn empty(status: ProbeStatus, error: Option<String>) -> Self {
61        Self { status, tool_count: None, server_name: None, server_version: None, error }
62    }
63    fn invalid(msg: impl Into<String>) -> Self {
64        Self::empty(ProbeStatus::InvalidConfig, Some(msg.into()))
65    }
66    fn unreachable(msg: impl Into<String>) -> Self {
67        Self::empty(ProbeStatus::Unreachable, Some(msg.into()))
68    }
69    fn handshake_failed(msg: impl Into<String>) -> Self {
70        Self::empty(ProbeStatus::HandshakeFailed, Some(msg.into()))
71    }
72    fn connected(server_name: Option<String>, server_version: Option<String>, tool_count: Option<usize>) -> Self {
73        Self { status: ProbeStatus::Connected, tool_count, server_name, server_version, error: None }
74    }
75}
76
77/// Probe a server given its `transport` (`"stdio"` | `"http"` | `"sse"` |
78/// any other free-text value from the pre-v11 field, treated as stdio for
79/// back-compat) and its raw `config` JSON blob — the exact object
80/// `agent_config::build_mcp_config_from_refs` merges into `.mcp.json`.
81pub async fn probe(transport: &str, config_json: &str) -> ProbeResult {
82    let config: Value = match serde_json::from_str(config_json) {
83        Ok(v) => v,
84        Err(e) => return ProbeResult::invalid(format!("config is not valid JSON: {e}")),
85    };
86
87    let is_http = matches!(transport, "http" | "sse" | "streamable-http") || config.get("url").is_some();
88
89    let outcome = if is_http {
90        timeout(PROBE_TIMEOUT, probe_http(&config)).await
91    } else {
92        timeout(PROBE_TIMEOUT, probe_stdio(&config)).await
93    };
94
95    match outcome {
96        Ok(result) => result,
97        Err(_) => ProbeResult::handshake_failed(format!("no response within {}s", PROBE_TIMEOUT.as_secs())),
98    }
99}
100
101fn initialize_request() -> Value {
102    json!({
103        "jsonrpc": "2.0",
104        "id": 1,
105        "method": "initialize",
106        "params": {
107            "protocolVersion": PROTOCOL_VERSION,
108            "capabilities": {},
109            "clientInfo": { "name": "agentmux-probe", "version": env!("CARGO_PKG_VERSION") },
110        },
111    })
112}
113
114fn tools_list_request() -> Value {
115    json!({ "jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {} })
116}
117
118// ── stdio ────────────────────────────────────────────────────────────────
119
120async fn probe_stdio(config: &Value) -> ProbeResult {
121    let command = match config.get("command").and_then(|v| v.as_str()) {
122        Some(c) if !c.is_empty() => c,
123        _ => return ProbeResult::invalid("stdio config missing a non-empty \"command\" string"),
124    };
125    let args: Vec<String> = config
126        .get("args")
127        .and_then(|v| v.as_array())
128        .map(|a| a.iter().filter_map(|x| x.as_str().map(str::to_string)).collect())
129        .unwrap_or_default();
130    let env: Vec<(String, String)> = config
131        .get("env")
132        .and_then(|v| v.as_object())
133        .map(|m| m.iter().filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string()))).collect())
134        .unwrap_or_default();
135
136    let mut cmd = Command::new(command);
137    // An arbitrary command from MCP server config — third-party by definition.
138    crate::backend::pane_env::sanitize_external_command(&mut cmd);
139    cmd.args(&args)
140        .envs(env)
141        .stdin(Stdio::piped())
142        .stdout(Stdio::piped())
143        .stderr(Stdio::null())
144        .kill_on_drop(true);
145    // CREATE_NO_WINDOW: console-flash suppression, see crates/common/src/cli.rs
146    #[cfg(windows)]
147    {
148        cmd.no_window();
149    }
150
151    let mut child = match cmd.spawn() {
152        Ok(c) => c,
153        Err(e) => return ProbeResult::unreachable(format!("failed to start \"{command}\": {e}")),
154    };
155
156    let mut stdin = match child.stdin.take() {
157        Some(s) => s,
158        None => return ProbeResult::unreachable("failed to open child stdin"),
159    };
160    let stdout = match child.stdout.take() {
161        Some(s) => s,
162        None => return ProbeResult::unreachable("failed to open child stdout"),
163    };
164    let mut reader = BufReader::new(stdout);
165
166    let result = run_stdio_handshake(&mut stdin, &mut reader).await;
167    let _ = child.start_kill(); // probe is ephemeral — never leave it running
168    result
169}
170
171async fn run_stdio_handshake<W, R>(stdin: &mut W, reader: &mut R) -> ProbeResult
172where
173    W: AsyncWrite + Unpin,
174    R: AsyncBufRead + Unpin,
175{
176    if let Err(e) = write_line(stdin, &initialize_request()).await {
177        return ProbeResult::unreachable(format!("failed to write initialize request: {e}"));
178    }
179    let init_resp = match read_json_line(reader).await {
180        Ok(Some(v)) => v,
181        Ok(None) => return ProbeResult::handshake_failed("process closed stdout before responding to initialize"),
182        Err(e) => return ProbeResult::handshake_failed(format!("failed to read initialize response: {e}")),
183    };
184    if let Some(err) = init_resp.get("error") {
185        return ProbeResult::handshake_failed(format!("server rejected initialize: {err}"));
186    }
187    let server_info = init_resp.pointer("/result/serverInfo").cloned().unwrap_or(Value::Null);
188    let server_name = server_info.get("name").and_then(|v| v.as_str()).map(str::to_string);
189    let server_version = server_info.get("version").and_then(|v| v.as_str()).map(str::to_string);
190
191    // notifications/initialized has no id and expects no response — best-effort.
192    let _ = write_line(stdin, &json!({"jsonrpc": "2.0", "method": "notifications/initialized"})).await;
193
194    if write_line(stdin, &tools_list_request()).await.is_err() {
195        // The handshake itself already succeeded even if tools/list couldn't be sent.
196        return ProbeResult::connected(server_name, server_version, None);
197    }
198    let tool_count = match read_json_line(reader).await {
199        Ok(Some(v)) => v.pointer("/result/tools").and_then(|t| t.as_array()).map(|a| a.len()),
200        _ => None,
201    };
202
203    ProbeResult::connected(server_name, server_version, tool_count)
204}
205
206async fn write_line<W: AsyncWrite + Unpin>(w: &mut W, v: &Value) -> std::io::Result<()> {
207    let s = serde_json::to_string(v).unwrap_or_default();
208    w.write_all(s.as_bytes()).await?;
209    w.write_all(b"\n").await?;
210    w.flush().await
211}
212
213/// Reads lines until one parses as JSON (skipping blank lines or stray
214/// non-JSON output, e.g. a server printing a startup banner to stdout),
215/// or the stream closes.
216async fn read_json_line<R: AsyncBufRead + Unpin>(r: &mut R) -> std::io::Result<Option<Value>> {
217    let mut line = String::new();
218    loop {
219        line.clear();
220        let n = r.read_line(&mut line).await?;
221        if n == 0 {
222            return Ok(None);
223        }
224        let trimmed = line.trim();
225        if trimmed.is_empty() {
226            continue;
227        }
228        if let Ok(v) = serde_json::from_str::<Value>(trimmed) {
229            return Ok(Some(v));
230        }
231    }
232}
233
234// ── http (Streamable HTTP transport) ────────────────────────────────────
235
236async fn probe_http(config: &Value) -> ProbeResult {
237    let url = match config.get("url").and_then(|v| v.as_str()) {
238        Some(u) if !u.is_empty() => u,
239        _ => return ProbeResult::invalid("http config missing a non-empty \"url\" string"),
240    };
241
242    let mut headers = reqwest::header::HeaderMap::new();
243    if let Some(h) = config.get("headers").and_then(|v| v.as_object()) {
244        for (k, v) in h {
245            let (Ok(name), Some(val)) = (reqwest::header::HeaderName::from_bytes(k.as_bytes()), v.as_str()) else {
246                continue;
247            };
248            if let Ok(hv) = reqwest::header::HeaderValue::from_str(val) {
249                headers.insert(name, hv);
250            }
251        }
252    }
253    let client = match reqwest::Client::builder().default_headers(headers).timeout(PROBE_TIMEOUT).build() {
254        Ok(c) => c,
255        Err(e) => return ProbeResult::invalid(format!("failed to build http client: {e}")),
256    };
257
258    let init_resp = match post_jsonrpc(&client, url, &initialize_request()).await {
259        Ok(v) => v,
260        Err(e) => return e,
261    };
262    if let Some(err) = init_resp.get("error") {
263        return ProbeResult::handshake_failed(format!("server rejected initialize: {err}"));
264    }
265    let server_info = init_resp.pointer("/result/serverInfo").cloned().unwrap_or(Value::Null);
266    let server_name = server_info.get("name").and_then(|v| v.as_str()).map(str::to_string);
267    let server_version = server_info.get("version").and_then(|v| v.as_str()).map(str::to_string);
268
269    let tool_count = match post_jsonrpc(&client, url, &tools_list_request()).await {
270        Ok(v) => v.pointer("/result/tools").and_then(|t| t.as_array()).map(|a| a.len()),
271        Err(_) => None, // reachability + handshake already proven above
272    };
273
274    ProbeResult::connected(server_name, server_version, tool_count)
275}
276
277/// POSTs one JSON-RPC request and returns the first JSON object in the
278/// response body — handling both a plain `application/json` body and a
279/// Streamable HTTP `text/event-stream` framing (`event: message\ndata:
280/// {...}\n\n`), since a probe only needs the first frame either way.
281async fn post_jsonrpc(client: &reqwest::Client, url: &str, req: &Value) -> Result<Value, ProbeResult> {
282    let resp = client
283        .post(url)
284        .header("Content-Type", "application/json")
285        .header("Accept", "application/json, text/event-stream")
286        .json(req)
287        .send()
288        .await
289        .map_err(|e| ProbeResult::unreachable(format!("request failed: {e}")))?;
290
291    if !resp.status().is_success() {
292        return Err(ProbeResult::handshake_failed(format!("server returned HTTP {}", resp.status())));
293    }
294    let body = resp.text().await.map_err(|e| ProbeResult::handshake_failed(format!("failed to read response body: {e}")))?;
295    let json_line = body.lines().find(|l| l.trim_start().starts_with('{')).unwrap_or(body.trim());
296    serde_json::from_str(json_line).map_err(|e| ProbeResult::handshake_failed(format!("non-JSON response: {e}")))
297}
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302
303    #[tokio::test]
304    async fn invalid_json_config_is_reported_as_invalid_config() {
305        let result = probe("stdio", "{not json").await;
306        assert!(matches!(result.status, ProbeStatus::InvalidConfig));
307        assert!(result.error.unwrap().contains("not valid JSON"));
308    }
309
310    #[tokio::test]
311    async fn stdio_config_missing_command_is_invalid() {
312        let result = probe("stdio", "{}").await;
313        assert!(matches!(result.status, ProbeStatus::InvalidConfig));
314        assert!(result.error.unwrap().contains("command"));
315    }
316
317    #[tokio::test]
318    async fn stdio_config_with_unresolvable_binary_is_unreachable() {
319        let config = json!({ "command": "definitely-not-a-real-agentmux-probe-binary-xyz" }).to_string();
320        let result = probe("stdio", &config).await;
321        assert!(matches!(result.status, ProbeStatus::Unreachable), "expected Unreachable, got {:?}", result.status);
322    }
323
324    #[tokio::test]
325    async fn http_config_missing_url_is_invalid() {
326        let result = probe("http", "{}").await;
327        assert!(matches!(result.status, ProbeStatus::InvalidConfig));
328        assert!(result.error.unwrap().contains("url"));
329    }
330
331    #[tokio::test]
332    async fn http_transport_is_inferred_from_url_field_even_without_explicit_kind() {
333        // transport is still the free-text pre-v11 field; a config that
334        // carries "url" should route to the http probe regardless of what
335        // (if anything) the transport string says.
336        let config = json!({ "url": "http://127.0.0.1:1/nonexistent" }).to_string();
337        let result = probe("stdio", &config).await;
338        assert!(matches!(result.status, ProbeStatus::Unreachable), "expected Unreachable, got {:?}", result.status);
339    }
340}
agentmux_srv/backend/mcp_probe.rs

agentmux_srv/backend/
mcp_probe.rs
