From e0326ecba4e3cfcecdec1b908e81d55b033722eb Mon Sep 17 00:00:00 2001 From: Jiahao <58540575+jiahao6635@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:06:23 +0800 Subject: [PATCH] fix(openai): base pooled WS breaker on opening results (#484) * fix(openai): base pooled WS breaker on opening results Keep foreground fallback independent of pooled opening health and fence late results by circuit generation. Refs #482 Co-authored-by: Codex * fix(openai): preserve concurrent breaker results while closed Co-authored-by: Codex --------- Co-authored-by: Codex Co-authored-by: Zyy --- .../openai/src/transport/websocket/breaker.rs | 176 ++++----- .../src/transport/websocket/coordinator.rs | 10 +- .../openai/tests/transport/latency.rs | 358 +++++++++++++++--- .../tests/transport/websocket/breaker.rs | 116 ++++++ .../openai/tests/transport/websocket/mod.rs | 1 + 5 files changed, 500 insertions(+), 161 deletions(-) create mode 100644 backend/crates/providers/openai/tests/transport/websocket/breaker.rs diff --git a/backend/crates/providers/openai/src/transport/websocket/breaker.rs b/backend/crates/providers/openai/src/transport/websocket/breaker.rs index 4a59c949f..eadc092d5 100644 --- a/backend/crates/providers/openai/src/transport/websocket/breaker.rs +++ b/backend/crates/providers/openai/src/transport/websocket/breaker.rs @@ -2,10 +2,7 @@ use std::{ collections::{HashMap, VecDeque}, - sync::{ - Arc, Mutex, - atomic::{AtomicBool, Ordering}, - }, + sync::{Arc, Mutex}, time::Duration, }; @@ -41,7 +38,7 @@ pub enum WebSocketOriginBreakerDecision { HalfOpenBusy, } -/// origin 级 WebSocket 快路径熔断器 +/// origin 级 WebSocket 建连熔断器 #[derive(Clone)] pub struct WebSocketOriginBreaker { inner: Arc>>, @@ -77,81 +74,90 @@ impl WebSocketOriginBreaker { .or_insert_with(CircuitState::closed); match state { - CircuitState::Closed { failures } => { + CircuitState::Closed { + generation, + failures, + } => { retain_window(failures, now, self.config.failure_window); WebSocketOriginBreakerDecision::Allowed(WebSocketOriginBreakerPermit::new( self.clone(), origin_key.to_string(), - None, + *generation, )) } CircuitState::Open { until } if now < *until => WebSocketOriginBreakerDecision::Open, CircuitState::Open { .. } => { - let probe_id = Uuid::new_v4(); - *state = CircuitState::HalfOpen { probe_id }; + let generation = Uuid::new_v4(); + *state = CircuitState::HalfOpen { generation }; WebSocketOriginBreakerDecision::Allowed(WebSocketOriginBreakerPermit::new( self.clone(), origin_key.to_string(), - Some(probe_id), + generation, )) } CircuitState::HalfOpen { .. } => WebSocketOriginBreakerDecision::HalfOpenBusy, } } - fn record_success(&self, origin_key: &str, probe_id: Option) { + fn record_success(&self, origin_key: &str, generation: Uuid) { let mut circuits = self .inner .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if permit_still_owns_state(circuits.get(origin_key), probe_id) { - circuits.insert(origin_key.to_string(), CircuitState::closed()); + let Some(state) = circuits.get_mut(origin_key) else { + return; + }; + if !permit_still_owns_state(Some(state), generation) { + return; + } + match state { + CircuitState::Closed { failures, .. } => { + // 保留同轮尚未完成的 opening，让其后续结果继续参与计数 + failures.clear(); + } + CircuitState::HalfOpen { .. } => *state = CircuitState::closed(), + CircuitState::Open { .. } => {} } } - fn record_failure(&self, origin_key: &str, probe_id: Option) { + fn record_failure(&self, origin_key: &str, generation: Uuid) { let now = Instant::now(); let mut circuits = self .inner .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if !permit_still_owns_state(circuits.get(origin_key), probe_id) { - return; - } - if probe_id.is_some() { - circuits.insert( - origin_key.to_string(), - CircuitState::Open { - until: now + self.config.open_duration, - }, - ); - return; - } - - let state = circuits - .entry(origin_key.to_string()) - .or_insert_with(CircuitState::closed); - let CircuitState::Closed { failures } = state else { + let Some(state) = circuits.get_mut(origin_key) else { return; }; - retain_window(failures, now, self.config.failure_window); - failures.push_back(now); - if failures.len() >= self.config.failure_threshold { - *state = CircuitState::Open { - until: now + self.config.open_duration, - }; + if !permit_still_owns_state(Some(state), generation) { + return; + } + match state { + CircuitState::HalfOpen { .. } => { + *state = CircuitState::Open { + until: now + self.config.open_duration, + }; + } + CircuitState::Closed { failures, .. } => { + retain_window(failures, now, self.config.failure_window); + failures.push_back(now); + if failures.len() >= self.config.failure_threshold { + *state = CircuitState::Open { + until: now + self.config.open_duration, + }; + } + } + CircuitState::Open { .. } => {} } } - fn cancel_probe(&self, origin_key: &str, probe_id: Option) { - let Some(probe_id) = probe_id else { - return; - }; + fn cancel_probe(&self, origin_key: &str, generation: Uuid) { let mut circuits = self .inner .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if permit_still_owns_state(circuits.get(origin_key), Some(probe_id)) { + if matches!(circuits.get(origin_key), Some(CircuitState::HalfOpen { generation: active }) if *active == generation) + { circuits.insert( origin_key.to_string(), CircuitState::Open { @@ -166,95 +172,71 @@ impl WebSocketOriginBreaker { pub struct WebSocketOriginBreakerPermit { breaker: WebSocketOriginBreaker, origin_key: String, - probe_id: Option, - degradation_recorded: Arc, + generation: Uuid, armed: bool, } impl WebSocketOriginBreakerPermit { - fn new(breaker: WebSocketOriginBreaker, origin_key: String, probe_id: Option) -> Self { + fn new(breaker: WebSocketOriginBreaker, origin_key: String, generation: Uuid) -> Self { Self { breaker, origin_key, - probe_id, - degradation_recorded: Arc::new(AtomicBool::new(false)), + generation, armed: true, } } - /// 创建可跨前台预算与后台 opening 生命周期共享的快路径观察器 - pub(crate) fn fast_path_reporter(&self) -> WebSocketOriginFastPathReporter { - WebSocketOriginFastPathReporter { - breaker: self.breaker.clone(), - origin_key: self.origin_key.clone(), - probe_id: self.probe_id, - degradation_recorded: Arc::clone(&self.degradation_recorded), - } - } - pub fn succeed(mut self) { - if !self.degradation_recorded.load(Ordering::Acquire) { - self.breaker.record_success(&self.origin_key, self.probe_id); - } + self.breaker + .record_success(&self.origin_key, self.generation); self.armed = false; } - pub fn fast_timeout(mut self) { - self.fast_path_reporter().missed(); - self.armed = false; + /// 非池化 opening 已被前台时限终止，按连接失败计数 + pub fn fast_timeout(self) { + self.fail(); } pub fn fail(mut self) { - self.record_failure_once(); + self.breaker + .record_failure(&self.origin_key, self.generation); self.armed = false; } /// opening 因账号驱逐或服务关闭而取消，不污染 origin 可用性 pub(crate) fn cancel(mut self) { - self.breaker.cancel_probe(&self.origin_key, self.probe_id); + self.breaker.cancel_probe(&self.origin_key, self.generation); self.armed = false; } - - fn record_failure_once(&self) { - if !self.degradation_recorded.swap(true, Ordering::AcqRel) { - self.breaker.record_failure(&self.origin_key, self.probe_id); - } - } } impl Drop for WebSocketOriginBreakerPermit { fn drop(&mut self) { if self.armed { - self.record_failure_once(); - } - } -} - -/// 前台等待结束后仍可报告同一次后台 opening 错过快路径预算 -pub(crate) struct WebSocketOriginFastPathReporter { - breaker: WebSocketOriginBreaker, - origin_key: String, - probe_id: Option, - degradation_recorded: Arc, -} - -impl WebSocketOriginFastPathReporter { - pub(crate) fn missed(&self) { - if !self.degradation_recorded.swap(true, Ordering::AcqRel) { - self.breaker.record_failure(&self.origin_key, self.probe_id); + self.breaker + .record_failure(&self.origin_key, self.generation); } } } enum CircuitState { - Closed { failures: VecDeque }, - Open { until: Instant }, - HalfOpen { probe_id: Uuid }, + Closed { + generation: Uuid, + failures: VecDeque, + }, + Open { + until: Instant, + }, + HalfOpen { + generation: Uuid, + }, } impl CircuitState { fn closed() -> Self { Self::Closed { + // 初始化或半开恢复时创建新代次，隔离之前熔断周期的迟到结果 + generation: Uuid::new_v4(), failures: VecDeque::new(), } } @@ -269,10 +251,14 @@ fn retain_window(timeouts: &mut VecDeque, now: Instant, window: Duratio } } -fn permit_still_owns_state(state: Option<&CircuitState>, probe_id: Option) -> bool { - match (state, probe_id) { - (Some(CircuitState::Closed { .. }), None) => true, - (Some(CircuitState::HalfOpen { probe_id: active }), Some(expected)) => *active == expected, +fn permit_still_owns_state(state: Option<&CircuitState>, generation: Uuid) -> bool { + match state { + Some( + CircuitState::Closed { + generation: active, .. + } + | CircuitState::HalfOpen { generation: active }, + ) => *active == generation, _ => false, } } diff --git a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs index 0b3197ea9..74afc9024 100644 --- a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs +++ b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs @@ -13,7 +13,6 @@ use tokio::{ use super::{ breaker::{ WebSocketOriginBreaker, WebSocketOriginBreakerDecision, WebSocketOriginBreakerPermit, - WebSocketOriginFastPathReporter, }, error::CodexWebSocketExchangeError, exchange::{ @@ -306,7 +305,6 @@ async fn prepare_pooled_websocket( struct PooledWebSocketConnectWaiter { started_at: tokio::time::Instant, receiver: oneshot::Receiver>, - fast_path_reporter: Option, } struct PooledWebSocketHandoff { @@ -324,9 +322,7 @@ impl PooledWebSocketConnectWaiter { match wait_for_fast_path(self.started_at, fast_path_budget, &mut self.receiver).await { WebSocketFastPath::Ready(received) => received, WebSocketFastPath::Missed => { - if let Some(reporter) = &self.fast_path_reporter { - reporter.missed(); - } + // 熔断许可仍归后台 opening 持有，前台停止等待不等于建连失败 return Ok(WebSocketFastPath::Missed); } }; @@ -347,9 +343,6 @@ fn start_pooled_websocket_connect( let task_key = connect_lease.key().clone(); let started_at = connect_lease.started_at(); let cancellation = connect_lease.cancellation_token(); - let fast_path_reporter = permit - .as_ref() - .map(WebSocketOriginBreakerPermit::fast_path_reporter); let keepalive = pool.keepalive(); let context = PumpLogContext::new( Some(task_key.account_id().to_owned()), @@ -437,7 +430,6 @@ fn start_pooled_websocket_connect( PooledWebSocketConnectWaiter { started_at, receiver, - fast_path_reporter, } } diff --git a/backend/crates/providers/openai/tests/transport/latency.rs b/backend/crates/providers/openai/tests/transport/latency.rs index 5e9b163b4..7e89a2d8b 100644 --- a/backend/crates/providers/openai/tests/transport/latency.rs +++ b/backend/crates/providers/openai/tests/transport/latency.rs @@ -117,7 +117,7 @@ async fn cold_websocket_should_fall_back_without_recording_a_successful_connect( pool.shutdown().await; } -#[tokio::test(start_paused = true)] +#[tokio::test] async fn downstream_websocket_new_chain_should_preserve_continuation_after_slow_opening() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -190,7 +190,7 @@ async fn downstream_websocket_new_chain_should_preserve_continuation_after_slow_ pool.shutdown().await; } -#[tokio::test(start_paused = true)] +#[tokio::test] async fn external_continuation_should_wait_for_a_cold_websocket() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -227,67 +227,102 @@ async fn external_continuation_should_wait_for_a_cold_websocket() { pool.shutdown().await; } -#[tokio::test(start_paused = true)] -async fn timed_out_websocket_should_finish_in_background_and_serve_the_next_request() { +#[tokio::test] +async fn slow_pooled_openings_should_remain_healthy_and_serve_subsequent_requests() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - let (websocket_ready_tx, websocket_ready_rx) = tokio::sync::oneshot::channel(); + let openings = WebSocketOriginBreakerConfig::default().failure_threshold + 1; + let (websocket_ready_tx, mut websocket_ready_rx) = tokio::sync::mpsc::channel(1); let server = tokio::spawn(async move { - let (websocket_stream, _) = listener.accept().await.unwrap(); - let websocket_server = tokio::spawn(async move { - tokio::time::sleep(Duration::from_secs(1)).await; - let mut websocket = accept_codex_test_websocket(websocket_stream).await; - websocket_ready_tx.send(()).unwrap(); - let _payload = websocket.next().await.unwrap().unwrap(); - websocket - .send(Message::Text( - completed_websocket_response("resp_background_ready", 2, 1).into(), - )) - .await - .unwrap(); - }); + for _ in 0..openings { + let (mut websocket_stream, _) = listener.accept().await.unwrap(); + let mut prefix = [0; 4]; + loop { + let read = websocket_stream.peek(&mut prefix).await.unwrap(); + assert!(read > 0, "opening socket closed before the request prefix"); + if read == prefix.len() { + break; + } + tokio::task::yield_now().await; + } + if prefix == *b"POST" { + read_http_request(&mut websocket_stream).await; + write_completed_sse_response(&mut websocket_stream).await; + continue; + } + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let ready = websocket_ready_tx.clone(); + let websocket_server = tokio::spawn(async move { + release_rx.await.unwrap(); + let mut websocket = accept_codex_test_websocket(websocket_stream).await; + ready.send(()).await.unwrap(); + // 回退的请求只发 HTTP，WS 的首个业务帧必须来自后续请求 + let payload = websocket.next().await.unwrap().unwrap(); + let payload: Value = serde_json::from_str(payload.to_text().unwrap()).unwrap(); + assert_eq!(payload["input"][0]["content"], "subsequent request"); + websocket + .send(Message::Text( + completed_websocket_response("resp_background_ready", 2, 1).into(), + )) + .await + .unwrap(); + }); - 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; - websocket_server.await.unwrap(); + 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; + release_tx.send(()).unwrap(); + websocket_server.await.unwrap(); + } }); - let pool = Arc::new(CodexWebSocketPool::default()); + 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 request = new_chain_request("conversation-background-ready"); + for index in 0..openings { + let request = new_chain_request(&format!("conversation-background-{index}")); + let first = backend + .create_response( + &request, + request_context("req_background_http", Some("chatgpt-account")), + ) + .await + .expect("foreground request should use HTTP after its fast-path budget"); + assert_eq!(first.transport, CodexBackendTransport::HttpSse); + assert_eq!(first.transport_metrics.ws_connect_ms, None); + assert_eq!( + first.transport_metrics.decision, + Some(CodexTransportDecision::Http2WebSocketBudgetExhausted) + ); + websocket_ready_rx.recv().await.unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + let mut subsequent = codex_request( + "gpt-5.5", + "be brief", + vec![json!({"role":"user","content":"subsequent request"})], + ); + subsequent.use_websocket = true; + subsequent.local_conversation_id = request.local_conversation_id.clone(); + let second = backend + .create_response( + &subsequent, + request_context("req_background_reuse", Some("chatgpt-account")), + ) + .await + .expect("next request should reuse the background websocket"); - let first = backend - .create_response( - &request, - request_context("req_background_http", Some("chatgpt-account")), - ) - .await - .expect("foreground request should use HTTP after its fast-path budget"); - websocket_ready_rx.await.unwrap(); - tokio::time::sleep(Duration::from_millis(20)).await; - let second = backend - .create_response( - &request, - request_context("req_background_reuse", Some("chatgpt-account")), - ) - .await - .expect("next request should reuse the background websocket"); + assert_eq!(second.transport, CodexBackendTransport::WebSocket); + assert_eq!( + second.transport_metrics.decision, + Some(CodexTransportDecision::ReusedWebSocket) + ); + assert!(second.body.contains("resp_background_ready")); + } server.await.unwrap(); - - assert_eq!(first.transport, CodexBackendTransport::HttpSse); - assert_eq!(first.transport_metrics.ws_connect_ms, None); - assert_eq!(second.transport, CodexBackendTransport::WebSocket); - assert_eq!( - second.transport_metrics.decision, - Some(CodexTransportDecision::ReusedWebSocket) - ); - assert!(second.body.contains("resp_background_ready")); pool.shutdown().await; } @@ -1265,6 +1300,201 @@ async fn origin_breaker_should_open_then_allow_only_one_half_open_probe() { )); } +#[tokio::test] +async fn slow_half_open_probe_should_remain_singleflight_and_recover_on_success() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let config = WebSocketOriginBreakerConfig::default(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let (ready_tx, ready_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + reject_websocket_openings(&listener, config.failure_threshold).await; + let (stream, _) = listener.accept().await.unwrap(); + let probe = tokio::spawn(async move { + release_rx.await.unwrap(); + let mut websocket = accept_codex_test_websocket(stream).await; + ready_tx.send(()).unwrap(); + let _payload = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + completed_websocket_response("resp_probe_reused", 2, 1).into(), + )) + .await + .unwrap(); + }); + // 探测未完成前，预算耗尽的原请求和其他新会话都只能发送 HTTP + for _ in 0..2 { + let (mut http, _) = listener.accept().await.unwrap(); + assert!( + read_http_request(&mut http) + .await + .starts_with("POST /codex/responses") + ); + write_completed_sse_response(&mut http).await; + } + probe.await.unwrap(); + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_codex_test_websocket(stream).await; + let _payload = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + completed_websocket_response("resp_after_recovery", 2, 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(Arc::clone(&pool)); + record_websocket_origin_failures(&backend, config.failure_threshold, "slow-half-open").await; + tokio::time::pause(); + tokio::time::advance(config.open_duration).await; + tokio::time::resume(); + + let probe_request = new_chain_request("slow-half-open-probe"); + let fallback = backend + .create_response( + &probe_request, + request_context("req_slow_probe", Some("chatgpt-account")), + ) + .await + .unwrap(); + assert_eq!( + fallback.transport_metrics.decision, + Some(CodexTransportDecision::Http2WebSocketBudgetExhausted) + ); + let suppressed = backend + .create_response( + &new_chain_request("other-cold-conversation"), + request_context("req_probe_busy", Some("chatgpt-account")), + ) + .await + .unwrap(); + assert_eq!( + suppressed.transport_metrics.decision, + Some(CodexTransportDecision::Http2BreakerOpen) + ); + release_tx.send(()).unwrap(); + ready_rx.await.unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + let reused = backend + .create_response( + &probe_request, + request_context("req_probe_reuse", Some("chatgpt-account")), + ) + .await + .unwrap(); + assert_eq!( + reused.transport_metrics.decision, + Some(CodexTransportDecision::ReusedWebSocket) + ); + let fresh = backend + .create_response( + &new_chain_request("after-probe-recovery"), + request_context("req_after_probe", Some("chatgpt-account")), + ) + .await + .unwrap(); + assert_eq!(fresh.transport, CodexBackendTransport::WebSocket); + assert_eq!( + fresh.transport_metrics.decision, + Some(CodexTransportDecision::ConnectedWebSocket) + ); + server.await.unwrap(); + pool.shutdown().await; +} + +#[tokio::test] +async fn slow_half_open_probe_should_reopen_after_the_hard_connect_timeout() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let config = WebSocketOriginBreakerConfig::default(); + let (closed_tx, closed_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + reject_websocket_openings(&listener, config.failure_threshold).await; + let (mut opening, _) = listener.accept().await.unwrap(); + assert!( + read_http_request(&mut opening) + .await + .starts_with("GET /codex/responses") + ); + let stalled = tokio::spawn(async move { + let mut byte = [0]; + assert_eq!(opening.read(&mut byte).await.unwrap(), 0); + closed_tx.send(()).unwrap(); + }); + for _ in 0..3 { + let (mut http, _) = listener.accept().await.unwrap(); + assert!( + read_http_request(&mut http) + .await + .starts_with("POST /codex/responses") + ); + write_completed_sse_response(&mut http).await; + } + stalled.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)); + record_websocket_origin_failures(&backend, config.failure_threshold, "hard-timeout").await; + tokio::time::pause(); + tokio::time::advance(config.open_duration).await; + tokio::time::resume(); + + let fallback = backend + .create_response( + &new_chain_request("hard-timeout-probe"), + request_context("req_hard_probe", Some("chatgpt-account")), + ) + .await + .unwrap(); + assert_eq!( + fallback.transport_metrics.decision, + Some(CodexTransportDecision::Http2WebSocketBudgetExhausted) + ); + let busy = backend + .create_response( + &new_chain_request("while-hard-probe-pending"), + request_context("req_hard_busy", Some("chatgpt-account")), + ) + .await + .unwrap(); + assert_eq!( + busy.transport_metrics.decision, + Some(CodexTransportDecision::Http2BreakerOpen) + ); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(15)).await; + tokio::time::resume(); + timeout(Duration::from_secs(1), closed_rx) + .await + .expect("15s hard timeout must close the opening socket") + .unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + let suppressed = backend + .create_response( + &new_chain_request("after-hard-probe-failure"), + request_context("req_after_hard_failure", Some("chatgpt-account")), + ) + .await + .unwrap(); + assert_eq!( + suppressed.transport_metrics.decision, + Some(CodexTransportDecision::Http2BreakerOpen) + ); + server.await.unwrap(); + pool.shutdown().await; +} + #[tokio::test] async fn cancelled_half_open_opening_should_allow_another_probe() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -1272,6 +1502,7 @@ async fn cancelled_half_open_opening_should_allow_another_probe() { let breaker_config = WebSocketOriginBreakerConfig::default(); let (opening_started_tx, opening_started_rx) = tokio::sync::oneshot::channel(); + let (opening_closed_tx, opening_closed_rx) = tokio::sync::oneshot::channel(); let server = tokio::spawn(async move { reject_websocket_openings(&listener, breaker_config.failure_threshold).await; @@ -1280,12 +1511,21 @@ async fn cancelled_half_open_opening_should_allow_another_probe() { assert!(request.starts_with("GET /codex/responses HTTP/1.1")); opening_started_tx.send(()).unwrap(); + let (mut http, _) = listener.accept().await.unwrap(); + assert!( + read_http_request(&mut http) + .await + .starts_with("POST /codex/responses") + ); + write_completed_sse_response(&mut http).await; + let mut byte = [0_u8; 1]; let read = timeout(Duration::from_secs(1), opening.read(&mut byte)) .await .expect("account eviction should close the half-open probe") .unwrap(); assert_eq!(read, 0); + opening_closed_tx.send(()).unwrap(); let (fresh_stream, _) = listener.accept().await.unwrap(); let mut websocket = accept_codex_test_websocket(fresh_stream).await; @@ -1318,23 +1558,27 @@ async fn cancelled_half_open_opening_should_allow_another_probe() { let attempt = tokio::spawn(async move { probe_backend .create_response( - &explicit_websocket_warmup_request("conversation-cancelled-half-open"), + &new_chain_request("conversation-cancelled-half-open"), request_context("req_cancelled_half_open", Some("chatgpt-account")), ) .await }); opening_started_rx.await.unwrap(); - pool.evict_account("chatgpt-account").await; - timeout(Duration::from_secs(1), attempt) + let fallback = attempt .await - .expect("cancelled request should finish") .unwrap() - .expect_err("cancelled half-open opening should fail the request"); + .expect("slow probe should fall back to HTTP"); + assert_eq!( + fallback.transport_metrics.decision, + Some(CodexTransportDecision::Http2WebSocketBudgetExhausted) + ); + pool.evict_account("chatgpt-account").await; + opening_closed_rx.await.unwrap(); let response = backend .create_response( - &explicit_websocket_warmup_request("conversation-next-half-open-probe"), + &new_chain_request("conversation-next-half-open-probe"), request_context("req_next_half_open_probe", Some("chatgpt-account")), ) .await @@ -1531,7 +1775,7 @@ async fn half_open_upstream_response_should_close_origin_breaker() { ); let response = backend .create_response( - &explicit_websocket_warmup_request("conversation-after-half-open-response"), + &new_chain_request("conversation-after-half-open-response"), request_context("req_after_half_open_response", Some("chatgpt-account")), ) .await diff --git a/backend/crates/providers/openai/tests/transport/websocket/breaker.rs b/backend/crates/providers/openai/tests/transport/websocket/breaker.rs new file mode 100644 index 000000000..74de83fc3 --- /dev/null +++ b/backend/crates/providers/openai/tests/transport/websocket/breaker.rs @@ -0,0 +1,116 @@ +//! 验证 WebSocket 熔断器的并发建连结果与跨轮次隔离 + +use std::time::Duration; + +use provider_openai::transport::websocket::{ + WebSocketOriginBreaker, WebSocketOriginBreakerConfig, WebSocketOriginBreakerDecision, + WebSocketOriginBreakerPermit, +}; + +const ORIGIN: &str = "https://example.test:443"; + +fn acquire(breaker: &WebSocketOriginBreaker) -> WebSocketOriginBreakerPermit { + let WebSocketOriginBreakerDecision::Allowed(permit) = breaker.try_acquire(ORIGIN) else { + panic!("expected an opening permit"); + }; + permit +} + +#[test] +fn closed_success_should_keep_counting_concurrent_failures() { + let config = WebSocketOriginBreakerConfig::default(); + let breaker = WebSocketOriginBreaker::with_config(config); + let success = acquire(&breaker); + let pending_failures = (0..config.failure_threshold) + .map(|_| acquire(&breaker)) + .collect::>(); + + success.succeed(); + for permit in pending_failures { + permit.fail(); + } + + assert!(matches!( + breaker.try_acquire(ORIGIN), + WebSocketOriginBreakerDecision::Open + )); +} + +#[test] +fn concurrent_closed_success_should_clear_previous_failures() { + let config = WebSocketOriginBreakerConfig::default(); + let breaker = WebSocketOriginBreaker::with_config(config); + let first_success = acquire(&breaker); + let pending_success = acquire(&breaker); + + first_success.succeed(); + for _ in 0..config.failure_threshold - 1 { + acquire(&breaker).fail(); + } + pending_success.succeed(); + for _ in 0..config.failure_threshold - 1 { + acquire(&breaker).fail(); + } + + acquire(&breaker).succeed(); +} + +#[tokio::test(start_paused = true)] +async fn stale_closed_permits_should_not_change_recovered_failure_history() { + for outcome in ["success", "failure", "drop"] { + let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig { + failure_threshold: 2, + failure_window: Duration::from_secs(30), + open_duration: Duration::from_secs(30), + }); + let stale = acquire(&breaker); + acquire(&breaker).fail(); + acquire(&breaker).fail(); + assert!(matches!( + breaker.try_acquire(ORIGIN), + WebSocketOriginBreakerDecision::Open + )); + + tokio::time::advance(Duration::from_secs(30)).await; + acquire(&breaker).succeed(); + acquire(&breaker).fail(); + match outcome { + "success" => stale.succeed(), + "failure" => stale.fail(), + _ => drop(stale), + } + + // 新代次的一次失败不能被旧成功清除，也不能被旧失败或 Drop 凑满阈值 + acquire(&breaker).fail(); + assert!(matches!( + breaker.try_acquire(ORIGIN), + WebSocketOriginBreakerDecision::Open + )); + } +} + +#[tokio::test(start_paused = true)] +async fn stale_closed_permits_should_not_finish_a_half_open_probe() { + for outcome in ["success", "failure", "drop"] { + let breaker = WebSocketOriginBreaker::with_config(WebSocketOriginBreakerConfig { + failure_threshold: 1, + failure_window: Duration::from_secs(30), + open_duration: Duration::from_secs(30), + }); + let stale = acquire(&breaker); + acquire(&breaker).fail(); + tokio::time::advance(Duration::from_secs(30)).await; + let probe = acquire(&breaker); + match outcome { + "success" => stale.succeed(), + "failure" => stale.fail(), + _ => drop(stale), + } + assert!(matches!( + breaker.try_acquire(ORIGIN), + WebSocketOriginBreakerDecision::HalfOpenBusy + )); + probe.succeed(); + acquire(&breaker).succeed(); + } +} diff --git a/backend/crates/providers/openai/tests/transport/websocket/mod.rs b/backend/crates/providers/openai/tests/transport/websocket/mod.rs index 835de8c84..83ceb1291 100644 --- a/backend/crates/providers/openai/tests/transport/websocket/mod.rs +++ b/backend/crates/providers/openai/tests/transport/websocket/mod.rs @@ -6,6 +6,7 @@ use provider_openai::transport::{ websocket::write_websocket_audit_artifact_for_dir, }; +mod breaker; mod diagnostics; mod fingerprint; mod metadata;