haematite-cli 0.21.8

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

use std::io;
use std::process::ExitCode;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, mpsc};

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

use super::lys_pass::LysAgent;
use super::remote_cli::Options;
use super::remote_http::{Client, OUTCOME_UNKNOWN, REQUEST_FAILED, UNKNOWN_MESSAGE};
use super::{remote_input, remote_wire};

pub(super) 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);
        }
    };
    match run_stdio(&Bridge::new(options)) {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("haematite_mcp_remote_io_failed: {}", error.kind());
            ExitCode::FAILURE
        }
    }
}

fn run_stdio(bridge: &Bridge) -> io::Result<()> {
    let output = Output(Arc::new(Mutex::new(io::stdout())));
    let active = Arc::new(AtomicBool::new(false));
    let (input, reader) = remote_input::spawn(bridge.clone(), output.clone(), Arc::clone(&active))?;
    bridge.relay(input, &output, &active)?;
    reader
        .join()
        .map_err(|_| io::Error::other("MCP input thread panicked"))?
}

#[derive(Clone)]
pub(super) struct Output(Arc<Mutex<io::Stdout>>);

impl Output {
    fn write(&self, message: &Value) -> io::Result<()> {
        let mut output = self
            .0
            .lock()
            .map_err(|_| io::Error::other("MCP output lock poisoned"))?;
        let result = remote_wire::write(&mut *output, message);
        drop(output);
        result
    }
}

#[derive(Clone, Default)]
pub(super) struct State {
    protocol: Option<String>,
    session: Option<HeaderValue>,
}

#[derive(Clone)]
pub(super) struct Bridge {
    url: Arc<str>,
    agent: Arc<Mutex<LysAgent>>,
    state: Arc<Mutex<State>>,
    client: Client,
}

impl Bridge {
    fn new(options: Options) -> Self {
        Self {
            client: Client::new(options.timeout),
            url: options.url.into(),
            agent: Arc::new(Mutex::new(options.agent)),
            state: Arc::new(Mutex::new(State::default())),
        }
    }

    fn relay(
        &self,
        input: mpsc::Receiver<io::Result<Vec<u8>>>,
        output: &Output,
        active: &AtomicBool,
    ) -> io::Result<()> {
        for line in input {
            let line = line?;
            let state = self.snapshot()?;
            active.store(true, Ordering::Release);
            let result = self.handle(&line, output, &state);
            active.store(false, Ordering::Release);
            result?;
        }
        Ok(())
    }

    pub(super) fn snapshot(&self) -> io::Result<State> {
        let state = self
            .state
            .lock()
            .map_err(|_| io::Error::other("MCP state lock poisoned"))?;
        let snapshot = state.clone();
        drop(state);
        Ok(snapshot)
    }

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

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

    fn send(
        &self,
        line: &[u8],
        state: &State,
        fresh: bool,
    ) -> Result<ureq::http::Response<ureq::Body>, String> {
        let mut agent = self.agent.lock().map_err(|_| REQUEST_FAILED.to_owned())?;
        // A pass Lys refused, or could not be asked for, is said by its own
        // name: nothing was sent, and the person needs to know which.
        let pass = if fresh {
            agent.ask_now()
        } else {
            agent.current_until_expiry()
        }
        .map_err(|refusal| refusal.code)?
        .to_owned();
        drop(agent);
        let mut bearer = HeaderValue::from_str(&format!("Bearer {pass}"))
            .map_err(|_| "haematite_mcp_remote_pass_invalid".to_owned())?;
        bearer.set_sensitive(true);
        self.client.send(
            &self.url,
            bearer,
            state.protocol.as_deref(),
            state.session.clone(),
            line,
        )
    }

    fn refused(output: &Output, request: &Value, code: &str) -> io::Result<()> {
        request.get("id").map_or_else(
            || {
                eprintln!("haematite_mcp_remote_notification_refused: {code}");
                Ok(())
            },
            |id| {
                let message = if code == OUTCOME_UNKNOWN {
                    UNKNOWN_MESSAGE
                } else {
                    "Remote MCP request refused"
                };
                output.write(&json!({"jsonrpc":"2.0", "id":id,
                    "error":{"code":-32000,"message":message,"data":{"code":code}}}))
            },
        )
    }
}