diff --git a/backend/crates/gateway-api/src/openai/responses/http.rs b/backend/crates/gateway-api/src/openai/responses/http.rs index a4f34bf88..49db85519 100644 --- a/backend/crates/gateway-api/src/openai/responses/http.rs +++ b/backend/crates/gateway-api/src/openai/responses/http.rs @@ -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) -> 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) -> 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, @@ -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 { diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs b/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs index 0b6bdd1c3..b61532f02 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/forward.rs @@ -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, +) -> 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, diff --git a/backend/crates/gateway-api/tests/openai/responses/http.rs b/backend/crates/gateway-api/tests/openai/responses/http.rs index b6a01bf01..ef12d5aeb 100644 --- a/backend/crates/gateway-api/tests/openai/responses/http.rs +++ b/backend/crates/gateway-api/tests/openai/responses/http.rs @@ -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 diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/connection.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/connection.rs index f74dc50a6..7631dfdee 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/connection.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/connection.rs @@ -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") diff --git a/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs b/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs index bc24434da..8795b1568 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs @@ -286,12 +286,20 @@ struct AtomicFailureTrace { struct AtomicFailureSession { trace: Arc, response_headers: Vec, + fail_before_first_event: bool, } impl ExecutionSession for AtomicFailureSession { fn next_event(&mut self) -> BoxFuture<'_, Result, 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, response_headers: Vec, + 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::::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!["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") diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index d5a7a54ad..91461dd6a 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -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, + 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) +} diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 779522c6c..86acc684f 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -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()); diff --git a/backend/crates/providers/openai/src/transport/client.rs b/backend/crates/providers/openai/src/transport/client.rs index 82c65e269..3bc7bbe5f 100644 --- a/backend/crates/providers/openai/src/transport/client.rs +++ b/backend/crates/providers/openai/src/transport/client.rs @@ -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, - /// 响应头里的最新 turn state。 + /// 响应头或 metadata 事件里最先确认的 turn state。 pub turn_state: Option, /// 上游透传的 `set-cookie` 列表。 pub set_cookie_headers: Vec, @@ -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()), diff --git a/backend/crates/providers/openai/src/transport/protocol/responses.rs b/backend/crates/providers/openai/src/transport/protocol/responses.rs index 0242d8a0c..18036033a 100644 --- a/backend/crates/providers/openai/src/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/src/transport/protocol/responses.rs @@ -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."; diff --git a/backend/crates/providers/openai/src/transport/protocol/websocket.rs b/backend/crates/providers/openai/src/transport/protocol/websocket.rs index caed6cedc..7e0e8daae 100644 --- a/backend/crates/providers/openai/src/transport/protocol/websocket.rs +++ b/backend/crates/providers/openai/src/transport/protocol/websocket.rs @@ -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 { - 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 连接寿命限制错误码。 diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index 0dc7e3f81..f3ad5fa0f 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -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) -> 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, diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs index d28912099..20dd4e255 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs @@ -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, /// 上游握手响应里的 `set-cookie` 列表。 pub set_cookie_headers: Vec, @@ -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, @@ -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 } diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs index f40589049..2d0cb9a3f 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs @@ -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, @@ -32,6 +31,7 @@ pub(super) enum ExchangeAction { pub(super) struct ReducedWebSocketEvent { pub(super) action: ExchangeAction, pub(super) diagnostic_event_type: Option, + pub(super) turn_state_update: Option, } 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, }) } diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs index 04f275405..7191eb2bd 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs @@ -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, }; diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index 86049234e..fb13fa8b6 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -461,6 +461,7 @@ fn contract_account_scope() -> Arc { "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()); } } diff --git a/backend/crates/providers/openai/tests/transport/protocol/websocket.rs b/backend/crates/providers/openai/tests/transport/protocol/websocket.rs index a29659d88..e834e3050 100644 --- a/backend/crates/providers/openai/tests/transport/protocol/websocket.rs +++ b/backend/crates/providers/openai/tests/transport/protocol/websocket.rs @@ -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] diff --git a/backend/crates/providers/openai/tests/transport/request.rs b/backend/crates/providers/openai/tests/transport/request.rs index 038e34b9f..943d29b3d 100644 --- a/backend/crates/providers/openai/tests/transport/request.rs +++ b/backend/crates/providers/openai/tests/transport/request.rs @@ -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] diff --git a/backend/crates/providers/openai/tests/transport/websocket/mod.rs b/backend/crates/providers/openai/tests/transport/websocket/mod.rs index 1b19be8b8..ad93e99e7 100644 --- a/backend/crates/providers/openai/tests/transport/websocket/mod.rs +++ b/backend/crates/providers/openai/tests/transport/websocket/mod.rs @@ -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::(&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();