rust_llm 2.1.1

//! Port of `lib/ruby_llm/transport/connection.rb`, `error_middleware.rb`, and the SSE handling in
//! `protocol/streaming.rb`.
//!
//! Retry rules match RubyLLM's Faraday setup: up to `max_retries` retries on rate limits, 5xx,
//! overloads, timeouts, and connection failures, with exponential backoff and jitter, honoring
//! `Retry-After`/`retry-after-ms`. A stream that has already delivered a chunk is never retried.

pub mod event_stream_parser;
pub mod websocket_connection;

pub use event_stream_parser::{EventStreamParser, SseEvent};
pub use websocket_connection::WebsocketConnection;

use event_stream_parser::{MAX_JSON_BODY_BYTES, StreamState, accumulate_failed_body};

use std::time::Duration;

use futures::StreamExt;
use serde_json::Value;

use crate::config::Config;
use crate::error::{Error, Result, error_for_status_message};
use crate::message::RawResponse;
use crate::providers::Provider;

#[derive(Clone)]
pub struct Connection {
    client: reqwest::Client,
    provider: Provider,
    config: std::sync::Arc<Config>,
}

impl std::fmt::Debug for Connection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Connection")
            .field("provider", &self.provider.slug())
            .finish()
    }
}

/// The incremental SSE parser the MCP transport shares.
pub(crate) type SseParser = EventStreamParser;

/// `Connection::Settings`: what a shared HTTP client is cached under. Credentials and the retry
/// settings are not in it: headers travel with each request, and the retry loop reads each
/// connection's own configuration. Pooled sockets belong to the tokio runtime that opened them,
/// so the runtime is part of the key, as Ruby's cache is per process.
#[derive(PartialEq)]
struct Settings {
    provider: Provider,
    api_base: Option<String>,
    timeout: Duration,
    proxy: Option<String>,
    runtime: tokio::runtime::Id,
}

/// `Connection::CACHE`.
static CLIENTS: crate::support::ProcessCache<Settings, reqwest::Client> =
    crate::support::ProcessCache::new(crate::support::ProcessCache::<(), ()>::LIMIT);

impl Connection {
    /// `Connection.new(provider, config)`: the HTTP client, and its pool of kept-alive sockets,
    /// is shared by every connection built from equal settings on this runtime (`15b48ed0`).
    pub fn new(provider: Provider, config: std::sync::Arc<Config>) -> Result<Connection> {
        let client = match tokio::runtime::Handle::try_current() {
            Ok(runtime) => {
                let settings = Settings {
                    provider,
                    api_base: provider.api_base(&config).ok(),
                    timeout: config.request_timeout,
                    proxy: config.http_proxy.clone(),
                    runtime: runtime.id(),
                };
                CLIENTS.try_fetch(settings, || basic(&config))?
            }
            Err(_) => basic(&config)?,
        };
        Ok(Connection {
            client,
            provider,
            config,
        })
    }

    /// A connection over an HTTP client you built yourself, the counterpart of RubyLLM's
    /// `config.faraday_adapter`: custom TLS roots, connectors, or pool settings. It is not cached.
    pub fn with_client(
        provider: Provider,
        config: std::sync::Arc<Config>,
        client: reqwest::Client,
    ) -> Connection {
        Connection {
            client,
            provider,
            config,
        }
    }

    pub fn client(&self) -> &reqwest::Client {
        &self.client
    }

    pub fn config(&self) -> &std::sync::Arc<Config> {
        &self.config
    }

    fn url(&self, path: &str) -> Result<String> {
        if path.starts_with("http://") || path.starts_with("https://") {
            return Ok(path.to_string());
        }
        let base = self.provider.api_base(&self.config)?;
        Ok(format!(
            "{}/{}",
            base.trim_end_matches('/'),
            path.trim_start_matches('/')
        ))
    }

    /// `additional_headers.merge(req.headers)`: request headers (a chat's `with_headers`, betas)
    /// are added, but the provider's own headers (auth, API version) win on a name clash, which
    /// is case-insensitive like Faraday's header hash.
    fn merged_headers(&self, extra: &[(String, String)]) -> Vec<(String, String)> {
        let provider = self.provider.headers(&self.config);
        let mut headers: Vec<(String, String)> = extra
            .iter()
            .filter(|(k, _)| !provider.iter().any(|(p, _)| p.eq_ignore_ascii_case(k)))
            .cloned()
            .collect();
        headers.extend(provider);
        headers
    }

    fn request(
        &self,
        url: &str,
        payload: &Value,
        extra: &[(String, String)],
    ) -> reqwest::RequestBuilder {
        let mut req = self.client.post(url).json(payload);
        for (k, v) in self.merged_headers(extra) {
            req = req.header(k, v);
        }
        req
    }

    /// faraday-retry's `calculate_sleep_amount`: a `Retry-After` longer than `retry_max_interval`
    /// means "don't retry" (`None`); otherwise wait the longer of it and the jittered backoff.
    fn backoff(&self, attempt: u32, retry_after: Option<f64>) -> Option<Duration> {
        let current = (self.config.retry_interval
            * self.config.retry_backoff_factor.powi(attempt as i32))
        .min(self.config.retry_max_interval);
        let interval = current
            + rand::random::<f64>()
                * self.config.retry_interval_randomness
                * self.config.retry_interval;
        match retry_after {
            Some(after) if after > self.config.retry_max_interval => None,
            Some(after) if after >= interval => Some(Duration::from_secs_f64(after)),
            _ => Some(Duration::from_secs_f64(interval)),
        }
    }

    /// Whether to retry after `error` on attempt `attempt`, and how long to wait first.
    fn retry_delay(
        &self,
        error: &Error,
        attempt: u32,
        retry_after: Option<f64>,
    ) -> Option<Duration> {
        if !error.retryable() || attempt >= self.config.max_retries {
            return None;
        }
        self.backoff(attempt, retry_after)
    }

    async fn send_with_retry(
        &self,
        url: &str,
        payload: &Value,
        extra: &[(String, String)],
        stream: bool,
        on_attempt: &mut (dyn FnMut(Option<&Error>) + Send),
    ) -> Result<reqwest::Response> {
        self.send_built(
            &|| self.request(url, payload, extra),
            true,
            stream,
            on_attempt,
        )
        .await
    }

    /// Sends the request `build` makes, retrying it under the rules above when `retry` is set.
    async fn send_built(
        &self,
        build: &(dyn Fn() -> reqwest::RequestBuilder + Send + Sync),
        retry: bool,
        stream: bool,
        on_attempt: &mut (dyn FnMut(Option<&Error>) + Send),
    ) -> Result<reqwest::Response> {
        // `instrument_request`: one `request.rust_llm` around the call and its retries.
        let mut event =
            crate::instrumentation::Event::start(&self.config, "request.rust_llm", || {
                let request = build().build().ok();
                crate::instrumentation::payload([
                    ("provider", self.provider.slug().into()),
                    (
                        "method",
                        request
                            .as_ref()
                            .map(|r| r.method().as_str().to_lowercase())
                            .into(),
                    ),
                    ("url", request.as_ref().map(|r| r.url().to_string()).into()),
                ])
            });
        let result = event
            .instrument(self.send_attempts(build, retry, stream, on_attempt))
            .await;
        let status = match &result {
            Ok(resp) => Some(resp.status().as_u16()),
            Err(e) => e.response().map(|r| r.status),
        };
        if let Some(status) = status {
            event.set("status", || status.into());
        }
        event.finish(result.as_ref().err());
        result
    }

    async fn send_attempts(
        &self,
        build: &(dyn Fn() -> reqwest::RequestBuilder + Send + Sync),
        retry: bool,
        stream: bool,
        on_attempt: &mut (dyn FnMut(Option<&Error>) + Send),
    ) -> Result<reqwest::Response> {
        let mut attempt = 0;
        let mut previous: Option<Error> = None;
        loop {
            on_attempt(previous.as_ref());
            let result = build().send().await;
            let (error, retry_after) = match result {
                Ok(resp) if resp.status().is_success() => return Ok(resp),
                Ok(resp) => {
                    let status = resp.status().as_u16();
                    let mut headers = header_pairs(&resp);
                    apply_retry_delay(Some(self.provider), status, &mut headers);
                    let body = if stream {
                        failed_stream_body(resp).await
                    } else {
                        resp.text().await.unwrap_or_default()
                    };
                    let message = self.provider.parse_error(&body);
                    (
                        error_for_status_message(status, &body, message),
                        retry_after_secs(&headers),
                    )
                }
                Err(e) => (transport_error(e), None),
            };
            let delay = if retry {
                self.retry_delay(&error, attempt, retry_after)
            } else {
                None
            };
            let Some(delay) = delay else {
                return Err(error);
            };
            tracing::debug!(
                provider = self.provider.slug(),
                attempt,
                "retrying after {error}"
            );
            tokio::time::sleep(delay).await;
            previous = Some(error);
            attempt += 1;
        }
    }

    /// POST JSON and parse a JSON response. `on_attempt` fires before every attempt so the usage
    /// ledger can record retries.
    pub async fn post(
        &self,
        path: &str,
        payload: &Value,
        extra: &[(String, String)],
        on_attempt: &mut (dyn FnMut(Option<&Error>) + Send),
    ) -> Result<RawResponse> {
        let url = self.url(path)?;
        let resp = self
            .send_with_retry(&url, payload, extra, false, on_attempt)
            .await?;
        let status = resp.status().as_u16();
        let headers = header_pairs(&resp);
        let text = resp
            .text()
            .await
            .map_err(|e| Error::ConnectionFailed(e.to_string()))?;
        if text.trim().is_empty() {
            return Err(Error::Api(
                "Provider returned an empty response body".into(),
                None,
            ));
        }
        let body: Value = serde_json::from_str(&text).unwrap_or(Value::String(text));
        Ok(RawResponse {
            status,
            headers,
            body,
            request_body: "".into(),
        })
    }

    /// `Connection#post` with a multipart body, for Files and Batches uploads. `form` builds a fresh
    /// body per attempt. Pass `idempotent: false` (`mark_non_idempotent`) for a request that creates
    /// server-side state: it is then never retried, since a retry after a lost response would
    /// create it twice.
    pub async fn post_multipart(
        &self,
        path: &str,
        form: impl Fn() -> reqwest::multipart::Form + Send + Sync,
        extra: &[(String, String)],
        idempotent: bool,
    ) -> Result<RawResponse> {
        let resp = self
            .send(reqwest::Method::POST, path, extra, idempotent, &|req| {
                req.multipart(form())
            })
            .await?;
        json_response(resp, Value::Null).await
    }

    /// `Connection#get`, parsing a JSON response. GETs retry like any idempotent request.
    pub async fn get(&self, path: &str, extra: &[(String, String)]) -> Result<RawResponse> {
        let resp = self
            .send(reqwest::Method::GET, path, extra, true, &|req| req)
            .await?;
        json_response(resp, Value::Null).await
    }

    /// `Connection#get` for a binary body, such as a file download.
    pub async fn get_bytes(&self, path: &str, extra: &[(String, String)]) -> Result<Vec<u8>> {
        let resp = self
            .send(reqwest::Method::GET, path, extra, true, &|req| req)
            .await?;
        Ok(resp
            .bytes()
            .await
            .map_err(|e| Error::ConnectionFailed(e.to_string()))?
            .to_vec())
    }

    /// `Connection#delete`, parsing a JSON response if there is one.
    pub async fn delete(&self, path: &str, extra: &[(String, String)]) -> Result<RawResponse> {
        let resp = self
            .send(reqwest::Method::DELETE, path, extra, true, &|req| req)
            .await?;
        let status = resp.status().as_u16();
        let headers = header_pairs(&resp);
        let text = resp
            .text()
            .await
            .map_err(|e| Error::ConnectionFailed(e.to_string()))?;
        let body = if text.trim().is_empty() {
            Value::Null
        } else {
            serde_json::from_str(&text).unwrap_or(Value::String(text))
        };
        Ok(RawResponse {
            status,
            headers,
            body,
            request_body: Default::default(),
        })
    }

    /// Any request with the provider's headers plus `extra`, with errors mapped like `post`.
    /// `body` adds the body to the builder; it runs once per attempt. `retry: false` sends once.
    pub async fn send(
        &self,
        method: reqwest::Method,
        path: &str,
        extra: &[(String, String)],
        retry: bool,
        body: &(dyn Fn(reqwest::RequestBuilder) -> reqwest::RequestBuilder + Send + Sync),
    ) -> Result<reqwest::Response> {
        self.send_tracked(method, path, extra, retry, body, &mut |_| {})
            .await
    }

    /// `send` with `on_attempt` firing before every attempt, like `post`, so the usage ledger can
    /// record retries of a request whose response is not JSON (speech audio, multipart uploads).
    pub async fn send_tracked(
        &self,
        method: reqwest::Method,
        path: &str,
        extra: &[(String, String)],
        retry: bool,
        body: &(dyn Fn(reqwest::RequestBuilder) -> reqwest::RequestBuilder + Send + Sync),
        on_attempt: &mut (dyn FnMut(Option<&Error>) + Send),
    ) -> Result<reqwest::Response> {
        let url = self.url(path)?;
        let build = || {
            let mut req = self.client.request(method.clone(), &url);
            for (k, v) in self.merged_headers(extra) {
                req = req.header(k, v);
            }
            body(req)
        };
        self.send_built(&build, retry, false, on_attempt).await
    }

    /// POST JSON and feed each server-sent event to `on_event`. Errors inside the stream
    /// (`event: error`, `{"error": ...}` data) are raised through the same status mapping, and are
    /// retried like HTTP errors as long as no event has reached `on_event` yet. Once the caller has
    /// seen a chunk, a failure is final: retrying would replay it into the same accumulator.
    pub async fn stream(
        &self,
        path: &str,
        payload: &Value,
        extra: &[(String, String)],
        on_attempt: &mut (dyn FnMut(Option<&Error>) + Send),
        on_event: &mut (dyn FnMut(SseEvent, Value) -> Result<()> + Send),
        streaming_error: fn(&str) -> Option<u16>,
    ) -> Result<RawResponse> {
        let url = self.url(path)?;
        let mut attempt = 0;
        let mut stream_error: Option<Error> = None;
        loop {
            let mut first = stream_error.take();
            let mut forward =
                |previous: Option<&Error>| on_attempt(previous.or(first.take().as_ref()));
            let resp = self
                .send_with_retry(&url, payload, extra, true, &mut forward)
                .await?;
            let status = resp.status().as_u16();
            let headers = header_pairs(&resp);
            let mut delivered = false;
            match self
                .read_stream(resp, on_event, streaming_error, &mut delivered)
                .await
            {
                Ok(()) => {
                    return Ok(RawResponse {
                        status,
                        headers,
                        body: Value::Null,
                        request_body: "".into(),
                    });
                }
                Err(error) if delivered => return Err(error),
                Err(error) => {
                    let Some(delay) = self.retry_delay(&error, attempt, None) else {
                        return Err(error);
                    };
                    tracing::debug!(
                        provider = self.provider.slug(),
                        attempt,
                        "retrying stream after {error}"
                    );
                    tokio::time::sleep(delay).await;
                    stream_error = Some(error);
                    attempt += 1;
                }
            }
        }
    }

    /// `process_stream_chunk` for each network read of a successful streaming response: a fresh
    /// `StreamState` per attempt, so nothing a failed attempt half-parsed reaches the next one.
    pub(crate) async fn read_stream(
        &self,
        resp: reqwest::Response,
        on_event: &mut (dyn FnMut(SseEvent, Value) -> Result<()> + Send),
        streaming_error: fn(&str) -> Option<u16>,
        delivered: &mut bool,
    ) -> Result<()> {
        let mut state = StreamState::new(MAX_JSON_BODY_BYTES);
        let mut stream = resp.bytes_stream();
        // `raise_stream_error`: the protocol's status for the error (500 by default).
        let raise = |text: &str| {
            error_for_status_message(
                streaming_error(text).unwrap_or(500),
                text,
                self.provider.parse_error(text),
            )
        };
        let mut handle = |event: SseEvent| -> Result<()> {
            if event.data == "[DONE]" {
                return Ok(());
            }
            let Ok(data) = serde_json::from_str::<Value>(&event.data) else {
                tracing::debug!("Failed to parse data chunk: {}", event.data);
                return Ok(());
            };
            let is_error = event.event.as_deref() == Some("error")
                || data.get("error").is_some()
                || data.get("type").and_then(Value::as_str) == Some("error");
            if is_error {
                return Err(raise(&event.data));
            }
            *delivered = true;
            // `stream_events`: `block.call(data) if data.is_a?(Hash)`; `data: true` is skipped
            // (after `handle_sse` has marked progress, as Ruby does).
            if !data.is_object() {
                return Ok(());
            }
            on_event(event, data)
        };
        while let Some(bytes) = stream.next().await {
            let bytes = bytes.map_err(transport_error)?;
            if let Some(error_body) = state.read(&bytes, &mut handle)? {
                return Err(raise(&error_body));
            }
        }
        for event in state.parser.finish() {
            handle(event)?;
        }
        Ok(())
    }
}

/// A `reqwest` failure as the error class Faraday raises for it: a timeout, or any other
/// connection failure (refused, reset, TLS, a body cut short), all of which are retried.
fn transport_error(error: reqwest::Error) -> Error {
    if error.is_timeout() {
        Error::Timeout(error.to_string())
    } else {
        Error::ConnectionFailed(error.to_string())
    }
}

/// `handle_failed_response`: a failed streaming response's body, read as it arrives and kept
/// only while it fits `MAX_JSON_BODY_BYTES`; past that it is dropped and the status alone picks
/// the error (`6f16112c`).
async fn failed_stream_body(resp: reqwest::Response) -> String {
    let mut buffer = Vec::new();
    let mut stream = resp.bytes_stream();
    while let Some(Ok(bytes)) = stream.next().await {
        if !accumulate_failed_body(&mut buffer, &bytes, MAX_JSON_BODY_BYTES) {
            return String::new();
        }
    }
    String::from_utf8_lossy(&buffer).into_owned()
}

/// Reads a JSON body the way `post` does: an empty body is an error, a non-JSON one a string.
/// The request is not kept (`release_request`), so `_request_body` is unused.
pub(crate) async fn json_response(
    resp: reqwest::Response,
    _request_body: Value,
) -> Result<RawResponse> {
    let status = resp.status().as_u16();
    let headers = header_pairs(&resp);
    let text = resp
        .text()
        .await
        .map_err(|e| Error::ConnectionFailed(e.to_string()))?;
    if text.trim().is_empty() {
        return Err(Error::Api(
            "Provider returned an empty response body".into(),
            None,
        ));
    }
    let body: Value = serde_json::from_str(&text).unwrap_or(Value::String(text));
    Ok(RawResponse {
        status,
        headers,
        body,
        request_body: "".into(),
    })
}

fn header_pairs(resp: &reqwest::Response) -> Vec<(String, String)> {
    resp.headers()
        .iter()
        .map(|(k, v)| {
            (
                k.as_str().to_string(),
                v.to_str().unwrap_or_default().to_string(),
            )
        })
        .collect()
}

/// `Transport::Connection.basic(config)`: a client with the configuration's `request_timeout` and
/// `http_proxy`. Provider connections build on it, and so do downloads outside a provider
/// (attachment URLs, hosted images and videos), each with the configuration it came from.
pub fn basic(config: &Config) -> Result<reqwest::Client> {
    let mut builder = reqwest::Client::builder().timeout(config.request_timeout);
    if let Some(proxy) = config
        .http_proxy
        .as_deref()
        .filter(|p| !p.trim().is_empty())
    {
        builder = builder
            .proxy(reqwest::Proxy::all(proxy).map_err(|e| Error::Configuration(e.to_string()))?);
    }
    builder
        .build()
        .map_err(|e| Error::Configuration(e.to_string()))
}

fn header<'a>(headers: &'a [(String, String)], name: &str) -> Option<&'a str> {
    headers
        .iter()
        .find(|(k, _)| k.eq_ignore_ascii_case(name))
        .map(|(_, v)| v.as_str())
}

/// `ErrorMiddleware#apply_retry_delay` (`transport/error_middleware.rb`): the retry loop only
/// reads the standard `Retry-After`, so on a failed response (status >= 400) without one, a
/// `retry-after-ms` delay (finite, not negative) or else, on a 429, the provider's own reset hint
/// (`Provider#retry_delay`) is written into `Retry-After` as seconds (`Float#to_s`).
pub fn apply_retry_delay(
    provider: Option<Provider>,
    status: u16,
    headers: &mut Vec<(String, String)>,
) {
    if status < 400 || header(headers, "retry-after").is_some() {
        return;
    }
    let millis = header(headers, "retry-after-ms")
        .and_then(|v| v.trim().parse::<f64>().ok())
        .filter(|ms| ms.is_finite() && *ms >= 0.0)
        .map(|ms| ms / 1000.0);
    let delay = millis.or_else(|| {
        if status == 429 {
            provider.and_then(|p| p.retry_delay(headers))
        } else {
            None
        }
    });
    if let Some(delay) = delay {
        headers.push(("Retry-After".into(), format!("{delay:?}")));
    }
}

/// faraday-retry's `calculate_retry_after` and `calculate_rate_limit_reset`: the longer of the
/// `Retry-After` and `RateLimit-Reset` waits, each an HTTP date or a number of seconds.
fn retry_after_secs(headers: &[(String, String)]) -> Option<f64> {
    ["retry-after", "ratelimit-reset"]
        .iter()
        .filter_map(|name| header(headers, name))
        .map(|v| match chrono::DateTime::parse_from_rfc2822(v.trim()) {
            Ok(at) => {
                (at.with_timezone(&chrono::Utc) - chrono::Utc::now()).num_milliseconds() as f64
                    / 1000.0
            }
            Err(_) => v.trim().parse::<f64>().unwrap_or(0.0),
        })
        .reduce(f64::max)
}