diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index d8dfad9621..d9e938befc 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -142,6 +142,8 @@ struct ModelClientState { /// This keeps startup preconnect task tracking and warmed-socket adoption in one lock so /// turn-time websocket setup observes a single, coherent state. preconnect: Mutex, + /// Session-scoped websocket v2 state shared across turns. + responses_websocket_v2: Mutex, } impl std::fmt::Debug for ModelClientState { @@ -167,6 +169,7 @@ impl std::fmt::Debug for ModelClientState { &self.disable_websockets.load(Ordering::Relaxed), ) .field("preconnect", &"") + .field("responses_websocket_v2", &"") .finish() } } @@ -186,7 +189,7 @@ struct CurrentClientSetup { /// This bundles the socket with optional sticky-routing state captured during /// handshake so they are taken and cleared atomically. struct PreconnectedWebSocket { - connection: ApiWebSocketConnection, + connection: Arc, turn_state: Option, } @@ -203,6 +206,13 @@ enum PreconnectState { Ready(PreconnectedWebSocket), } +#[derive(Default)] +struct ResponsesWebsocketV2State { + connection: Option>, + websocket_last_items: Vec, + websocket_last_response_id: Option, +} + /// A session-scoped client for model-provider API calls. /// /// This holds configuration and state that should be shared across turns within a Codex session @@ -240,7 +250,7 @@ pub struct ModelClient { /// contract and can cause routing bugs. pub struct ModelClientSession { client: ModelClient, - connection: Option, + connection: Option>, websocket_last_items: Vec, websocket_last_response_id: Option, websocket_last_response_id_rx: Option>, @@ -289,6 +299,7 @@ impl ModelClient { beta_features_header, disable_websockets: AtomicBool::new(false), preconnect: Mutex::new(PreconnectState::Idle), + responses_websocket_v2: Mutex::new(ResponsesWebsocketV2State::default()), }), } } @@ -493,6 +504,63 @@ impl ModelClient { self.state.enable_responses_websockets_v2 } + fn with_shared_v2_state( + &self, + action: impl FnOnce(&mut ResponsesWebsocketV2State) -> R, + ) -> R { + let mut state = self + .state + .responses_websocket_v2 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + action(&mut state) + } + + fn shared_v2_connection(&self) -> Option> { + self.with_shared_v2_state(|state| state.connection.clone()) + } + + fn set_shared_v2_connection(&self, connection: Option>) { + self.with_shared_v2_state(|state| { + state.connection = connection; + }); + } + + fn shared_v2_request_chain_state(&self) -> (Vec, Option) { + self.with_shared_v2_state(|state| { + ( + state.websocket_last_items.clone(), + state.websocket_last_response_id.clone(), + ) + }) + } + + fn set_shared_v2_request_chain_state( + &self, + websocket_last_items: Vec, + response_id: String, + ) { + self.with_shared_v2_state(|state| { + state.websocket_last_items = websocket_last_items; + state.websocket_last_response_id = Some(response_id); + }); + } + + fn clear_shared_v2_request_chain_state(&self) { + self.with_shared_v2_state(|state| { + state.websocket_last_items.clear(); + state.websocket_last_response_id = None; + }); + } + + fn clear_shared_v2_state(&self) { + self.with_shared_v2_state(|state| { + state.connection = None; + state.websocket_last_items.clear(); + state.websocket_last_response_id = None; + }); + } + /// Returns whether websocket transport has been permanently disabled for this session. /// /// Once set by fallback activation, subsequent turns must stay on HTTP transport. @@ -612,7 +680,7 @@ impl ModelClient { return; } *state = PreconnectState::Ready(PreconnectedWebSocket { - connection, + connection: Arc::new(connection), turn_state, }); } @@ -875,6 +943,17 @@ impl ModelClientSession { options: &ApiResponsesOptions, ) -> ResponsesWsRequest { let responses_websockets_v2_enabled = self.client.responses_websockets_v2_enabled(); + if responses_websockets_v2_enabled + && self.websocket_last_items.is_empty() + && self.websocket_last_response_id.is_none() + && self.websocket_last_response_id_rx.is_none() + { + let (websocket_last_items, websocket_last_response_id) = + self.client.shared_v2_request_chain_state(); + self.websocket_last_items = websocket_last_items; + self.websocket_last_response_id = websocket_last_response_id; + } + let incremental_items = self.get_incremental_items(&api_prompt.input); if let Some(append_items) = incremental_items { if responses_websockets_v2_enabled @@ -920,6 +999,11 @@ impl ModelClientSession { turn_metadata_header: Option<&str>, options: &ApiResponsesOptions, ) -> std::result::Result<&ApiWebSocketConnection, ApiError> { + let responses_websockets_v2_enabled = self.client.responses_websockets_v2_enabled(); + if responses_websockets_v2_enabled && self.connection.is_none() { + self.connection = self.client.shared_v2_connection(); + } + // Prefer the session-level preconnect slot before creating a new websocket. if self.connection.is_none() { if let Some(preconnected) = self.try_use_preconnected_websocket() { @@ -942,24 +1026,36 @@ impl ModelClientSession { self.websocket_last_items.clear(); self.websocket_last_response_id = None; self.websocket_last_response_id_rx = None; + if responses_websockets_v2_enabled { + self.client.clear_shared_v2_request_chain_state(); + } let turn_state = options .turn_state .clone() .unwrap_or_else(|| Arc::clone(&self.turn_state)); - let new_conn = self - .client - .connect_websocket( - otel_manager, - api_provider, - api_auth, - Some(turn_state), - turn_metadata_header, - ) - .await?; + let new_conn = Arc::new( + self.client + .connect_websocket( + otel_manager, + api_provider, + api_auth, + Some(turn_state), + turn_metadata_header, + ) + .await?, + ); + if responses_websockets_v2_enabled { + self.client + .set_shared_v2_connection(Some(Arc::clone(&new_conn))); + } self.connection = Some(new_conn); + } else if responses_websockets_v2_enabled && let Some(connection) = self.connection.as_ref() + { + self.client + .set_shared_v2_connection(Some(Arc::clone(connection))); } - self.connection.as_ref().ok_or(ApiError::Stream( + self.connection.as_deref().ok_or(ApiError::Stream( "websocket connection is unavailable".to_string(), )) } @@ -1136,16 +1232,33 @@ impl ModelClientSession { .stream_request(request) .await .map_err(map_api_error)?; + let responses_websockets_v2_enabled = self.client.responses_websockets_v2_enabled(); + if responses_websockets_v2_enabled { + // Keep chain state completion-driven: an interrupted request should not + // leave a mixed {new input, old response_id} pair for the next turn. + self.client.clear_shared_v2_request_chain_state(); + } self.websocket_last_items = api_prompt.input.clone(); + let mut completed_input = Some(self.websocket_last_items.clone()); let (last_response_id_sender, last_response_id_receiver) = oneshot::channel(); self.websocket_last_response_id_rx = Some(last_response_id_receiver); let mut last_response_id_sender = Some(last_response_id_sender); + let client = self.client.clone(); let stream_result = stream_result.inspect(move |event| { if let Ok(ResponseEvent::Completed { response_id, .. }) = event && !response_id.is_empty() - && let Some(sender) = last_response_id_sender.take() { - let _ = sender.send(response_id.clone()); + if let Some(sender) = last_response_id_sender.take() { + let _ = sender.send(response_id.clone()); + } + if responses_websockets_v2_enabled + && let Some(completed_input) = completed_input.take() + { + client.set_shared_v2_request_chain_state( + completed_input, + response_id.clone(), + ); + } } }); @@ -1241,6 +1354,7 @@ impl ModelClientSession { self.connection = None; self.websocket_last_items.clear(); + self.client.clear_shared_v2_state(); self.client.clear_preconnect(); } activated diff --git a/codex-rs/core/tests/suite/turn_state.rs b/codex-rs/core/tests/suite/turn_state.rs index 833c2aad9b..df939f01cc 100644 --- a/codex-rs/core/tests/suite/turn_state.rs +++ b/codex-rs/core/tests/suite/turn_state.rs @@ -1,6 +1,7 @@ #![allow(clippy::expect_used, clippy::unwrap_used)] use anyhow::Result; +use codex_core::features::Feature; use core_test_support::responses::WebSocketConnectionConfig; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -123,3 +124,108 @@ async fn websocket_turn_state_persists_within_turn_and_resets_after() -> Result< server.shutdown().await; Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn websocket_v2_reuses_connection_across_turns() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_websocket_server_with_headers(vec![WebSocketConnectionConfig { + requests: vec![ + vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ], + vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "done"), + ev_completed("resp-2"), + ], + ], + response_headers: Vec::new(), + accept_delay: None, + }]) + .await; + + let builder = test_codex(); + let test = builder + .with_config(|config| { + config.features.enable(Feature::ResponsesWebsocketsV2); + }) + .build_with_websocket_server(&server) + .await?; + test.submit_turn("first turn").await?; + test.submit_turn("second turn").await?; + + let handshakes = server.handshakes(); + assert_eq!(handshakes.len(), 1); + let requests = server.single_connection(); + assert_eq!(requests.len(), 2); + let second_request = requests[1].body_json(); + assert_eq!( + second_request["previous_response_id"].as_str(), + Some("resp-1") + ); + + server.shutdown().await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn websocket_v2_reconnect_after_turn_boundary_does_not_replay_turn_state() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_websocket_server_with_headers(vec![ + WebSocketConnectionConfig { + requests: vec![vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]], + response_headers: vec![(TURN_STATE_HEADER.to_string(), "ts-1".to_string())], + accept_delay: None, + }, + WebSocketConnectionConfig { + requests: vec![vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "done"), + ev_completed("resp-2"), + ]], + response_headers: Vec::new(), + accept_delay: None, + }, + ]) + .await; + + let builder = test_codex(); + let test = builder + .with_config(|config| { + config.features.enable(Feature::ResponsesWebsocketsV2); + }) + .build_with_websocket_server(&server) + .await?; + test.submit_turn("first turn").await?; + test.submit_turn("second turn").await?; + + let handshakes = server.handshakes(); + assert_eq!(handshakes.len(), 2); + assert_eq!(handshakes[0].header(TURN_STATE_HEADER), None); + assert_eq!(handshakes[1].header(TURN_STATE_HEADER), None); + + let connections = server.connections(); + assert_eq!(connections.len(), 2); + let second_connection = connections + .get(1) + .expect("second websocket connection should exist"); + let second_request = second_connection + .first() + .expect("second websocket connection should have a request") + .body_json(); + assert_eq!( + second_request["previous_response_id"].as_str(), + Some("resp-1") + ); + + server.shutdown().await; + Ok(()) +}