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 <noreply@openai.com>

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<String>,
     codex_window_id: Option<String>,
     parent_thread_id: Option<String>,
+    downstream_websocket_connection_id: Option<String>,
 
     conversation_id: Option<String>,
     session_id: Option<String>,
@@ -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<String>,
+    ) -> Self {
+        self.downstream_websocket_connection_id = Some(connection_id.into());
+        self
+    }
+
     fn apply_subagent(&self, body: &mut Map<String, Value>) {
         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",

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();

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<Vec<String>>,
 }
 
 struct AtomicFailureSession {
@@ -365,6 +366,19 @@ impl ExecutionService for AtomicFailureExecution {
         request: StartExecution,
     ) -> BoxFuture<'_, Result<StartedExecution, GatewayError>> {
         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();

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)
     }
 
     /// 获取后端模型目录条目。

backend/crates/providers/openai/src/transport/protocol/responses.rs

@@ -82,6 +82,8 @@ pub struct CodexResponsesRequest {
     pub memgen_request: Option<String>,
     /// Codex window ID。
     pub codex_window_id: Option<String>,
+    /// 代理分配的下游 WebSocket 连接 ID，仅用于隔离本地连接池通道。
+    pub downstream_websocket_connection_id: Option<String>,
     /// 父线程 ID。
     pub parent_thread_id: Option<String>,
     /// 已知 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,
         }

backend/crates/providers/openai/src/transport/request.rs

@@ -511,6 +511,9 @@ fn apply_protocol_context(request: &mut CodexResponsesRequest, context: &Map<Str
         .or_else(|| request.include_timing_metrics.take());
     request.codex_window_id =
         context_string(context, "codex_window_id").or_else(|| request.codex_window_id.take());
+    request.downstream_websocket_connection_id =
+        context_string(context, "downstream_websocket_connection_id")
+            .or_else(|| request.downstream_websocket_connection_id.take());
     request.parent_thread_id =
         context_string(context, "parent_thread_id").or_else(|| request.parent_thread_id.take());
     let prompt_cache_key = request.prompt_cache_key().map(ToOwned::to_owned);

backend/crates/providers/openai/src/transport/websocket/pool/state.rs

@@ -23,6 +23,7 @@ pub struct CodexWebSocketPoolKey {
     account_id: String,
     conversation_id: String,
     connection_profile: String,
+    downstream_connection_id: String,
 }
 
 impl CodexWebSocketPoolKey {
@@ -37,6 +38,7 @@ impl CodexWebSocketPoolKey {
             account_id: account_id.into(),
             conversation_id: conversation_id.into(),
             connection_profile: String::new(),
+            downstream_connection_id: String::new(),
         }
     }
 
@@ -46,6 +48,15 @@ impl CodexWebSocketPoolKey {
         self
     }
 
+    /// 隔离同一逻辑会话中由不同下游 WebSocket 驱动的并发响应链。
+    pub(crate) fn with_downstream_connection_id(
+        mut self,
+        connection_id: impl Into<String>,
+    ) -> 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(),
         ])
     }
 

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(

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();