agent/codex
2
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
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/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(())