From c8dc2e293a8720ee678a0e532fcf330c8f786e5c Mon Sep 17 00:00:00 2001 From: zyycn Date: Thu, 16 Jul 2026 01:46:33 +0800 Subject: [PATCH] fix(transport): eliminate websocket fallback race Notify foreground and shared waiters before awaiting pool cleanup so an observed WebSocket opening error cannot be replaced by a fast-path timeout and HTTP fallback. Drive the half-open breaker test through explicit state transitions instead of timing sleeps. Co-authored-by: Codex --- .../upstream/openai/transport/websocket.rs | 5 +- .../openai/transport/websocket_pool.rs | 3 +- .../upstream/openai/transport/latency.rs | 52 ++++++++----------- 3 files changed, 28 insertions(+), 32 deletions(-) diff --git a/backend/src/upstream/openai/transport/websocket.rs b/backend/src/upstream/openai/transport/websocket.rs index d29ce4d54..eb519158f 100644 --- a/backend/src/upstream/openai/transport/websocket.rs +++ b/backend/src/upstream/openai/transport/websocket.rs @@ -666,8 +666,8 @@ fn start_pooled_websocket_connect( biased; _ = cancellation.cancelled() => { permit.cancel(); - connect_lease.failed().await; let _ = sender.send(Err(CodexWebSocketExchangeError::SharedConnectFailed)); + connect_lease.failed().await; tracing::info!( account_id = task_key.account_id(), conversation_id_hash = task_key.conversation_id_hash(), @@ -710,8 +710,9 @@ fn start_pooled_websocket_connect( } Err(error) => { let error_message = error.to_string(); - connect_lease.failed().await; + // 先交付 opening 原始错误，避免连接池清理侵占前台 fast-path 预算。 let foreground_waiting = sender.send(Err(error)).is_ok(); + connect_lease.failed().await; tracing::warn!( account_id = task_key.account_id(), conversation_id_hash = task_key.conversation_id_hash(), diff --git a/backend/src/upstream/openai/transport/websocket_pool.rs b/backend/src/upstream/openai/transport/websocket_pool.rs index cba54c8f8..b60937913 100644 --- a/backend/src/upstream/openai/transport/websocket_pool.rs +++ b/backend/src/upstream/openai/transport/websocket_pool.rs @@ -657,9 +657,10 @@ impl WebSocketPoolConnectLease { pub(crate) async fn failed(mut self) { self.cancellation.cancel(); - self.pool.fail_connect(&self.key, self.id).await; + // 先唤醒共享等待者，pool mutex 清理不得延迟前台 transport 决策。 self.outcome .send_replace(WebSocketPoolConnectOutcome::Failed); + self.pool.fail_connect(&self.key, self.id).await; self.armed = false; } } diff --git a/backend/tests/upstream/openai/transport/latency.rs b/backend/tests/upstream/openai/transport/latency.rs index e9d876f97..7cd43d2f4 100644 --- a/backend/tests/upstream/openai/transport/latency.rs +++ b/backend/tests/upstream/openai/transport/latency.rs @@ -18,6 +18,18 @@ fn new_chain_request(conversation_id: &str) -> CodexResponsesRequest { request } +fn explicit_websocket_warmup_request(conversation_id: &str) -> CodexResponsesRequest { + let mut body = serde_json::Map::new(); + body.insert("model".to_string(), json!("gpt-5.5")); + body.insert("input".to_string(), json!([])); + body.insert("generate".to_string(), json!(false)); + body.insert("store".to_string(), json!(false)); + let mut request = CodexResponsesRequest::from_body(body); + request.use_websocket = true; + request.local_conversation_id = Some(conversation_id.to_string()); + request +} + #[tokio::test] async fn cold_websocket_should_fall_back_without_recording_a_successful_connect() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -326,14 +338,7 @@ async fn store_false_warmup_should_never_fall_back_to_http() { crate::support::wire_profile::test_wire_profile(), ) .with_websocket_pool(Arc::new(CodexWebSocketPool::default())); - let mut body = serde_json::Map::new(); - body.insert("model".to_string(), json!("gpt-5.5")); - body.insert("input".to_string(), json!([])); - body.insert("generate".to_string(), json!(false)); - body.insert("store".to_string(), json!(false)); - let mut request = CodexResponsesRequest::from_body(body); - request.use_websocket = true; - request.local_conversation_id = Some("conversation-warmup-required".to_string()); + let request = explicit_websocket_warmup_request("conversation-warmup-required"); let error = backend .create_response( @@ -784,21 +789,18 @@ async fn fast_path_miss_and_late_failure_should_count_as_one_breaker_failure() { async fn half_open_upstream_response_should_close_origin_breaker() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); + let origin_key = format!("http://{addr}"); let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig { failure_threshold: 1, failure_window: Duration::from_secs(1), - open_duration: Duration::from_millis(20), + open_duration: Duration::ZERO, }); + let WebSocketOriginBreakerDecision::Allowed(permit) = breaker.try_acquire(&origin_key) else { + panic!("closed breaker should grant the initial permit"); + }; + permit.fast_timeout(); + let server = tokio::spawn(async move { - let (mut stalled_websocket, _) = listener.accept().await.unwrap(); - let opening = read_http_request(&mut stalled_websocket).await; - assert!(opening.starts_with("GET /codex/responses HTTP/1.1")); - - let (mut http, _) = listener.accept().await.unwrap(); - let request = read_http_request(&mut http).await; - assert!(request.starts_with("POST /codex/responses HTTP/1.1")); - write_completed_sse_response(&mut http).await; - let (mut probe, _) = listener.accept().await.unwrap(); let opening = read_http_request(&mut probe).await; assert!(opening.starts_with("GET /codex/responses HTTP/1.1")); @@ -811,21 +813,13 @@ async fn half_open_upstream_response_should_close_origin_breaker() { }); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), - format!("http://{addr}"), + origin_key.clone(), crate::support::wire_profile::test_wire_profile(), ) .with_websocket_pool(Arc::new(CodexWebSocketPool::default())) - .with_websocket_fast_path_budget(Duration::from_millis(100)) + .with_websocket_fast_path_budget(Duration::from_secs(1)) .with_websocket_origin_breaker(breaker.clone()); - backend - .create_response( - &new_chain_request("conversation-breaker-open"), - request_context("req_breaker_open", Some("chatgpt-account")), - ) - .await - .expect("first timeout should use HTTP"); - tokio::time::sleep(Duration::from_millis(50)).await; let error = backend .create_response( &new_chain_request("conversation-breaker-probe"), @@ -844,7 +838,7 @@ async fn half_open_upstream_response_should_close_origin_breaker() { } if status == reqwest::StatusCode::UNAUTHORIZED ); assert!(matches!( - breaker.try_acquire(&format!("http://{addr}")), + breaker.try_acquire(&origin_key), WebSocketOriginBreakerDecision::Allowed(_) )); }