1use 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
40pub const PROTOCOL_VERSION: &str = "2025-03-26";
44
45pub const CLIENT_NAME: &str = "omgbase-sync";
47
48pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60);
51
52const EXIT_GRACE: Duration = Duration::from_secs(5);
55
56#[derive(Clone, Debug, PartialEq, Eq)]
60pub enum EngineSpec {
61 Http {
64 url: String,
65 headers: Vec<(String, String)>,
66 },
67 Stdio { command: String, args: Vec<String> },
69}
70
71#[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
78pub 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#[derive(Clone, Debug, PartialEq)]
113pub struct ToolResult {
114 pub body: Value,
115 pub is_error: bool,
116}
117
118impl ToolResult {
119 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
149enum Transport {
152 Stdio(StdioTransport),
153 #[cfg(feature = "http")]
154 Http(HttpTransport),
155}
156
157pub 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 pub fn connect(spec: &EngineSpec) -> Result<Self> {
178 Self::connect_with(spec, DEFAULT_TIMEOUT)
179 }
180
181 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 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 #[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 #[must_use]
249 pub fn protocol_version(&self) -> &str {
250 &self.protocol_version
251 }
252
253 #[must_use]
255 pub fn server_info(&self) -> &Value {
256 &self.server_info
257 }
258
259 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 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 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 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 pub fn call(&mut self, name: &str, args: Value) -> Result<Value> {
346 self.call_tool(name, args)?.into_result(name)
347 }
348
349 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
369fn 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
386fn 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
421fn 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 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
529struct 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 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#[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 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 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#[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 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 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 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 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 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}