[codex] Send request-scoped turn state over WebSocket (#27996)

## Context

Turn state is scoped to one logical turn, but the WebSocket path
currently exchanges it through upgrade headers, which are scoped to the
physical connection. A connection may be reused across turns, so its
handshake cannot represent the turn lifecycle reliably.

## Change

Exchange turn state on each WebSocket response request instead:

- send an established value in `response.create.client_metadata`
- read the returned value from the existing `response.metadata` event
- retain the first value in the turn-scoped `ModelClientSession`
`OnceLock`
- start the next logical turn without state, even when it reuses the
same WebSocket connection

This gives WebSocket requests the same first-value-wins contract as the
existing HTTP path.

## Test plan

Integration coverage verifies that:

- WebSocket replays returned state on same-turn follow-ups
- later response metadata does not replace the first value
- state resets at the logical turn boundary without requiring a
reconnect

CI validates the full change.

## Stack

This is 1/2. #28002 builds on this request-scoped transport to carry
established state through compact requests.
This commit is contained in:
Ahmed Ibrahim
2026-06-13 00:44:39 -07:00
committed by GitHub
Unverified
parent 9d938a46d9
commit 640d61b121
5 changed files with 140 additions and 70 deletions
@@ -215,6 +215,7 @@ impl ResponsesWebsocketConnection {
&self,
request: ResponsesWsRequest,
connection_reused: bool,
turn_state: Option<Arc<OnceLock<String>>>,
) -> Result<ResponseStream, ApiError> {
let (tx_event, rx_event) =
mpsc::channel::<std::result::Result<ResponseEvent, ApiError>>(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<Arc<dyn WebsocketTelemetry>>,
connection_reused: bool,
turn_state: Option<&OnceLock<String>>,
) -> Result<(), ApiError> {
let mut last_server_model: Option<String> = 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" {
+24 -2
View File
@@ -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<String> {
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<Vec<ModelVerification>> {
if self.kind() != "response.metadata" {
return None;
@@ -220,6 +231,17 @@ fn header_openai_model_value_from_json(value: &Value) -> Option<String> {
})
}
fn header_turn_state_value_from_json(value: &Value) -> Option<String> {
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<Vec<ModelVerification>> {
let verifications = value
.as_array()