agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
## 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
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-api/src/endpoint/responses_websocket.rs
+8
View File
@@ -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" {
codex-rs/codex-api/src/sse/responses.rs
+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()
codex-rs/core/src/client.rs
+18 -34
View File
@@ -806,13 +806,10 @@ impl ModelClient {
api_provider: codex_api::Provider,
api_auth: SharedAuthProvider,
responses_metadata: &CodexResponsesMetadata,
turn_state: Option<Arc<OnceLock<String>>>,
auth_context: AuthRequestTelemetryContext,
request_route_telemetry: RequestRouteTelemetry,
) -> std::result::Result<ApiWebSocketConnection, ApiError> {
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<OnceLock<String>>>,
) -> 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,
}
codex-rs/core/src/client_tests.rs
+1 -1
View File
@@ -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!(
codex-rs/core/tests/suite/turn_state.rs
+89 -33
View File
@@ -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<_>>(),
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<_>>(),
vec![json!(null), json!("ts-1"), json!("ts-1")]
);
assert_eq!(handshakes[2].header(TURN_STATE_HEADER), None);
server.shutdown().await;
Ok(())