reuse encoded Responses request bodies (#28327)

## 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`
2026-07-01 00:31:56 +08:00 · 2026-06-15 18:11:26 +01:00
parent bed60e3746
commit 495da45643
6 changed files with 306 additions and 78 deletions
@@ -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<T: HttpTransport> ResponsesClient<T> {
            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<T: HttpTransport> ResponsesClient<T> {
            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<T: HttpTransport> ResponsesClient<T> {
        extra_headers: HeaderMap,
        compression: Compression,
        turn_state: Option<Arc<OnceLock<String>>>,
+    ) -> Result<ResponseStream, ApiError> {
+        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<Arc<OnceLock<String>>>,
    ) -> Result<ResponseStream, ApiError> {
        let request_compression = match compression {
            Compression::None => RequestCompression::None,
@@ -128,7 +148,7 @@ impl<T: HttpTransport> ResponsesClient<T> {

        let stream_response = self
            .session
-            .stream_with(
+            .stream_encoded_json_with(
                Method::POST,
                Self::path(),
                extra_headers,
@@ -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<T: HttpTransport> EndpointSession<T> {
        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<T: HttpTransport> EndpointSession<T> {
    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<T: HttpTransport> EndpointSession<T> {
    }

    #[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<C>(
+    pub(crate) async fn stream_encoded_json_with<C>(
        &self,
        method: Method,
        path: &str,
        extra_headers: HeaderMap,
-        body: Option<Value>,
+        body: Option<EncodedJsonBody>,
        configure: C,
    ) -> Result<StreamResponse, ApiError>
    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(),
@@ -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<Mutex<Vec<Request>>>,
@@ -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<Mutex<i64>>,
+    state: Arc<Mutex<FlakyTransportState>>,
 }

 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<StreamResponse, TransportError> {
-        let mut attempts = self
+    async fn stream(&self, req: Request) -> Result<StreamResponse, TransportError> {
+        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());
@@ -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;
@@ -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<Bytes>,
+    prepared: bool,
+}
+
+impl EncodedJsonBody {
+    /// Serializes `value` into a reusable JSON body.
+    pub fn encode<T: Serialize + ?Sized>(value: &T) -> Result<Self, serde_json::Error> {
+        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<Self, String> {
+        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<PreparedRequestBody, String> {
-        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<PreparedRequestBody, String> {
+        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)]
@@ -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!("<raw body: {} bytes>", body.len()),
        None => String::new(),
    }