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 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
77pub 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
185pub 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
212fn 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 first_presented: bool,
346 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 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
428pub(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}