From b05a1ffb8edc6cb6c0a390985630a93b1d4cdfff Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Fri, 12 Jun 2026 18:21:48 +0000 Subject: [PATCH] Make turn state replaceable --- codex-rs/codex-api/src/endpoint/responses.rs | 6 +-- .../src/endpoint/responses_websocket.rs | 10 ++-- codex-rs/codex-api/src/sse/responses.rs | 8 +-- codex-rs/core/src/client.rs | 41 ++++++-------- codex-rs/core/tests/suite/turn_state.rs | 54 +++++++++++++++++++ 5 files changed, 83 insertions(+), 36 deletions(-) diff --git a/codex-rs/codex-api/src/endpoint/responses.rs b/codex-rs/codex-api/src/endpoint/responses.rs index cc1be2846a..6f74a909ce 100644 --- a/codex-rs/codex-api/src/endpoint/responses.rs +++ b/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 { @@ -35,7 +35,7 @@ pub struct ResponsesOptions { pub session_source: Option, pub extra_headers: HeaderMap, pub compression: Compression, - pub turn_state: Option>>, + pub turn_state: Option>>, } impl ResponsesClient { @@ -119,7 +119,7 @@ impl ResponsesClient { body: Value, extra_headers: HeaderMap, compression: Compression, - turn_state: Option>>, + turn_state: Option>>, ) -> Result { let request_compression = match compression { Compression::None => RequestCompression::None, diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 44cb6a544b..33fe870aee 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/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>>, + turn_state: Option>>, telemetry: Option>, ) -> Result { let ws_url = self @@ -436,7 +436,7 @@ fn merge_request_headers( async fn connect_websocket( url: Url, headers: HeaderMap, - turn_state: Option>>, + turn_state: Option>>, ) -> Result<(WsStream, StatusCode, bool, Option, Option), 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), diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index ab1be640df..defa07833d 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/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>, - turn_state: Option>>, + turn_state: Option>>, ) -> 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::>(1600); tokio::spawn(async move { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 6c4058bb01..4ae9ee712e 100644 --- a/codex-rs/core/src/client.rs +++ b/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>, + turn_state: Arc>, } #[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>>, + turn_state: Option>>, auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, ) -> std::result::Result { @@ -892,7 +881,7 @@ impl ModelClient { async fn build_websocket_headers( &self, responses_metadata: &CodexResponsesMetadata, - turn_state: Option<&Arc>>, + turn_state: Option<&Arc>>, ) -> 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>>, + turn_state: Option<&Arc>>, ) -> 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 } diff --git a/codex-rs/core/tests/suite/turn_state.rs b/codex-rs/core/tests/suite/turn_state.rs index 7a930af606..8bebf327bb 100644 --- a/codex-rs/core/tests/suite/turn_state.rs +++ b/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(()));