diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 44cb6a544..559a8faf0 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -215,6 +215,7 @@ impl ResponsesWebsocketConnection { &self, request: ResponsesWsRequest, connection_reused: bool, + turn_state: Option>>, ) -> Result { let (tx_event, rx_event) = mpsc::channel::>(1600); @@ -264,6 +265,7 @@ impl ResponsesWebsocketConnection { idle_timeout, telemetry, connection_reused, + turn_state.as_deref(), ) .await }; @@ -631,6 +633,7 @@ async fn run_websocket_response_stream( idle_timeout: Duration, telemetry: Option>, connection_reused: bool, + turn_state: Option<&OnceLock>, ) -> Result<(), ApiError> { let mut last_server_model: Option = None; send_websocket_request( @@ -682,6 +685,11 @@ async fn run_websocket_response_stream( continue; } }; + if let Some(response_turn_state) = event.turn_state() + && let Some(turn_state) = turn_state + { + let _ = turn_state.set(response_turn_state); + } let model_verifications = event.model_verifications(); let turn_moderation_metadata = event.turn_moderation_metadata(); if event.kind() == "codex.rate_limits" { diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index ab1be640d..a451f67af 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/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"; @@ -56,8 +57,8 @@ pub fn spawn_response_stream( 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()) + .get(X_CODEX_TURN_STATE_HEADER) + .and_then(|value| value.to_str().ok()) { let _ = turn_state.set(header_value.to_string()); } @@ -184,6 +185,16 @@ impl ResponsesStreamEvent { } } + pub(crate) fn turn_state(&self) -> Option { + if self.kind() != "response.metadata" { + return None; + } + + self.headers + .as_ref() + .and_then(header_turn_state_value_from_json) + } + pub(crate) fn model_verifications(&self) -> Option> { if self.kind() != "response.metadata" { return None; @@ -220,6 +231,17 @@ fn header_openai_model_value_from_json(value: &Value) -> Option { }) } +fn header_turn_state_value_from_json(value: &Value) -> Option { + let headers = value.as_object()?; + headers.iter().find_map(|(name, value)| { + if name.eq_ignore_ascii_case(X_CODEX_TURN_STATE_HEADER) { + json_value_as_string(value) + } else { + None + } + }) +} + fn model_verifications_from_json_value(value: &Value) -> Option> { let verifications = value .as_array() diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 6c4058bb0..94ddc036a 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -806,13 +806,10 @@ impl ModelClient { api_provider: codex_api::Provider, api_auth: SharedAuthProvider, responses_metadata: &CodexResponsesMetadata, - turn_state: Option>>, auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, ) -> std::result::Result { - let headers = self - .build_websocket_headers(responses_metadata, turn_state.as_ref()) - .await; + let headers = self.build_websocket_headers(responses_metadata).await; let websocket_telemetry = ModelClientSession::build_websocket_telemetry( session_telemetry, auth_context, @@ -826,7 +823,7 @@ impl ModelClient { ApiWebSocketResponsesClient::new(api_provider, api_auth).connect( headers, codex_login::default_client::default_headers(), - turn_state, + /*turn_state*/ None, Some(websocket_telemetry), ), ) @@ -886,16 +883,14 @@ impl ModelClient { } /// Builds websocket handshake headers for both prewarm and turn-time reconnect. - /// - /// Callers should pass the current turn-state lock when available so sticky-routing state is - /// replayed on reconnect within the same turn. async fn build_websocket_headers( &self, responses_metadata: &CodexResponsesMetadata, - turn_state: Option<&Arc>>, ) -> ApiHeaderMap { - let mut headers = - build_responses_headers(self.state.beta_features_header.as_deref(), turn_state); + let mut headers = build_responses_headers( + self.state.beta_features_header.as_deref(), + /*turn_state*/ None, + ); if let Ok(header_value) = HeaderValue::from_str(&responses_metadata.thread_id) { headers.insert("x-client-request-id", header_value); } @@ -1088,7 +1083,6 @@ impl ModelClientSession { client_setup.api_provider, client_setup.api_auth, responses_metadata, - Some(Arc::clone(&self.turn_state)), auth_context, RequestRouteTelemetry::for_endpoint(RESPONSES_ENDPOINT), ) @@ -1120,7 +1114,6 @@ impl ModelClientSession { api_provider, api_auth, responses_metadata, - options, auth_context, request_route_telemetry, } = params; @@ -1133,10 +1126,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 +1133,6 @@ impl ModelClientSession { api_provider, api_auth, responses_metadata, - Some(turn_state), auth_context, request_route_telemetry, ) @@ -1346,15 +1334,6 @@ impl ModelClientSession { client_setup.api_auth.as_ref(), pending_retry, ); - let compression = self.responses_request_compression(client_setup.auth.as_ref()); - - let options = self - .build_responses_options( - responses_metadata, - compression, - model_info.use_responses_lite, - ) - .await; let request = self.client.build_responses_request( &client_setup.api_provider, prompt, @@ -1364,12 +1343,15 @@ impl ModelClientSession { service_tier.clone(), responses_metadata, )?; + let mut client_metadata = self + .client + .build_ws_client_metadata(responses_metadata, model_info.use_responses_lite); + if let Some(turn_state) = self.turn_state.get() { + client_metadata.insert(X_CODEX_TURN_STATE_HEADER.to_string(), turn_state.clone()); + } let mut ws_payload = ResponseCreateWsRequest { client_metadata: response_create_client_metadata( - Some(self.client.build_ws_client_metadata( - responses_metadata, - model_info.use_responses_lite, - )), + Some(client_metadata), request_trace.as_ref(), ), ..ResponseCreateWsRequest::from(&request) @@ -1384,7 +1366,6 @@ impl ModelClientSession { api_provider: client_setup.api_provider, api_auth: client_setup.api_auth, responses_metadata, - options: &options, auth_context: request_auth_context, request_route_telemetry: RequestRouteTelemetry::for_endpoint( RESPONSES_ENDPOINT, @@ -1441,7 +1422,11 @@ impl ModelClientSession { )) })?; let stream_result = websocket_connection - .stream_request(ws_request, self.websocket_session.connection_reused()) + .stream_request( + ws_request, + self.websocket_session.connection_reused(), + Some(Arc::clone(&self.turn_state)), + ) .await .map_err(|err| { let response_debug_context = @@ -1920,7 +1905,6 @@ struct WebsocketConnectParams<'a> { api_provider: codex_api::Provider, api_auth: SharedAuthProvider, responses_metadata: &'a CodexResponsesMetadata, - options: &'a ApiResponsesOptions, auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 18cd20223..29f397d5c 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -585,7 +585,7 @@ async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() ); let headers = model_client - .build_websocket_headers(&responses_metadata, /*turn_state*/ None) + .build_websocket_headers(&responses_metadata) .await; assert_eq!( diff --git a/codex-rs/core/tests/suite/turn_state.rs b/codex-rs/core/tests/suite/turn_state.rs index 7a930af60..6a9c70714 100644 --- a/codex-rs/core/tests/suite/turn_state.rs +++ b/codex-rs/core/tests/suite/turn_state.rs @@ -16,6 +16,7 @@ use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; use serde_json::Value; +use serde_json::json; const TURN_STATE_HEADER: &str = "x-codex-turn-state"; @@ -91,56 +92,111 @@ async fn responses_turn_state_persists_within_turn_and_resets_after() -> Result< async fn websocket_turn_state_persists_within_turn_and_resets_after() -> Result<()> { skip_if_no_network!(Ok(())); - let call_id = "ws-shell-turn-state"; - // First connection delivers turn_state; second (same turn) must send it; third (new turn) must not. - let server = start_websocket_server_with_headers(vec![ - WebSocketConnectionConfig { - requests: vec![vec![ + let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig { + requests: vec![ + vec![ + json!({ + "type": "response.metadata", + "headers": {(TURN_STATE_HEADER): "ts-1"}, + }), ev_response_created("resp-1"), ev_reasoning_item("rsn-1", &["thinking"], &[]), - ev_shell_command_call(call_id, "echo websocket"), + ev_shell_command_call("ws-shell-turn-state", "echo websocket"), ev_completed("resp-1"), - ]], - response_headers: vec![(TURN_STATE_HEADER.to_string(), "ts-1".to_string())], - accept_delay: None, - close_after_requests: true, - }, - WebSocketConnectionConfig { - requests: vec![vec![ + ], + vec![ ev_response_created("resp-2"), ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), - ]], - response_headers: Vec::new(), - accept_delay: None, - close_after_requests: true, - }, - WebSocketConnectionConfig { - requests: vec![vec![ + ], + vec![ ev_response_created("resp-3"), ev_assistant_message("msg-2", "done"), ev_completed("resp-3"), - ]], - response_headers: Vec::new(), - accept_delay: None, - close_after_requests: true, - }, - ]) + ], + ], + response_headers: Vec::new(), + accept_delay: None, + close_after_requests: false, + }]) .await; let mut builder = test_codex(); let test = builder.build_with_websocket_server(&server).await?; + // Phase 1: the first response mints state for its same-turn tool follow-up. test.submit_turn("run the echo command").await?; - test.submit_turn("second turn").await?; + // Phase 2: the follow-up replays that state on the same physical connection. + // Phase 3: the next logical turn reuses the connection but starts with empty state. + test.submit_turn("start another turn").await?; - let handshakes = server.handshakes(); - assert_eq!(handshakes.len(), 3); - assert_eq!(handshakes[0].header(TURN_STATE_HEADER), None); + assert_eq!(server.handshakes().len(), 1); + let requests = server.single_connection(); + assert_eq!(requests.len(), 3); assert_eq!( - handshakes[1].header(TURN_STATE_HEADER), - Some("ts-1".to_string()) + requests + .iter() + .map(|request| request.body_json()["client_metadata"][TURN_STATE_HEADER].clone()) + .collect::>(), + vec![json!(null), json!("ts-1"), json!(null)] + ); + + server.shutdown().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn websocket_turn_state_is_stable_within_turn() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig { + requests: vec![ + vec![ + json!({ + "type": "response.metadata", + "headers": {(TURN_STATE_HEADER): "ts-1"}, + }), + ev_response_created("resp-1"), + ev_shell_command_call("ws-shell-1", "echo one"), + ev_completed("resp-1"), + ], + vec![ + json!({ + "type": "response.metadata", + "headers": {(TURN_STATE_HEADER): "ts-2"}, + }), + ev_response_created("resp-2"), + ev_shell_command_call("ws-shell-2", "echo two"), + ev_completed("resp-2"), + ], + vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-3"), + ], + ], + response_headers: Vec::new(), + accept_delay: None, + close_after_requests: false, + }]) + .await; + let mut builder = test_codex(); + let test = builder.build_with_websocket_server(&server).await?; + + // Phase 1: the initial request starts empty and receives the first metadata value. + // Phase 2: the first tool follow-up replays it and ignores a later value. + // Phase 3: the second follow-up sends the original value on the same connection. + test.submit_turn("run two echo commands").await?; + + assert_eq!(server.handshakes().len(), 1); + let requests = server.single_connection(); + assert_eq!(requests.len(), 3); + assert_eq!( + requests + .iter() + .map(|request| request.body_json()["client_metadata"][TURN_STATE_HEADER].clone()) + .collect::>(), + vec![json!(null), json!("ts-1"), json!("ts-1")] ); - assert_eq!(handshakes[2].header(TURN_STATE_HEADER), None); server.shutdown().await; Ok(())