Skip to main content

omgbase_sync/
mcp_client.rs

1//! The MCP engine client (`spec/sync/README.md` §6; the reference's
2//! `McpEngineClient` in `@omgbase/sync`): reaches "the omgbase side" over
3//! the Model Context Protocol, so the [`Coordinator`](crate::Coordinator)
4//! runs unchanged against a remote engine, and the CLI's `--server` mode
5//! calls any catalog tool (`spec/surface` §4) and renders the result it
6//! would have computed locally.
7//!
8//! Two transports, one JSON-RPC 2.0 client on top:
9//!
10//! - **stdio** — spawn a command (`omg mcp -C /vault`, `omgbase mcp …`) and
11//!   speak newline-delimited JSON-RPC over its stdin/stdout (the framing of
12//!   [`omgbase mcp`](https://github.com/omgbase/omgbase) and of the SDK's
13//!   `StdioServerTransport`); stderr is inherited.
14//! - **Streamable HTTP** (feature `http`) — `POST` each message to the URL
15//!   with `Accept: application/json, text/event-stream`, read the reply from
16//!   a JSON body or an SSE stream, round-trip `Mcp-Session-Id`, and send
17//!   `MCP-Protocol-Version` once negotiated; extra headers ride on every
18//!   request (auth beyond a secret path); `DELETE` ends the session.
19//!
20//! The one rule for a `--server` value lives in [`parse_engine_spec`]: an
21//! `http(s)://` value is an HTTP endpoint (used verbatim, base path and
22//! all), anything else is whitespace-split and spawned — a URL is never
23//! spawned. Synchronous like the rest of the crate: one request in flight,
24//! a bounded wait per reply.
25
26use std::collections::BTreeMap;
27use std::io::{BufRead, BufReader, Write};
28use std::process::{Child, ChildStdin, Command, Stdio};
29use std::sync::mpsc::{Receiver, RecvTimeoutError, channel};
30use std::thread::JoinHandle;
31use std::time::{Duration, Instant};
32
33use omgbase_format::hash::{hex, sha256};
34use omgbase_store::{ChangesPage, CommitDigest, DeleteOutcome, DigestRevision, ObserveOutcome};
35use serde_json::{Map, Value, json};
36
37use crate::engine::{DocBytes, EngineClient, FileBytes};
38use crate::error::{Error, Result};
39
40/// The MCP protocol revision this client proposes in `initialize` (the
41/// server may answer with the one it speaks; the reply is what is sent as
42/// `MCP-Protocol-Version` afterwards).
43pub const PROTOCOL_VERSION: &str = "2025-03-26";
44
45/// `clientInfo.name` (the reference's).
46pub const CLIENT_NAME: &str = "omgbase-sync";
47
48/// How long one request may wait for its reply (the SDK's default request
49/// timeout).
50pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60);
51
52/// After stdin EOF a spawned server gets this long to exit before it is
53/// killed.
54const EXIT_GRACE: Duration = Duration::from_secs(5);
55
56// ---- the spec ---------------------------------------------------------------------------
57
58/// Where a `--server` value points.
59#[derive(Clone, Debug, PartialEq, Eq)]
60pub enum EngineSpec {
61    /// An http(s) endpoint reached over Streamable HTTP, with extra request
62    /// headers (`-H "Name: value"`, in flag order).
63    Http {
64        url: String,
65        headers: Vec<(String, String)>,
66    },
67    /// A command to spawn and speak stdio to.
68    Stdio { command: String, args: Vec<String> },
69}
70
71/// Whether `s` starts with `http://` or `https://` (case-insensitive).
72#[must_use]
73pub fn is_http_url(s: &str) -> bool {
74    let lower: String = s.chars().take(8).collect::<String>().to_ascii_lowercase();
75    lower.starts_with("http://") || lower.starts_with("https://")
76}
77
78/// Classify a raw `--server` value (the reference's `parseEngineSpec`): an
79/// `http(s)://` value is an HTTP endpoint (verbatim; `headers` ride along);
80/// anything else is split on whitespace and spawned. Errors: an empty
81/// value, or headers for a command (they have nowhere to go).
82pub fn parse_engine_spec(raw: &str, headers: &[(String, String)]) -> Result<EngineSpec> {
83    let spec = raw.trim();
84    if is_http_url(spec) {
85        return Ok(EngineSpec::Http {
86            url: spec.to_owned(),
87            headers: headers.to_vec(),
88        });
89    }
90    if !headers.is_empty() {
91        return Err(Error::Other(
92            "extra headers only apply to an http(s) server url, not a command to spawn".to_owned(),
93        ));
94    }
95    let mut argv = spec.split_whitespace().map(str::to_owned);
96    let Some(command) = argv.next() else {
97        return Err(Error::Other(
98            "server spec is empty — expected a command to spawn (e.g. \"omg mcp -C /vault\") or an http(s) URL".to_owned(),
99        ));
100    };
101    Ok(EngineSpec::Stdio {
102        command,
103        args: argv.collect(),
104    })
105}
106
107// ---- tool results ----------------------------------------------------------------------
108
109/// One `tools/call` outcome: the JSON the tool's text content held, and
110/// whether the server flagged it (`isError`; the body is then the error
111/// envelope of `spec/surface` §4 — `{ error, message, data?, retriable }`).
112#[derive(Clone, Debug, PartialEq)]
113pub struct ToolResult {
114    pub body: Value,
115    pub is_error: bool,
116}
117
118impl ToolResult {
119    /// The body, or the envelope as an [`Error::Other`] carrying the code
120    /// and message (the reference's `tool <name> failed: <code> — <message>`).
121    pub fn into_result(self, tool: &str) -> Result<Value> {
122        if !self.is_error {
123            return Ok(self.body);
124        }
125        let code = self
126            .body
127            .get("error")
128            .and_then(Value::as_str)
129            .unwrap_or("error");
130        let message = self
131            .body
132            .get("message")
133            .and_then(Value::as_str)
134            .map_or_else(|| clip(&self.body.to_string(), 200), str::to_owned);
135        Err(Error::Other(format!(
136            "tool {tool} failed: {code} — {message}"
137        )))
138    }
139}
140
141fn clip(s: &str, max: usize) -> String {
142    if s.chars().count() <= max {
143        s.to_owned()
144    } else {
145        s.chars().take(max).collect()
146    }
147}
148
149// ---- the client -------------------------------------------------------------------------
150
151enum Transport {
152    Stdio(StdioTransport),
153    #[cfg(feature = "http")]
154    Http(HttpTransport),
155}
156
157/// A connected MCP client: `initialize` done, ready for `tools/call`.
158pub struct McpEngineClient {
159    transport: Transport,
160    next_id: i64,
161    timeout: Duration,
162    protocol_version: String,
163    server_info: Value,
164}
165
166impl std::fmt::Debug for McpEngineClient {
167    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
168        f.debug_struct("McpEngineClient")
169            .field("protocol_version", &self.protocol_version)
170            .field("server_info", &self.server_info)
171            .finish_non_exhaustive()
172    }
173}
174
175impl McpEngineClient {
176    /// Connect over the transport `spec` selects and run the handshake.
177    pub fn connect(spec: &EngineSpec) -> Result<Self> {
178        Self::connect_with(spec, DEFAULT_TIMEOUT)
179    }
180
181    /// [`Self::connect`] with a per-request timeout.
182    pub fn connect_with(spec: &EngineSpec, timeout: Duration) -> Result<Self> {
183        match spec {
184            EngineSpec::Stdio { command, args } => Self::connect_stdio(command, args, timeout),
185            #[cfg(feature = "http")]
186            EngineSpec::Http { url, headers } => Self::connect_http(url, headers, timeout),
187            #[cfg(not(feature = "http"))]
188            EngineSpec::Http { url, .. } => Err(Error::Other(format!(
189                "cannot reach {url}: this build has no HTTP transport (omgbase-sync feature `http`)"
190            ))),
191        }
192    }
193
194    /// Spawn `command args…` (stderr inherited) and connect over its stdio.
195    pub fn connect_stdio(command: &str, args: &[String], timeout: Duration) -> Result<Self> {
196        let transport = StdioTransport::spawn(command, args)?;
197        Self::handshake(Transport::Stdio(transport), timeout)
198    }
199
200    /// Connect to a Streamable HTTP endpoint; `headers` are sent on every
201    /// request.
202    #[cfg(feature = "http")]
203    pub fn connect_http(
204        url: &str,
205        headers: &[(String, String)],
206        timeout: Duration,
207    ) -> Result<Self> {
208        let transport = HttpTransport::new(url, headers, timeout);
209        Self::handshake(Transport::Http(transport), timeout)
210    }
211
212    fn handshake(transport: Transport, timeout: Duration) -> Result<Self> {
213        let mut client = Self {
214            transport,
215            next_id: 0,
216            timeout,
217            protocol_version: PROTOCOL_VERSION.to_owned(),
218            server_info: Value::Null,
219        };
220        let init = client.request(
221            "initialize",
222            json!({
223                "protocolVersion": PROTOCOL_VERSION,
224                "capabilities": {},
225                "clientInfo": { "name": CLIENT_NAME, "version": env!("CARGO_PKG_VERSION") },
226            }),
227        );
228        let init = match init {
229            Ok(v) => v,
230            Err(e) => {
231                let _ = client.close();
232                return Err(e);
233            }
234        };
235        if let Some(v) = init.get("protocolVersion").and_then(Value::as_str) {
236            client.protocol_version = v.to_owned();
237            #[cfg(feature = "http")]
238            if let Transport::Http(h) = &mut client.transport {
239                h.protocol_version = Some(v.to_owned());
240            }
241        }
242        client.server_info = init.get("serverInfo").cloned().unwrap_or(Value::Null);
243        client.notify("notifications/initialized", Value::Null)?;
244        Ok(client)
245    }
246
247    /// The protocol revision the server answered with.
248    #[must_use]
249    pub fn protocol_version(&self) -> &str {
250        &self.protocol_version
251    }
252
253    /// The server's `serverInfo` (`{ name, version }`).
254    #[must_use]
255    pub fn server_info(&self) -> &Value {
256        &self.server_info
257    }
258
259    /// One JSON-RPC request; the `result`, or the server's error as
260    /// [`Error::Other`].
261    pub fn request(&mut self, method: &str, params: Value) -> Result<Value> {
262        self.next_id += 1;
263        let id = json!(self.next_id);
264        let mut msg = json!({ "jsonrpc": "2.0", "id": id, "method": method });
265        if !params.is_null() {
266            msg["params"] = params;
267        }
268        let reply = match &mut self.transport {
269            Transport::Stdio(t) => t.request(&msg, &id, self.timeout)?,
270            #[cfg(feature = "http")]
271            Transport::Http(t) => t.post(&msg, Some(&id))?.ok_or_else(|| {
272                Error::Other(format!(
273                    "{method}: the server accepted the request without a reply"
274                ))
275            })?,
276        };
277        if let Some(err) = reply.get("error") {
278            let code = err.get("code").and_then(Value::as_i64).unwrap_or(0);
279            let message = err
280                .get("message")
281                .and_then(Value::as_str)
282                .unwrap_or("unknown error");
283            return Err(Error::Other(format!(
284                "{method} failed: {message} (JSON-RPC {code})"
285            )));
286        }
287        reply
288            .get("result")
289            .cloned()
290            .ok_or_else(|| Error::Other(format!("{method}: reply without result")))
291    }
292
293    /// One JSON-RPC notification (no reply).
294    pub fn notify(&mut self, method: &str, params: Value) -> Result<()> {
295        let mut msg = json!({ "jsonrpc": "2.0", "method": method });
296        if !params.is_null() {
297            msg["params"] = params;
298        }
299        match &mut self.transport {
300            Transport::Stdio(t) => t.send(&msg),
301            #[cfg(feature = "http")]
302            Transport::Http(t) => t.post(&msg, None).map(|_| ()),
303        }
304    }
305
306    /// `tools/list`: the catalog as the server describes it.
307    pub fn list_tools(&mut self) -> Result<Vec<Value>> {
308        let res = self.request("tools/list", Value::Null)?;
309        Ok(res
310            .get("tools")
311            .and_then(Value::as_array)
312            .cloned()
313            .unwrap_or_default())
314    }
315
316    /// `tools/call`: the tool's JSON body and its error flag. A reply whose
317    /// text content is not JSON is an error (the reference's
318    /// `tool <name> returned non-JSON`).
319    pub fn call_tool(&mut self, name: &str, args: Value) -> Result<ToolResult> {
320        let res = self.request("tools/call", json!({ "name": name, "arguments": args }))?;
321        let text = res
322            .get("content")
323            .and_then(Value::as_array)
324            .and_then(|items| {
325                items
326                    .iter()
327                    .find(|c| c.get("type").and_then(Value::as_str) == Some("text"))
328            })
329            .and_then(|c| c.get("text"))
330            .and_then(Value::as_str)
331            .unwrap_or("");
332        let body: Value = serde_json::from_str(text).map_err(|_| {
333            Error::Other(format!(
334                "tool {name} returned non-JSON: {}",
335                clip(text, 200)
336            ))
337        })?;
338        Ok(ToolResult {
339            body,
340            is_error: res.get("isError").and_then(Value::as_bool).unwrap_or(false),
341        })
342    }
343
344    /// [`Self::call_tool`] with an error result turned into an error.
345    pub fn call(&mut self, name: &str, args: Value) -> Result<Value> {
346        self.call_tool(name, args)?.into_result(name)
347    }
348
349    /// End the session: stdio — stdin EOF, wait, kill after the grace
350    /// period; HTTP — `DELETE` with the session id. Idempotent.
351    pub fn close(&mut self) -> Result<()> {
352        match &mut self.transport {
353            Transport::Stdio(t) => t.close(),
354            #[cfg(feature = "http")]
355            Transport::Http(t) => {
356                t.delete_session();
357                Ok(())
358            }
359        }
360    }
361}
362
363impl Drop for McpEngineClient {
364    fn drop(&mut self) {
365        let _ = self.close();
366    }
367}
368
369// ---- the engine seam --------------------------------------------------------------------
370
371fn str_field(v: &Value, key: &str) -> String {
372    v.get(key)
373        .and_then(Value::as_str)
374        .unwrap_or_default()
375        .to_owned()
376}
377
378fn opt_str_field(v: &Value, key: &str) -> Option<String> {
379    v.get(key).and_then(Value::as_str).map(str::to_owned)
380}
381
382fn bool_field(v: &Value, key: &str) -> bool {
383    v.get(key).and_then(Value::as_bool).unwrap_or(false)
384}
385
386/// An `observe_many` member as the wire carries it (`{ _id, _path, rev,
387/// commitId, … }` — the document's identity under its engine names,
388/// `spec/surface` §3, 5.0; the `observe*` ARGUMENTS keep `path` /
389/// `files[].path`, a file event's path being a role) → the store's outcome.
390/// Hashes are not on the wire (`old_hash_hex` is `None`, `new_hash_hex`
391/// empty); the coordinator reads `echo`/`conflicted`/`path` only.
392fn observe_outcome(v: &Value) -> ObserveOutcome {
393    let dispositions: BTreeMap<String, u64> = v
394        .get("dispositions")
395        .and_then(Value::as_array)
396        .map(|a| {
397            a.iter()
398                .filter_map(|d| {
399                    Some((
400                        d.get("kind")?.as_str()?.to_owned(),
401                        d.get("count")?.as_u64()?,
402                    ))
403                })
404                .collect()
405        })
406        .unwrap_or_default();
407    ObserveOutcome {
408        path: str_field(v, "_path"),
409        doc_id: str_field(v, "_id"),
410        rev: opt_str_field(v, "rev"),
411        commit_id: opt_str_field(v, "commitId"),
412        converged: bool_field(v, "converged"),
413        echo: bool_field(v, "echo"),
414        conflicted: bool_field(v, "conflicted"),
415        dispositions,
416        old_hash_hex: None,
417        new_hash_hex: String::new(),
418    }
419}
420
421/// A `changes_since` page as the wire carries it (`camelCase`; a revision is
422/// `{ _id, _path, _hash }` — `spec/sync` §6, 3.0).
423fn changes_page(v: &Value, cursor: i64) -> ChangesPage {
424    let digests = v
425        .get("digests")
426        .and_then(Value::as_array)
427        .map(|a| {
428            a.iter()
429                .map(|d| CommitDigest {
430                    commit: str_field(d, "commit"),
431                    seq: d.get("seq").and_then(Value::as_i64).unwrap_or(0),
432                    ts: str_field(d, "ts"),
433                    origin: str_field(d, "origin"),
434                    actor: opt_str_field(d, "actor"),
435                    summary: str_field(d, "summary"),
436                    revisions: d
437                        .get("revisions")
438                        .and_then(Value::as_array)
439                        .map(|rs| {
440                            rs.iter()
441                                .map(|r| DigestRevision {
442                                    doc: str_field(r, "_id"),
443                                    path: str_field(r, "_path"),
444                                    content_hash: str_field(r, "_hash"),
445                                })
446                                .collect()
447                        })
448                        .unwrap_or_default(),
449                })
450                .collect()
451        })
452        .unwrap_or_default();
453    ChangesPage {
454        digests,
455        cursor: v.get("cursor").and_then(Value::as_i64).unwrap_or(cursor),
456        truncated: bool_field(v, "truncated"),
457        head: v.get("head").and_then(Value::as_i64).unwrap_or(0),
458    }
459}
460
461impl EngineClient for McpEngineClient {
462    fn observe_many(&mut self, files: &[FileBytes]) -> Result<Vec<ObserveOutcome>> {
463        let files: Vec<Value> = files
464            .iter()
465            .map(|f| json!({ "path": f.path, "content": f.content }))
466            .collect();
467        let out = self.call("observe_many", json!({ "files": files }))?;
468        let items = out.as_array().ok_or_else(|| {
469            Error::Other(format!(
470                "observe_many returned {} instead of an array",
471                clip(&out.to_string(), 200)
472            ))
473        })?;
474        Ok(items.iter().map(observe_outcome).collect())
475    }
476
477    fn observe_delete(&mut self, path: &str) -> Result<DeleteOutcome> {
478        let out = self.call("observe_delete", json!({ "path": path }))?;
479        Ok(DeleteOutcome {
480            path: opt_str_field(&out, "_path").unwrap_or_else(|| path.to_owned()),
481            doc_id: opt_str_field(&out, "_id"),
482            old_hash_hex: None,
483        })
484    }
485
486    fn changes_since(
487        &mut self,
488        cursor: i64,
489        limit: Option<usize>,
490        origin: Option<&str>,
491    ) -> Result<ChangesPage> {
492        let mut args = Map::new();
493        args.insert("cursor".to_owned(), json!(cursor));
494        if let Some(l) = limit {
495            args.insert("limit".to_owned(), json!(l));
496        }
497        if let Some(o) = origin {
498            args.insert("origin".to_owned(), json!(o));
499        }
500        let out = self.call("changes_since", Value::Object(args))?;
501        Ok(changes_page(&out, cursor))
502    }
503
504    /// `docs_read { _path }` — the argument addresses the row's `_path`
505    /// (`spec/surface` §4, 5.0: a bare `path` is refused, no tolerance) — and
506    /// `_raw` is the document's reconstructed bytes; a tool error
507    /// (`doc_missing`, …) reads as absent, as the reference treats it. The
508    /// hash is computed here: the wire carries the bytes, not the store's
509    /// `file_hash`.
510    fn read_doc(&mut self, path: &str) -> Result<Option<DocBytes>> {
511        let res = self.call_tool("docs_read", json!({ "_path": path }))?;
512        if res.is_error {
513            return Ok(None);
514        }
515        let Some(content) = res.body.get("_raw").and_then(Value::as_str) else {
516            return Ok(None);
517        };
518        Ok(Some(DocBytes {
519            content: content.to_owned(),
520            content_hash: hex(&sha256(content.as_bytes())),
521        }))
522    }
523
524    fn close(&mut self) -> Result<()> {
525        McpEngineClient::close(self)
526    }
527}
528
529// ---- stdio ------------------------------------------------------------------------------
530
531/// A spawned server and the line reader over its stdout.
532struct StdioTransport {
533    command: String,
534    child: Option<Child>,
535    stdin: Option<ChildStdin>,
536    lines: Receiver<String>,
537    reader: Option<JoinHandle<()>>,
538}
539
540impl StdioTransport {
541    fn spawn(command: &str, args: &[String]) -> Result<Self> {
542        let spawn_err = |message: String| {
543            Error::Other(format!(
544                "cannot start the MCP server `{}`: {message}",
545                std::iter::once(command)
546                    .chain(args.iter().map(String::as_str))
547                    .collect::<Vec<_>>()
548                    .join(" ")
549            ))
550        };
551        let mut child = Command::new(command)
552            .args(args)
553            .stdin(Stdio::piped())
554            .stdout(Stdio::piped())
555            .stderr(Stdio::inherit())
556            .spawn()
557            .map_err(|e| spawn_err(e.to_string()))?;
558        let stdin = child
559            .stdin
560            .take()
561            .ok_or_else(|| spawn_err("no stdin pipe".to_owned()))?;
562        let stdout = child
563            .stdout
564            .take()
565            .ok_or_else(|| spawn_err("no stdout pipe".to_owned()))?;
566        let (tx, lines) = channel();
567        let reader = std::thread::Builder::new()
568            .name("omgbase-mcp-client".to_owned())
569            .spawn(move || {
570                for line in BufReader::new(stdout).lines() {
571                    let Ok(line) = line else { break };
572                    if line.trim().is_empty() {
573                        continue;
574                    }
575                    if tx.send(line).is_err() {
576                        break;
577                    }
578                }
579            })
580            .map_err(|e| spawn_err(format!("cannot spawn the reader thread: {e}")))?;
581        Ok(Self {
582            command: command.to_owned(),
583            child: Some(child),
584            stdin: Some(stdin),
585            lines,
586            reader: Some(reader),
587        })
588    }
589
590    fn send(&mut self, msg: &Value) -> Result<()> {
591        let stdin = self
592            .stdin
593            .as_mut()
594            .ok_or_else(|| Error::Other(format!("MCP server `{}` is closed", self.command)))?;
595        serde_json::to_writer(&mut *stdin, msg)
596            .map_err(std::io::Error::other)
597            .and_then(|()| stdin.write_all(b"\n"))
598            .and_then(|()| stdin.flush())
599            .map_err(|e| Error::Other(format!("MCP server `{}`: write failed: {e}", self.command)))
600    }
601
602    /// Send and wait for the reply carrying `id`; every other line (a
603    /// notification, a stray reply) is skipped.
604    fn request(&mut self, msg: &Value, id: &Value, timeout: Duration) -> Result<Value> {
605        self.send(msg)?;
606        let method = msg.get("method").and_then(Value::as_str).unwrap_or("?");
607        let deadline = Instant::now() + timeout;
608        loop {
609            let remaining = deadline.saturating_duration_since(Instant::now());
610            let line = self.lines.recv_timeout(remaining).map_err(|e| match e {
611                RecvTimeoutError::Timeout => Error::Other(format!(
612                    "MCP server `{}`: no reply to {method} within {timeout:?}",
613                    self.command
614                )),
615                RecvTimeoutError::Disconnected => Error::Other(format!(
616                    "MCP server `{}` exited while waiting for {method}",
617                    self.command
618                )),
619            })?;
620            let Ok(reply) = serde_json::from_str::<Value>(&line) else {
621                return Err(Error::Other(format!(
622                    "MCP server `{}` wrote non-JSON to stdout: {}",
623                    self.command,
624                    clip(&line, 200)
625                )));
626            };
627            if reply.get("id") == Some(id) {
628                return Ok(reply);
629            }
630        }
631    }
632
633    fn close(&mut self) -> Result<()> {
634        drop(self.stdin.take());
635        if let Some(mut child) = self.child.take() {
636            let deadline = Instant::now() + EXIT_GRACE;
637            loop {
638                match child.try_wait() {
639                    Ok(Some(_)) => break,
640                    Ok(None) if Instant::now() < deadline => {
641                        std::thread::sleep(Duration::from_millis(20));
642                    }
643                    Ok(None) => {
644                        let _ = child.kill();
645                        let _ = child.wait();
646                        break;
647                    }
648                    Err(_) => break,
649                }
650            }
651        }
652        if let Some(r) = self.reader.take() {
653            let _ = r.join();
654        }
655        Ok(())
656    }
657}
658
659// ---- Streamable HTTP --------------------------------------------------------------------
660
661#[cfg(feature = "http")]
662struct HttpTransport {
663    agent: ureq::Agent,
664    url: String,
665    headers: Vec<(String, String)>,
666    session_id: Option<String>,
667    protocol_version: Option<String>,
668}
669
670#[cfg(feature = "http")]
671impl HttpTransport {
672    fn new(url: &str, headers: &[(String, String)], timeout: Duration) -> Self {
673        let config = ureq::Agent::config_builder()
674            .http_status_as_error(false)
675            .timeout_global(Some(timeout))
676            .build();
677        Self {
678            agent: ureq::Agent::new_with_config(config),
679            url: url.to_owned(),
680            headers: headers.to_vec(),
681            session_id: None,
682            protocol_version: None,
683        }
684    }
685
686    fn fail(&self, what: impl std::fmt::Display) -> Error {
687        Error::Other(format!("MCP server {}: {what}", self.url))
688    }
689
690    fn common_headers<B>(&self, mut req: ureq::RequestBuilder<B>) -> ureq::RequestBuilder<B> {
691        for (k, v) in &self.headers {
692            req = req.header(k.as_str(), v.as_str());
693        }
694        if let Some(sid) = &self.session_id {
695            req = req.header("mcp-session-id", sid.as_str());
696        }
697        if let Some(pv) = &self.protocol_version {
698            req = req.header("mcp-protocol-version", pv.as_str());
699        }
700        req
701    }
702
703    /// `POST` one message. With `want_id`, the reply carrying that id from
704    /// a JSON body or an SSE stream; `None` for a notification (or a 202).
705    fn post(&mut self, msg: &Value, want_id: Option<&Value>) -> Result<Option<Value>> {
706        let req = self
707            .agent
708            .post(&self.url)
709            .header("content-type", "application/json")
710            .header("accept", "application/json, text/event-stream");
711        let req = self.common_headers(req);
712        let mut res = req.send(msg.to_string()).map_err(|e| self.fail(e))?;
713        if let Some(sid) = res
714            .headers()
715            .get("mcp-session-id")
716            .and_then(|v| v.to_str().ok())
717        {
718            self.session_id = Some(sid.to_owned());
719        }
720        let status = res.status().as_u16();
721        if status >= 400 {
722            let text = res.body_mut().read_to_string().unwrap_or_default();
723            return Err(self.fail(format!(
724                "HTTP {status} posting to endpoint: {}",
725                clip(text.trim(), 300)
726            )));
727        }
728        let Some(id) = want_id else {
729            return Ok(None);
730        };
731        if status == 202 {
732            return Ok(None);
733        }
734        let media = res
735            .headers()
736            .get("content-type")
737            .and_then(|v| v.to_str().ok())
738            .map(|ct| {
739                ct.split(';')
740                    .next()
741                    .unwrap_or("")
742                    .trim()
743                    .to_ascii_lowercase()
744            })
745            .unwrap_or_default();
746        match media.as_str() {
747            "application/json" => {
748                let text = res.body_mut().read_to_string().map_err(|e| self.fail(e))?;
749                let body: Value = serde_json::from_str(&text)
750                    .map_err(|e| self.fail(format!("non-JSON body: {e}")))?;
751                let found = match body {
752                    Value::Array(items) => items.into_iter().find(|m| m.get("id") == Some(id)),
753                    other if other.get("id") == Some(id) => Some(other),
754                    _ => None,
755                };
756                found
757                    .map(Some)
758                    .ok_or_else(|| self.fail("the JSON body carried no reply to the request"))
759            }
760            "text/event-stream" => {
761                let reader = BufReader::new(res.body_mut().as_reader());
762                let found = sse_find(reader, id).map_err(|e| self.fail(e))?;
763                found
764                    .map(Some)
765                    .ok_or_else(|| self.fail("the SSE stream ended without a reply to the request"))
766            }
767            other => Err(self.fail(format!("unexpected content type: {other:?}"))),
768        }
769    }
770
771    /// `DELETE` the session (best effort; a 405 means the server keeps no
772    /// sessions).
773    fn delete_session(&mut self) {
774        if self.session_id.is_none() {
775            return;
776        }
777        let req = self.agent.delete(&self.url);
778        let req = self.common_headers(req);
779        let _ = req.call();
780        self.session_id = None;
781    }
782}
783
784/// Read SSE events from `reader` until one's `data` is the JSON-RPC message
785/// carrying `id` (returned), or the stream ends (`None`). Events are blank
786/// line separated; `data:` lines of one event join with `\n`; other fields
787/// (`event`, `id`, `retry`) and comments are skipped.
788#[cfg(feature = "http")]
789fn sse_find(reader: impl BufRead, id: &Value) -> std::io::Result<Option<Value>> {
790    let mut data: Vec<String> = Vec::new();
791    let mut lines = reader.lines();
792    loop {
793        let line = match lines.next() {
794            Some(l) => Some(l?),
795            None => None,
796        };
797        let end_of_event = line.as_deref().is_none_or(|l| l.is_empty());
798        if end_of_event {
799            if !data.is_empty() {
800                let payload = data.join("\n");
801                data.clear();
802                if let Ok(msg) = serde_json::from_str::<Value>(&payload) {
803                    let hit = match &msg {
804                        Value::Array(items) => {
805                            items.iter().find(|m| m.get("id") == Some(id)).cloned()
806                        }
807                        m if m.get("id") == Some(id) => Some(msg.clone()),
808                        _ => None,
809                    };
810                    if hit.is_some() {
811                        return Ok(hit);
812                    }
813                }
814            }
815            if line.is_none() {
816                return Ok(None);
817            }
818            continue;
819        }
820        let line = line.unwrap_or_default();
821        if line.starts_with(':') {
822            continue;
823        }
824        let (field, value) = match line.split_once(':') {
825            Some((f, v)) => (f, v.strip_prefix(' ').unwrap_or(v)),
826            None => (line.as_str(), ""),
827        };
828        if field == "data" {
829            data.push(value.to_owned());
830        }
831    }
832}
833
834#[cfg(test)]
835mod tests {
836    use super::*;
837
838    #[test]
839    fn engine_specs_follow_the_one_rule() {
840        let h = |k: &str, v: &str| (k.to_owned(), v.to_owned());
841        assert_eq!(
842            parse_engine_spec("  https://host/k/secret/mcp ", &[h("X-A", "1")]).unwrap(),
843            EngineSpec::Http {
844                url: "https://host/k/secret/mcp".into(),
845                headers: vec![h("X-A", "1")],
846            }
847        );
848        assert_eq!(
849            parse_engine_spec("HTTP://host/mcp", &[]).unwrap(),
850            EngineSpec::Http {
851                url: "HTTP://host/mcp".into(),
852                headers: vec![],
853            }
854        );
855        assert_eq!(
856            parse_engine_spec("omg   mcp -C /vault", &[]).unwrap(),
857            EngineSpec::Stdio {
858                command: "omg".into(),
859                args: vec!["mcp".into(), "-C".into(), "/vault".into()],
860            }
861        );
862        assert!(
863            parse_engine_spec("omg mcp", &[h("X", "y")]).is_err(),
864            "headers need a url"
865        );
866        assert!(parse_engine_spec("   ", &[]).is_err(), "an empty spec");
867        assert!(!is_http_url("httpx://nope") && !is_http_url("http:/one-slash"));
868    }
869
870    #[test]
871    fn tool_results_turn_envelopes_into_errors() {
872        let ok = ToolResult {
873            body: json!({ "items": [] }),
874            is_error: false,
875        };
876        assert_eq!(ok.into_result("docs_list").unwrap(), json!({ "items": [] }));
877        let bad = ToolResult {
878            body: json!({ "error": "doc_missing", "message": "no document for a.md", "retriable": false }),
879            is_error: true,
880        };
881        let msg = bad.into_result("docs_read").unwrap_err().to_string();
882        assert_eq!(
883            msg,
884            "tool docs_read failed: doc_missing — no document for a.md"
885        );
886    }
887
888    #[test]
889    fn wire_shapes_map_to_store_outcomes() {
890        let o = observe_outcome(&json!({
891            "_id": "d_0", "_path": "a.md", "rev": "r_0", "commitId": "c_0",
892            "converged": true, "echo": false, "conflicted": false,
893            "dispositions": [{ "kind": "same", "count": 2 }, { "kind": "edited", "count": 1 }],
894        }));
895        assert_eq!((o.doc_id.as_str(), o.path.as_str()), ("d_0", "a.md"));
896        assert_eq!(o.dispositions["same"], 2);
897        assert!(o.converged && !o.echo);
898        let echo = observe_outcome(
899            &json!({ "_id": "d_0", "_path": "a.md", "rev": null, "commitId": null, "converged": true, "echo": true, "conflicted": false, "dispositions": [] }),
900        );
901        assert!(echo.echo && echo.rev.is_none());
902        let page = changes_page(
903            &json!({
904                "digests": [{ "commit": "c_1", "seq": 1, "ts": "t", "origin": "api", "actor": null, "summary": "s",
905                              "revisions": [{ "_id": "d_0", "_path": "a.md", "_hash": "ab" }] }],
906                "cursor": 1, "truncated": false, "head": 1,
907            }),
908            0,
909        );
910        assert_eq!(page.digests[0].revisions[0].content_hash, "ab");
911        assert_eq!(page.cursor, 1);
912        assert_eq!(
913            changes_page(&json!({}), 7).cursor,
914            7,
915            "an empty page keeps the input cursor"
916        );
917    }
918
919    /// A `sh` MCP server: replies to `initialize` and to `tools/call` by id,
920    /// swallows notifications, writes one stray notification first.
921    fn sh_server() -> String {
922        r##"
923printf '%s\n' '{"jsonrpc":"2.0","method":"notifications/message","params":{"level":"info","data":"hi"}}'
924while IFS= read -r line; do
925  case "$line" in
926    *'"id":'*) ;;
927    *) continue ;;
928  esac
929  id="${line#*\"id\":}"; id="${id%%,*}"
930  case "$line" in
931    *'"initialize"'*) printf '%s\n' "{\"jsonrpc\":\"2.0\",\"id\":$id,\"result\":{\"protocolVersion\":\"2024-11-05\",\"capabilities\":{\"tools\":{}},\"serverInfo\":{\"name\":\"sh\",\"version\":\"0\"}}}" ;;
932    *'"docs_list"'*) printf '%s\n' "{\"jsonrpc\":\"2.0\",\"id\":$id,\"result\":{\"content\":[{\"type\":\"text\",\"text\":\"{\\\"items\\\":[]}\"}]}}" ;;
933    *'"docs_read"'*) printf '%s\n' "{\"jsonrpc\":\"2.0\",\"id\":$id,\"result\":{\"isError\":true,\"content\":[{\"type\":\"text\",\"text\":\"{\\\"error\\\":\\\"doc_missing\\\",\\\"message\\\":\\\"nope\\\",\\\"retriable\\\":false}\"}]}}" ;;
934    *'"tools/list"'*) printf '%s\n' "{\"jsonrpc\":\"2.0\",\"id\":$id,\"result\":{\"tools\":[{\"name\":\"docs_list\"}]}}" ;;
935    *) printf '%s\n' "{\"jsonrpc\":\"2.0\",\"id\":$id,\"error\":{\"code\":-32601,\"message\":\"method not found\"}}" ;;
936  esac
937done
938"##
939        .to_owned()
940    }
941
942    #[test]
943    fn stdio_handshake_calls_and_closes() {
944        let spec = EngineSpec::Stdio {
945            command: "sh".into(),
946            args: vec!["-c".into(), sh_server()],
947        };
948        let mut c = McpEngineClient::connect_with(&spec, Duration::from_secs(10)).unwrap();
949        assert_eq!(c.protocol_version(), "2024-11-05");
950        assert_eq!(c.server_info()["name"], "sh");
951        assert_eq!(c.list_tools().unwrap()[0]["name"], "docs_list");
952        assert_eq!(
953            c.call("docs_list", json!({})).unwrap(),
954            json!({ "items": [] })
955        );
956        let r = c.call_tool("docs_read", json!({ "_path": "x" })).unwrap();
957        assert!(r.is_error);
958        assert_eq!(r.body["error"], "doc_missing");
959        assert!(
960            c.read_doc("x.md").unwrap().is_none(),
961            "a tool error reads as absent"
962        );
963        let err = c
964            .request("resources/list", Value::Null)
965            .unwrap_err()
966            .to_string();
967        assert!(
968            err.contains("method not found") && err.contains("-32601"),
969            "{err}"
970        );
971        c.close().unwrap();
972        c.close().unwrap();
973        assert!(c.call("docs_list", json!({})).is_err(), "closed");
974    }
975
976    #[test]
977    fn a_missing_command_and_a_silent_server_fail_cleanly() {
978        let spec = EngineSpec::Stdio {
979            command: "definitely-not-an-mcp-server-xyz".into(),
980            args: vec![],
981        };
982        let err = McpEngineClient::connect(&spec).unwrap_err().to_string();
983        assert!(err.contains("cannot start the MCP server"), "{err}");
984        // A server that never answers: the bounded wait, then the error.
985        let spec = EngineSpec::Stdio {
986            command: "sh".into(),
987            args: vec!["-c".into(), "cat >/dev/null".into()],
988        };
989        let err = McpEngineClient::connect_with(&spec, Duration::from_millis(300))
990            .unwrap_err()
991            .to_string();
992        assert!(err.contains("no reply to initialize"), "{err}");
993        // A server that exits at once.
994        let spec = EngineSpec::Stdio {
995            command: "sh".into(),
996            args: vec!["-c".into(), "exit 0".into()],
997        };
998        let err = McpEngineClient::connect_with(&spec, Duration::from_secs(5))
999            .unwrap_err()
1000            .to_string();
1001        assert!(err.contains("exited while waiting"), "{err}");
1002    }
1003
1004    #[cfg(feature = "http")]
1005    mod http {
1006        use super::*;
1007        use std::io::Read;
1008        use std::net::TcpListener;
1009
1010        /// A one-connection-at-a-time HTTP server that answers `n` requests:
1011        /// `initialize` → JSON with `Mcp-Session-Id: s1`; a notification →
1012        /// 202; `tools/call` → an SSE stream (a comment, an unrelated event,
1013        /// then the reply); `DELETE` → 200. Records every request's headers.
1014        fn serve(n: usize) -> (String, std::thread::JoinHandle<Vec<String>>) {
1015            let listener = TcpListener::bind("127.0.0.1:0").unwrap();
1016            let url = format!("http://{}/k/secret/mcp", listener.local_addr().unwrap());
1017            let h = std::thread::spawn(move || {
1018                let mut seen = Vec::new();
1019                for _ in 0..n {
1020                    let (mut s, _) = listener.accept().unwrap();
1021                    let mut buf = Vec::new();
1022                    let mut tmp = [0u8; 4096];
1023                    let (head, body_len) = loop {
1024                        let k = s.read(&mut tmp).unwrap();
1025                        buf.extend_from_slice(&tmp[..k]);
1026                        if let Some(pos) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
1027                            let head = String::from_utf8_lossy(&buf[..pos]).into_owned();
1028                            let len = head
1029                                .lines()
1030                                .find_map(|l| {
1031                                    let (k, v) = l.split_once(':')?;
1032                                    k.eq_ignore_ascii_case("content-length")
1033                                        .then(|| v.trim().parse::<usize>().ok())?
1034                                })
1035                                .unwrap_or(0);
1036                            buf.drain(..pos + 4);
1037                            break (head, len);
1038                        }
1039                    };
1040                    while buf.len() < body_len {
1041                        let k = s.read(&mut tmp).unwrap();
1042                        buf.extend_from_slice(&tmp[..k]);
1043                    }
1044                    let body = String::from_utf8_lossy(&buf[..body_len]).into_owned();
1045                    seen.push(format!("{head}\n\n{body}"));
1046                    let req: Value = serde_json::from_str(&body).unwrap_or(Value::Null);
1047                    let method = req.get("method").and_then(Value::as_str).unwrap_or("");
1048                    let id = req.get("id").cloned().unwrap_or(Value::Null);
1049                    let reply = if head.starts_with("DELETE") {
1050                        "HTTP/1.1 200 OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
1051                            .to_owned()
1052                    } else if method == "initialize" {
1053                        let b = json!({ "jsonrpc": "2.0", "id": id, "result": { "protocolVersion": "2025-03-26", "capabilities": {}, "serverInfo": { "name": "tcp", "version": "1" } } }).to_string();
1054                        format!(
1055                            "HTTP/1.1 200 OK\r\nContent-Type: application/json; charset=utf-8\r\nMcp-Session-Id: s1\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{b}",
1056                            b.len()
1057                        )
1058                    } else if id.is_null() {
1059                        "HTTP/1.1 202 Accepted\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
1060                            .to_owned()
1061                    } else {
1062                        let b = json!({ "jsonrpc": "2.0", "id": id, "result": { "content": [{ "type": "text", "text": "{\"items\":[{\"path\":\"a.md\"}]}" }] } }).to_string();
1063                        let sse = format!(
1064                            ": keep-alive\n\nevent: message\ndata: {{\"jsonrpc\":\"2.0\",\"method\":\"notifications/message\",\"params\":{{}}}}\n\nid: 7\nevent: message\ndata: {b}\n\n"
1065                        );
1066                        format!(
1067                            "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{sse}",
1068                            sse.len()
1069                        )
1070                    };
1071                    s.write_all(reply.as_bytes()).unwrap();
1072                }
1073                seen
1074            });
1075            (url, h)
1076        }
1077
1078        #[test]
1079        fn streamable_http_round_trips_session_headers_json_and_sse() {
1080            let (url, server) = serve(4);
1081            let spec = EngineSpec::Http {
1082                url: url.clone(),
1083                headers: vec![("X-Auth".to_owned(), "t0k".to_owned())],
1084            };
1085            let mut c = McpEngineClient::connect_with(&spec, Duration::from_secs(10)).unwrap();
1086            assert_eq!(c.server_info()["name"], "tcp");
1087            let out = c.call("docs_list", json!({})).unwrap();
1088            assert_eq!(out["items"][0]["path"], "a.md");
1089            c.close().unwrap();
1090            let seen = server.join().unwrap();
1091            assert_eq!(seen.len(), 4, "initialize, initialized, tools/call, DELETE");
1092            let lower: Vec<String> = seen.iter().map(|s| s.to_ascii_lowercase()).collect();
1093            assert!(
1094                lower[0].starts_with("post /k/secret/mcp "),
1095                "the url is used verbatim"
1096            );
1097            assert!(lower[0].contains("accept: application/json, text/event-stream"));
1098            assert!(
1099                lower[0].contains("x-auth: t0k"),
1100                "extra headers on every request"
1101            );
1102            assert!(
1103                !lower[0].contains("mcp-session-id"),
1104                "no session before initialize"
1105            );
1106            assert!(
1107                lower[1].contains("mcp-session-id: s1"),
1108                "the session id rides after"
1109            );
1110            assert!(lower[1].contains("mcp-protocol-version: 2025-03-26"));
1111            assert!(
1112                lower[2].contains("\"method\":\"tools/call\"") && lower[2].contains("x-auth: t0k")
1113            );
1114            assert!(lower[3].starts_with("delete ") && lower[3].contains("mcp-session-id: s1"));
1115        }
1116
1117        #[test]
1118        fn http_errors_are_reported() {
1119            let listener = TcpListener::bind("127.0.0.1:0").unwrap();
1120            let url = format!("http://{}/mcp", listener.local_addr().unwrap());
1121            std::thread::spawn(move || {
1122                let (mut s, _) = listener.accept().unwrap();
1123                let mut tmp = [0u8; 4096];
1124                let _ = s.read(&mut tmp);
1125                s.write_all(b"HTTP/1.1 401 Unauthorized\r\nContent-Type: text/plain\r\nContent-Length: 6\r\nConnection: close\r\n\r\nnot ok").unwrap();
1126            });
1127            let spec = EngineSpec::Http {
1128                url,
1129                headers: vec![],
1130            };
1131            let err = McpEngineClient::connect_with(&spec, Duration::from_secs(5))
1132                .unwrap_err()
1133                .to_string();
1134            assert!(err.contains("HTTP 401") && err.contains("not ok"), "{err}");
1135        }
1136
1137        #[test]
1138        fn sse_parsing_finds_the_reply() {
1139            let stream = ": hello\n\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"n\"}\n\ndata: {\"jsonrpc\":\"2.0\",\ndata: \"id\":3,\"result\":{}}\n\n";
1140            let found = sse_find(stream.as_bytes(), &json!(3)).unwrap().unwrap();
1141            assert_eq!(found["result"], json!({}));
1142            assert!(sse_find(stream.as_bytes(), &json!(4)).unwrap().is_none());
1143            // No trailing blank line: the final event still counts.
1144            let tail = "data: {\"id\":1,\"result\":1}";
1145            assert_eq!(
1146                sse_find(tail.as_bytes(), &json!(1)).unwrap().unwrap()["result"],
1147                1
1148            );
1149        }
1150    }
1151}