cookie_agent_models/catalog/
transport.rs

1use std::{future::Future, pin::Pin, time::Duration};
2
3use futures_util::{Stream, StreamExt as _};
4use reqwest::{Client, header};
5use thiserror::Error;
6
7use super::{CATALOG_MAX_BYTES, MODELS_DEV_CATALOG_URL};
8
9pub const MODELS_DEV_USER_AGENT: &str = "cookie-agent/0.1 catalog-cache/1";
10
11pub type CatalogBodyStream =
12    Pin<Box<dyn Stream<Item = Result<Vec<u8>, CatalogTransportError>> + Send>>;
13pub type CatalogTransportFuture<'a> = Pin<
14    Box<dyn Future<Output = Result<CatalogTransportResponse, CatalogTransportError>> + Send + 'a>,
15>;
16
17/// The only request shape accepted by the catalog transport boundary.
18#[derive(Clone, Debug, Eq, PartialEq)]
19pub struct CatalogRequest {
20    pub url: &'static str,
21    pub user_agent: &'static str,
22    pub accept: &'static str,
23    pub accept_encoding: &'static str,
24    pub if_none_match: Option<String>,
25}
26
27impl CatalogRequest {
28    pub(crate) fn fixed(etag: Option<String>) -> Self {
29        Self {
30            url: MODELS_DEV_CATALOG_URL,
31            user_agent: MODELS_DEV_USER_AGENT,
32            accept: "application/json",
33            accept_encoding: "identity",
34            if_none_match: etag,
35        }
36    }
37}
38
39/// Streamed transport response. Tests inject this without live network access.
40pub struct CatalogTransportResponse {
41    pub status: u16,
42    pub content_type: Option<String>,
43    pub content_encoding: Option<String>,
44    pub content_length: Option<u64>,
45    pub etag: Option<String>,
46    pub body: CatalogBodyStream,
47}
48
49impl CatalogTransportResponse {
50    #[must_use]
51    pub fn from_bytes(status: u16, bytes: Vec<u8>) -> Self {
52        let length = bytes.len() as u64;
53        Self {
54            status,
55            content_type: Some("application/json".to_owned()),
56            content_encoding: None,
57            content_length: Some(length),
58            etag: None,
59            body: Box::pin(futures_util::stream::once(async move { Ok(bytes) })),
60        }
61    }
62
63    #[must_use]
64    pub fn not_modified() -> Self {
65        Self {
66            status: 304,
67            content_type: None,
68            content_encoding: None,
69            content_length: Some(0),
70            etag: None,
71            body: Box::pin(futures_util::stream::empty()),
72        }
73    }
74}
75
76/// Injectable fixed catalog transport.
77pub trait CatalogTransport: Send + Sync {
78    fn fetch(&self, request: CatalogRequest) -> CatalogTransportFuture<'_>;
79}
80
81/// Production rustls client with redirects and automatic content coding disabled.
82#[derive(Clone, Debug)]
83pub struct HttpCatalogTransport {
84    client: Client,
85}
86
87impl HttpCatalogTransport {
88    pub fn new() -> Result<Self, CatalogTransportError> {
89        let client = Client::builder()
90            .redirect(reqwest::redirect::Policy::none())
91            .connect_timeout(Duration::from_secs(5))
92            .timeout(Duration::from_secs(15))
93            .user_agent(MODELS_DEV_USER_AGENT)
94            .build()
95            .map_err(|_| CatalogTransportError::ClientBuild)?;
96        Ok(Self { client })
97    }
98}
99
100impl CatalogTransport for HttpCatalogTransport {
101    fn fetch(&self, request: CatalogRequest) -> CatalogTransportFuture<'_> {
102        Box::pin(async move {
103            if request.url != MODELS_DEV_CATALOG_URL
104                || request.user_agent != MODELS_DEV_USER_AGENT
105                || request.accept != "application/json"
106                || request.accept_encoding != "identity"
107            {
108                return Err(CatalogTransportError::InvalidRequest);
109            }
110            let mut builder = self
111                .client
112                .get(MODELS_DEV_CATALOG_URL)
113                .header(header::ACCEPT, request.accept)
114                .header(header::ACCEPT_ENCODING, request.accept_encoding);
115            if let Some(etag) = request.if_none_match {
116                let value = header::HeaderValue::from_str(&etag)
117                    .map_err(|_| CatalogTransportError::InvalidEtag)?;
118                builder = builder.header(header::IF_NONE_MATCH, value);
119            }
120            let response = builder
121                .send()
122                .await
123                .map_err(|_| CatalogTransportError::RequestFailed)?;
124            let status = response.status().as_u16();
125            let headers = response.headers();
126            let content_type = header_string(headers, header::CONTENT_TYPE)?;
127            let content_encoding = header_string(headers, header::CONTENT_ENCODING)?;
128            let etag = header_string(headers, header::ETAG)?;
129            let content_length = match headers.get(header::CONTENT_LENGTH) {
130                Some(value) => Some(
131                    value
132                        .to_str()
133                        .map_err(|_| CatalogTransportError::InvalidHeaders)?
134                        .parse::<u64>()
135                        .map_err(|_| CatalogTransportError::InvalidHeaders)?,
136                ),
137                None => None,
138            };
139            if content_length.is_some_and(|length| length > CATALOG_MAX_BYTES as u64) {
140                return Err(CatalogTransportError::BodyTooLarge);
141            }
142            let body = response.bytes_stream().map(|chunk| {
143                chunk
144                    .map(|bytes| bytes.to_vec())
145                    .map_err(|_| CatalogTransportError::BodyReadFailed)
146            });
147            Ok(CatalogTransportResponse {
148                status,
149                content_type,
150                content_encoding,
151                content_length,
152                etag,
153                body: Box::pin(body),
154            })
155        })
156    }
157}
158
159fn header_string(
160    headers: &header::HeaderMap,
161    name: header::HeaderName,
162) -> Result<Option<String>, CatalogTransportError> {
163    headers
164        .get(name)
165        .map(|value| {
166            value
167                .to_str()
168                .map(str::to_owned)
169                .map_err(|_| CatalogTransportError::InvalidHeaders)
170        })
171        .transpose()
172}
173
174/// Body-free transport errors.
175#[derive(Clone, Copy, Debug, Error, Eq, PartialEq)]
176pub enum CatalogTransportError {
177    #[error("catalog transport client could not be constructed")]
178    ClientBuild,
179    #[error("catalog transport request shape is invalid")]
180    InvalidRequest,
181    #[error("catalog cache ETag is invalid")]
182    InvalidEtag,
183    #[error("catalog response headers are invalid")]
184    InvalidHeaders,
185    #[error("catalog request failed")]
186    RequestFailed,
187    #[error("catalog response body read failed")]
188    BodyReadFailed,
189    #[error("catalog response body exceeds the byte limit")]
190    BodyTooLarge,
191}
cookie_agent_models/catalog/transport.rs

cookie_agent_models/catalog/
transport.rs
