From cf0e84c54881bae2a7562bded9e29542bce27a2f Mon Sep 17 00:00:00 2001 From: Zyy Date: Fri, 4 Sep 2026 11:50:34 +0800 Subject: [PATCH] fix(openai): isolate websocket continuation lanes Include the downstream connection identity in pooled upstream WebSocket keys so concurrent clients cannot contend for connection-local response state. Co-authored-by: Codex --- .../src/openai/responses/request.rs | 18 +++ .../src/openai/responses/websocket/mod.rs | 2 + .../tests/openai/responses/websocket/mod.rs | 24 +++- .../openai/src/transport/client_sse.rs | 10 +- .../src/transport/protocol/responses.rs | 3 + .../providers/openai/src/transport/request.rs | 3 + .../src/transport/websocket/pool/state.rs | 12 ++ .../openai/tests/transport/request.rs | 29 ++++ .../openai/tests/transport/websocket_pool.rs | 132 ++++++++++++++++++ 9 files changed, 228 insertions(+), 5 deletions(-) diff --git a/backend/crates/gateway-api/src/openai/responses/request.rs b/backend/crates/gateway-api/src/openai/responses/request.rs index 132c08e66..3aeaceeb9 100644 --- a/backend/crates/gateway-api/src/openai/responses/request.rs +++ b/backend/crates/gateway-api/src/openai/responses/request.rs @@ -16,6 +16,7 @@ const OPENAI_PROTOCOL: &str = "openai"; const OPENAI_SUBAGENT_KEY: &str = "x-openai-subagent"; const CODEX_TURN_METADATA_KEY: &str = "x-codex-turn-metadata"; const PASSTHROUGH_HEADERS_CONTEXT_KEY: &str = "opaque_request_headers"; +const DOWNSTREAM_WEBSOCKET_CONNECTION_ID_CONTEXT_KEY: &str = "downstream_websocket_connection_id"; /// Responses 请求进入共享解码内核时的下游传输来源。 #[derive(Clone, Copy)] @@ -35,6 +36,7 @@ pub struct OpenAiRequestHeaders { include_timing_metrics: Option, codex_window_id: Option, parent_thread_id: Option, + downstream_websocket_connection_id: Option, conversation_id: Option, session_id: Option, @@ -60,6 +62,7 @@ impl OpenAiRequestHeaders { include_timing_metrics: header_string(headers, "x-responsesapi-include-timing-metrics"), codex_window_id: header_string(headers, "x-codex-window-id"), parent_thread_id: header_string(headers, "x-codex-parent-thread-id"), + downstream_websocket_connection_id: None, conversation_id: header_string(headers, "conversation-id") .or_else(|| header_string(headers, "conversation_id")), session_id: header_string(headers, "session-id") @@ -74,6 +77,16 @@ impl OpenAiRequestHeaders { } } + /// 绑定代理为当前下游 WebSocket 分配的连接身份。 + #[must_use] + pub(super) fn with_downstream_websocket_connection_id( + mut self, + connection_id: impl Into, + ) -> Self { + self.downstream_websocket_connection_id = Some(connection_id.into()); + self + } + fn apply_subagent(&self, body: &mut Map) { if let Some(subagent) = self.subagent.as_deref() { inject_subagent_metadata(body, subagent); @@ -109,6 +122,11 @@ impl OpenAiRequestHeaders { "parent_thread_id", self.parent_thread_id.as_ref(), ); + insert_protocol_context( + &mut context, + DOWNSTREAM_WEBSOCKET_CONNECTION_ID_CONTEXT_KEY, + self.downstream_websocket_connection_id.as_ref(), + ); insert_protocol_context( &mut context, "conversation_id", diff --git a/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs index a436cb75c..8759dd169 100644 --- a/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs +++ b/backend/crates/gateway-api/src/openai/responses/websocket/mod.rs @@ -143,6 +143,8 @@ async fn serve_responses_websocket(socket: WebSocket, session: ResponsesWebSocke ); let _connection_guard = connection_guard; let cancellation = lifecycle.cancellation(); + let request_headers = + request_headers.with_downstream_websocket_connection_id(connection_id.clone()); let mut connection = ResponsesWebSocketConnection::new(socket, connection_id, cancellation); let mut request_count = 0_u64; let mut replay = ConnectionReplaySnapshot::default(); 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 d1a74c81c..bc24434da 100644 --- a/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/websocket/mod.rs @@ -2,7 +2,7 @@ mod connection; mod protocol; use std::sync::{ - Arc, + Arc, Mutex, atomic::{AtomicBool, AtomicUsize, Ordering}, }; use std::time::{Duration, SystemTime}; @@ -280,6 +280,7 @@ struct AtomicFailureTrace { next_calls: AtomicUsize, committed: AtomicBool, finalized: AtomicBool, + downstream_websocket_connection_ids: Mutex>, } struct AtomicFailureSession { @@ -365,6 +366,19 @@ impl ExecutionService for AtomicFailureExecution { request: StartExecution, ) -> BoxFuture<'_, Result> { self.trace.starts.fetch_add(1, Ordering::AcqRel); + if let Operation::Generate(generate) = &request.operation + && let Some(connection_id) = generate + .protocol_payload() + .context() + .get("downstream_websocket_connection_id") + .and_then(Value::as_str) + { + self.trace + .downstream_websocket_connection_ids + .lock() + .expect("downstream connection trace lock") + .push(connection_id.to_owned()); + } Box::pin(async move { Ok(StartedExecution { request_id: gateway_core::engine::ModelRequestId::new("req_ws_atomic") @@ -510,6 +524,14 @@ async fn websocket_atomic_upstream_failure_batch_should_be_forwarded_once() { assert!(trace.committed.load(Ordering::Acquire)); assert_eq!(trace.next_calls.load(Ordering::Acquire), 2); assert!(trace.finalized.load(Ordering::Acquire)); + { + let downstream_connection_ids = trace + .downstream_websocket_connection_ids + .lock() + .expect("downstream connection trace lock"); + assert_eq!(downstream_connection_ids.len(), 1); + assert!(downstream_connection_ids[0].starts_with("ws_")); + } socket.close(None).await.expect("close WebSocket"); server.abort(); diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index 076edaa0f..7b81e1a8d 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -682,10 +682,12 @@ impl CodexBackendClient { .local_conversation_id .as_deref() .or(request.previous_response_id())?; - Some( - CodexWebSocketPoolKey::new(&self.base_url, account_id, conversation_id) - .with_connection_profile(connection_profile), - ) + let mut key = CodexWebSocketPoolKey::new(&self.base_url, account_id, conversation_id) + .with_connection_profile(connection_profile); + if let Some(connection_id) = request.downstream_websocket_connection_id.as_deref() { + key = key.with_downstream_connection_id(connection_id); + } + Some(key) } /// 获取后端模型目录条目。 diff --git a/backend/crates/providers/openai/src/transport/protocol/responses.rs b/backend/crates/providers/openai/src/transport/protocol/responses.rs index cdf452e22..0242d8a0c 100644 --- a/backend/crates/providers/openai/src/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/src/transport/protocol/responses.rs @@ -82,6 +82,8 @@ pub struct CodexResponsesRequest { pub memgen_request: Option, /// Codex window ID。 pub codex_window_id: Option, + /// 代理分配的下游 WebSocket 连接 ID，仅用于隔离本地连接池通道。 + pub downstream_websocket_connection_id: Option, /// 父线程 ID。 pub parent_thread_id: Option, /// 已知 previous response 的持久化范围，仅用于本地 transport 校验。 @@ -519,6 +521,7 @@ impl CodexResponsesRequest { responses_lite: None, memgen_request: None, codex_window_id: None, + downstream_websocket_connection_id: None, parent_thread_id: None, previous_response_scope: None, } diff --git a/backend/crates/providers/openai/src/transport/request.rs b/backend/crates/providers/openai/src/transport/request.rs index 6f94ea075..0dc7e3f81 100644 --- a/backend/crates/providers/openai/src/transport/request.rs +++ b/backend/crates/providers/openai/src/transport/request.rs @@ -511,6 +511,9 @@ fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map, + ) -> Self { + self.downstream_connection_id = connection_id.into(); + self + } + pub(crate) fn account_id(&self) -> &str { &self.account_id } @@ -60,6 +71,7 @@ impl CodexWebSocketPoolKey { self.account_id.as_str(), self.conversation_id.as_str(), self.connection_profile.as_str(), + self.downstream_connection_id.as_str(), ]) } diff --git a/backend/crates/providers/openai/tests/transport/request.rs b/backend/crates/providers/openai/tests/transport/request.rs index 72702c558..038e34b9f 100644 --- a/backend/crates/providers/openai/tests/transport/request.rs +++ b/backend/crates/providers/openai/tests/transport/request.rs @@ -260,6 +260,35 @@ fn header_context_should_win_over_body_topline_aliases() { ); } +#[test] +fn encoder_should_preserve_downstream_websocket_connection_identity_outside_wire_body() { + let payload = ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("client-model")), + ("input".to_owned(), json!("prompt")), + ]), + ) + .expect("OpenAI payload") + .with_context(Map::from_iter([( + "downstream_websocket_connection_id".to_owned(), + Value::String("ws_downstream_a".to_owned()), + )])); + let request = GenerateRequest::from_protocol_payload(payload); + + let encoded = encode_generate_request(&request, "gpt-test").expect("encode"); + + assert_eq!( + encoded.downstream_websocket_connection_id.as_deref(), + Some("ws_downstream_a") + ); + assert!( + !encoded + .body() + .contains_key("downstream_websocket_connection_id") + ); +} + #[test] fn body_topline_alias_should_only_fill_an_absent_header_context() { let payload = ProtocolPayload::json_object( diff --git a/backend/crates/providers/openai/tests/transport/websocket_pool.rs b/backend/crates/providers/openai/tests/transport/websocket_pool.rs index 4f571007f..a1c01a13c 100644 --- a/backend/crates/providers/openai/tests/transport/websocket_pool.rs +++ b/backend/crates/providers/openai/tests/transport/websocket_pool.rs @@ -71,6 +71,138 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and assert_eq!(accepted_connections.load(Ordering::SeqCst), 1); } +#[tokio::test] +async fn websocket_pool_should_isolate_concurrent_downstream_lanes_for_one_conversation() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let accepted_connections = Arc::new(AtomicUsize::new(0)); + let accepted_connections_for_server = Arc::clone(&accepted_connections); + let (first_accepted_tx, first_accepted_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (first_stream, _) = listener.accept().await.unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut first = accept_codex_test_websocket(first_stream).await; + first_accepted_tx.send(()).unwrap(); + + let (second_stream, _) = timeout(Duration::from_secs(2), listener.accept()) + .await + .expect("the second downstream lane should open its own upstream websocket") + .unwrap(); + accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); + let mut second = accept_codex_test_websocket(second_stream).await; + + let (first_warmup, second_warmup) = tokio::join!(first.next(), second.next()); + first_warmup.unwrap().unwrap(); + second_warmup.unwrap().unwrap(); + first + .send(Message::Text( + completed_websocket_response("resp_lane_a", 1, 0).into(), + )) + .await + .unwrap(); + second + .send(Message::Text( + completed_websocket_response("resp_lane_b", 1, 0).into(), + )) + .await + .unwrap(); + + let _first_continuation = first.next().await.unwrap().unwrap(); + first + .send(Message::Text( + completed_websocket_response("resp_lane_a_next", 2, 1).into(), + )) + .await + .unwrap(); + let _second_continuation = second.next().await.unwrap().unwrap(); + second + .send(Message::Text( + completed_websocket_response("resp_lane_b_next", 2, 1).into(), + )) + .await + .unwrap(); + first.close(None).await.unwrap(); + second.close(None).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(Arc::clone(&pool)); + let warmup = |connection_id: &str| { + let mut body = Map::new(); + body.insert("model".to_owned(), json!("gpt-5.5")); + body.insert("input".to_owned(), json!([])); + body.insert("generate".to_owned(), json!(false)); + body.insert("store".to_owned(), json!(false)); + let mut request = CodexResponsesRequest::from_body(body); + request.use_websocket = true; + request.local_conversation_id = Some("conversation-lanes".to_owned()); + request.downstream_websocket_connection_id = Some(connection_id.to_owned()); + request + }; + + let first_backend = backend.clone(); + let first_request = warmup("ws_downstream_a"); + let first = tokio::spawn(async move { + first_backend + .create_response( + &first_request, + request_context("req_lane_a_warmup", Some("chatgpt-account")), + ) + .await + }); + first_accepted_rx.await.unwrap(); + let second = backend + .create_response( + &warmup("ws_downstream_b"), + request_context("req_lane_b_warmup", Some("chatgpt-account")), + ) + .await + .expect("the second downstream lane should not contend with the first warmup"); + let first = first.await.unwrap().unwrap(); + + let mut first_continuation = pooled_websocket_request("conversation-lanes"); + first_continuation.downstream_websocket_connection_id = Some("ws_downstream_a".to_owned()); + first_continuation.set_previous_response_id(Some("resp_lane_a".to_owned())); + first_continuation.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal); + let first_continuation = backend + .create_response( + &first_continuation, + request_context("req_lane_a_next", Some("chatgpt-account")), + ) + .await + .expect("the first lane should retain its connection-local response owner"); + + let mut second_continuation = pooled_websocket_request("conversation-lanes"); + second_continuation.downstream_websocket_connection_id = Some("ws_downstream_b".to_owned()); + second_continuation.set_previous_response_id(Some("resp_lane_b".to_owned())); + second_continuation.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal); + let second_continuation = backend + .create_response( + &second_continuation, + request_context("req_lane_b_next", Some("chatgpt-account")), + ) + .await + .expect("the second lane should retain its connection-local response owner"); + server.await.unwrap(); + + assert_eq!(first.websocket_pool_decision.unwrap().kind(), "new"); + assert_eq!(second.websocket_pool_decision.unwrap().kind(), "new"); + assert_eq!( + first_continuation.websocket_pool_decision.unwrap().kind(), + "reuse" + ); + assert_eq!( + second_continuation.websocket_pool_decision.unwrap().kind(), + "reuse" + ); + assert_eq!(accepted_connections.load(Ordering::SeqCst), 2); + pool.shutdown().await; +} + #[tokio::test] async fn exact_continuation_reuses_owning_socket_after_connection_profile_updates() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();