mirror of
https://github.com/openai/codex.git
synced 2026-05-04 21:32:21 +03:00
Add feature for optional request compression (#8767)
Adds a new feature `enable_request_compression` that will compress using zstd requests to the codex-backend. Currently only enabled for codex-backend so only enabled for openai providers when using chatgpt::auth even when the feature is enabled Added a new info log line too for evaluating the compression ratio and overhead off compressing before requesting. You can enable with `RUST_LOG=$RUST_LOG,codex_client::transport=info` ``` 2026-01-06T00:09:48.272113Z INFO codex_client::transport: Compressed request body with zstd pre_compression_bytes=28914 post_compression_bytes=11485 compression_duration_ms=0 ```
This commit is contained in:
@@ -104,6 +104,13 @@ impl CodexRequestBuilder {
|
||||
self.map(|builder| builder.json(value))
|
||||
}
|
||||
|
||||
pub fn body<B>(self, body: B) -> Self
|
||||
where
|
||||
B: Into<reqwest::Body>,
|
||||
{
|
||||
self.map(|builder| builder.body(body))
|
||||
}
|
||||
|
||||
pub async fn send(self) -> Result<Response, reqwest::Error> {
|
||||
let headers = trace_headers();
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ pub use crate::default_client::CodexRequestBuilder;
|
||||
pub use crate::error::StreamError;
|
||||
pub use crate::error::TransportError;
|
||||
pub use crate::request::Request;
|
||||
pub use crate::request::RequestCompression;
|
||||
pub use crate::request::Response;
|
||||
pub use crate::retry::RetryOn;
|
||||
pub use crate::retry::RetryPolicy;
|
||||
|
||||
@@ -5,12 +5,20 @@ use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use std::time::Duration;
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub enum RequestCompression {
|
||||
#[default]
|
||||
None,
|
||||
Zstd,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Request {
|
||||
pub method: Method,
|
||||
pub url: String,
|
||||
pub headers: HeaderMap,
|
||||
pub body: Option<Value>,
|
||||
pub compression: RequestCompression,
|
||||
pub timeout: Option<Duration>,
|
||||
}
|
||||
|
||||
@@ -21,6 +29,7 @@ impl Request {
|
||||
url,
|
||||
headers: HeaderMap::new(),
|
||||
body: None,
|
||||
compression: RequestCompression::None,
|
||||
timeout: None,
|
||||
}
|
||||
}
|
||||
@@ -29,6 +38,11 @@ impl Request {
|
||||
self.body = serde_json::to_value(body).ok();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_compression(mut self, compression: RequestCompression) -> Self {
|
||||
self.compression = compression;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
|
||||
@@ -2,6 +2,7 @@ use crate::default_client::CodexHttpClient;
|
||||
use crate::default_client::CodexRequestBuilder;
|
||||
use crate::error::TransportError;
|
||||
use crate::request::Request;
|
||||
use crate::request::RequestCompression;
|
||||
use crate::request::Response;
|
||||
use async_trait::async_trait;
|
||||
use bytes::Bytes;
|
||||
@@ -41,18 +42,70 @@ impl ReqwestTransport {
|
||||
}
|
||||
|
||||
fn build(&self, req: Request) -> Result<CodexRequestBuilder, TransportError> {
|
||||
let mut builder = self
|
||||
.client
|
||||
.request(
|
||||
Method::from_bytes(req.method.as_str().as_bytes()).unwrap_or(Method::GET),
|
||||
&req.url,
|
||||
)
|
||||
.headers(req.headers);
|
||||
if let Some(timeout) = req.timeout {
|
||||
let Request {
|
||||
method,
|
||||
url,
|
||||
mut headers,
|
||||
body,
|
||||
compression,
|
||||
timeout,
|
||||
} = req;
|
||||
|
||||
let mut builder = self.client.request(
|
||||
Method::from_bytes(method.as_str().as_bytes()).unwrap_or(Method::GET),
|
||||
&url,
|
||||
);
|
||||
|
||||
if let Some(timeout) = timeout {
|
||||
builder = builder.timeout(timeout);
|
||||
}
|
||||
if let Some(body) = req.body {
|
||||
builder = builder.json(&body);
|
||||
|
||||
if let Some(body) = body {
|
||||
if compression != RequestCompression::None {
|
||||
if headers.contains_key(http::header::CONTENT_ENCODING) {
|
||||
return Err(TransportError::Build(
|
||||
"request compression was requested but content-encoding is already set"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let json = serde_json::to_vec(&body)
|
||||
.map_err(|err| TransportError::Build(err.to_string()))?;
|
||||
let pre_compression_bytes = json.len();
|
||||
let compression_start = std::time::Instant::now();
|
||||
let (compressed, content_encoding) = match compression {
|
||||
RequestCompression::None => unreachable!("guarded by compression != None"),
|
||||
RequestCompression::Zstd => (
|
||||
zstd::stream::encode_all(std::io::Cursor::new(json), 3)
|
||||
.map_err(|err| TransportError::Build(err.to_string()))?,
|
||||
http::HeaderValue::from_static("zstd"),
|
||||
),
|
||||
};
|
||||
let post_compression_bytes = compressed.len();
|
||||
let compression_duration = compression_start.elapsed();
|
||||
|
||||
// Ensure the server knows to unpack the request body.
|
||||
headers.insert(http::header::CONTENT_ENCODING, content_encoding);
|
||||
if !headers.contains_key(http::header::CONTENT_TYPE) {
|
||||
headers.insert(
|
||||
http::header::CONTENT_TYPE,
|
||||
http::HeaderValue::from_static("application/json"),
|
||||
);
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
pre_compression_bytes,
|
||||
post_compression_bytes,
|
||||
compression_duration_ms = compression_duration.as_millis(),
|
||||
"Compressed request body with zstd"
|
||||
);
|
||||
|
||||
builder = builder.headers(headers).body(compressed);
|
||||
} else {
|
||||
builder = builder.headers(headers).json(&body);
|
||||
}
|
||||
} else {
|
||||
builder = builder.headers(headers);
|
||||
}
|
||||
Ok(builder)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user