From 10877f09fbc7793d64ce830674d8969f647eb08e Mon Sep 17 00:00:00 2001
From: Jiaming Zhang <jm@openai.com>
Date: Thu, 7 May 2026 18:04:26 -0700
Subject: [PATCH] codex: inline attestation header insertion

---
 codex-rs/core/src/client.rs       | 26 ++++++++++++--------------
 codex-rs/core/src/client_tests.rs | 18 +++++++++---------
 2 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs
index d328e143de..789515a355 100644
--- a/codex-rs/core/src/client.rs
+++ b/codex-rs/core/src/client.rs
@@ -498,7 +498,9 @@ impl ModelClient {
             Some(self.state.session_id.to_string()),
             Some(self.state.thread_id.to_string()),
         ));
-        self.extend_attestation_header_for(&mut extra_headers).await;
+        if let Some(header_value) = self.generate_attestation_header_for().await {
+            extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
+        }
         let client =
             ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth)
                 .with_telemetry(Some(request_telemetry));
@@ -520,7 +522,9 @@ impl ModelClient {
         // Create the media call over HTTP first, then retain matching auth so realtime can attach
         // the server-side control WebSocket to the call id from that HTTP response.
         let client_setup = self.current_client_setup().await?;
-        self.extend_attestation_header_for(&mut extra_headers).await;
+        if let Some(header_value) = self.generate_attestation_header_for().await {
+            extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
+        }
         let mut sideband_headers = extra_headers.clone();
         sideband_headers.extend(sideband_websocket_auth_headers(
             client_setup.api_auth.as_ref(),
@@ -900,7 +904,9 @@ impl ModelClient {
         }
         headers.extend(build_session_headers(Some(session_id), Some(thread_id)));
         headers.extend(self.build_responses_identity_headers());
-        self.extend_attestation_header_for(&mut headers).await;
+        if let Some(header_value) = self.generate_attestation_header_for().await {
+            headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
+        }
         headers.insert(
             OPENAI_BETA_HEADER,
             HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE),
@@ -969,9 +975,9 @@ impl ModelClientSession {
                     turn_metadata_header.as_ref(),
                 );
                 headers.extend(self.client.build_responses_identity_headers());
-                self.client
-                    .extend_attestation_header_for(&mut headers)
-                    .await;
+                if let Some(header_value) = self.client.generate_attestation_header_for().await {
+                    headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
+                }
                 headers
             },
             compression,
@@ -1671,14 +1677,6 @@ fn build_responses_headers(
     headers
 }
 
-impl ModelClient {
-    async fn extend_attestation_header_for(&self, headers: &mut ApiHeaderMap) {
-        if let Some(header_value) = self.generate_attestation_header_for().await {
-            headers.insert(X_OAI_ATTESTATION_HEADER, header_value);
-        }
-    }
-}
-
 fn subagent_header_value(session_source: &SessionSource) -> Option<String> {
     match session_source {
         SessionSource::SubAgent(subagent_source) => match subagent_source {
diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs
index 4370bd009a..901eba32bb 100644
--- a/codex-rs/core/src/client_tests.rs
+++ b/codex-rs/core/src/client_tests.rs
@@ -573,17 +573,17 @@ async fn non_chatgpt_codex_endpoints_omit_attestation_generation() {
     let (model_client, attestation_calls) = model_client_with_counting_attestation(false);
     let mut response_headers = http::HeaderMap::new();
 
-    model_client
-        .extend_attestation_header_for(&mut response_headers)
-        .await;
+    if let Some(header_value) = model_client.generate_attestation_header_for().await {
+        response_headers.insert(crate::attestation::X_OAI_ATTESTATION_HEADER, header_value);
+    }
     let mut compaction_headers = http::HeaderMap::new();
-    model_client
-        .extend_attestation_header_for(&mut compaction_headers)
-        .await;
+    if let Some(header_value) = model_client.generate_attestation_header_for().await {
+        compaction_headers.insert(crate::attestation::X_OAI_ATTESTATION_HEADER, header_value);
+    }
     let mut realtime_headers = http::HeaderMap::new();
-    model_client
-        .extend_attestation_header_for(&mut realtime_headers)
-        .await;
+    if let Some(header_value) = model_client.generate_attestation_header_for().await {
+        realtime_headers.insert(crate::attestation::X_OAI_ATTESTATION_HEADER, header_value);
+    }
 
     assert_eq!(
         response_headers.get(crate::attestation::X_OAI_ATTESTATION_HEADER),