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 {