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(), }