haematite_cli/mcp/
stdio.rs

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

haematite_cli/mcp/
stdio.rs
