From 495da4564337099f064b25f7e2f00436f56bc076 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 15 Jun 2026 18:11:26 +0100 Subject: [PATCH] reuse encoded Responses request bodies (#28327) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why Responses HTTP requests were converted from `ResponsesApiRequest` into a full `serde_json::Value`. `EndpointSession` then deep-cloned that value for each retry, and the transport serialized and compressed it again before every send. Large histories make those copies expensive. Retry attempts should reuse the same immutable request bytes. ## What - Serialize standard Responses requests directly into a ref-counted `EncodedJsonBody`. - Preserve the Azure path that attaches item IDs before encoding. - Prepare JSON, compression, and derived content headers once before the retry loop. - Clone the prepared request per attempt so body clones only bump the `Bytes` reference count. - Keep auth inside the retry loop. Signing auth sees the exact final headers and body bytes that the transport sends. - Preserve request-body TRACE output. With TRACE plus compression, retain the original JSON bytes for logging; normal requests keep only the final wire bytes. - Leave non-Responses endpoint bodies on the existing `Value` path. ## Performance A temporary release-mode measurement used a 10 MiB JSON body and 10 retry preparations: - old `Value` clone + serialize path: 30 ms total - prepared shared-byte path: less than 1 ms total That is about 3 ms avoided per retry for this payload on the test machine. Each retry also stops allocating another request-sized JSON tree and serialized buffer. Without TRACE, compressed requests retain only the final compressed wire bytes. ## Validation - `just test -p codex-client` — 28 passed - `just test -p codex-api` — 125 passed - `just fix -p codex-client` - `just fix -p codex-api` --- codex-rs/codex-api/src/endpoint/responses.rs | 30 ++- codex-rs/codex-api/src/endpoint/session.rs | 22 +- codex-rs/codex-api/tests/clients.rs | 116 ++++++++-- codex-rs/codex-client/src/lib.rs | 1 + codex-rs/codex-client/src/request.rs | 212 ++++++++++++++----- codex-rs/codex-client/src/transport.rs | 3 + 6 files changed, 306 insertions(+), 78 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/responses.rs b/codex-rs/codex-api/src/endpoint/responses.rs index cc1be2846..ad79adf33 100644 --- a/codex-rs/codex-api/src/endpoint/responses.rs +++ b/codex-rs/codex-api/src/endpoint/responses.rs @@ -11,6 +11,7 @@ use crate::requests::headers::insert_header; use crate::requests::headers::subagent_header; use crate::sse::spawn_response_stream; use crate::telemetry::SseTelemetry; +use codex_client::EncodedJsonBody; use codex_client::HttpTransport; use codex_client::RequestCompression; use codex_client::RequestTelemetry; @@ -81,11 +82,16 @@ impl ResponsesClient { turn_state, } = options; - let mut body = serde_json::to_value(&request) - .map_err(|e| ApiError::Stream(format!("failed to encode responses request: {e}")))?; - if request.store && self.session.provider().is_azure_responses_endpoint() { + let body = if request.store && self.session.provider().is_azure_responses_endpoint() { + let mut body = serde_json::to_value(&request).map_err(|e| { + ApiError::Stream(format!("failed to encode responses request: {e}")) + })?; attach_item_ids(&mut body, &request.input); + EncodedJsonBody::encode(&body) + } else { + EncodedJsonBody::encode(&request) } + .map_err(|e| ApiError::Stream(format!("failed to encode responses request: {e}")))?; let mut headers = extra_headers; if let Some(ref thread_id) = thread_id { @@ -96,7 +102,8 @@ impl ResponsesClient { insert_header(&mut headers, "x-openai-subagent", &subagent); } - self.stream(body, headers, compression, turn_state).await + self.stream_encoded(body, headers, compression, turn_state) + .await } fn path() -> &'static str { @@ -120,6 +127,19 @@ impl ResponsesClient { extra_headers: HeaderMap, compression: Compression, turn_state: Option>>, + ) -> Result { + let body = EncodedJsonBody::encode(&body) + .map_err(|e| ApiError::Stream(format!("failed to encode responses request: {e}")))?; + self.stream_encoded(body, extra_headers, compression, turn_state) + .await + } + + async fn stream_encoded( + &self, + body: EncodedJsonBody, + extra_headers: HeaderMap, + compression: Compression, + turn_state: Option>>, ) -> Result { let request_compression = match compression { Compression::None => RequestCompression::None, @@ -128,7 +148,7 @@ impl ResponsesClient { let stream_response = self .session - .stream_with( + .stream_encoded_json_with( Method::POST, Self::path(), extra_headers, diff --git a/codex-rs/codex-api/src/endpoint/session.rs b/codex-rs/codex-api/src/endpoint/session.rs index 132c3abd9..7849225b5 100644 --- a/codex-rs/codex-api/src/endpoint/session.rs +++ b/codex-rs/codex-api/src/endpoint/session.rs @@ -2,6 +2,7 @@ use crate::auth::SharedAuthProvider; use crate::error::ApiError; use crate::provider::Provider; use crate::telemetry::run_with_request_telemetry; +use codex_client::EncodedJsonBody; use codex_client::HttpTransport; use codex_client::Request; use codex_client::RequestBody; @@ -49,12 +50,12 @@ impl EndpointSession { method: &Method, path: &str, extra_headers: &HeaderMap, - body: Option<&Value>, + body: Option<&RequestBody>, ) -> Request { let mut req = self.provider.build_request(method.clone(), path); req.headers.extend(extra_headers.clone()); if let Some(body) = body { - req.body = Some(RequestBody::Json(body.clone())); + req.body = Some(body.clone()); } req } @@ -87,6 +88,7 @@ impl EndpointSession { where C: Fn(&mut Request), { + let body = body.map(RequestBody::Json); let make_request = || { let mut req = self.make_request(&method, path, &extra_headers, body.as_ref()); configure(&mut req); @@ -112,27 +114,27 @@ impl EndpointSession { } #[instrument( - name = "endpoint_session.stream_with", + name = "endpoint_session.stream_encoded_json_with", level = "info", skip_all, fields(http.method = %method, api.path = path) )] - pub(crate) async fn stream_with( + pub(crate) async fn stream_encoded_json_with( &self, method: Method, path: &str, extra_headers: HeaderMap, - body: Option, + body: Option, configure: C, ) -> Result where C: Fn(&mut Request), { - let make_request = || { - let mut req = self.make_request(&method, path, &extra_headers, body.as_ref()); - configure(&mut req); - req - }; + let body = body.map(RequestBody::EncodedJson); + let mut request = self.make_request(&method, path, &extra_headers, body.as_ref()); + configure(&mut request); + let request = request.into_prepared().map_err(TransportError::Build)?; + let make_request = || request.clone(); let stream = run_with_request_telemetry( self.provider.retry.to_policy(), diff --git a/codex-rs/codex-api/tests/clients.rs b/codex-rs/codex-api/tests/clients.rs index 34284634f..ae34accb6 100644 --- a/codex-rs/codex-api/tests/clients.rs +++ b/codex-rs/codex-api/tests/clients.rs @@ -36,6 +36,13 @@ fn assert_path_ends_with(requests: &[Request], suffix: &str) { ); } +fn request_body_bytes(request: &Request) -> &[u8] { + let Some(RequestBody::EncodedJson(body)) = request.body.as_ref() else { + panic!("expected a prepared request body"); + }; + body.as_bytes() +} + #[derive(Debug, Default, Clone)] struct RecordingState { stream_requests: Arc>>, @@ -138,9 +145,15 @@ fn provider(name: &str) -> Provider { } } +#[derive(Debug, Default)] +struct FlakyTransportState { + attempts: i64, + requests: Vec<(RequestBody, HeaderMap, codex_client::RequestCompression)>, +} + #[derive(Clone)] struct FlakyTransport { - state: Arc>, + state: Arc>, } impl Default for FlakyTransport { @@ -152,15 +165,23 @@ impl Default for FlakyTransport { impl FlakyTransport { fn new() -> Self { Self { - state: Arc::new(Mutex::new(0)), + state: Arc::new(Mutex::new(FlakyTransportState::default())), } } fn attempts(&self) -> i64 { - *self - .state + self.state .lock() .unwrap_or_else(|err| panic!("mutex poisoned: {err}")) + .attempts + } + + fn requests(&self) -> Vec<(RequestBody, HeaderMap, codex_client::RequestCompression)> { + self.state + .lock() + .unwrap_or_else(|err| panic!("mutex poisoned: {err}")) + .requests + .clone() } } @@ -225,14 +246,20 @@ impl HttpTransport for FlakyTransport { Err(TransportError::Build("execute should not run".to_string())) } - async fn stream(&self, _req: Request) -> Result { - let mut attempts = self + async fn stream(&self, req: Request) -> Result { + let Some(body) = req.body.clone() else { + panic!("request should have a body"); + }; + let mut state = self .state .lock() .unwrap_or_else(|err| panic!("mutex poisoned: {err}")); - *attempts += 1; + state.attempts += 1; + state + .requests + .push((body, req.headers.clone(), req.compression)); - if *attempts == 1 { + if state.attempts == 1 { return Err(TransportError::Network("first attempt fails".to_string())); } @@ -272,6 +299,51 @@ async fn responses_client_uses_responses_path() -> Result<()> { Ok(()) } +#[tokio::test] +async fn responses_client_stream_request_preserves_exact_json_body() -> Result<()> { + let state = RecordingState::default(); + let transport = RecordingTransport::new(state.clone()); + let client = ResponsesClient::new(transport, provider("openai"), Arc::new(NoAuth)); + let request = ResponsesApiRequest { + model: "gpt-test".into(), + instructions: "Say hi".into(), + input: vec![ResponseItem::Message { + id: Some("msg_1".into()), + role: "user".into(), + content: vec![ContentItem::InputText { text: "hi".into() }], + phase: None, + }], + tools: Vec::new(), + tool_choice: "auto".into(), + parallel_tool_calls: false, + reasoning: None, + store: false, + stream: true, + include: Vec::new(), + service_tier: None, + prompt_cache_key: None, + text: None, + client_metadata: None, + }; + let expected = serde_json::to_vec(&request)?; + + let _stream = client + .stream_request(request, ResponsesOptions::default()) + .await?; + + let requests = state.take_stream_requests(); + assert_eq!(requests.len(), 1); + let prepared = requests[0] + .prepare_body_for_send() + .expect("body should prepare"); + assert_eq!(prepared.body.as_deref(), Some(expected.as_slice())); + assert_eq!( + prepared.headers.get(http::header::CONTENT_TYPE), + Some(&HeaderValue::from_static("application/json")) + ); + Ok(()) +} + #[tokio::test] async fn streaming_client_adds_auth_headers() -> Result<()> { let state = RecordingState::default(); @@ -342,12 +414,30 @@ async fn streaming_client_retries_on_transport_error() -> Result<()> { .stream_request( request, ResponsesOptions { - compression: Compression::None, + compression: Compression::Zstd, ..Default::default() }, ) .await?; assert_eq!(transport.attempts(), 2); + let requests = transport.requests(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0], requests[1]); + let RequestBody::EncodedJson(first_body) = &requests[0].0 else { + panic!("expected an encoded JSON body"); + }; + let RequestBody::EncodedJson(second_body) = &requests[1].0 else { + panic!("expected an encoded JSON body"); + }; + assert_eq!( + first_body.as_bytes().as_ptr(), + second_body.as_bytes().as_ptr() + ); + assert_eq!( + requests[0].1.get(http::header::CONTENT_ENCODING), + Some(&HeaderValue::from_static("zstd")) + ); + assert_eq!(requests[0].2, codex_client::RequestCompression::None); Ok(()) } @@ -485,11 +575,9 @@ async fn azure_default_store_attaches_ids_and_headers() -> Result<()> { Some("present") ); - let input_id = req - .body - .as_ref() - .and_then(RequestBody::json) - .and_then(|body| body.get("input")) + let body: serde_json::Value = serde_json::from_slice(request_body_bytes(req))?; + let input_id = body + .get("input") .and_then(|input| input.get(0)) .and_then(|item| item.get("id")) .and_then(|id| id.as_str()); diff --git a/codex-rs/codex-client/src/lib.rs b/codex-rs/codex-client/src/lib.rs index 0f503fb3e..e7c5edfc7 100644 --- a/codex-rs/codex-client/src/lib.rs +++ b/codex-rs/codex-client/src/lib.rs @@ -25,6 +25,7 @@ pub use crate::default_client::CodexHttpClient; pub use crate::default_client::CodexRequestBuilder; pub use crate::error::StreamError; pub use crate::error::TransportError; +pub use crate::request::EncodedJsonBody; pub use crate::request::PreparedRequestBody; pub use crate::request::Request; pub use crate::request::RequestBody; diff --git a/codex-rs/codex-client/src/request.rs b/codex-rs/codex-client/src/request.rs index 5fc076627..6d50f7cb5 100644 --- a/codex-rs/codex-client/src/request.rs +++ b/codex-rs/codex-client/src/request.rs @@ -6,6 +6,38 @@ use serde::Serialize; use serde_json::Value; use std::time::Duration; +/// A JSON request body serialized once into reference-counted bytes. +/// +/// Clones share the encoded allocation. Internally, the body can also hold the +/// final compressed wire bytes while retaining the original JSON only when +/// request-body trace logging is enabled. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EncodedJsonBody { + bytes: Bytes, + trace_bytes: Option, + prepared: bool, +} + +impl EncodedJsonBody { + /// Serializes `value` into a reusable JSON body. + pub fn encode(value: &T) -> Result { + serde_json::to_vec(value).map(|bytes| Self { + bytes: Bytes::from(bytes), + trace_bytes: None, + prepared: false, + }) + } + + /// Returns the encoded bytes currently stored by this body. + pub fn as_bytes(&self) -> &[u8] { + &self.bytes + } + + pub(crate) fn trace_bytes(&self) -> &[u8] { + self.trace_bytes.as_ref().unwrap_or(&self.bytes) + } +} + #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub enum RequestCompression { #[default] @@ -16,6 +48,7 @@ pub enum RequestCompression { #[derive(Debug, Clone, PartialEq, Eq)] pub enum RequestBody { Json(Value), + EncodedJson(EncodedJsonBody), Raw(Bytes), } @@ -23,7 +56,7 @@ impl RequestBody { pub fn json(&self) -> Option<&Value> { match self { Self::Json(value) => Some(value), - Self::Raw(_) => None, + Self::EncodedJson(_) | Self::Raw(_) => None, } } } @@ -77,13 +110,51 @@ impl Request { self } + /// Prepares the body once and stores the exact bytes that will be sent. + /// + /// Cloning the returned request shares the body bytes, so retry attempts do + /// not repeat JSON serialization or compression. Request-signing auth also + /// sees the same final headers and bytes that the transport will send. + pub fn into_prepared(mut self) -> Result { + let is_json = matches!( + self.body, + Some(RequestBody::Json(_) | RequestBody::EncodedJson(_)) + ); + let trace_bytes = if self.compression != RequestCompression::None + && tracing::enabled!(target: "codex_client::transport", tracing::Level::TRACE) + { + match self.body.as_ref() { + Some(RequestBody::Json(body)) => Some(Bytes::from( + serde_json::to_vec(body).map_err(|err| err.to_string())?, + )), + Some(RequestBody::EncodedJson(body)) => Some(body.bytes.clone()), + Some(RequestBody::Raw(_)) | None => None, + } + } else { + None + }; + let prepared = self.prepare_body_for_send()?; + self.headers = prepared.headers; + self.body = match (is_json, prepared.body) { + (true, Some(bytes)) => Some(RequestBody::EncodedJson(EncodedJsonBody { + bytes, + trace_bytes, + prepared: true, + })), + (false, Some(body)) => Some(RequestBody::Raw(body)), + (_, None) => None, + }; + self.compression = RequestCompression::None; + Ok(self) + } + /// Convert the request body into the exact bytes that will be sent. /// /// Auth schemes such as AWS SigV4 need to sign the final body bytes, including /// compression and content headers. Calling this method does not mutate the /// request. pub fn prepare_body_for_send(&self) -> Result { - let mut headers = self.headers.clone(); + let headers = self.headers.clone(); match self.body.as_ref() { Some(RequestBody::Raw(raw_body)) => { if self.compression != RequestCompression::None { @@ -95,60 +166,76 @@ impl Request { }) } Some(RequestBody::Json(body)) => { - let json = serde_json::to_vec(&body).map_err(|err| err.to_string())?; - let bytes = if self.compression != RequestCompression::None { - if headers.contains_key(http::header::CONTENT_ENCODING) { - return Err( - "request compression was requested but content-encoding is already set" - .to_string(), - ); - } - - let pre_compression_bytes = json.len(); - let compression_start = std::time::Instant::now(); - let (compressed, content_encoding) = match self.compression { - RequestCompression::None => unreachable!("guarded by compression != None"), - RequestCompression::Zstd => ( - zstd::stream::encode_all(std::io::Cursor::new(json), 3) - .map_err(|err| err.to_string())?, - HeaderValue::from_static("zstd"), - ), - }; - let post_compression_bytes = compressed.len(); - let compression_duration = compression_start.elapsed(); - - headers.insert(http::header::CONTENT_ENCODING, content_encoding); - - tracing::debug!( - pre_compression_bytes, - post_compression_bytes, - compression_duration_ms = compression_duration.as_millis(), - "Compressed request body with zstd" - ); - - compressed - } else { - json - }; - - if !headers.contains_key(http::header::CONTENT_TYPE) { - headers.insert( - http::header::CONTENT_TYPE, - HeaderValue::from_static("application/json"), - ); - } - - Ok(PreparedRequestBody { - headers, - body: Some(Bytes::from(bytes)), - }) + let body = EncodedJsonBody::encode(body).map_err(|err| err.to_string())?; + self.prepare_encoded_json(headers, &body) } + Some(RequestBody::EncodedJson(body)) => self.prepare_encoded_json(headers, body), None => Ok(PreparedRequestBody { headers, body: None, }), } } + + fn prepare_encoded_json( + &self, + mut headers: HeaderMap, + body: &EncodedJsonBody, + ) -> Result { + if body.prepared { + return Ok(PreparedRequestBody { + headers, + body: Some(body.bytes.clone()), + }); + } + + let bytes = if self.compression != RequestCompression::None { + if headers.contains_key(http::header::CONTENT_ENCODING) { + return Err( + "request compression was requested but content-encoding is already set" + .to_string(), + ); + } + + let pre_compression_bytes = body.bytes.len(); + let compression_start = std::time::Instant::now(); + let (compressed, content_encoding) = match self.compression { + RequestCompression::None => unreachable!("guarded by compression != None"), + RequestCompression::Zstd => ( + zstd::stream::encode_all(std::io::Cursor::new(body.as_bytes()), 3) + .map_err(|err| err.to_string())?, + HeaderValue::from_static("zstd"), + ), + }; + let post_compression_bytes = compressed.len(); + let compression_duration = compression_start.elapsed(); + + headers.insert(http::header::CONTENT_ENCODING, content_encoding); + + tracing::debug!( + pre_compression_bytes, + post_compression_bytes, + compression_duration_ms = compression_duration.as_millis(), + "Compressed request body with zstd" + ); + + Bytes::from(compressed) + } else { + body.bytes.clone() + }; + + if !headers.contains_key(http::header::CONTENT_TYPE) { + headers.insert( + http::header::CONTENT_TYPE, + HeaderValue::from_static("application/json"), + ); + } + + Ok(PreparedRequestBody { + headers, + body: Some(bytes), + }) + } } #[cfg(test)] @@ -205,6 +292,33 @@ mod tests { "request compression was requested but content-encoding is already set" ); } + + #[test] + fn into_prepared_stores_compressed_body_for_reuse() { + let body = + EncodedJsonBody::encode(&json!({"model": "test-model"})).expect("JSON should encode"); + let mut request = + Request::new(Method::POST, "https://example.com/v1/responses".to_string()) + .with_compression(RequestCompression::Zstd); + request.body = Some(RequestBody::EncodedJson(body)); + let request = request.into_prepared().expect("body should prepare"); + let Some(RequestBody::EncodedJson(body)) = request.body.as_ref() else { + panic!("expected an encoded JSON body"); + }; + let decompressed = zstd::stream::decode_all(std::io::Cursor::new(body.as_bytes())) + .expect("body should decompress"); + + assert_eq!(decompressed, br#"{"model":"test-model"}"#); + assert_eq!(request.compression, RequestCompression::None); + assert_eq!( + request.headers.get(http::header::CONTENT_ENCODING), + Some(&HeaderValue::from_static("zstd")) + ); + assert_eq!( + request.headers.get(http::header::CONTENT_TYPE), + Some(&HeaderValue::from_static("application/json")) + ); + } } #[derive(Debug, Clone)] diff --git a/codex-rs/codex-client/src/transport.rs b/codex-rs/codex-client/src/transport.rs index 00b84a614..708b73bf2 100644 --- a/codex-rs/codex-client/src/transport.rs +++ b/codex-rs/codex-client/src/transport.rs @@ -85,6 +85,9 @@ impl ReqwestTransport { fn request_body_for_trace(req: &Request) -> String { match req.body.as_ref() { Some(RequestBody::Json(body)) => body.to_string(), + Some(RequestBody::EncodedJson(body)) => { + String::from_utf8_lossy(body.trace_bytes()).into_owned() + } Some(RequestBody::Raw(body)) => format!("", body.len()), None => String::new(), }