Skip to main content

codex_http_client/
request.rs

1use bytes::Bytes;
2use http::HeaderMap;
3use http::HeaderValue;
4use http::Method;
5use serde::Serialize;
6use serde_json::Value;
7use std::time::Duration;
8
9/// A JSON request body serialized once into reference-counted bytes.
10///
11/// Clones share the encoded allocation. Internally, the body can also hold the
12/// final compressed wire bytes while retaining the original JSON only when
13/// request-body trace logging is enabled.
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct EncodedJsonBody {
16    bytes: Bytes,
17    trace_bytes: Option<Bytes>,
18    prepared: bool,
19}
20
21impl EncodedJsonBody {
22    /// Serializes `value` into a reusable JSON body.
23    pub fn encode<T: Serialize + ?Sized>(value: &T) -> Result<Self, serde_json::Error> {
24        serde_json::to_vec(value).map(|bytes| Self {
25            bytes: Bytes::from(bytes),
26            trace_bytes: None,
27            prepared: false,
28        })
29    }
30
31    /// Returns the encoded bytes currently stored by this body.
32    pub fn as_bytes(&self) -> &[u8] {
33        &self.bytes
34    }
35
36    pub(crate) fn trace_bytes(&self) -> &[u8] {
37        self.trace_bytes.as_ref().unwrap_or(&self.bytes)
38    }
39}
40
41#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
42pub enum RequestCompression {
43    #[default]
44    None,
45    Zstd,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub enum RequestBody {
50    Json(Value),
51    EncodedJson(EncodedJsonBody),
52    Raw(Bytes),
53}
54
55impl RequestBody {
56    pub fn json(&self) -> Option<&Value> {
57        match self {
58            Self::Json(value) => Some(value),
59            Self::EncodedJson(_) | Self::Raw(_) => None,
60        }
61    }
62}
63
64#[derive(Debug, Clone, PartialEq, Eq)]
65pub struct PreparedRequestBody {
66    pub headers: HeaderMap,
67    pub body: Option<Bytes>,
68}
69
70impl PreparedRequestBody {
71    pub fn body_bytes(&self) -> Bytes {
72        self.body.clone().unwrap_or_default()
73    }
74}
75
76#[derive(Debug, Clone)]
77pub struct Request {
78    pub method: Method,
79    pub url: String,
80    pub headers: HeaderMap,
81    pub body: Option<RequestBody>,
82    pub compression: RequestCompression,
83    pub timeout: Option<Duration>,
84}
85
86impl Request {
87    pub fn new(method: Method, url: String) -> Self {
88        Self {
89            method,
90            url,
91            headers: HeaderMap::new(),
92            body: None,
93            compression: RequestCompression::None,
94            timeout: None,
95        }
96    }
97
98    pub fn with_json<T: Serialize>(mut self, body: &T) -> Self {
99        self.body = serde_json::to_value(body).ok().map(RequestBody::Json);
100        self
101    }
102
103    pub fn with_raw_body(mut self, body: impl Into<Bytes>) -> Self {
104        self.body = Some(RequestBody::Raw(body.into()));
105        self
106    }
107
108    pub fn with_compression(mut self, compression: RequestCompression) -> Self {
109        self.compression = compression;
110        self
111    }
112
113    /// Prepares the body once and stores the exact bytes that will be sent.
114    ///
115    /// Cloning the returned request shares the body bytes, so retry attempts do
116    /// not repeat JSON serialization or compression. Request-signing auth also
117    /// sees the same final headers and bytes that the transport will send.
118    pub fn into_prepared(mut self) -> Result<Self, String> {
119        let is_json = matches!(
120            self.body,
121            Some(RequestBody::Json(_) | RequestBody::EncodedJson(_))
122        );
123        let trace_bytes = if self.compression != RequestCompression::None
124            && tracing::enabled!(target: "codex_http_client::transport", tracing::Level::TRACE)
125        {
126            match self.body.as_ref() {
127                Some(RequestBody::Json(body)) => Some(Bytes::from(
128                    serde_json::to_vec(body).map_err(|err| err.to_string())?,
129                )),
130                Some(RequestBody::EncodedJson(body)) => Some(body.bytes.clone()),
131                Some(RequestBody::Raw(_)) | None => None,
132            }
133        } else {
134            None
135        };
136        let prepared = self.prepare_body_for_send()?;
137        self.headers = prepared.headers;
138        self.body = match (is_json, prepared.body) {
139            (true, Some(bytes)) => Some(RequestBody::EncodedJson(EncodedJsonBody {
140                bytes,
141                trace_bytes,
142                prepared: true,
143            })),
144            (false, Some(body)) => Some(RequestBody::Raw(body)),
145            (_, None) => None,
146        };
147        self.compression = RequestCompression::None;
148        Ok(self)
149    }
150
151    /// Convert the request body into the exact bytes that will be sent.
152    ///
153    /// Auth schemes such as AWS SigV4 need to sign the final body bytes, including
154    /// compression and content headers. Calling this method does not mutate the
155    /// request.
156    pub fn prepare_body_for_send(&self) -> Result<PreparedRequestBody, String> {
157        let headers = self.headers.clone();
158        match self.body.as_ref() {
159            Some(RequestBody::Raw(raw_body)) => {
160                if self.compression != RequestCompression::None {
161                    return Err("request compression cannot be used with raw bodies".to_string());
162                }
163                Ok(PreparedRequestBody {
164                    headers,
165                    body: Some(raw_body.clone()),
166                })
167            }
168            Some(RequestBody::Json(body)) => {
169                let body = EncodedJsonBody::encode(body).map_err(|err| err.to_string())?;
170                self.prepare_encoded_json(headers, &body)
171            }
172            Some(RequestBody::EncodedJson(body)) => self.prepare_encoded_json(headers, body),
173            None => Ok(PreparedRequestBody {
174                headers,
175                body: None,
176            }),
177        }
178    }
179
180    fn prepare_encoded_json(
181        &self,
182        mut headers: HeaderMap,
183        body: &EncodedJsonBody,
184    ) -> Result<PreparedRequestBody, String> {
185        if body.prepared {
186            return Ok(PreparedRequestBody {
187                headers,
188                body: Some(body.bytes.clone()),
189            });
190        }
191
192        let bytes = if self.compression != RequestCompression::None {
193            if headers.contains_key(http::header::CONTENT_ENCODING) {
194                return Err(
195                    "request compression was requested but content-encoding is already set"
196                        .to_string(),
197                );
198            }
199
200            let pre_compression_bytes = body.bytes.len();
201            let compression_start = std::time::Instant::now();
202            let (compressed, content_encoding) = match self.compression {
203                RequestCompression::None => unreachable!("guarded by compression != None"),
204                RequestCompression::Zstd => (
205                    zstd::stream::encode_all(std::io::Cursor::new(body.as_bytes()), 3)
206                        .map_err(|err| err.to_string())?,
207                    HeaderValue::from_static("zstd"),
208                ),
209            };
210            let post_compression_bytes = compressed.len();
211            let compression_duration = compression_start.elapsed();
212
213            headers.insert(http::header::CONTENT_ENCODING, content_encoding);
214
215            tracing::debug!(
216                pre_compression_bytes,
217                post_compression_bytes,
218                compression_duration_ms = compression_duration.as_millis(),
219                "Compressed request body with zstd"
220            );
221
222            Bytes::from(compressed)
223        } else {
224            body.bytes.clone()
225        };
226
227        if !headers.contains_key(http::header::CONTENT_TYPE) {
228            headers.insert(
229                http::header::CONTENT_TYPE,
230                HeaderValue::from_static("application/json"),
231            );
232        }
233
234        Ok(PreparedRequestBody {
235            headers,
236            body: Some(bytes),
237        })
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use http::HeaderValue;
245    use pretty_assertions::assert_eq;
246    use serde_json::json;
247
248    #[test]
249    fn prepare_body_for_send_serializes_json_and_sets_content_type() {
250        let request = Request::new(Method::POST, "https://example.com/v1/responses".to_string())
251            .with_json(&json!({"model": "test-model"}));
252
253        let prepared = request
254            .prepare_body_for_send()
255            .expect("body should prepare");
256
257        assert_eq!(
258            prepared.body,
259            Some(Bytes::from_static(br#"{"model":"test-model"}"#))
260        );
261        assert_eq!(
262            prepared
263                .headers
264                .get(http::header::CONTENT_TYPE)
265                .and_then(|value| value.to_str().ok()),
266            Some("application/json")
267        );
268        assert_eq!(
269            request.body,
270            Some(RequestBody::Json(json!({"model": "test-model"})))
271        );
272        assert_eq!(request.compression, RequestCompression::None);
273    }
274
275    #[test]
276    fn prepare_body_for_send_rejects_existing_content_encoding_when_compressing() {
277        let mut request =
278            Request::new(Method::POST, "https://example.com/v1/responses".to_string())
279                .with_json(&json!({"model": "test-model"}))
280                .with_compression(RequestCompression::Zstd);
281        request.headers.insert(
282            http::header::CONTENT_ENCODING,
283            HeaderValue::from_static("gzip"),
284        );
285
286        let err = request
287            .prepare_body_for_send()
288            .expect_err("conflicting content-encoding should fail");
289
290        assert_eq!(
291            err,
292            "request compression was requested but content-encoding is already set"
293        );
294    }
295
296    #[test]
297    fn into_prepared_stores_compressed_body_for_reuse() {
298        let body =
299            EncodedJsonBody::encode(&json!({"model": "test-model"})).expect("JSON should encode");
300        let mut request =
301            Request::new(Method::POST, "https://example.com/v1/responses".to_string())
302                .with_compression(RequestCompression::Zstd);
303        request.body = Some(RequestBody::EncodedJson(body));
304        let request = request.into_prepared().expect("body should prepare");
305        let Some(RequestBody::EncodedJson(body)) = request.body.as_ref() else {
306            panic!("expected an encoded JSON body");
307        };
308        let decompressed = zstd::stream::decode_all(std::io::Cursor::new(body.as_bytes()))
309            .expect("body should decompress");
310
311        assert_eq!(decompressed, br#"{"model":"test-model"}"#);
312        assert_eq!(request.compression, RequestCompression::None);
313        assert_eq!(
314            request.headers.get(http::header::CONTENT_ENCODING),
315            Some(&HeaderValue::from_static("zstd"))
316        );
317        assert_eq!(
318            request.headers.get(http::header::CONTENT_TYPE),
319            Some(&HeaderValue::from_static("application/json"))
320        );
321    }
322}
323
324#[derive(Debug, Clone)]
325pub struct Response {
326    pub status: http::StatusCode,
327    pub headers: HeaderMap,
328    pub body: Bytes,
329}