Round-trip turn state through compact

codex-rs/codex-api/src/endpoint/compact.rs

@@ -3,6 +3,7 @@ use crate::common::CompactionInput;
 use crate::endpoint::session::EndpointSession;
 use crate::error::ApiError;
 use crate::provider::Provider;
+use crate::sse::responses::capture_turn_state;
 use codex_client::HttpTransport;
 use codex_client::RequestTelemetry;
 use codex_protocol::models::ResponseItem;
@@ -11,6 +12,7 @@ use http::Method;
 use serde::Deserialize;
 use serde_json::to_value;
 use std::sync::Arc;
+use std::sync::Mutex;
 use std::time::Duration;
 
 pub struct CompactClient<T: HttpTransport> {
@@ -39,6 +41,7 @@ impl<T: HttpTransport> CompactClient<T> {
         body: serde_json::Value,
         extra_headers: HeaderMap,
         request_timeout: Duration,
+        turn_state: Option<&Mutex<String>>,
     ) -> Result<Vec<ResponseItem>, ApiError> {
         let resp = self
             .session
@@ -52,6 +55,9 @@ impl<T: HttpTransport> CompactClient<T> {
                 },
             )
             .await?;
+        if let Some(turn_state) = turn_state {
+            capture_turn_state(&resp.headers, turn_state);
+        }
         let parsed: CompactHistoryResponse =
             serde_json::from_slice(&resp.body).map_err(|e| ApiError::Stream(e.to_string()))?;
         Ok(parsed.output)
@@ -62,10 +68,12 @@ impl<T: HttpTransport> CompactClient<T> {
         input: &CompactionInput<'_>,
         extra_headers: HeaderMap,
         request_timeout: Duration,
+        turn_state: Option<&Mutex<String>>,
     ) -> Result<Vec<ResponseItem>, ApiError> {
         let body = to_value(input)
             .map_err(|e| ApiError::Stream(format!("failed to encode compaction input: {e}")))?;
-        self.compact(body, extra_headers, request_timeout).await
+        self.compact(body, extra_headers, request_timeout, turn_state)
+            .await
     }
 }
 

codex-rs/codex-api/src/sse/responses.rs

@@ -23,6 +23,7 @@ use tracing::debug;
 use tracing::trace;
 
 const X_REASONING_INCLUDED_HEADER: &str = "x-reasoning-included";
+const X_CODEX_TURN_STATE_HEADER: &str = "x-codex-turn-state";
 const OPENAI_MODEL_HEADER: &str = "openai-model";
 const REQUEST_ID_HEADER: &str = "x-request-id";
 const TRUSTED_ACCESS_FOR_CYBER_VERIFICATION: &str = "trusted_access_for_cyber";
@@ -53,15 +54,8 @@ pub fn spawn_response_stream(
         .get(REQUEST_ID_HEADER)
         .and_then(|value| value.to_str().ok())
         .map(str::to_string);
-    if let Some(turn_state) = turn_state.as_ref()
-        && let Some(header_value) = stream_response
-            .headers
-            .get("x-codex-turn-state")
-            .and_then(|v| v.to_str().ok())
-    {
-        *turn_state
-            .lock()
-            .unwrap_or_else(std::sync::PoisonError::into_inner) = header_value.to_string();
+    if let Some(turn_state) = turn_state.as_deref() {
+        capture_turn_state(&stream_response.headers, turn_state);
     }
     let (tx_event, rx_event) = mpsc::channel::<Result<ResponseEvent, ApiError>>(1600);
     tokio::spawn(async move {
@@ -88,6 +82,17 @@ pub fn spawn_response_stream(
     }
 }
 
+pub(crate) fn capture_turn_state(headers: &http::HeaderMap, turn_state: &Mutex<String>) {
+    if let Some(header_value) = headers
+        .get(X_CODEX_TURN_STATE_HEADER)
+        .and_then(|value| value.to_str().ok())
+    {
+        *turn_state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner) = header_value.to_string();
+    }
+}
+
 #[derive(Debug, Deserialize)]
 #[allow(dead_code)]
 struct Error {