haematite-cli 0.21.7

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use ureq::http::{HeaderValue, Response, header};
use ureq::unversioned::resolver::DefaultResolver;
use ureq::unversioned::transport::{
    Buffers, ConnectionDetails, Connector, DefaultConnector, NextTimeout, Transport,
};

pub(super) const REQUEST_FAILED: &str = "haematite_mcp_remote_request_failed";
pub(super) const OUTCOME_UNKNOWN: &str = "haematite_mcp_remote_outcome_unknown";
pub(super) const UNKNOWN_MESSAGE: &str =
    "the request reached the server; whether it was carried out is not known; read before retrying";

#[derive(Clone)]
pub(super) struct Client {
    config: ureq::config::Config,
}

impl Client {
    pub fn new(timeout: Duration) -> 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(None)
            .timeout_per_call(None)
            .timeout_resolve(Some(timeout))
            .timeout_connect(Some(timeout))
            .timeout_send_request(Some(timeout))
            .timeout_send_body(Some(timeout))
            .timeout_recv_response(None)
            .timeout_recv_body(None)
            .build();
        Self { config }
    }

    pub fn send(
        &self,
        url: &str,
        bearer: HeaderValue,
        protocol: Option<&str>,
        session: Option<HeaderValue>,
        line: &[u8],
    ) -> Result<Response<ureq::Body>, String> {
        let attempted = Arc::new(AtomicBool::new(false));
        let connector = ObservedConnector {
            attempted: Arc::clone(&attempted),
        };
        let client =
            ureq::Agent::with_parts(self.config.clone(), connector, DefaultResolver::default());
        let mut request = client
            .post(url)
            .header(header::AUTHORIZATION, bearer)
            .header(header::CONTENT_TYPE, "application/json")
            .header(header::ACCEPT, "application/json, text/event-stream");
        if let Some(version) = protocol {
            request = request.header("MCP-Protocol-Version", version);
        }
        if let Some(session) = session {
            request = request.header("Mcp-Session-Id", session);
        }
        request.send(line).map_err(|_| {
            if attempted.load(Ordering::Acquire) {
                OUTCOME_UNKNOWN
            } else {
                REQUEST_FAILED
            }
            .to_owned()
        })
    }
}

#[derive(Debug)]
struct ObservedConnector {
    attempted: Arc<AtomicBool>,
}

impl Connector for ObservedConnector {
    type Out = ObservedTransport;

    fn connect(
        &self,
        details: &ConnectionDetails,
        chained: Option<()>,
    ) -> Result<Option<Self::Out>, ureq::Error> {
        // The default connector completes TLS before exposing HTTP transport.
        Ok(DefaultConnector::default()
            .connect(details, chained)?
            .map(|inner| ObservedTransport {
                inner,
                attempted: Arc::clone(&self.attempted),
            }))
    }
}

#[derive(Debug)]
struct ObservedTransport {
    inner: Box<dyn Transport>,
    attempted: Arc<AtomicBool>,
}

impl Transport for ObservedTransport {
    fn buffers(&mut self) -> &mut dyn Buffers {
        self.inner.buffers()
    }

    fn transmit_output(&mut self, amount: usize, timeout: NextTimeout) -> Result<(), ureq::Error> {
        if amount != 0 {
            // A failed write may have sent bytes before failing. It cannot prove refusal.
            self.attempted.store(true, Ordering::Release);
        }
        self.inner.transmit_output(amount, timeout)
    }

    fn await_input(&mut self, timeout: NextTimeout) -> Result<bool, ureq::Error> {
        self.inner.await_input(timeout)
    }

    fn is_open(&mut self) -> bool {
        self.inner.is_open()
    }

    fn is_tls(&self) -> bool {
        self.inner.is_tls()
    }
}