Skip to main content

haematite_cli/mcp/
stdio.rs

1use std::fmt;
2use std::io::{self, BufRead, Write};
3use std::num::{NonZeroU32, NonZeroU64};
4use std::path::PathBuf;
5use std::process::ExitCode;
6use std::time::Duration;
7
8use serde::Deserialize;
9use serde_json::value::RawValue;
10
11use super::lys_pass::{LysAgent, PassSource, Proof};
12use super::{adapter, rpc};
13use crate::client::{self, Client, Reply};
14use crate::serve::{cli::Config, protocol};
15
16const USAGE: &str = "usage: haematite mcp (--socket <path> | --config <file.json>) [--timeout-ms N] [--pass-file <path> | --lys-agent <issuer origin> <agent id> ((--credential-file | --key-file) <absolute path> | --run-pass-env <variable>)]";
17const DEFAULT_TIMEOUT_MS: NonZeroU64 = NonZeroU64::MIN.saturating_add(29_999);
18
19#[derive(Debug)]
20pub enum RelayError {
21    Input(io::Error),
22    ReceiptOutput {
23        cause: io::Error,
24        response: Vec<u8>,
25    },
26    ResultEncoding {
27        refusal: rpc::Refusal,
28        result: adapter::ToolResult,
29    },
30    Encoding(serde_json::Error),
31    Adapter(adapter::EncodeError),
32    Arguments(&'static str),
33    ConfigRead(io::Error),
34    ConfigDecode(serde_json::Error),
35    Runtime(crate::serve::runtime::Refusal),
36    SocketCap(client::Ended),
37    CapNotRepresentable,
38}
39
40impl fmt::Display for RelayError {
41    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
42        // Diagnostics name the stage without exposing request or record text.
43        let code = match self {
44            Self::Input(_) => "mcp_input_read_failed",
45            Self::ReceiptOutput { .. } => "mcp_output_write_failed",
46            Self::ResultEncoding { .. } | Self::Encoding(_) => "mcp_response_encoding_failed",
47            Self::Adapter(_) => "mcp_result_encoding_failed",
48            Self::Arguments(why) => {
49                return write!(formatter, "mcp_arguments_invalid: {why}; {USAGE}");
50            }
51            Self::ConfigRead(_) => "mcp_config_read_failed",
52            Self::ConfigDecode(_) => "mcp_config_invalid",
53            Self::Runtime(error) => return fmt::Display::fmt(error, formatter),
54            Self::SocketCap(_) => "mcp_socket_cap_unavailable",
55            Self::CapNotRepresentable => "mcp_frame_cap_unrepresentable",
56        };
57        formatter.write_str(code)
58    }
59}
60
61impl std::error::Error for RelayError {
62    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
63        match self {
64            Self::Input(error) | Self::ConfigRead(error) => Some(error),
65            Self::Encoding(error) | Self::ConfigDecode(error) => Some(error),
66            Self::Adapter(error) => Some(error),
67            Self::Runtime(error) => Some(error),
68            Self::ReceiptOutput { cause, .. } => Some(cause),
69            Self::ResultEncoding { .. }
70            | Self::Arguments(_)
71            | Self::SocketCap(_)
72            | Self::CapNotRepresentable => None,
73        }
74    }
75}
76
77/// Serves one input message at a time and flushes each complete response.
78///
79/// # Errors
80///
81/// Returns a named input, output or encoding failure. Output errors stop the
82/// loop immediately; no later request is read or sent.
83pub fn relay(
84    mut input: impl BufRead,
85    mut output: impl Write,
86    sender: &mut impl adapter::Sender,
87    cap: NonZeroU32,
88) -> Result<(), RelayError> {
89    let byte_limit = usize::try_from(cap.get()).map_err(|_| RelayError::CapNotRepresentable)?;
90    let codec = rpc::RpcCodec::default();
91    while let Some(line) = read_line_bounded(&mut input, byte_limit).map_err(RelayError::Input)? {
92        if let Some(mut response) = answer(&codec, sender, &line, byte_limit)? {
93            response.push(b'\n');
94            if let Err(cause) = output.write_all(&response).and_then(|()| output.flush()) {
95                return Err(RelayError::ReceiptOutput { cause, response });
96            }
97        }
98    }
99    Ok(())
100}
101
102pub(super) fn answer(
103    codec: &rpc::RpcCodec,
104    sender: &mut impl adapter::Sender,
105    line: &[u8],
106    byte_limit: usize,
107) -> Result<Option<Vec<u8>>, RelayError> {
108    let prepared = codec
109        .decode(line, byte_limit)
110        .and_then(|request| codec.prepare(request));
111    match prepared {
112        Err(refusal) => refusal.encode().map_err(RelayError::Encoding),
113        Ok(rpc::Prepared::NoResponse) => Ok(None),
114        Ok(rpc::Prepared::Response(response)) => Ok(Some(response)),
115        Ok(rpc::Prepared::ToolCall { id, call }) => {
116            let mut result = adapter::dispatch(sender, &call).map_err(RelayError::Adapter)?;
117            compact_result(&mut result)?;
118            match rpc::encode_result(&id, &result) {
119                Ok(response) => Ok(Some(response)),
120                Err(refusal) => Err(RelayError::ResultEncoding { refusal, result }),
121            }
122        }
123    }
124}
125
126fn compact_result(result: &mut adapter::ToolResult) -> Result<(), RelayError> {
127    if result.is_error {
128        return Ok(());
129    }
130    let mut body = result.structured_content.get().as_bytes().to_vec();
131    compact_line(&mut body);
132    result.structured_content = serde_json::from_slice(&body).map_err(RelayError::Encoding)?;
133    let [adapter::TextContent::Text { text }] = &mut result.content;
134    text.clear();
135    text.push_str(result.structured_content.get());
136    Ok(())
137}
138
139pub(super) fn read_line_bounded(
140    reader: &mut impl BufRead,
141    cap: usize,
142) -> io::Result<Option<Vec<u8>>> {
143    let mut line = Vec::new();
144    let retained = cap.saturating_add(1);
145    let mut consumed_any = false;
146    loop {
147        let available = reader.fill_buf()?;
148        if available.is_empty() {
149            return Ok(consumed_any.then_some(line));
150        }
151        consumed_any = true;
152        let newline = available.iter().position(|byte| *byte == b'\n');
153        let before = newline.unwrap_or(available.len());
154        let take = before.min(retained.saturating_sub(line.len()));
155        line.extend_from_slice(&available[..take]);
156        reader.consume(before + usize::from(newline.is_some()));
157        if newline.is_some() {
158            return Ok(Some(line));
159        }
160    }
161}
162
163fn compact_line(frame: &mut Vec<u8>) {
164    let mut in_string = false;
165    let mut escaped = false;
166    frame.retain(|byte| {
167        if in_string {
168            if escaped {
169                escaped = false;
170            } else if *byte == b'\\' {
171                escaped = true;
172            } else if *byte == b'"' {
173                in_string = false;
174            }
175            true
176        } else if *byte == b'"' {
177            in_string = true;
178            true
179        } else {
180            !matches!(*byte, b' ' | b'\t' | b'\r' | b'\n')
181        }
182    });
183}
184
185/// Runs the child-process transport. Standard output contains protocol only.
186pub fn run(args: impl Iterator<Item = String>) -> ExitCode {
187    match run_inner(args) {
188        Ok(()) => ExitCode::SUCCESS,
189        Err(error) => {
190            eprintln!("haematite mcp: {error}");
191            ExitCode::FAILURE
192        }
193    }
194}
195
196fn run_inner(args: impl Iterator<Item = String>) -> Result<(), RelayError> {
197    let (mut sender, cap) = Options::parse(args)?.sender()?;
198    relay(io::stdin().lock(), io::stdout().lock(), &mut sender, cap)
199}
200
201enum Target {
202    Socket(PathBuf),
203    Config(PathBuf),
204}
205
206struct Options {
207    target: Target,
208    timeout: Duration,
209    pass: Option<PassSource>,
210}
211
212/// The value after a flag: present, not empty and not itself a flag.
213fn flag_value(
214    args: &mut impl Iterator<Item = String>,
215    why: &'static str,
216) -> Result<String, RelayError> {
217    args.next()
218        .filter(|value| !value.is_empty() && !value.starts_with("--"))
219        .ok_or(RelayError::Arguments(why))
220}
221
222impl Options {
223    fn parse(mut args: impl Iterator<Item = String>) -> Result<Self, RelayError> {
224        let mut target = None;
225        let mut timeout_ms = None;
226        let mut pass_file = None;
227        let mut lys_agent = None;
228        let mut credential_file = None;
229        while let Some(arg) = args.next() {
230            match arg.as_str() {
231                "--socket" | "--config" if target.is_none() => {
232                    let path = flag_value(&mut args, "target requires a path")?;
233                    let path = PathBuf::from(path);
234                    target = Some(if arg == "--socket" {
235                        Target::Socket(path)
236                    } else {
237                        Target::Config(path)
238                    });
239                }
240                "--timeout-ms" if timeout_ms.is_none() => {
241                    let value = args
242                        .next()
243                        .ok_or(RelayError::Arguments("timeout requires milliseconds"))?;
244                    timeout_ms = Some(value.parse::<NonZeroU64>().map_err(|_| {
245                        RelayError::Arguments("timeout requires positive integer milliseconds")
246                    })?);
247                }
248                "--pass-file" if pass_file.is_none() => {
249                    let path = flag_value(&mut args, "pass file requires a path")?;
250                    pass_file = Some(PathBuf::from(path));
251                }
252                "--lys-agent" if lys_agent.is_none() => {
253                    let words = (args.next(), args.next());
254                    let (Some(issuer), Some(agent)) = words else {
255                        return Err(RelayError::Arguments(
256                            "lys agent requires an issuer origin and an agent id",
257                        ));
258                    };
259                    if issuer.starts_with("--") || agent.starts_with("--") {
260                        return Err(RelayError::Arguments(
261                            "lys agent requires an issuer origin and an agent id",
262                        ));
263                    }
264                    lys_agent = Some((issuer, agent));
265                }
266                "--credential-file" if credential_file.is_none() => {
267                    let path = flag_value(&mut args, "credential file requires a path")?;
268                    credential_file = Some(Proof::Grant(PathBuf::from(path)));
269                }
270                "--key-file" if credential_file.is_none() => {
271                    let path = flag_value(&mut args, "key file requires a path")?;
272                    credential_file = Some(Proof::Key(PathBuf::from(path)));
273                }
274                "--run-pass-env" if credential_file.is_none() => {
275                    let name = flag_value(&mut args, "run pass requires a variable name")?;
276                    credential_file = Some(Proof::RunPass(name));
277                }
278                _ => return Err(RelayError::Arguments("unknown or repeated argument")),
279            }
280        }
281        let target = target.ok_or(RelayError::Arguments(
282            "give exactly one of socket and config",
283        ))?;
284        let timeout = Duration::from_millis(timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS).get());
285        let pass = match (pass_file, lys_agent, credential_file) {
286            (None, None, None) => None,
287            (Some(path), None, None) => Some(PassSource::File(path)),
288            (None, Some((issuer, agent)), Some(credential)) => Some(PassSource::Lys(
289                LysAgent::new(&issuer, &agent, credential, timeout)
290                    .map_err(RelayError::Arguments)?,
291            )),
292            (Some(_), Some(_), _) | (Some(_), None, Some(_)) => {
293                return Err(RelayError::Arguments(
294                    "give at most one of pass file and lys agent",
295                ));
296            }
297            (None, Some(_), None) | (None, None, Some(_)) => {
298                return Err(RelayError::Arguments(
299                    "lys agent and a credential or key file are given together",
300                ));
301            }
302        };
303        Ok(Self {
304            target,
305            timeout,
306            pass,
307        })
308    }
309
310    fn sender(self) -> Result<(NativeSender, NonZeroU32), RelayError> {
311        let (socket, first, cap) = match self.target {
312            Target::Config(path) => {
313                let bytes = std::fs::read(path).map_err(RelayError::ConfigRead)?;
314                let config: Config =
315                    serde_json::from_slice(&bytes).map_err(RelayError::ConfigDecode)?;
316                let endpoint = config.endpoint(false).map_err(RelayError::Runtime)?;
317                (endpoint.socket, None, config.max_frame)
318            }
319            Target::Socket(socket) => {
320                let first =
321                    Client::connect(&socket, self.timeout).map_err(RelayError::SocketCap)?;
322                let cap = first.max_frame();
323                (socket, Some(first), cap)
324            }
325        };
326        Ok((
327            NativeSender {
328                socket,
329                timeout: self.timeout,
330                first,
331                first_presented: false,
332                pass: self.pass,
333            },
334            cap,
335        ))
336    }
337}
338
339pub(super) struct NativeSender {
340    socket: PathBuf,
341    timeout: Duration,
342    first: Option<Client>,
343    /// The first connection already presented the pass, so the first call
344    /// on it does not present it again.
345    first_presented: bool,
346    /// Where the Lys pass presented as `session.pass` on every connection,
347    /// before the tool's verb, comes from: a file read afresh each time, or
348    /// Lys's agent pass route, asked only when the held pass is near expiry.
349    pass: Option<PassSource>,
350}
351
352impl NativeSender {
353    #[cfg(test)]
354    pub(super) const fn new(socket: PathBuf, timeout: Duration) -> Self {
355        Self {
356            socket,
357            timeout,
358            first: None,
359            first_presented: false,
360            pass: None,
361        }
362    }
363
364    /// A sender for the HTTP MCP: `first` is a connection on which the
365    /// caller's own `pass` was already presented and accepted, and every
366    /// further connection presents the same pass before its verb.
367    pub(super) const fn presented(
368        socket: PathBuf,
369        timeout: Duration,
370        first: Client,
371        pass: String,
372    ) -> Self {
373        Self {
374            socket,
375            timeout,
376            first: Some(first),
377            first_presented: true,
378            pass: Some(PassSource::Held(pass)),
379        }
380    }
381}
382
383#[derive(Deserialize)]
384struct WireOperation {
385    method: String,
386    params: Option<Box<RawValue>>,
387}
388
389impl adapter::Sender for NativeSender {
390    fn send(
391        &mut self,
392        request: protocol::Operation,
393    ) -> Result<Box<RawValue>, adapter::ServiceRefusal> {
394        let encode_failure = |error: serde_json::Error| adapter::ServiceRefusal {
395            code: "mcp_client_request_encoding_failed".into(),
396            message: error.to_string(),
397            receipt: None,
398        };
399        let raw = serde_json::value::to_raw_value(&request).map_err(encode_failure)?;
400        let operation: WireOperation = serde_json::from_str(raw.get()).map_err(encode_failure)?;
401        let (mut client, presented) = match self.first.take() {
402            Some(first) => (first, std::mem::take(&mut self.first_presented)),
403            None => (
404                Client::connect(&self.socket, self.timeout).map_err(client_refusal)?,
405                false,
406            ),
407        };
408        if !presented && let Some(source) = &mut self.pass {
409            source.present(|pass| present_pass(&mut client, pass))?;
410        }
411        match client
412            .call(&operation.method, operation.params.as_deref())
413            .map_err(client_refusal)?
414        {
415            Reply::Answered(raw) => Ok(raw),
416            Reply::Refused(refusal) => Err(adapter::ServiceRefusal {
417                code: refusal.code,
418                message: refusal.message,
419                receipt: refusal
420                    .verified
421                    .as_ref()
422                    .map(|receipt| receipt.raw().to_owned()),
423            }),
424        }
425    }
426}
427
428/// Presents `pass` on `client`; a pass the store refuses ends the call
429/// there (after the one renewal a Lys source makes), and the verb is never
430/// sent.
431pub(super) fn present_pass(client: &mut Client, pass: &str) -> Result<(), adapter::ServiceRefusal> {
432    let pass =
433        serde_json::value::to_raw_value(&serde_json::json!({ "pass": pass })).map_err(|error| {
434            adapter::ServiceRefusal {
435                code: "mcp_pass_unreadable".into(),
436                message: format!("the pass could not be written: {error}"),
437                receipt: None,
438            }
439        })?;
440    match client
441        .call("session.pass", Some(&pass))
442        .map_err(client_refusal)?
443    {
444        Reply::Answered(_) => Ok(()),
445        Reply::Refused(refusal) => Err(adapter::ServiceRefusal {
446            code: refusal.code,
447            message: refusal.message,
448            receipt: None,
449        }),
450    }
451}
452
453fn client_refusal(ended: client::Ended) -> adapter::ServiceRefusal {
454    let (code, message, receipt) = match ended {
455        client::Ended::Usage(message) => ("mcp_client_usage", message, None),
456        client::Ended::NotAsked(message) => ("mcp_client_not_asked", message, None),
457        client::Ended::Unknown(message) => ("mcp_client_outcome_unknown", message, None),
458        client::Ended::Received { why, receipt } => (
459            "mcp_client_complete_receipt",
460            why,
461            Some(receipt.raw().to_owned()),
462        ),
463    };
464    adapter::ServiceRefusal {
465        code: code.into(),
466        message,
467        receipt,
468    }
469}
470
471#[cfg(test)]
472#[path = "stdio_tests.rs"]
473mod tests;
474
475#[cfg(test)]
476mod runtime_endpoint_tests {
477    use super::{Options, Target};
478    use crate::serve::cli::endpoint_tests::with_fixture;
479    use std::time::Duration;
480
481    #[test]
482    fn stdio_configuration_uses_the_services_default_endpoint()
483    -> Result<(), Box<dyn std::error::Error>> {
484        with_fixture(|fixture| {
485            let expected = fixture.config()?.endpoint(false)?.socket;
486            let options = Options {
487                target: Target::Config(fixture.config_path.clone()),
488                timeout: Duration::from_secs(1),
489                pass: None,
490            };
491            let (sender, cap) = options.sender()?;
492            assert_eq!(sender.socket, expected);
493            assert_eq!(cap.get(), 65_536);
494            assert!(sender.first.is_none());
495            Ok(())
496        })
497    }
498}