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#[derive(Debug, Clone, PartialEq, Eq)]
15pub struct EncodedJsonBody {
16 bytes: Bytes,
17 trace_bytes: Option<Bytes>,
18 prepared: bool,
19}
20
21impl EncodedJsonBody {
22 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 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 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 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}