rig_core/driver/

http_transport.rs

//! The HTTP transport: any [`HttpClientExt`] sends [`Encoded`] requests and
//! delivers their replies as [`WireFrame`]s. It owns every HTTP concern:
//! headers, framing, status and content-type checks, request ids and the
//! observation of the exchange.

use std::sync::Arc;

use bytes::Bytes;
use futures::StreamExt;

use super::{Exchange, Opened, Opening, Received, Transport};
use crate::error::ProviderError;
use crate::http_client::framing::{Framer, Framing};
use crate::http_client::{self, HttpClientExt};
use crate::observe::{AdapterContext, AdapterErrorBoundary, AdapterSlot};
use crate::wasm_compat::{WasmCompatSend, WasmCompatSync};
use crate::wire::{Body, Encoded, Mode, Projector, Wire, WireFrame};

impl<W, H> Transport<W> for H
where
    W: Wire<Payload = Encoded, Frame = WireFrame>,
    H: HttpClientExt + Clone + WasmCompatSend + WasmCompatSync + 'static,
{
    fn send(&self, payload: Encoded, exchange: Exchange) -> Opening<WireFrame> {
        let Encoded {
            mut request,
            framing,
            request_id_header,
            relaxed_content_type,
            route,
            project,
            analysis_only,
        } = payload;
        let Exchange {
            mode,
            observation,
            extensions,
        } = exchange;
        request.extensions_mut().extend(extensions);
        // A whole reply has nothing to stream: it is read as a unary one is,
        // so its document is a fact of the reply in both modes.
        let streamed = framing.framer().filter(|_| mode == Mode::Streaming);
        if streamed.is_some() && matches!(request.body(), Body::Multipart(_)) {
            return Opening::failed(ProviderError::request(
                "a multipart request cannot open a streamed reply",
            ));
        }
        accept_header(&mut request, framing);
        // Errors need the actual path; observations use the declared
        // template to group attempts independently of concrete URLs.
        let path = request.uri().path().to_owned();
        let declared = route.map_or_else(|| path.clone(), str::to_owned);
        let exchange = HttpExchange {
            framing,
            request_id_header,
            relaxed_content_type,
            path,
            observation: observation.map(|context| (context, AdapterSlot::default())),
            project,
        };
        let http = self.clone();
        let slot = exchange.slot().cloned();
        let opening = slot.clone();
        Opening::new(async move {
            // Unpolled replies must not report transport attempts.
            exchange.install(&request, &declared);
            let mut opened = match streamed {
                Some(framer) => match byte_request(request) {
                    Ok(request) => exchange.streaming(&http, request, framer).await,
                    Err(error) => Opened::failed(error),
                },
                None => exchange.unary(&http, request).await,
            };
            opened.slot = slot;
            opened.analysis_only = analysis_only;
            Ok(opened)
        })
        .with_slot(opening)
    }
}

/// What one request's exchange needs besides the client.
struct HttpExchange {
    framing: Framing,
    request_id_header: Option<&'static str>,
    relaxed_content_type: bool,
    path: String,
    observation: Option<(AdapterContext, AdapterSlot)>,
    project: Option<Projector>,
}

impl HttpExchange {
    fn slot(&self) -> Option<&AdapterSlot> {
        self.observation.as_ref().map(|(_, slot)| slot)
    }

    fn project(&self, payload: &[u8]) {
        project(self.slot(), self.project, payload);
    }

    fn install<B>(&self, request: &http::Request<B>, declared: &str) {
        if let Some((context, slot)) = &self.observation {
            slot.install(context.attempt_for(request, declared));
        }
    }

    fn failed(self, error: ProviderError, request_id: Option<String>) -> Opened<WireFrame> {
        Opened::failed(error)
            .with_request_id(request_id)
            .with_route(self.path)
    }

    /// Send a buffered request and frame its whole body. Every frame is
    /// projected as the driver reads it.
    async fn unary<H: HttpClientExt>(
        self,
        http: &H,
        request: http::Request<Body>,
    ) -> Opened<WireFrame> {
        let sent = match send(http, request, self.request_id_header, self.slot()).await {
            Ok(sent) => sent,
            Err(error) => {
                // The reply the failure carries is still the provider's:
                // project its facts before reporting the ending.
                if let Some(body) = error.provider_response_body() {
                    self.project(body.as_bytes());
                }
                return self.failed(error, None);
            }
        };
        // The status said success, but an SSE framer over a body that is not
        // an event stream yields no frames at all, which would fold to a
        // contentless success. The reply the provider actually sent is the
        // error.
        if let Some(rejected) =
            wrong_content_type(&sent.headers, self.framing, self.relaxed_content_type)
        {
            let error = ProviderError::from_transport_error(rejected)
                .with_provider_status(Some(sent.status))
                .with_provider_request_id(sent.provider_request_id.clone())
                .with_response_headers(Some(sent.headers.clone()));
            self.project(&sent.body);
            return self.failed(error, sent.provider_request_id);
        }
        // Frame the reply the way a stream is framed: a wire whose unary
        // reply is an event stream would otherwise hand every projector a
        // document it cannot parse.
        let document = serde_json::from_slice(&sent.body).ok();
        let frames = self.framing.split(&sent.body);
        let slot = self.slot().cloned();
        let projector = self.project;
        let frames = futures::stream::iter(frames).map(move |frame| {
            project(slot.as_ref(), projector, frame.as_bytes());
            Ok(frame)
        });
        Opened {
            document,
            ..Opened::new(frames)
                .with_request_id(sent.provider_request_id)
                .with_http(sent.status, sent.headers)
                .with_route(self.path)
        }
    }

    /// Open a streamed reply and frame its body as it arrives. Each frame is
    /// projected only when the driver reads it, and the bytes left unframed
    /// are recorded when the body ends.
    async fn streaming<H: HttpClientExt>(
        self,
        http: &H,
        request: http::Request<Vec<u8>>,
        mut framer: Framer,
    ) -> Opened<WireFrame> {
        let response = match http.send_streaming(request).await {
            // Custom transports may return rejected responses directly; preserve
            // their status, headers, and bounded body in the error.
            Ok(response) if response.status() != http::StatusCode::OK => {
                Err(reject_response(response).await)
            }
            Ok(response) => {
                match wrong_content_type(
                    response.headers(),
                    self.framing,
                    self.relaxed_content_type,
                ) {
                    Some(error) => Err(error),
                    None => Ok(response),
                }
            }
            other => other,
        };
        let response = match response {
            Ok(response) => response,
            Err(error) => {
                if let Some(slot) = self.slot() {
                    slot.error_boundary(AdapterErrorBoundary::from_http(&error));
                    if let Some(status) = error.non_success_status() {
                        slot.response_with_headers(status, error.non_success_headers());
                    }
                    if let Some(body) = error.non_success_body() {
                        self.project(body.as_bytes());
                    }
                }
                let request_id = error
                    .non_success_headers()
                    .and_then(|headers| request_id_from(headers, self.request_id_header));
                let error = ProviderError::from_transport_error(error)
                    .with_provider_request_id(request_id.clone());
                return self.failed(error, request_id);
            }
        };
        if let Some(slot) = self.slot() {
            slot.response_with_headers(response.status(), Some(response.headers()));
        }
        let request_id = request_id_from(response.headers(), self.request_id_header);
        let status = response.status();
        let headers = response.headers().clone();
        let slot = self.slot().cloned();
        let Self {
            path,
            project: projector,
            ..
        } = self;
        let mut body = response.into_body();
        let received = Arc::new(Received::new());
        let marked = Arc::clone(&received);
        let frames = async_stream::stream! {
            while let Some(chunk) = body.next().await {
                let chunk = match chunk {
                    Ok(chunk) => {
                        marked.mark();
                        chunk
                    }
                    Err(error) => {
                        if let Some(slot) = &slot {
                            slot.error_boundary(AdapterErrorBoundary::Transport);
                        }
                        yield Err(ProviderError::from_transport_error(error));
                        return;
                    }
                };
                for frame in framer.push(&chunk) {
                    project(slot.as_ref(), projector, frame.as_bytes());
                    yield Ok(frame);
                }
            }
            if let Some(slot) = &slot {
                slot.unframed(framer.pending());
            }
            if let Some(frame) = framer.finish() {
                project(slot.as_ref(), projector, frame.as_bytes());
                yield Ok(frame);
            }
        };
        Opened {
            received: Some(received),
            ..Opened::new(frames)
                .with_request_id(request_id)
                .with_http(status, headers)
                .with_route(path)
        }
    }
}

/// Project a payload's observation facts onto the attempt, when the call is
/// observed and the wire reads any.
fn project(slot: Option<&AdapterSlot>, project: Option<Projector>, payload: &[u8]) {
    if let (Some(slot), Some(project)) = (slot, project) {
        slot.project(|sink| project(payload, sink));
    }
}

/// What one send learned about its reply.
struct Sent {
    status: http::StatusCode,
    headers: http::HeaderMap,
    body: Bytes,
    provider_request_id: Option<String>,
}

/// Sends a buffered request, preserving non-success response details and IDs.
async fn send<H>(
    http: &H,
    request: http::Request<Body>,
    request_id_header: Option<&'static str>,
    observation: Option<&AdapterSlot>,
) -> Result<Sent, ProviderError>
where
    H: HttpClientExt,
{
    let (parts, body) = request.into_parts();
    let response = match body {
        Body::Bytes(bytes) => {
            http.send::<_, Bytes>(http::Request::from_parts(parts, bytes))
                .await
        }
        Body::Multipart(form) => {
            http.send_multipart::<Bytes>(http::Request::from_parts(parts, form))
                .await
        }
    };
    let response = match response {
        Ok(response) => response,
        // A transport that reports the non-success reply as an error: the
        // reply is the provider's, so it funnels to a preserved provider
        // response with the id read off its headers and the headers
        // themselves; a response-less failure stays a transport error.
        Err(error) => {
            if let Some(observation) = observation
                && let Some(status) = error.non_success_status()
            {
                observation.response_with_headers(status, error.non_success_headers());
            }
            let request_id = error
                .non_success_headers()
                .and_then(|headers| request_id_from(headers, request_id_header));
            return Err(
                ProviderError::from_transport_error(error).with_provider_request_id(request_id)
            );
        }
    };

    // Take the reply apart before awaiting the body: the headers are then
    // owned, so preserving them onto an error costs no clone and every
    // error path below can afford them.
    let (parts, body) = response.into_parts();
    let status = parts.status;
    if let Some(observation) = observation {
        observation.response_with_headers(status, Some(&parts.headers));
    }
    let provider_request_id = request_id_from(&parts.headers, request_id_header);
    let body = body.await.map_err(ProviderError::from_transport_error)?;

    if !status.is_success() {
        return Err(
            ProviderError::from_http_response(status, String::from_utf8_lossy(&body))
                .with_provider_request_id(provider_request_id)
                .with_response_headers(Some(parts.headers)),
        );
    }
    Ok(Sent {
        status,
        headers: parts.headers,
        body,
        provider_request_id,
    })
}

/// The provider's transport request id, when it names such a header and the
/// reply carries a non-empty value.
pub(crate) fn request_id_from(headers: &http::HeaderMap, header: Option<&str>) -> Option<String> {
    header.and_then(|header| {
        headers
            .get(header)
            .and_then(|value| value.to_str().ok())
            .filter(|value| !value.is_empty())
            .map(str::to_string)
    })
}

/// Defaults byte-body requests to `application/json`, including bodyless GETs.
/// Preserves an explicitly supplied content type.
fn content_type(request: &mut http::Request<Body>) {
    if matches!(request.body(), Body::Bytes(_)) {
        request
            .headers_mut()
            .entry(http::header::CONTENT_TYPE)
            .or_insert(http::HeaderValue::from_static("application/json"));
    }
}

/// Add `Accept: text/event-stream` to an SSE request, without overriding an
/// `Accept` the wire set itself.
fn accept_header(request: &mut http::Request<Body>, framing: Framing) {
    content_type(request);
    if framing == Framing::Sse {
        request
            .headers_mut()
            .entry("Accept")
            .or_insert(http::HeaderValue::from_static("text/event-stream"));
    }
}

/// Whether a reply's content type is not the event stream an SSE wire asked
/// for: the framer would silently produce no frames, which reads as
/// truncation rather than as the wrong endpoint. A reply that names no
/// content type at all is accepted only by a wire that opted in.
///
/// The one predicate both paths ask, so a unary reply and a streamed one
/// cannot disagree about what the provider sent.
fn wrong_content_type(
    headers: &http::HeaderMap,
    framing: Framing,
    relaxed: bool,
) -> Option<http_client::Error> {
    if framing != Framing::Sse {
        return None;
    }
    let Some(content_type) = headers.get(&http::header::CONTENT_TYPE) else {
        return (!relaxed)
            .then(|| http_client::Error::InvalidContentType(http::HeaderValue::from_static("")));
    };
    let event_stream = content_type
        .to_str()
        .ok()
        .and_then(|value| value.parse::<mime::Mime>().ok())
        .is_some_and(|mime_type| {
            matches!(
                (mime_type.type_(), mime_type.subtype()),
                (mime::TEXT, mime::EVENT_STREAM)
            )
        });
    (!event_stream).then(|| http_client::Error::InvalidContentType(content_type.clone()))
}

/// A request whose body is bytes. Multipart replies are never streamed.
fn byte_request(request: http::Request<Body>) -> Result<http::Request<Vec<u8>>, ProviderError> {
    let (parts, body) = request.into_parts();
    match body {
        Body::Bytes(bytes) => Ok(http::Request::from_parts(parts, bytes)),
        Body::Multipart(_) => Err(ProviderError::request(
            "a multipart request cannot open a streamed reply",
        )),
    }
}

/// Bytes of a rejected reply's body kept on the error; a reply longer than
/// this is cut there.
const REJECTED_BODY_LIMIT: usize = 1 << 20;

/// Chunks read off a rejected reply before giving up on it, so a transport
/// that keeps yielding empty chunks cannot hold the opener.
const REJECTED_CHUNK_LIMIT: usize = 4096;

/// Turn a reply the driver will not stream (any status but 200, a 204
/// included: a status is a status) into the non-success error, reading the
/// body to its end (bounded in bytes and chunks) so the provider's payload
/// and the transport's headers ride on the error.
async fn reject_response(
    response: http::Response<crate::http_client::BoxedStream>,
) -> http_client::Error {
    let status = response.status();
    let headers = response.headers().clone();
    let mut body = response.into_body();
    let mut bytes: Vec<u8> = Vec::new();
    let mut chunks = 0usize;
    while let Some(chunk) = body.next().await {
        chunks += 1;
        if let Ok(chunk) = chunk {
            let room = REJECTED_BODY_LIMIT.saturating_sub(bytes.len());
            bytes.extend_from_slice(chunk.get(..chunk.len().min(room)).unwrap_or_default());
        }
        if bytes.len() >= REJECTED_BODY_LIMIT || chunks >= REJECTED_CHUNK_LIMIT {
            break;
        }
    }
    http_client::Error::InvalidStatusCodeWithDetails {
        status,
        body: String::from_utf8_lossy(&bytes).into_owned(),
        headers,
    }
}