haematite-cli 0.21.8

use std::io::{self, BufRead, Write};
use std::num::{NonZeroU32, NonZeroU64};
use std::path::PathBuf;
use std::process::ExitCode;
use std::time::Duration;

use serde::Deserialize;
use serde_json::value::RawValue;

use super::lys_pass::{LysAgent, PassSource, Proof};
use super::{adapter, rpc};
use crate::client::{self, Client, Reply};
use crate::serve::{cli::Config, protocol};

const 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>)]";
const DEFAULT_TIMEOUT_MS: NonZeroU64 = NonZeroU64::MIN.saturating_add(29_999);

#[path = "stdio_error.rs"]
mod error;
pub use error::RelayError;

/// Serves one input message at a time and flushes each complete response.
///
/// # Errors
///
/// Returns a named input, output or encoding failure. Output errors stop the
/// loop immediately; no later request is read or sent.
pub fn relay(
    mut input: impl BufRead,
    mut output: impl Write,
    sender: &mut impl adapter::Sender,
    cap: NonZeroU32,
) -> Result<(), RelayError> {
    let byte_limit = usize::try_from(cap.get()).map_err(|_| RelayError::CapNotRepresentable)?;
    let codec = rpc::RpcCodec::default();
    while let Some(line) = read_line_bounded(&mut input, byte_limit).map_err(RelayError::Input)? {
        if let Some(mut response) = answer(&codec, sender, &line, byte_limit)? {
            response.push(b'\n');
            if let Err(cause) = output.write_all(&response).and_then(|()| output.flush()) {
                return Err(RelayError::ReceiptOutput { cause, response });
            }
        }
    }
    Ok(())
}

pub(super) fn answer(
    codec: &rpc::RpcCodec,
    sender: &mut impl adapter::Sender,
    line: &[u8],
    byte_limit: usize,
) -> Result<Option<Vec<u8>>, RelayError> {
    let prepared = codec
        .decode(line, byte_limit)
        .and_then(|request| codec.prepare(request));
    match prepared {
        Err(refusal) => refusal.encode().map_err(RelayError::Encoding),
        Ok(rpc::Prepared::NoResponse) => Ok(None),
        Ok(rpc::Prepared::Response(response)) => Ok(Some(response)),
        Ok(rpc::Prepared::ToolCall { id, call }) => {
            let mut result = adapter::dispatch(sender, &call).map_err(RelayError::Adapter)?;
            compact_result(&mut result)?;
            match rpc::encode_result(&id, &result) {
                Ok(response) => Ok(Some(response)),
                Err(refusal) => Err(RelayError::ResultEncoding { refusal, result }),
            }
        }
    }
}

fn compact_result(result: &mut adapter::ToolResult) -> Result<(), RelayError> {
    if result.is_error {
        return Ok(());
    }
    let mut body = result.structured_content.get().as_bytes().to_vec();
    compact_line(&mut body);
    result.structured_content = serde_json::from_slice(&body).map_err(RelayError::Encoding)?;
    let [adapter::TextContent::Text { text }] = &mut result.content;
    text.clear();
    text.push_str(result.structured_content.get());
    Ok(())
}

pub(super) fn read_line_bounded(
    reader: &mut impl BufRead,
    cap: usize,
) -> io::Result<Option<Vec<u8>>> {
    let mut line = Vec::new();
    let retained = cap.saturating_add(1);
    let mut consumed_any = false;
    loop {
        let available = reader.fill_buf()?;
        if available.is_empty() {
            return Ok(consumed_any.then_some(line));
        }
        consumed_any = true;
        let newline = available.iter().position(|byte| *byte == b'\n');
        let before = newline.unwrap_or(available.len());
        let take = before.min(retained.saturating_sub(line.len()));
        line.extend_from_slice(&available[..take]);
        reader.consume(before + usize::from(newline.is_some()));
        if newline.is_some() {
            return Ok(Some(line));
        }
    }
}

fn compact_line(frame: &mut Vec<u8>) {
    let mut in_string = false;
    let mut escaped = false;
    frame.retain(|byte| {
        if in_string {
            if escaped {
                escaped = false;
            } else if *byte == b'\\' {
                escaped = true;
            } else if *byte == b'"' {
                in_string = false;
            }
            true
        } else if *byte == b'"' {
            in_string = true;
            true
        } else {
            !matches!(*byte, b' ' | b'\t' | b'\r' | b'\n')
        }
    });
}

/// Runs the child-process transport. Standard output contains protocol only.
pub fn run(args: impl Iterator<Item = String>) -> ExitCode {
    match run_inner(args) {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("haematite mcp: {error}");
            ExitCode::FAILURE
        }
    }
}

fn run_inner(args: impl Iterator<Item = String>) -> Result<(), RelayError> {
    let (mut sender, cap) = Options::parse(args)?.sender()?;
    relay(io::stdin().lock(), io::stdout().lock(), &mut sender, cap)
}

enum Target {
    Socket(PathBuf),
    Config(PathBuf),
}

struct Options {
    target: Target,
    timeout: Duration,
    pass: Option<PassSource>,
}

/// The value after a flag: present, not empty and not itself a flag.
fn flag_value(
    args: &mut impl Iterator<Item = String>,
    why: &'static str,
) -> Result<String, RelayError> {
    args.next()
        .filter(|value| !value.is_empty() && !value.starts_with("--"))
        .ok_or(RelayError::Arguments(why))
}

impl Options {
    fn parse(mut args: impl Iterator<Item = String>) -> Result<Self, RelayError> {
        let mut target = None;
        let mut timeout_ms = None;
        let mut pass_file = None;
        let mut lys_agent = None;
        let mut credential_file = None;
        while let Some(arg) = args.next() {
            match arg.as_str() {
                "--socket" | "--config" if target.is_none() => {
                    let path = flag_value(&mut args, "target requires a path")?;
                    let path = PathBuf::from(path);
                    target = Some(if arg == "--socket" {
                        Target::Socket(path)
                    } else {
                        Target::Config(path)
                    });
                }
                "--timeout-ms" if timeout_ms.is_none() => {
                    let value = args
                        .next()
                        .ok_or(RelayError::Arguments("timeout requires milliseconds"))?;
                    timeout_ms = Some(value.parse::<NonZeroU64>().map_err(|_| {
                        RelayError::Arguments("timeout requires positive integer milliseconds")
                    })?);
                }
                "--pass-file" if pass_file.is_none() => {
                    let path = flag_value(&mut args, "pass file requires a path")?;
                    pass_file = Some(PathBuf::from(path));
                }
                "--lys-agent" if lys_agent.is_none() => {
                    let words = (args.next(), args.next());
                    let (Some(issuer), Some(agent)) = words else {
                        return Err(RelayError::Arguments(
                            "lys agent requires an issuer origin and an agent id",
                        ));
                    };
                    if issuer.starts_with("--") || agent.starts_with("--") {
                        return Err(RelayError::Arguments(
                            "lys agent requires an issuer origin and an agent id",
                        ));
                    }
                    lys_agent = Some((issuer, agent));
                }
                "--credential-file" if credential_file.is_none() => {
                    let path = flag_value(&mut args, "credential file requires a path")?;
                    credential_file = Some(Proof::Grant(PathBuf::from(path)));
                }
                "--key-file" if credential_file.is_none() => {
                    let path = flag_value(&mut args, "key file requires a path")?;
                    credential_file = Some(Proof::Key(PathBuf::from(path)));
                }
                "--launch-proof-file" if credential_file.is_none() => {
                    let path = flag_value(&mut args, "launch proof file requires a path")?;
                    credential_file = Some(Proof::Launch(PathBuf::from(path)));
                }
                "--run-pass-env" if credential_file.is_none() => {
                    let name = flag_value(&mut args, "run pass requires a variable name")?;
                    credential_file = Some(Proof::RunPass(name));
                }
                _ => return Err(RelayError::Arguments("unknown or repeated argument")),
            }
        }
        let target = target.ok_or(RelayError::Arguments(
            "give exactly one of socket and config",
        ))?;
        let timeout = Duration::from_millis(timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS).get());
        let pass = match (pass_file, lys_agent, credential_file) {
            (None, None, None) => None,
            (Some(path), None, None) => Some(PassSource::File(path)),
            (None, Some((issuer, agent)), Some(credential)) => Some(PassSource::Lys(
                LysAgent::new(&issuer, &agent, credential, timeout)
                    .map_err(RelayError::Arguments)?,
            )),
            (Some(_), Some(_), _) | (Some(_), None, Some(_)) => {
                return Err(RelayError::Arguments(
                    "give at most one of pass file and lys agent",
                ));
            }
            (None, Some(_), None) | (None, None, Some(_)) => {
                return Err(RelayError::Arguments(
                    "lys agent and exactly one proof source are given together",
                ));
            }
        };
        Ok(Self {
            target,
            timeout,
            pass,
        })
    }

    fn sender(self) -> Result<(NativeSender, NonZeroU32), RelayError> {
        let (socket, first, cap) = match self.target {
            Target::Config(path) => {
                let bytes = std::fs::read(path).map_err(RelayError::ConfigRead)?;
                let config: Config =
                    serde_json::from_slice(&bytes).map_err(RelayError::ConfigDecode)?;
                let endpoint = config.endpoint(false).map_err(RelayError::Runtime)?;
                (endpoint.socket, None, config.max_frame)
            }
            Target::Socket(socket) => {
                let first =
                    Client::connect(&socket, self.timeout).map_err(RelayError::SocketCap)?;
                let cap = first.max_frame();
                (socket, Some(first), cap)
            }
        };
        Ok((
            NativeSender {
                socket,
                timeout: self.timeout,
                first,
                first_presented: false,
                pass: self.pass,
            },
            cap,
        ))
    }
}

pub(super) struct NativeSender {
    socket: PathBuf,
    timeout: Duration,
    first: Option<Client>,
    /// The first connection already presented the pass, so the first call
    /// on it does not present it again.
    first_presented: bool,
    /// Where the Lys pass presented as `session.pass` on every connection,
    /// before the tool's verb, comes from: a file read afresh each time, or
    /// Lys's agent pass route, asked only when the held pass is near expiry.
    pass: Option<PassSource>,
}

impl NativeSender {
    #[cfg(test)]
    pub(super) const fn new(socket: PathBuf, timeout: Duration) -> Self {
        Self {
            socket,
            timeout,
            first: None,
            first_presented: false,
            pass: None,
        }
    }

    /// A sender for the HTTP MCP: `first` is a connection on which the
    /// caller's own `pass` was already presented and accepted, and every
    /// further connection presents the same pass before its verb.
    pub(super) const fn presented(
        socket: PathBuf,
        timeout: Duration,
        first: Client,
        pass: String,
    ) -> Self {
        Self {
            socket,
            timeout,
            first: Some(first),
            first_presented: true,
            pass: Some(PassSource::Held(pass)),
        }
    }
}

#[derive(Deserialize)]
struct WireOperation {
    method: String,
    params: Option<Box<RawValue>>,
}

impl adapter::Sender for NativeSender {
    fn send(
        &mut self,
        request: protocol::Operation,
    ) -> Result<Box<RawValue>, adapter::ServiceRefusal> {
        let encode_failure = |error: serde_json::Error| adapter::ServiceRefusal {
            code: "mcp_client_request_encoding_failed".into(),
            message: error.to_string(),
            receipt: None,
        };
        let raw = serde_json::value::to_raw_value(&request).map_err(encode_failure)?;
        let operation: WireOperation = serde_json::from_str(raw.get()).map_err(encode_failure)?;
        let (mut client, presented) = match self.first.take() {
            Some(first) => (first, std::mem::take(&mut self.first_presented)),
            None => (
                Client::connect(&self.socket, self.timeout).map_err(client_refusal)?,
                false,
            ),
        };
        if !presented && let Some(source) = &mut self.pass {
            source.present(|pass| present_pass(&mut client, pass))?;
        }
        match client
            .call(&operation.method, operation.params.as_deref())
            .map_err(client_refusal)?
        {
            Reply::Answered(raw) => Ok(raw),
            Reply::Refused(refusal) => Err(adapter::ServiceRefusal {
                code: refusal.code,
                message: refusal.message,
                receipt: refusal
                    .verified
                    .as_ref()
                    .map(|receipt| receipt.raw().to_owned()),
            }),
        }
    }
}

/// Presents `pass` on `client`; a pass the store refuses ends the call
/// there (after the one renewal a Lys source makes), and the verb is never
/// sent.
pub(super) fn present_pass(client: &mut Client, pass: &str) -> Result<(), adapter::ServiceRefusal> {
    let pass =
        serde_json::value::to_raw_value(&serde_json::json!({ "pass": pass })).map_err(|error| {
            adapter::ServiceRefusal {
                code: "mcp_pass_unreadable".into(),
                message: format!("the pass could not be written: {error}"),
                receipt: None,
            }
        })?;
    match client
        .call("session.pass", Some(&pass))
        .map_err(client_refusal)?
    {
        Reply::Answered(_) => Ok(()),
        Reply::Refused(refusal) => Err(adapter::ServiceRefusal {
            code: refusal.code,
            message: refusal.message,
            receipt: None,
        }),
    }
}

fn client_refusal(ended: client::Ended) -> adapter::ServiceRefusal {
    let (code, message, receipt) = match ended {
        client::Ended::Usage(message) => ("mcp_client_usage", message, None),
        client::Ended::NotAsked(message) => ("mcp_client_not_asked", message, None),
        client::Ended::Unknown(message) => ("mcp_client_outcome_unknown", message, None),
        client::Ended::Received { why, receipt } => (
            "mcp_client_complete_receipt",
            why,
            Some(receipt.raw().to_owned()),
        ),
    };
    adapter::ServiceRefusal {
        code: code.into(),
        message,
        receipt,
    }
}

#[cfg(test)]
#[path = "stdio_tests.rs"]
mod tests;

#[cfg(test)]
#[path = "launch_proof_tests.rs"]
mod launch_proof_tests;

#[cfg(test)]
mod runtime_endpoint_tests {
    use super::{Options, Target};
    use crate::serve::cli::endpoint_tests::with_fixture;
    use std::time::Duration;

    #[test]
    fn stdio_configuration_uses_the_services_default_endpoint()
    -> Result<(), Box<dyn std::error::Error>> {
        with_fixture(|fixture| {
            let expected = fixture.config()?.endpoint(false)?.socket;
            let options = Options {
                target: Target::Config(fixture.config_path.clone()),
                timeout: Duration::from_secs(1),
                pass: None,
            };
            let (sender, cap) = options.sender()?;
            assert_eq!(sender.socket, expected);
            assert_eq!(cap.get(), 65_536);
            assert!(sender.first.is_none());
            Ok(())
        })
    }
}