Make turn state replaceable

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

@@ -20,7 +20,7 @@ use http::HeaderValue;
 use http::Method;
 use serde_json::Value;
 use std::sync::Arc;
-use std::sync::OnceLock;
+use std::sync::Mutex;
 use tracing::instrument;
 
 pub struct ResponsesClient<T: HttpTransport> {
@@ -35,7 +35,7 @@ pub struct ResponsesOptions {
     pub session_source: Option<SessionSource>,
     pub extra_headers: HeaderMap,
     pub compression: Compression,
-    pub turn_state: Option<Arc<OnceLock<String>>>,
+    pub turn_state: Option<Arc<Mutex<String>>>,
 }
 
 impl<T: HttpTransport> ResponsesClient<T> {
@@ -119,7 +119,7 @@ impl<T: HttpTransport> ResponsesClient<T> {
         body: Value,
         extra_headers: HeaderMap,
         compression: Compression,
-        turn_state: Option<Arc<OnceLock<String>>>,
+        turn_state: Option<Arc<Mutex<String>>>,
     ) -> Result<ResponseStream, ApiError> {
         let request_compression = match compression {
             Compression::None => RequestCompression::None,

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

@@ -21,7 +21,7 @@ use serde::Deserialize;
 use serde_json::Value;
 use serde_json::map::Map as JsonMap;
 use std::sync::Arc;
-use std::sync::OnceLock;
+use std::sync::Mutex as StdMutex;
 use std::time::Duration;
 use tokio::net::TcpStream;
 use tokio::sync::Mutex;
@@ -335,7 +335,7 @@ impl ResponsesWebsocketClient {
         &self,
         extra_headers: HeaderMap,
         default_headers: HeaderMap,
-        turn_state: Option<Arc<OnceLock<String>>>,
+        turn_state: Option<Arc<StdMutex<String>>>,
         telemetry: Option<Arc<dyn WebsocketTelemetry>>,
     ) -> Result<ResponsesWebsocketConnection, ApiError> {
         let ws_url = self
@@ -436,7 +436,7 @@ fn merge_request_headers(
 async fn connect_websocket(
     url: Url,
     headers: HeaderMap,
-    turn_state: Option<Arc<OnceLock<String>>>,
+    turn_state: Option<Arc<StdMutex<String>>>,
 ) -> Result<(WsStream, StatusCode, bool, Option<String>, Option<String>), ApiError> {
     ensure_rustls_crypto_provider();
     info!("connecting to websocket: {url}");
@@ -493,7 +493,9 @@ async fn connect_websocket(
             .get(X_CODEX_TURN_STATE_HEADER)
             .and_then(|value| value.to_str().ok())
     {
-        let _ = turn_state.set(header_value.to_string());
+        *turn_state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner) = header_value.to_string();
     }
     Ok((
         WsStream::new(stream),

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

@@ -14,7 +14,7 @@ use futures::StreamExt;
 use serde::Deserialize;
 use serde_json::Value;
 use std::sync::Arc;
-use std::sync::OnceLock;
+use std::sync::Mutex;
 use std::time::Duration;
 use tokio::sync::mpsc;
 use tokio::time::Instant;
@@ -31,7 +31,7 @@ pub fn spawn_response_stream(
     stream_response: StreamResponse,
     idle_timeout: Duration,
     telemetry: Option<Arc<dyn SseTelemetry>>,
-    turn_state: Option<Arc<OnceLock<String>>>,
+    turn_state: Option<Arc<Mutex<String>>>,
 ) -> ResponseStream {
     let rate_limit_snapshots = parse_all_rate_limits(&stream_response.headers);
     let models_etag = stream_response
@@ -59,7 +59,9 @@ pub fn spawn_response_stream(
             .get("x-codex-turn-state")
             .and_then(|v| v.to_str().ok())
     {
-        let _ = turn_state.set(header_value.to_string());
+        *turn_state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner) = header_value.to_string();
     }
     let (tx_event, rx_event) = mpsc::channel::<Result<ResponseEvent, ApiError>>(1600);
     tokio::spawn(async move {

codex-rs/core/src/client.rs

@@ -26,7 +26,6 @@
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::sync::Mutex as StdMutex;
-use std::sync::OnceLock;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 
@@ -238,17 +237,7 @@ pub struct ModelClient {
 pub struct ModelClientSession {
     client: ModelClient,
     websocket_session: WebsocketSession,
-    /// Turn state for sticky routing.
-    ///
-    /// This is an `OnceLock` that stores the turn state value received from the server
-    /// on turn start via the `x-codex-turn-state` response header. Once set, this value
-    /// should be sent back to the server in the `x-codex-turn-state` request header for
-    /// all subsequent requests within the same turn to maintain sticky routing.
-    ///
-    /// This is a contract between the client and server: we receive it at turn start,
-    /// keep sending it unchanged between turn requests (e.g., for retries, incremental
-    /// appends, or continuation requests), and must not send it between different turns.
-    turn_state: Arc<OnceLock<String>>,
+    turn_state: Arc<StdMutex<String>>,
 }
 
 #[derive(Debug, Clone)]
@@ -375,7 +364,7 @@ impl ModelClient {
         ModelClientSession {
             client: self.clone(),
             websocket_session: self.take_cached_websocket_session(),
-            turn_state: Arc::new(OnceLock::new()),
+            turn_state: Arc::new(StdMutex::new(String::new())),
         }
     }
 
@@ -806,7 +795,7 @@ impl ModelClient {
         api_provider: codex_api::Provider,
         api_auth: SharedAuthProvider,
         responses_metadata: &CodexResponsesMetadata,
-        turn_state: Option<Arc<OnceLock<String>>>,
+        turn_state: Option<Arc<StdMutex<String>>>,
         auth_context: AuthRequestTelemetryContext,
         request_route_telemetry: RequestRouteTelemetry,
     ) -> std::result::Result<ApiWebSocketConnection, ApiError> {
@@ -892,7 +881,7 @@ impl ModelClient {
     async fn build_websocket_headers(
         &self,
         responses_metadata: &CodexResponsesMetadata,
-        turn_state: Option<&Arc<OnceLock<String>>>,
+        turn_state: Option<&Arc<StdMutex<String>>>,
     ) -> ApiHeaderMap {
         let mut headers =
             build_responses_headers(self.state.beta_features_header.as_deref(), turn_state);
@@ -1133,10 +1122,6 @@ impl ModelClientSession {
             self.websocket_session.last_request = None;
             self.websocket_session.last_response_rx = None;
             self.websocket_session.last_response_from_untraced_warmup = false;
-            let turn_state = options
-                .turn_state
-                .clone()
-                .unwrap_or_else(|| Arc::clone(&self.turn_state));
             let new_conn = match self
                 .client
                 .connect_websocket(
@@ -1144,7 +1129,7 @@ impl ModelClientSession {
                     api_provider,
                     api_auth,
                     responses_metadata,
-                    Some(turn_state),
+                    options.turn_state.clone(),
                     auth_context,
                     request_route_telemetry,
                 )
@@ -1656,7 +1641,7 @@ fn stamp_ws_stream_request_start_ms(request: &mut ResponsesWsRequest) {
 /// - `x-codex-turn-state`: sticky routing token captured earlier in the turn.
 fn build_responses_headers(
     beta_features_header: Option<&str>,
-    turn_state: Option<&Arc<OnceLock<String>>>,
+    turn_state: Option<&Arc<StdMutex<String>>>,
 ) -> ApiHeaderMap {
     let mut headers = ApiHeaderMap::new();
     if let Some(value) = beta_features_header
@@ -1665,11 +1650,15 @@ fn build_responses_headers(
     {
         headers.insert("x-codex-beta-features", header_value);
     }
-    if let Some(turn_state) = turn_state
-        && let Some(state) = turn_state.get()
-        && let Ok(header_value) = HeaderValue::from_str(state)
-    {
-        headers.insert(X_CODEX_TURN_STATE_HEADER, header_value);
+    if let Some(turn_state) = turn_state {
+        let state = turn_state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if !state.is_empty()
+            && let Ok(header_value) = HeaderValue::from_str(&state)
+        {
+            headers.insert(X_CODEX_TURN_STATE_HEADER, header_value);
+        }
     }
     headers
 }

codex-rs/core/tests/suite/turn_state.rs

@@ -87,6 +87,60 @@ async fn responses_turn_state_persists_within_turn_and_resets_after() -> Result<
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn responses_turn_state_updates_are_replayed_within_turn() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = start_mock_server().await;
+    let responses = vec![
+        sse_response(sse(vec![
+            ev_response_created("resp-1"),
+            ev_shell_command_call("shell-1", "echo one"),
+            ev_completed("resp-1"),
+        ]))
+        .insert_header(TURN_STATE_HEADER, "ts-1"),
+        sse_response(sse(vec![
+            ev_response_created("resp-2"),
+            ev_shell_command_call("shell-2", "echo two"),
+            ev_completed("resp-2"),
+        ]))
+        .insert_header(TURN_STATE_HEADER, "ts-2"),
+        sse_response(sse(vec![
+            ev_response_created("resp-3"),
+            ev_assistant_message("msg-1", "done"),
+            ev_completed("resp-3"),
+        ])),
+        sse_response(sse(vec![
+            ev_response_created("resp-4"),
+            ev_assistant_message("msg-2", "done"),
+            ev_completed("resp-4"),
+        ])),
+    ];
+    let request_log = mount_response_sequence(&server, responses).await;
+    let test = test_codex().build(&server).await?;
+
+    // Phase 1: the first response mints state for the first same-turn follow-up.
+    // Phase 2: the next response replaces it before the second follow-up.
+    test.submit_turn("run two shell commands").await?;
+    // Phase 3: a new logical turn starts without state from the completed turn.
+    test.submit_turn("start a new turn").await?;
+
+    let requests = request_log.requests();
+    assert_eq!(requests.len(), 4);
+    assert_eq!(requests[0].header(TURN_STATE_HEADER), None);
+    assert_eq!(
+        requests[1].header(TURN_STATE_HEADER).as_deref(),
+        Some("ts-1")
+    );
+    assert_eq!(
+        requests[2].header(TURN_STATE_HEADER).as_deref(),
+        Some("ts-2")
+    );
+    assert_eq!(requests[3].header(TURN_STATE_HEADER), None);
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn websocket_turn_state_persists_within_turn_and_resets_after() -> Result<()> {
     skip_if_no_network!(Ok(()));