haematite-cli 0.21.3

//! A request-driven stdio bridge to the public HTTP MCP, using a Lys agent pass.

use std::io::{self, BufRead, Write};
use std::process::ExitCode;

use serde_json::{Value, json};
use ureq::http::{HeaderValue, header};

use super::remote_cli::Options;
use super::remote_wire;

const MESSAGE_BYTES: usize = 16 * 1024 * 1024;

/// Selects remote HTTP when explicitly requested, or the native service transport.
pub fn dispatch(args: impl Iterator<Item = String>) -> ExitCode {
    let args: Vec<_> = args.collect();
    if args.iter().any(|arg| arg == "--remote") {
        run(args.into_iter())
    } else {
        super::stdio::run(args.into_iter())
    }
}

fn run(args: impl Iterator<Item = String>) -> ExitCode {
    let options = match Options::parse(args) {
        Ok(options) => options,
        Err(error) => {
            eprintln!("{error}");
            return ExitCode::from(78);
        }
    };
    let mut bridge = Bridge::new(options);
    match bridge.relay(io::stdin().lock(), &mut io::stdout().lock()) {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("haematite_mcp_remote_io_failed: {}", error.kind());
            ExitCode::FAILURE
        }
    }
}

struct Bridge {
    options: Options,
    client: ureq::Agent,
    protocol: Option<String>,
    session: Option<HeaderValue>,
}

impl Bridge {
    fn new(options: Options) -> Self {
        let config = ureq::Agent::config_builder()
            .proxy(None)
            .http_status_as_error(false)
            .max_redirects(0)
            .max_redirects_will_error(false)
            .max_idle_connections(0)
            .timeout_global(Some(options.timeout))
            .timeout_resolve(Some(options.timeout))
            .build();
        Self {
            options,
            client: ureq::Agent::new_with_config(config),
            protocol: None,
            session: None,
        }
    }

    fn relay(&mut self, mut input: impl BufRead, output: &mut impl Write) -> io::Result<()> {
        while let Some(line) = super::stdio::read_line_bounded(&mut input, MESSAGE_BYTES)? {
            let request: Value = match serde_json::from_slice(&line) {
                Ok(Value::Object(object)) if line.len() <= MESSAGE_BYTES => Value::Object(object),
                _ => {
                    remote_wire::write(
                        output,
                        &json!({"jsonrpc":"2.0", "id":null,
                        "error":{"code":-32700,"message":"haematite_mcp_remote_request_invalid"}}),
                    )?;
                    continue;
                }
            };
            let initialize_id = (request["method"] == "initialize").then(|| request["id"].clone());
            match self.post(&line) {
                Ok(mut response) => {
                    if let Some(session) = response.headers().get("mcp-session-id") {
                        let mut session = session.clone();
                        session.set_sensitive(true);
                        self.session = Some(session);
                    }
                    let outcome = remote_wire::response(&mut response, |message| {
                        if initialize_id.as_ref() == message.get("id")
                            && let Some(version) = message["result"]["protocolVersion"].as_str()
                        {
                            self.protocol = Some(version.to_owned());
                        }
                        remote_wire::write(output, &message)
                    });
                    match outcome {
                        Ok(()) => {}
                        Err(remote_wire::Error::Output(error)) => return Err(error),
                        Err(remote_wire::Error::Response(code)) => {
                            Self::refused(output, &request, &code)?;
                        }
                    }
                }
                Err(code) => Self::refused(output, &request, &code)?,
            }
        }
        Ok(())
    }

    fn post(&mut self, line: &[u8]) -> Result<ureq::http::Response<ureq::Body>, String> {
        let mut response = self.send(line, false)?;
        if response.status().as_u16() == 401 {
            // Only a definitive authentication refusal permits repeating a request.
            drop(response);
            response = self.send(line, true)?;
        }
        Ok(response)
    }

    fn send(
        &mut self,
        line: &[u8],
        fresh: bool,
    ) -> Result<ureq::http::Response<ureq::Body>, String> {
        let pass = if fresh {
            self.options.agent.ask_now()
        } else {
            self.options.agent.current_until_expiry()
        }
        .map_err(|refusal| refusal.code)?;
        let mut bearer = HeaderValue::from_str(&format!("Bearer {pass}"))
            .map_err(|_| "haematite_mcp_remote_pass_invalid".to_owned())?;
        bearer.set_sensitive(true);
        let mut request = self
            .client
            .post(&self.options.url)
            .header(header::AUTHORIZATION, bearer)
            .header(header::CONTENT_TYPE, "application/json")
            .header(header::ACCEPT, "application/json, text/event-stream");
        if let Some(version) = &self.protocol {
            request = request.header("MCP-Protocol-Version", version);
        }
        if let Some(session) = &self.session {
            request = request.header("Mcp-Session-Id", session.clone());
        }
        request
            .send(line)
            .map_err(|_| "haematite_mcp_remote_request_failed".to_owned())
    }

    fn refused(output: &mut impl Write, request: &Value, code: &str) -> io::Result<()> {
        request.get("id").map_or_else(
            || {
                // Notifications have no JSON-RPC answer, but a failure must be visible.
                eprintln!("haematite_mcp_remote_notification_refused: {code}");
                Ok(())
            },
            |id| {
                remote_wire::write(
                    output,
                    &json!({"jsonrpc":"2.0", "id":id,
            "error":{"code":-32000,"message":"Remote MCP request refused", "data":{"code":code}}}),
                )
            },
        )
    }
}