starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 08:06:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): preserve websocket turn state on failure

Browse Source 
Bridge both metadata event names, project sticky state per response.create, and return captured state on pre-terminal failures without replaying ambiguous requests.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-04 16:51:01 +08:00
co-authored by Codex
parent bcccebc9fa
commit 3554ba13db
18 changed files with 702 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/gateway-api/src/openai/responses/http.rs
+14 -5
View File
@@ -14,8 +14,8 @@ use axum::{
response::{IntoResponse, Response},
};
use futures::stream;
use gateway_core::engine::CommitRequirement;
use gateway_core::engine::execution::{ClientTransport, ExecutionSession, StartedExecution};
use gateway_core::engine::{CommitRequirement, EngineError};
use gateway_core::error::{GatewayError, GatewayErrorKind};
use gateway_core::event::{ProviderEvent, ProviderResponseHeader};
use gateway_core::lifecycle::ConnectionGuard;
@@ -152,7 +152,8 @@ pub async fn collect_execution_response(session: Box<dyn ExecutionSession>) -> R
let events = match session.collect_uncommitted().await {
Ok(events) => events,
Err(error) => {
let response = engine_error_response(&error);
let response_headers = session.response_headers().to_vec();
let response = engine_error_response_with_headers(&error, &response_headers);
return execution.record_response_status(response).await;
}
};
@@ -178,7 +179,7 @@ pub async fn collect_execution_response(session: Box<dyn ExecutionSession>) -> R
.await
{
execution.cancel_and_finalize().await;
let response = engine_error_response(&error);
let response = engine_error_response_with_headers(&error, &response_headers);
return execution.record_response_status(response).await;
}
if !session.is_finalized() {
@@ -246,6 +247,13 @@ fn apply_response_headers(
response
}
fn engine_error_response_with_headers(
error: &EngineError,
response_headers: &[ProviderResponseHeader],
) -> Response {
apply_response_headers(engine_error_response(error), response_headers)
}
/// 编码首个 SSE frame 后提交下游，再持续驱动同一执行会话。
pub async fn stream_execution_response(
session: Box<dyn ExecutionSession>,
@@ -267,7 +275,8 @@ pub async fn stream_execution_response(
return response;
}
Err(error) => {
let response = engine_error_response(&error);
let response_headers = session.response_headers().to_vec();
let response = engine_error_response_with_headers(&error, &response_headers);
return execution.record_response_status(response).await;
}
};
@@ -299,7 +308,7 @@ pub async fn stream_execution_response(
.await
{
execution.cancel_and_finalize().await;
let response = engine_error_response(&error);
let response = engine_error_response_with_headers(&error, &response_headers);
return execution.record_response_status(response).await;
}
let Some(session) = execution.into_session() else {
backend/crates/gateway-api/src/openai/responses/websocket/forward.rs
+26 -2
View File
@@ -77,8 +77,8 @@ pub(super) async fn forward_execution(
return send_gateway_error(connection, &error, &request_id).await;
}
ActiveInput::Event(Err(error)) => {
let error = gateway_error_from_engine(&error);
return send_gateway_error(connection, &error, &request_id).await;
return send_initial_engine_error(connection, &mut execution, &error, &request_id)
.await;
}
ActiveInput::Control => continue,
ActiveInput::Disconnect => return ForwardOutcome::Disconnect,
@@ -271,6 +271,30 @@ pub(super) async fn forward_execution(
}
}
async fn send_initial_engine_error(
connection: &mut ResponsesWebSocketConnection,
execution: &mut PendingExecution,
error: &EngineError,
request_id: &Arc<str>,
) -> ForwardOutcome {
let response_headers = execution
.session_mut()
.map(|session| session.response_headers().to_vec())
.unwrap_or_default();
if !response_headers.is_empty()
&& !send_metadata(
connection,
request_id,
response_metadata_event(request_id, &response_headers),
)
.await
{
return ForwardOutcome::Disconnect;
}
let error = gateway_error_from_engine(error);
send_gateway_error(connection, &error, request_id).await
}
fn commit_connection_replay(
replay: &mut ConnectionReplaySnapshot,
encoder: &OpenAiResponsesEncoder,
backend/crates/gateway-api/tests/openai/responses/http.rs
+54 -1
View File
@@ -893,6 +893,41 @@ async fn streaming_rate_limit_before_first_frame_should_persist_the_returned_429
assert_eq!(trace.client_statuses(), vec![429]);
}
#[tokio::test]
async fn streaming_failure_before_first_frame_should_return_safe_observed_headers() {
let trace = Arc::new(Trace::default());
let session = FakeSession::streaming(
Arc::clone(&trace),
vec![NextStep::Error(EngineError::Provider(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::Ambiguous,
)))],
)
.with_response_headers(vec![
ProviderResponseHeader::new(
"x-codex-turn-state",
Bytes::from_static(b"turn-state-before-close"),
),
ProviderResponseHeader::new(
"authorization",
Bytes::from_static(b"must-not-cross-boundary"),
),
]);
let response = stream_execution_response(Box::new(session), None).await;
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
assert_eq!(
response
.headers()
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok()),
Some("turn-state-before-close")
);
assert!(response.headers().get("authorization").is_none());
assert_eq!(trace.client_statuses(), vec![502]);
}
#[tokio::test]
async fn streaming_upstream_http_failure_before_first_frame_should_preserve_raw_response() {
let trace = Arc::new(Trace::default());
@@ -979,11 +1014,29 @@ async fn buffered_upstream_http_failure_should_preserve_raw_response() {
raw_body.clone(),
));
let session = FakeSession::buffered(Arc::clone(&trace), vec![started(), completed()])
.with_collect_error(EngineError::Provider(error));
.with_collect_error(EngineError::Provider(error))
.with_response_headers(vec![
ProviderResponseHeader::new(
"x-codex-turn-state",
Bytes::from_static(b"turn-state-before-close"),
),
ProviderResponseHeader::new(
"authorization",
Bytes::from_static(b"must-not-cross-boundary"),
),
]);
let response = collect_execution_response(Box::new(session)).await;
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(
response
.headers()
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok()),
Some("turn-state-before-close")
);
assert!(response.headers().get("authorization").is_none());
assert_eq!(
to_bytes(response.into_body(), usize::MAX)
.await
backend/crates/gateway-api/tests/openai/responses/websocket/connection.rs
+1
View File
@@ -428,6 +428,7 @@ async fn idle_connection_reaches_the_official_limit_without_starting_an_executio
client: authenticated_client("sk_ws_atomic"),
trace: Arc::clone(&trace),
response_headers: Vec::new(),
fail_before_first_event: false,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs
+99 -1
View File
@@ -286,12 +286,20 @@ struct AtomicFailureTrace {
struct AtomicFailureSession {
trace: Arc<AtomicFailureTrace>,
response_headers: Vec<ProviderResponseHeader>,
fail_before_first_event: bool,
}
impl ExecutionSession for AtomicFailureSession {
fn next_event(&mut self) -> BoxFuture<'_, Result<Option<CoordinatedEvent>, EngineError>> {
Box::pin(async move {
match self.trace.next_calls.fetch_add(1, Ordering::AcqRel) {
let next_call = self.trace.next_calls.fetch_add(1, Ordering::AcqRel);
if self.fail_before_first_event && next_call == 0 {
return Err(EngineError::Provider(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::Ambiguous,
)));
}
match next_call {
0 => Ok(Some(atomic_failure_batch())),
1 => {
self.trace.finalized.store(true, Ordering::Release);
@@ -341,6 +349,7 @@ struct AtomicFailureExecution {
client: AuthenticatedClient,
trace: Arc<AtomicFailureTrace>,
response_headers: Vec<ProviderResponseHeader>,
fail_before_first_event: bool,
}
impl ExecutionService for AtomicFailureExecution {
@@ -388,6 +397,7 @@ impl ExecutionService for AtomicFailureExecution {
session: Box::new(AtomicFailureSession {
trace: Arc::clone(&self.trace),
response_headers: self.response_headers.clone(),
fail_before_first_event: self.fail_before_first_event,
}),
})
})
@@ -453,6 +463,7 @@ async fn websocket_atomic_upstream_failure_batch_should_be_forwarded_once() {
client: authenticated_client("sk_ws_atomic"),
trace: Arc::clone(&trace),
response_headers: Vec::new(),
fail_before_first_event: false,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
@@ -562,6 +573,7 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
ProviderResponseHeader::new("content-type", Bytes::from_static(b"application/private")),
ProviderResponseHeader::new("x-request-id", Bytes::from_static(b"req_upstream")),
],
fail_before_first_event: false,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
@@ -646,6 +658,91 @@ async fn websocket_response_metadata_should_preserve_ordinary_headers_without_bl
server.abort();
}
#[tokio::test]
async fn websocket_failure_before_first_event_should_return_observed_metadata_then_error() {
let trace = Arc::new(AtomicFailureTrace::default());
let execution = Arc::new(AtomicFailureExecution {
client: authenticated_client("sk_ws_atomic"),
trace: Arc::clone(&trace),
response_headers: vec![
ProviderResponseHeader::new(
"x-codex-turn-state",
Bytes::from_static(b"turn-state-before-close"),
),
ProviderResponseHeader::new(
"authorization",
Bytes::from_static(b"must-not-cross-boundary"),
),
],
fail_before_first_event: true,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket test server");
let address = listener.local_addr().expect("test server address");
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("serve WebSocket test router");
});
let mut request = format!("ws://{address}/v1/responses")
.into_client_request()
.expect("WebSocket request");
request.headers_mut().insert(
AUTHORIZATION,
"Bearer sk_ws_atomic".parse().expect("authorization"),
);
let (mut socket, response) = connect_async(request).await.expect("upgrade WebSocket");
assert_eq!(response.status(), StatusCode::SWITCHING_PROTOCOLS);
socket
.send(ClientMessage::Text(
json!({
"type": "response.create",
"model": "model-a",
"input": "hello"
})
.to_string()
.into(),
))
.await
.expect("send response.create");
let mut messages = Vec::<Value>::new();
while messages.len() < 2 {
let message = tokio::time::timeout(Duration::from_secs(2), socket.next())
.await
.expect("WebSocket response timeout")
.expect("WebSocket remains open")
.expect("valid WebSocket frame");
let ClientMessage::Text(text) = message else {
continue;
};
messages.push(serde_json::from_str(&text).expect("JSON WebSocket event"));
}
assert_eq!(
messages
.iter()
.filter_map(|message| message.get("type").and_then(Value::as_str))
.collect::<Vec<_>>(),
vec!["response.metadata", "error"]
);
let headers = messages[0]
.get("headers")
.and_then(Value::as_object)
.expect("response metadata headers");
assert_eq!(
headers.get("x-codex-turn-state"),
Some(&json!("turn-state-before-close"))
);
assert!(!headers.contains_key("authorization"));
assert!(!trace.committed.load(Ordering::Acquire));
socket.close(None).await.expect("close WebSocket");
server.abort();
}
#[tokio::test]
async fn websocket_response_create_should_not_have_a_private_16_mib_frame_limit() {
let trace = Arc::new(AtomicFailureTrace::default());
@@ -653,6 +750,7 @@ async fn websocket_response_create_should_not_have_a_private_16_mib_frame_limit(
client: authenticated_client("sk_ws_atomic"),
trace,
response_headers: Vec::new(),
fail_before_first_event: false,
});
let app = api_router(execution).await;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
backend/crates/providers/openai/src/provider/execution.rs
+44 -26
View File
@@ -732,14 +732,23 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
);
}
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
if !updates.is_empty() {
let rate_limits_changed = if updates.is_empty() {
false
} else {
passive_quota_observation.observe(&updates);
let update_headers = rate_limit_update_headers(&updates);
if observation_state.merge_rate_limit_headers(&update_headers)
&& let Some(observation) = observation_state.observation()
{
yield ProviderEvent::observation(observation);
}
observation_state.merge_rate_limit_headers(&update_headers)
};
let turn_state_changed = merge_turn_state_update(
turn_state_updates.as_ref(),
&mut session_capture,
&mut observation_state,
)
.await;
if (rate_limits_changed || turn_state_changed)
&& let Some(observation) = observation_state.observation()
{
yield ProviderEvent::observation(observation);
}
if allows_account_state_mutation {
synchronize_passive_quota(
@@ -763,16 +772,12 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
passive_quota_observation.observe(&updates);
observation_state.merge_rate_limit_headers(&rate_limit_update_headers(&updates))
};
let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref()
&& let Some(turn_state) = updates.lock().await.take()
{
if let Some(capture) = session_capture.as_mut() {
capture.turn_state = Some(turn_state.clone());
}
observation_state.merge_client_header("x-codex-turn-state", &turn_state)
} else {
false
};
let turn_state_changed = merge_turn_state_update(
turn_state_updates.as_ref(),
&mut session_capture,
&mut observation_state,
)
.await;
let first_event_changed =
observation_state.observe_stream_chunk(&chunk, output_started_at);
let chunk_len = chunk.len();
@@ -949,16 +954,12 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
apply_failure(&failure_context, &active_account, failure)
.await;
}
let turn_state_changed = if let Some(updates) = turn_state_updates.as_ref()
&& let Some(turn_state) = updates.lock().await.take()
{
if let Some(capture) = session_capture.as_mut() {
capture.turn_state = Some(turn_state.clone());
}
observation_state.merge_client_header("x-codex-turn-state", &turn_state)
} else {
false
};
let turn_state_changed = merge_turn_state_update(
turn_state_updates.as_ref(),
&mut session_capture,
&mut observation_state,
)
.await;
attach_openai_session_update(&mut events, &mut session_capture);
let completed = events
.iter()
@@ -1024,3 +1025,20 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream {
}
})
}
async fn merge_turn_state_update(
updates: Option<&CodexTurnStateUpdate>,
session_capture: &mut Option<OpenAiSessionCapture>,
observation_state: &mut OpenAiResponseObservationState,
) -> bool {
let Some(updates) = updates else {
return false;
};
let Some(turn_state) = updates.lock().await.take() else {
return false;
};
if let Some(capture) = session_capture.as_mut() {
capture.turn_state = Some(turn_state.clone());
}
observation_state.merge_client_header("x-codex-turn-state", &turn_state)
}
backend/crates/providers/openai/src/provider/mod.rs
+14 -9
View File
@@ -93,7 +93,7 @@ use crate::transport::{
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
CodexBackendJsonResponse, CodexBackendStreamingResponse, CodexBackendTransport,
CodexClientError, CodexRateLimitUpdates, CodexRequestContext, CodexResponseMetadata,
CodexTransportMetrics, CodexUpstreamDiagnostics, CodexWebSocketPool,
CodexTransportMetrics, CodexTurnStateUpdate, CodexUpstreamDiagnostics, CodexWebSocketPool,
WebSocketConnectionPreference, endpoint_url,
};
@@ -320,15 +320,20 @@ impl Provider for CodexProvider {
if let Some(identity) = &self.session_identity {
identity.prepare_local_conversation(&mut upstream_request);
}
if previous_session.as_ref().is_some_and(|state| {
same_client_turn(
state.client_turn_id.as_deref(),
if let Some(previous_session) = previous_session.as_ref() {
upstream_request.turn_state = if same_client_turn(
previous_session.client_turn_id.as_deref(),
upstream_request.client_turn_id.as_deref(),
)
}) {
upstream_request.turn_state = previous_session
.as_ref()
.and_then(|state| state.turn_state.clone());
) {
// 客户端可能携带 metadata-close 失败后才拿到的新状态；同一
// turn 内显式回传值比最后一次成功响应保存的状态更新。
upstream_request
.turn_state
.take()
.or_else(|| previous_session.turn_state.clone())
} else {
None
};
}
let session_affinity =
derive_codex_session_affinity(&upstream_request, context.client_api_key_ref());
backend/crates/providers/openai/src/transport/client.rs
+23 -20
View File
@@ -24,7 +24,9 @@ use serde_json::{Value, map::Map};
use thiserror::Error;
use uuid::Uuid;
use crate::transport::protocol::responses::{CodexResponsesRequest, TransportRequirement};
use crate::transport::protocol::responses::{
CodexResponsesRequest, TransportRequirement, X_CODEX_TURN_STATE_CLIENT_METADATA_KEY,
};
use super::diagnostics::{CodexUpstreamDiagnostics, CodexUpstreamFailure, CodexUpstreamSendPhase};
use super::response_meta::CodexResponseMetadata;
@@ -546,7 +548,7 @@ pub struct CodexBackendStreamingResponse {
pub transport: CodexBackendTransport,
/// WebSocket 响应所绑定的连接；HTTP transport 为 `None`。
pub websocket_connection_id: Option<Uuid>,
/// 响应头里的最新 turn state。
/// 响应头或 metadata 事件里最先确认的 turn state。
pub turn_state: Option<String>,
/// 上游透传的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
@@ -760,32 +762,33 @@ pub(super) async fn read_error_response_body(
pub(super) fn websocket_upstream_request(request: &CodexResponsesRequest) -> CodexResponsesRequest {
let mut request = request.clone();
project_responses_lite_to_ws_metadata(&mut request);
stamp_ws_stream_request_start_ms(&mut request);
project_websocket_client_metadata(&mut request);
request
}
fn project_responses_lite_to_ws_metadata(request: &mut CodexResponsesRequest) {
let Some(responses_lite) = request.responses_lite.clone() else {
return;
};
let mut metadata = match request.client_metadata() {
Some(Value::Object(metadata)) => metadata.clone(),
None => Map::new(),
Some(_) => return,
};
metadata
.entry(WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY.to_string())
.or_insert(Value::String(responses_lite));
request.set_client_metadata(Some(Value::Object(metadata)));
}
fn stamp_ws_stream_request_start_ms(request: &mut CodexResponsesRequest) {
fn project_websocket_client_metadata(request: &mut CodexResponsesRequest) {
let mut metadata = match request.client_metadata() {
Some(Value::Object(metadata)) => metadata.clone(),
None => Map::new(),
Some(_) => return,
};
if let Some(responses_lite) = request.responses_lite.clone() {
metadata
.entry(WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY.to_owned())
.or_insert(Value::String(responses_lite));
}
let turn_state = request.turn_state.clone();
match turn_state {
Some(turn_state) => {
metadata.insert(
X_CODEX_TURN_STATE_CLIENT_METADATA_KEY.to_owned(),
Value::String(turn_state),
);
}
None => {
metadata.remove(X_CODEX_TURN_STATE_CLIENT_METADATA_KEY);
}
}
metadata.insert(
X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY.to_string(),
Value::String(now_unix_timestamp_millis().to_string()),
backend/crates/providers/openai/src/transport/protocol/responses.rs
+2
View File
@@ -10,6 +10,8 @@ use serde_json::{Map, Value};
/// 官方 Codex 客户端据此触发完整历史重放的稳定错误码。
pub(crate) const PREVIOUS_RESPONSE_NOT_FOUND_CODE: &str = "previous_response_not_found";
/// Responses WebSocket 用于回传同一 turn 不透明状态的官方 client metadata 键。
pub(crate) const X_CODEX_TURN_STATE_CLIENT_METADATA_KEY: &str = "x-codex-turn-state";
/// 本地生成 history unavailable 错误时使用的官方提示文本。
pub(crate) const PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE: &str =
"Previous response was not found. Retrying the full request.";
backend/crates/providers/openai/src/transport/protocol/websocket.rs
+7 -13
View File
@@ -110,9 +110,9 @@ fn is_internal_websocket_event(event: &str) -> bool {
event == "codex.rate_limits"
}
/// 提取 `response.metadata` 帧中的字符串响应头。
/// 提取 Responses WebSocket metadata 帧中的字符串响应头。
pub fn websocket_metadata_headers(value: &Value) -> Vec<(String, String)> {
if value.get("type").and_then(Value::as_str) != Some("response.metadata") {
if !is_websocket_metadata_event(websocket_event_type(value)) {
return Vec::new();
}
value
@@ -124,24 +124,18 @@ pub fn websocket_metadata_headers(value: &Value) -> Vec<(String, String)> {
.collect()
}
/// 从 `response.metadata` 帧中提取 `x-codex-turn-state`。
/// 从 Responses WebSocket metadata 帧中提取 `x-codex-turn-state`。
pub fn websocket_metadata_turn_state(value: &Value) -> Option<String> {
if value.get("type").and_then(Value::as_str) != Some("response.metadata") {
return None;
}
websocket_metadata_headers(value)
.into_iter()
.find_map(|(name, value)| {
name.eq_ignore_ascii_case("x-codex-turn-state")
.then_some(value)
})
.or_else(|| {
value
.pointer("/metadata/turn_state")
.or_else(|| value.get("turn_state"))
.and_then(Value::as_str)
.map(ToString::to_string)
})
}
fn is_websocket_metadata_event(event: Option<&str>) -> bool {
matches!(event, Some("response.metadata" | "codex.response.metadata"))
}
/// 上游 WebSocket 连接寿命限制错误码。
backend/crates/providers/openai/src/transport/request.rs
+20 -9
View File
@@ -7,9 +7,12 @@ use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::transport::protocol::responses::CodexResponsesRequest;
use crate::transport::protocol::responses::{
CodexResponsesRequest, X_CODEX_TURN_STATE_CLIENT_METADATA_KEY,
};
const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers";
const TURN_ID_CLIENT_METADATA_KEY: &str = "turn_id";
const THREAD_SPAWN_SUBAGENT_KIND: &str = "thread_spawn";
const THREAD_SPAWN_CONVERSATION_PREFIX: &str = "thread-spawn:";
const UNSUPPORTED_CODEX_RESPONSES_FIELDS: &[&str] = &["max_output_tokens", "temperature"];
@@ -144,8 +147,16 @@ struct ExtractedRequestContext {
impl ExtractedRequestContext {
fn from_body(body: &Map<String, Value>) -> Self {
let client_metadata = body.get("client_metadata").and_then(Value::as_object);
Self {
turn_state: body_string(body, "turnState"),
// 官方 downstream WebSocket 无法逐帧更新 HTTP header，因此把
// response.metadata 返回的 turn state 放回下一帧 client_metadata。
// Provider 仍会在账号/turn 归属确定后决定是否允许复用该状态。
turn_state: body_string(body, "turnState").or_else(|| {
client_metadata.and_then(|metadata| {
string_value(metadata.get(X_CODEX_TURN_STATE_CLIENT_METADATA_KEY))
})
}),
turn_metadata: body_string(body, "turnMetadata"),
beta_features: body_string(body, "betaFeatures"),
version: body_string(body, "version"),
@@ -156,15 +167,15 @@ impl ExtractedRequestContext {
session_id: body_string(body, "session_id"),
thread_id: body_string(body, "thread_id"),
client_request_id: body_string(body, "x-client-request-id"),
turn_id: body_string(body, "turn_id"),
turn_id: body_string(body, "turn_id").or_else(|| {
client_metadata
.and_then(|metadata| string_value(metadata.get(TURN_ID_CLIENT_METADATA_KEY)))
}),
// Responses Lite 在 WebSocket body 中的这个键是官方 header 投影，
// 不是普通上下文字段的 metadata 回退。
responses_lite: body
.get("client_metadata")
.and_then(Value::as_object)
.and_then(|metadata| {
string_value(metadata.get(WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY))
}),
responses_lite: client_metadata.and_then(|metadata| {
string_value(metadata.get(WS_REQUEST_HEADER_RESPONSES_LITE_CLIENT_METADATA_KEY))
}),
// Memory consolidation 只由官方请求头提供；body 中同名字段不参与
// transport 事实提取。
memgen_request: None,
backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs
+7 -2
View File
@@ -31,7 +31,7 @@ pub struct CodexWebSocketStreamingExchange {
pub(crate) websocket_connection_id: Uuid,
/// 由 WebSocket 事件转换出的 live SSE 字节流。
pub body: CodexWebSocketSseStream,
/// 上游 metadata 帧中的最新 turn state。
/// 上游为本次响应返回的首个 turn state。
pub turn_state: Option<String>,
/// 上游握手响应里的 `set-cookie` 列表。
pub set_cookie_headers: Vec<String>,
@@ -39,7 +39,7 @@ pub struct CodexWebSocketStreamingExchange {
pub rate_limit_headers: Vec<(String, String)>,
/// 上游内部 `codex.rate_limits` 事件里的结构化动态更新。
pub rate_limit_updates: CodexWebSocketRateLimitUpdates,
/// 上游内部 metadata 事件里的动态 turn state。
/// 上游内部 metadata 事件里的首个动态 turn state。
pub turn_state_update: CodexWebSocketTurnStateUpdate,
/// WebSocket 连接池决策。
pub pool_decision: Option<WebSocketPoolDecision>,
@@ -63,5 +63,10 @@ pub(super) fn reusable_websocket_metadata(
mut metadata: CodexWebSocketConnectionMetadata,
) -> CodexWebSocketConnectionMetadata {
metadata.rate_limit_headers.clear();
metadata.turn_state = None;
metadata
.response_metadata
.client_headers
.retain(|(name, _)| !name.eq_ignore_ascii_case("x-codex-turn-state"));
metadata
}
backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs
+10 -7
View File
@@ -21,7 +21,6 @@ pub(in crate::transport::websocket) enum WebSocketTerminalKind {
pub(super) enum ExchangeAction {
RateLimits(events::ParsedRateLimits),
TurnState(String),
Forward {
frame: String,
terminal: Option<WebSocketTerminalKind>,
@@ -32,6 +31,7 @@ pub(super) enum ExchangeAction {
pub(super) struct ReducedWebSocketEvent {
pub(super) action: ExchangeAction,
pub(super) diagnostic_event_type: Option<String>,
pub(super) turn_state_update: Option<String>,
}
pub(super) fn reduce_websocket_event(
@@ -45,6 +45,7 @@ pub(super) fn reduce_websocket_event(
return Ok(ReducedWebSocketEvent {
action: ExchangeAction::Ignore,
diagnostic_event_type: None,
turn_state_update: None,
});
};
let diagnostic_event_type = diagnostic_event_type(websocket_event_type(&value));
@@ -54,6 +55,7 @@ pub(super) fn reduce_websocket_event(
return Ok(ReducedWebSocketEvent {
action: ExchangeAction::RateLimits(parsed),
diagnostic_event_type,
turn_state_update: None,
});
}
@@ -61,13 +63,13 @@ pub(super) fn reduce_websocket_event(
&mut metadata.response_metadata,
websocket_metadata_headers(&value),
);
if let Some(turn_state) = websocket_metadata_turn_state(&value) {
let turn_state_update = websocket_metadata_turn_state(&value).and_then(|turn_state| {
if metadata.turn_state.is_some() {
return None;
}
metadata.turn_state = Some(turn_state.clone());
return Ok(ReducedWebSocketEvent {
action: ExchangeAction::TurnState(turn_state),
diagnostic_event_type,
});
}
Some(turn_state)
});
let event = websocket_event_type(&value);
if event == Some("response.completed")
@@ -89,6 +91,7 @@ pub(super) fn reduce_websocket_event(
Ok(ReducedWebSocketEvent {
action,
diagnostic_event_type,
turn_state_update,
})
}
backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs
+6 -4
View File
@@ -271,15 +271,17 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
if let Some(event_type) = reduced.diagnostic_event_type {
last_event_type = Some(event_type);
}
if let Some(turn_state) = reduced.turn_state_update {
let mut pending = turn_state_update.lock().await;
if pending.is_none() {
*pending = Some(turn_state);
}
}
let (frame, terminal) = match reduced.action {
ExchangeAction::RateLimits(rate_limits) => {
rate_limit_updates.lock().await.push(rate_limits);
continue;
}
ExchangeAction::TurnState(turn_state) => {
*turn_state_update.lock().await = Some(turn_state);
continue;
}
ExchangeAction::Forward { frame, terminal } => (frame, terminal),
ExchangeAction::Ignore => continue,
};
backend/crates/providers/openai/tests/provider/contract.rs
+116 -10
View File
@@ -461,6 +461,7 @@ fn contract_account_scope() -> Arc<FrozenAccountScope> {
"acct_websocket_cooldown",
"acct_websocket_fast_path",
"acct_websocket_busy_replay",
"acct_websocket_metadata_close",
"acct_websocket_turn_state",
]
.into_iter()
@@ -2519,6 +2520,96 @@ async fn websocket_turn_state_metadata_is_exposed_through_response_observation()
assert!(observed_turn_state);
}
#[tokio::test]
async fn websocket_turn_state_metadata_survives_close_1000_without_changing_failure_semantics() {
const ACCOUNT_ID: &str = "acct_websocket_metadata_close";
const TURN_STATE: &str = "turn-state-before-close";
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, ACCOUNT_ID).await;
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket listener");
let base_url = format!(
"http://{}",
listener.local_addr().expect("listener address")
);
let server = tokio::spawn(async move {
let (stream, _) = listener
.accept()
.await
.expect("accept WebSocket connection");
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = websocket
.next()
.await
.expect("WebSocket request")
.expect("valid WebSocket request");
websocket
.send(Message::Text(
json!({
"type": "codex.response.metadata",
"headers": {"x-codex-turn-state": TURN_STATE}
})
.to_string()
.into(),
))
.await
.expect("send response metadata");
websocket
.close(Some(CloseFrame {
code: CloseCode::Normal,
reason: "".into(),
}))
.await
.expect("close before terminal response");
});
let operation = Operation::Generate(generate_with_persisted_session_context(
ACCOUNT_ID,
"conversation-metadata-close",
"session-metadata-close",
"turn-metadata-close",
));
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", operation),
context("req_websocket_metadata_close", CancellationToken::new()),
)
.await
.expect("prepare WebSocket provider stream");
let mut observed_turn_state = false;
let error = loop {
match stream.next().await {
Some(Ok(event)) => {
if let Some(observation) = event.response_observation() {
observed_turn_state |= observation.client_headers().iter().any(|header| {
header.name().eq_ignore_ascii_case("x-codex-turn-state")
&& header.value().as_ref() == TURN_STATE.as_bytes()
});
}
}
Some(Err(error)) => break error,
None => panic!("metadata close must surface a provider error"),
}
};
server.await.expect("WebSocket server");
assert!(observed_turn_state);
assert_eq!(error.send_state(), UpstreamSendState::Ambiguous);
assert_eq!(error.pre_delivery_retry(), None);
assert_eq!(
error.upstream_code().map(|code| code.as_str()),
Some("websocket_close_1000")
);
assert_eq!(
error.diagnostic().map(|diagnostic| diagnostic.as_str()),
Some(
"OpenAI WebSocket closed before a terminal response (close code 1000); last event type: codex.response.metadata"
)
);
}
#[tokio::test]
async fn disabled_account_is_excluded_from_normal_scheduling() {
let store = Arc::new(MemoryAccountStore::default());
@@ -2955,30 +3046,45 @@ async fn matching_turn_id_should_restore_previous_turn_state() {
}
#[tokio::test]
async fn matching_turn_id_should_prefer_the_latest_provider_state_over_a_stale_client_echo() {
async fn matching_turn_id_should_prefer_an_explicit_client_echo_over_saved_provider_state() {
let request = capture_turn_state_request(
"req_replace_stale_client_turn_state",
"req_use_client_turn_state",
Some("turn-same"),
Some("turn-same"),
Some("stale-client-turn-state"),
Some("client-turn-state"),
)
.await;
assert_eq!(
captured_header_values(&request, "x-codex-turn-state"),
vec![b"previous-turn-state".to_vec()]
vec![b"client-turn-state".to_vec()]
);
}
#[tokio::test]
async fn new_or_unidentified_turn_should_not_restore_previous_turn_state() {
for (request_id, previous_turn_id, current_turn_id) in [
("req_new_turn_state", Some("turn-old"), Some("turn-new")),
("req_unidentified_turn_state", None, Some("turn-new")),
("req_missing_current_turn_state", Some("turn-old"), None),
for (request_id, previous_turn_id, current_turn_id, client_turn_state) in [
(
"req_new_turn_state",
Some("turn-old"),
Some("turn-new"),
Some("stale-client-turn-state"),
),
("req_unidentified_turn_state", None, Some("turn-new"), None),
(
"req_missing_current_turn_state",
Some("turn-old"),
None,
None,
),
] {
let request =
capture_turn_state_request(request_id, previous_turn_id, current_turn_id, None).await;
let request = capture_turn_state_request(
request_id,
previous_turn_id,
current_turn_id,
client_turn_state,
)
.await;
assert!(captured_header_values(&request, "x-codex-turn-state").is_empty());
}
}
backend/crates/providers/openai/tests/transport/protocol/websocket.rs
+42 -8
View File
@@ -155,6 +155,7 @@ fn websocket_event_to_sse_should_forward_public_events_and_strip_internal_events
);
assert!(websocket_event_to_sse_frame(r#"{"type":"codex.rate_limits"}"#).is_none());
assert!(websocket_event_to_sse_frame(r#"{"type":"response.metadata"}"#).is_none());
assert!(websocket_event_to_sse_frame(r#"{"type":"codex.response.metadata"}"#).is_some());
assert!(websocket_event_to_sse_frame(r#"{"response":{}}"#).is_none());
assert!(websocket_event_to_sse_frame("not-json").is_none());
}
@@ -178,15 +179,48 @@ fn websocket_event_to_sse_should_add_missing_previous_response_recovery_code() {
#[test]
fn websocket_metadata_turn_state_should_accept_case_insensitive_header() {
let event = json!({
"type": "response.metadata",
"headers": {"X-Codex-Turn-State": ["turn-from-metadata"]}
});
for (event, expected) in [
(
json!({
"type": "response.metadata",
"headers": {"X-Codex-Turn-State": ["turn-from-legacy-metadata"]}
}),
"turn-from-legacy-metadata",
),
(
json!({
"type": "codex.response.metadata",
"headers": {"X-Codex-Turn-State": "turn-from-codex-metadata"}
}),
"turn-from-codex-metadata",
),
] {
assert_eq!(
websocket_metadata_turn_state(&event).as_deref(),
Some(expected)
);
}
}
assert_eq!(
websocket_metadata_turn_state(&event).as_deref(),
Some("turn-from-metadata")
);
#[test]
fn websocket_metadata_turn_state_should_ignore_unverified_shapes() {
for event in [
json!({
"type": "codex.response.metadata",
"headers": {"x-codex-turn-state": {"future": "shape"}}
}),
json!({
"type": "codex.response.metadata",
"metadata": {"turn_state": "nested-value"},
"turn_state": "top-level-value"
}),
json!({
"type": "future.response.metadata",
"headers": {"x-codex-turn-state": "unverified-event"}
}),
] {
assert_eq!(websocket_metadata_turn_state(&event), None);
}
}
#[test]
backend/crates/providers/openai/tests/transport/request.rs
+32
View File
@@ -227,6 +227,26 @@ fn encoder_should_ignore_legacy_context_aliases_and_metadata_fallbacks() {
);
}
#[test]
fn encoder_should_extract_official_websocket_context_projection() {
let request = request(Map::from_iter([
("model".to_owned(), json!("client-model")),
("input".to_owned(), json!("prompt")),
(
"client_metadata".to_owned(),
json!({
"x-codex-turn-state": "metadata-turn-state",
"turn_id": "metadata-turn-id"
}),
),
]));
let encoded = encode_generate_request(&request, "gpt-test").expect("encode");
assert_eq!(encoded.turn_state.as_deref(), Some("metadata-turn-state"));
assert_eq!(encoded.client_turn_id.as_deref(), Some("metadata-turn-id"));
}
#[test]
fn header_context_should_win_over_body_topline_aliases() {
let payload = ProtocolPayload::json_object(
@@ -236,6 +256,13 @@ fn header_context_should_win_over_body_topline_aliases() {
("input".to_owned(), json!("prompt")),
("turnState".to_owned(), json!("body-turn-state")),
("turnMetadata".to_owned(), json!("body-turn-metadata")),
(
"client_metadata".to_owned(),
json!({
"x-codex-turn-state": "metadata-turn-state",
"turn_id": "metadata-turn-id"
}),
),
]),
)
.expect("OpenAI payload")
@@ -248,6 +275,10 @@ fn header_context_should_win_over_body_topline_aliases() {
"turn_metadata".to_owned(),
Value::String("header-turn-metadata".to_owned()),
),
(
"turn_id".to_owned(),
Value::String("header-turn-id".to_owned()),
),
]));
let request = GenerateRequest::from_protocol_payload(payload);
@@ -258,6 +289,7 @@ fn header_context_should_win_over_body_topline_aliases() {
encoded.turn_metadata.as_deref(),
Some("header-turn-metadata")
);
assert_eq!(encoded.client_turn_id.as_deref(), Some("header-turn-id"));
}
#[test]
backend/crates/providers/openai/tests/transport/websocket/mod.rs
+185
View File
@@ -914,6 +914,191 @@ async fn websocket_execute_response_create_request_should_capture_internal_metad
);
}
#[tokio::test]
async fn websocket_codex_metadata_should_capture_first_turn_state_without_changing_forwarding() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
json!({
"type": "codex.response.metadata",
"headers": {"x-codex-turn-state": "turn-first"}
})
.to_string()
.into(),
))
.await
.unwrap();
websocket
.send(Message::Text(
json!({
"type": "response.metadata",
"headers": {"x-codex-turn-state": ["turn-second"]}
})
.to_string()
.into(),
))
.await
.unwrap();
websocket
.send(Message::Text(
completed_websocket_response("resp_metadata_alias", 1, 1).into(),
))
.await
.unwrap();
});
let prepared = prepared_websocket_request(&format!("http://{addr}"));
let response = execute_response_create_request(&prepared)
.await
.expect("metadata aliases should preserve a completed response");
server.await.unwrap();
assert_eq!(response.turn_state.as_deref(), Some("turn-first"));
assert!(response.body.contains("event: codex.response.metadata"));
assert!(!response.body.contains("event: response.metadata\n"));
assert!(response.body.contains("event: response.completed"));
}
#[tokio::test]
async fn reused_websocket_should_not_leak_turn_state_into_the_next_response() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _first = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
json!({
"type": "codex.response.metadata",
"headers": {"x-codex-turn-state": "turn-from-first-response"}
})
.to_string()
.into(),
))
.await
.unwrap();
websocket
.send(Message::Text(
completed_websocket_response("resp_turn_state_first", 1, 1).into(),
))
.await
.unwrap();
let _second = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
completed_websocket_response("resp_turn_state_second", 1, 1).into(),
))
.await
.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
test_wire_profile(),
)
.with_websocket_pool(pool);
let request = pooled_websocket_request("conversation-turn-state-reset");
let first = backend
.create_response(
&request,
request_context("req_turn_state_first", Some("chatgpt-account")),
)
.await
.expect("first response should complete");
let second = backend
.create_response(
&request,
request_context("req_turn_state_second", Some("chatgpt-account")),
)
.await
.expect("second response should reuse the WebSocket");
server.await.unwrap();
assert_eq!(
first.turn_state.as_deref(),
Some("turn-from-first-response")
);
assert!(
second
.websocket_pool_decision
.is_some_and(WebSocketPoolDecision::is_reuse)
);
assert_eq!(second.turn_state, None);
assert!(
second
.response_metadata
.client_headers
.iter()
.all(|(name, _)| { !name.eq_ignore_ascii_case("x-codex-turn-state") })
);
}
#[tokio::test]
async fn websocket_request_should_project_trusted_turn_state_and_remove_stale_metadata() {
for (case, turn_state) in [("trusted", Some("trusted-turn-state")), ("cleared", None)] {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let message = websocket.next().await.unwrap().unwrap();
let payload = serde_json::from_str::<Value>(&message.into_text().unwrap())
.expect("response.create payload");
websocket
.send(Message::Text(
completed_websocket_response("resp_turn_state_projection", 1, 1).into(),
))
.await
.unwrap();
payload
});
let mut body = codex_request_body("gpt-5.5", "be brief", Vec::new());
body.insert("generate".to_owned(), json!(false));
body.insert("store".to_owned(), json!(false));
let mut request = CodexResponsesRequest::from_body(body);
request.turn_state = turn_state.map(str::to_owned);
request.local_conversation_id = Some(format!("turn-state-projection-{case}"));
request.set_client_metadata(Some(json!({
"x-codex-turn-state": "stale-client-value",
"keep": "preserved"
})));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
test_wire_profile(),
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))));
let mut context = request_context(case, Some("chatgpt-account"));
context.turn_state = turn_state;
let response = backend
.create_response(&request, context)
.await
.unwrap_or_else(|error| panic!("required WebSocket request should complete: {error}"));
let payload = server.await.unwrap();
assert_eq!(response.transport, CodexBackendTransport::WebSocket);
assert_eq!(
payload.pointer("/client_metadata/keep"),
Some(&json!("preserved"))
);
assert_eq!(
payload.pointer("/client_metadata/x-codex-turn-state"),
turn_state.map(|turn_state| json!(turn_state)).as_ref(),
"case {case}"
);
}
}
#[tokio::test]
async fn codex_backend_client_should_not_reuse_websocket_rate_limit_headers() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();