haematite-cli 0.21.7

use std::io::{self, BufRead, BufReader, Write};

use serde_json::Value;

use super::remote_http::OUTCOME_UNKNOWN;

const ANSWER_BYTES: u64 = 16 * 1024 * 1024;

pub(super) enum Error {
    Response(String),
    Output(io::Error),
}

pub(super) fn write(output: &mut impl Write, message: &Value) -> io::Result<()> {
    serde_json::to_writer(&mut *output, message).map_err(io::Error::other)?;
    output.write_all(b"\n")?;
    output.flush()
}

fn unknown() -> Error {
    Error::Response(OUTCOME_UNKNOWN.into())
}

fn invalid() -> Error {
    Error::Response("haematite_mcp_remote_response_invalid".into())
}

pub(super) fn response(
    response: &mut ureq::http::Response<ureq::Body>,
    mut emit: impl FnMut(Value) -> io::Result<()>,
) -> Result<(), Error> {
    let status = response.status().as_u16();
    if status == 202 {
        return Ok(());
    }
    let kind = response
        .headers()
        .get("content-type")
        .and_then(|value| value.to_str().ok())
        .unwrap_or_default()
        .split(';')
        .next()
        .unwrap_or_default()
        .trim()
        .to_ascii_lowercase();
    if status != 200 {
        let bytes = response
            .body_mut()
            .with_config()
            .limit(ANSWER_BYTES)
            .read_to_vec()
            .map_err(|_| unknown())?;
        let value: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
        let code = [
            &value["code"],
            &value["error"],
            &value["refusal"],
            &value["error"]["data"]["code"],
        ]
        .into_iter()
        .find_map(Value::as_str)
        .filter(|code| {
            !code.is_empty()
                && code.len() <= 96
                && code
                    .bytes()
                    .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_')
        })
        .map_or_else(
            || format!("haematite_mcp_remote_http_{status}"),
            str::to_owned,
        );
        return Err(Error::Response(code));
    }
    match kind.as_str() {
        "application/json" => {
            let bytes = response
                .body_mut()
                .with_config()
                .limit(ANSWER_BYTES)
                .read_to_vec()
                .map_err(|_| unknown())?;
            let message = serde_json::from_slice(&bytes).map_err(|_| invalid())?;
            emit(message).map_err(Error::Output)
        }
        "text/event-stream" => {
            let reader = response
                .body_mut()
                .with_config()
                .limit(ANSWER_BYTES)
                .reader();
            events(BufReader::new(reader), emit)
        }
        _ => Err(invalid()),
    }
}

fn events(
    reader: impl BufRead,
    mut emit: impl FnMut(Value) -> io::Result<()>,
) -> Result<(), Error> {
    let mut data = String::new();
    for line in reader.lines() {
        let line = line.map_err(|_| unknown())?;
        if line.is_empty() {
            if !data.is_empty() {
                let message = serde_json::from_str(&data).map_err(|_| invalid())?;
                emit(message).map_err(Error::Output)?;
                data.clear();
            }
        } else if let Some(value) = line.strip_prefix("data:") {
            data.push_str(value.strip_prefix(' ').unwrap_or(value));
            data.push('\n');
        }
    }
    if !data.is_empty() {
        let message = serde_json::from_str(&data).map_err(|_| invalid())?;
        emit(message).map_err(Error::Output)?;
    }
    Ok(())
}

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