From 75f9d2a290e1f810565e30a437aafe83aa06a123 Mon Sep 17 00:00:00 2001 From: Zyy Date: Tue, 25 Aug 2026 11:22:07 +0800 Subject: [PATCH] fix(openai): remove websocket pool slot caps Co-authored-by: Codex --- backend/crates/providers/openai/src/config.rs | 13 +- .../src/transport/websocket/pool/mod.rs | 68 ++---- .../src/transport/websocket/pool/state.rs | 29 --- .../crates/providers/openai/tests/config.rs | 4 - .../openai/tests/transport/headers.rs | 6 +- .../openai/tests/transport/latency.rs | 10 +- .../providers/openai/tests/transport/mod.rs | 2 - .../openai/tests/transport/websocket.rs | 6 +- .../openai/tests/transport/websocket_pool.rs | 218 +++++------------- 9 files changed, 94 insertions(+), 262 deletions(-) diff --git a/backend/crates/providers/openai/src/config.rs b/backend/crates/providers/openai/src/config.rs index c92455008..2332353e0 100644 --- a/backend/crates/providers/openai/src/config.rs +++ b/backend/crates/providers/openai/src/config.rs @@ -155,8 +155,6 @@ impl CodexApiConfig { pub struct CodexWebSocketPoolSettings { pub enabled: bool, pub max_age_ms: u64, - pub max_per_account: usize, - pub max_total: usize, pub max_connecting: usize, pub initial_event_timeout_ms: u64, } @@ -166,8 +164,6 @@ impl Default for CodexWebSocketPoolSettings { Self { enabled: true, max_age_ms: 55 * 60 * 1000, - max_per_account: 8, - max_total: 64, max_connecting: 8, initial_event_timeout_ms: 20_000, } @@ -176,12 +172,7 @@ impl Default for CodexWebSocketPoolSettings { impl CodexWebSocketPoolSettings { fn validate(&self) -> Result<(), OpenAiConfigError> { - if self.max_age_ms == 0 - || self.max_per_account == 0 - || self.max_total == 0 - || self.max_connecting == 0 - || self.max_connecting > self.max_total - { + if self.max_age_ms == 0 || self.max_connecting == 0 { return Err(OpenAiConfigError::InvalidField("openai.ws_pool")); } Ok(()) @@ -191,8 +182,6 @@ impl CodexWebSocketPoolSettings { CodexWebSocketPoolConfig { enabled: self.enabled, max_age: Duration::from_millis(self.max_age_ms), - max_per_account: self.max_per_account, - max_total: self.max_total, max_connecting: self.max_connecting, initial_event_timeout: (self.initial_event_timeout_ms != 0) .then(|| Duration::from_millis(self.initial_event_timeout_ms)), diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs b/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs index 88217440e..261438425 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs @@ -10,13 +10,13 @@ use std::{ time::Duration, }; -use tokio::{sync::Semaphore, time::Instant}; +use tokio::sync::Semaphore; use tokio_util::{sync::CancellationToken, task::TaskTracker}; use uuid::Uuid; use self::state::{ - WebSocketPoolConnecting, WebSocketPoolSlot, WebSocketPoolState, account_slot_count, - close_pooled_connection, close_pooled_connections, take_lru_idle_connection, + WebSocketPoolConnecting, WebSocketPoolSlot, WebSocketPoolState, close_pooled_connection, + close_pooled_connections, }; use super::pump::PumpKeepalive; @@ -31,8 +31,6 @@ pub(crate) use self::{ }, }; -const DEFAULT_MAX_PER_ACCOUNT: usize = 8; -const DEFAULT_MAX_TOTAL: usize = 64; const DEFAULT_MAX_CONNECTING: usize = 16; const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55); const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25); @@ -64,10 +62,6 @@ pub struct CodexWebSocketPoolConfig { pub enabled: bool, /// 单个 socket 的最大生命周期。 pub max_age: Duration, - /// 单个账号允许占用的最大池 slot 数。 - pub max_per_account: usize, - /// 所有账号合计允许占用的最大池 slot 数。 - pub max_total: usize, /// 所有账号合计允许并发执行的 opening 数。 pub max_connecting: usize, /// 后台维护间隔；`None` 表示不启动后台任务。 @@ -87,8 +81,6 @@ impl Default for CodexWebSocketPoolConfig { Self { enabled: true, max_age: DEFAULT_MAX_AGE, - max_per_account: DEFAULT_MAX_PER_ACCOUNT, - max_total: DEFAULT_MAX_TOTAL, max_connecting: DEFAULT_MAX_CONNECTING, maintenance_interval: Some(DEFAULT_MAINTENANCE_INTERVAL), ping_interval: Some(DEFAULT_PING_INTERVAL), @@ -113,10 +105,9 @@ impl CodexWebSocketPoolConfig { } impl CodexWebSocketPool { - /// 构造连接池策略和状态。 - pub fn new(max_per_account: usize, max_age: Duration) -> Self { + /// 构造不限制累计 slot 数的连接池策略和状态。 + pub fn new(max_age: Duration) -> Self { Self::with_config(CodexWebSocketPoolConfig { - max_per_account: max_per_account.max(1), max_age, maintenance_interval: None, ping_interval: None, @@ -218,43 +209,25 @@ impl CodexWebSocketPool { let acquire = if required_response_id.is_some() { WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::ContinuationNotFound) - } else if self.config.max_per_account == 0 - || account_slot_count(&state.slots, key.account_id()) >= self.config.max_per_account - { - WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap) } else { let connect_permit = self.connect_semaphore.clone().try_acquire_owned().ok(); match connect_permit { Some(connect_permit) => { - let has_total_capacity = if self.config.max_total == 0 { - false - } else if state.slots.len() < self.config.max_total { - true - } else if let Some(connection) = take_lru_idle_connection(&mut state) { - connections_to_close.push(connection); - true - } else { - false - }; - if has_total_capacity { - let lease = WebSocketPoolConnectLease::reserve( - self.clone(), - key.clone(), - connect_permit, - ); - state.slots.insert( - key, - WebSocketPoolSlot::Connecting(WebSocketPoolConnecting { - id: lease.id, - started_at: lease.started_at, - outcome: lease.outcome.clone(), - cancellation: lease.cancellation.clone(), - }), - ); - WebSocketPoolAcquire::Connect(lease) - } else { - WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap) - } + let lease = WebSocketPoolConnectLease::reserve( + self.clone(), + key.clone(), + connect_permit, + ); + state.slots.insert( + key, + WebSocketPoolSlot::Connecting(WebSocketPoolConnecting { + id: lease.id, + started_at: lease.started_at, + outcome: lease.outcome.clone(), + cancellation: lease.cancellation.clone(), + }), + ); + WebSocketPoolAcquire::Connect(lease) } None => WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Cap), } @@ -292,7 +265,6 @@ impl CodexWebSocketPool { key.clone(), WebSocketPoolSlot::Idle { connection: Box::new(connection), - last_used_at: Instant::now(), }, ); } diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/state.rs b/backend/crates/providers/openai/src/transport/websocket/pool/state.rs index 87c03c7ee..cd4e9fba3 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/state.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/state.rs @@ -108,7 +108,6 @@ impl WebSocketContinuationState { pub(super) enum WebSocketPoolSlot { Idle { connection: Box, - last_used_at: Instant, }, Busy(WebSocketPoolReservation), Connecting(WebSocketPoolConnecting), @@ -138,34 +137,6 @@ pub(super) struct WebSocketPoolReservation { pub(super) latest_response_id: Option, } -pub(super) fn account_slot_count( - slots: &HashMap, - account_id: &str, -) -> usize { - slots - .keys() - .filter(|key| key.account_id() == account_id) - .count() -} - -pub(super) fn take_lru_idle_connection( - state: &mut WebSocketPoolState, -) -> Option { - let key = state - .slots - .iter() - .filter_map(|(key, slot)| match slot { - WebSocketPoolSlot::Idle { last_used_at, .. } => Some((key, *last_used_at)), - WebSocketPoolSlot::Busy(_) | WebSocketPoolSlot::Connecting(_) => None, - }) - .min_by_key(|(_, last_used_at)| *last_used_at) - .map(|(key, _)| key.clone())?; - let Some(WebSocketPoolSlot::Idle { connection, .. }) = state.slots.remove(&key) else { - return None; - }; - Some(*connection) -} - pub(super) async fn close_pooled_connection(connection: PooledWebSocketConnection) { connection.websocket.close().await; } diff --git a/backend/crates/providers/openai/tests/config.rs b/backend/crates/providers/openai/tests/config.rs index df9a10646..873d87953 100644 --- a/backend/crates/providers/openai/tests/config.rs +++ b/backend/crates/providers/openai/tests/config.rs @@ -101,8 +101,6 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() { config.api.base_url.as_str(), config.ws_pool.enabled, config.ws_pool.max_age_ms, - config.ws_pool.max_per_account, - config.ws_pool.max_total, config.ws_pool.max_connecting, config.ws_pool.initial_event_timeout_ms, config.quota.refresh_interval_minutes, @@ -115,8 +113,6 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() { true, 3_300_000, 8, - 64, - 8, 20_000, 15, true, diff --git a/backend/crates/providers/openai/tests/transport/headers.rs b/backend/crates/providers/openai/tests/transport/headers.rs index d147e7db6..6bf7c1864 100644 --- a/backend/crates/providers/openai/tests/transport/headers.rs +++ b/backend/crates/providers/openai/tests/transport/headers.rs @@ -265,7 +265,7 @@ async fn backend_websocket_should_forward_context_headers_and_preserve_payload_f format!("http://{address}"), test_wire_profile(), ) - .with_websocket_pool(Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)))); + .with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)))); let response = backend .create_response( @@ -480,7 +480,7 @@ async fn backend_websocket_should_drop_only_unrepresentable_opaque_header_values format!("http://{address}"), test_wire_profile(), ) - .with_websocket_pool(Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)))); + .with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)))); client .create_response( @@ -932,7 +932,7 @@ async fn websocket_should_keep_an_exact_chain_while_new_connections_adopt_the_la }); let profile = test_wire_profile(); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let client = CodexBackendClient::new( reqwest::Client::builder() .no_proxy() diff --git a/backend/crates/providers/openai/tests/transport/latency.rs b/backend/crates/providers/openai/tests/transport/latency.rs index 68b11e42f..cb9bd8191 100644 --- a/backend/crates/providers/openai/tests/transport/latency.rs +++ b/backend/crates/providers/openai/tests/transport/latency.rs @@ -300,7 +300,7 @@ async fn account_eviction_should_cancel_a_background_websocket_opening() { .await .unwrap(); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -354,7 +354,7 @@ async fn pool_shutdown_should_cancel_and_join_a_background_websocket_opening() { .unwrap(); assert_eq!(read, 0); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -733,7 +733,7 @@ async fn concurrent_multi_conversation_continuations_should_select_each_exact_so tokio::join!(first_continuation, second_continuation); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -1190,7 +1190,7 @@ async fn cancelled_half_open_opening_should_allow_another_probe() { .await .unwrap(); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -1325,7 +1325,7 @@ async fn fast_path_miss_and_late_failure_should_count_as_one_breaker_failure() { write_completed_sse_response(&mut second_http).await; } }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), diff --git a/backend/crates/providers/openai/tests/transport/mod.rs b/backend/crates/providers/openai/tests/transport/mod.rs index fc5caaf0e..08866502b 100644 --- a/backend/crates/providers/openai/tests/transport/mod.rs +++ b/backend/crates/providers/openai/tests/transport/mod.rs @@ -495,8 +495,6 @@ fn websocket_pool_config_for_tests( CodexWebSocketPoolConfig { enabled: true, max_age: Duration::from_mins(1), - max_per_account: 8, - max_total: 64, max_connecting: 16, maintenance_interval, ping_interval, diff --git a/backend/crates/providers/openai/tests/transport/websocket.rs b/backend/crates/providers/openai/tests/transport/websocket.rs index 0bb262cf5..c48cdeb53 100644 --- a/backend/crates/providers/openai/tests/transport/websocket.rs +++ b/backend/crates/providers/openai/tests/transport/websocket.rs @@ -126,7 +126,7 @@ async fn codex_backend_client_should_decode_permessage_deflate_context_takeover_ websocket.send(Message::Text(payload.into())).await.unwrap(); } }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -932,7 +932,7 @@ async fn codex_backend_client_should_not_reuse_websocket_rate_limit_headers() { .unwrap(); websocket.close(None).await.unwrap(); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -1914,7 +1914,7 @@ async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_ let mut request = codex_request("gpt-5.5", "be brief", Vec::new()); request.set_previous_response_id(Some("resp_previous".to_string())); request.previous_response_scope = Some(PreviousResponseScope::Persisted); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), diff --git a/backend/crates/providers/openai/tests/transport/websocket_pool.rs b/backend/crates/providers/openai/tests/transport/websocket_pool.rs index 41a56f71e..f27805213 100644 --- a/backend/crates/providers/openai/tests/transport/websocket_pool.rs +++ b/backend/crates/providers/openai/tests/transport/websocket_pool.rs @@ -35,7 +35,7 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and } websocket.close(None).await.unwrap(); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -97,7 +97,7 @@ async fn exact_continuation_reuses_owning_socket_after_connection_profile_update format!("http://{addr}"), profile.clone(), ) - .with_websocket_pool(Arc::new(CodexWebSocketPool::new(1, Duration::from_mins(1)))); + .with_websocket_pool(Arc::new(CodexWebSocketPool::new(Duration::from_mins(1)))); let first_request = pooled_websocket_request("wire-profile-continuation"); let first = backend .create_response( @@ -401,7 +401,7 @@ async fn codex_backend_client_stream_should_keep_reused_socket_after_structural_ } #[tokio::test] -async fn codex_backend_client_stream_should_use_http_when_pool_is_at_cap() { +async fn codex_backend_client_stream_should_use_http_when_connecting_limit_is_exhausted() { 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)); @@ -414,7 +414,7 @@ async fn codex_backend_client_stream_should_use_http_when_pool_is_at_cap() { write_completed_sse_response(&mut first_stream).await; }); let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { - max_per_account: 0, + max_connecting: 0, initial_event_timeout: Some(Duration::from_millis(30)), ..websocket_pool_config_for_tests(None, None, None) })); @@ -432,7 +432,7 @@ async fn codex_backend_client_stream_should_use_http_when_pool_is_at_cap() { request_context("req_structural_bypass", Some("chatgpt-account")), ) .await - .expect("pool cap should select HTTP before sending payload"); + .expect("exhausted connection-opening limit should select HTTP before sending payload"); assert_eq!(response.transport, CodexBackendTransport::HttpSse); assert_eq!( response.transport_metrics.decision, @@ -472,7 +472,7 @@ async fn codex_backend_client_should_not_reuse_pooled_websocket_across_local_acc websocket.close(None).await.unwrap(); } }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -841,140 +841,32 @@ async fn websocket_pool_shutdown_should_cancel_and_join_active_stream() { } #[tokio::test] -async fn websocket_pool_should_use_http_for_new_keys_after_account_cap() { +async fn websocket_pool_should_keep_sixty_fifth_conversation_on_same_account_on_websocket() { + const REQUEST_COUNT: usize = 65; + 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 server = tokio::spawn(async move { - let (first_stream, _) = listener.accept().await.unwrap(); - accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); - let mut first_websocket = accept_codex_test_websocket(first_stream).await; - let _first_message = first_websocket.next().await.unwrap().unwrap(); - first_websocket - .send(Message::Text( - completed_websocket_response("resp_cap_first", 2, 1).into(), - )) - .await - .unwrap(); - - for _ in 0..2 { - let (mut stream, _) = listener.accept().await.unwrap(); + let mut websockets = Vec::with_capacity(REQUEST_COUNT); + for index in 0..REQUEST_COUNT { + let (stream, _) = listener.accept().await.unwrap(); accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); - let request = read_http_request(&mut stream).await; - assert!(request.starts_with("POST /codex/responses HTTP/1.1")); - write_completed_sse_response(&mut stream).await; + let mut websocket = accept_codex_test_websocket(stream).await; + let _message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + completed_websocket_response(&format!("resp_same_account_{index}"), 2, 1) + .into(), + )) + .await + .unwrap(); + websockets.push(websocket); } - drop(first_websocket); - }); - let pool = Arc::new(CodexWebSocketPool::new(1, 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 first_request = pooled_websocket_request("conversation-cap-one"); - let second_request = pooled_websocket_request("conversation-cap-two"); - - let first = backend - .create_response( - &first_request, - request_context("req_cap_first", Some("chatgpt-account")), - ) - .await - .expect("first capped websocket response should succeed"); - let second = backend - .create_response( - &second_request, - request_context("req_cap_second", Some("chatgpt-account")), - ) - .await - .expect("new key over account cap should use HTTP"); - let third = backend - .create_response( - &second_request, - request_context("req_cap_third", Some("chatgpt-account")), - ) - .await - .expect("capped key should keep using HTTP instead of opening one-shot sockets"); - server.await.unwrap(); - - assert!(first.body.contains("resp_cap_first")); - assert!(second.body.contains("response.completed")); - assert!(third.body.contains("response.completed")); - assert_eq!(first.websocket_pool_decision.unwrap().kind(), "new"); - assert_eq!( - second.transport_metrics.decision, - Some(CodexTransportDecision::Http2PoolUnavailable) - ); - assert_eq!( - third.transport_metrics.decision, - Some(CodexTransportDecision::Http2PoolUnavailable) - ); - assert_eq!(accepted_connections.load(Ordering::SeqCst), 3); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn websocket_pool_should_bound_concurrent_openings_across_accounts() { - let (websockets, http) = concurrent_pool_transport_counts(8, 2).await; - - assert_eq!((websockets, http), (2, 6)); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn websocket_pool_should_bound_total_slots_across_accounts() { - let (websockets, http) = concurrent_pool_transport_counts(3, 8).await; - - assert_eq!((websockets, http), (3, 5)); -} - -#[tokio::test] -async fn websocket_pool_should_evict_lru_idle_connection_at_total_cap() { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { - let (first_stream, _) = listener.accept().await.unwrap(); - let mut first = accept_codex_test_websocket(first_stream).await; - let _ = first.next().await.unwrap().unwrap(); - first - .send(Message::Text( - completed_websocket_response("resp_global_lru_first", 2, 1).into(), - )) - .await - .unwrap(); - - let (second_stream, _) = listener.accept().await.unwrap(); - let mut second = accept_codex_test_websocket(second_stream).await; - let _ = second.next().await.unwrap().unwrap(); - second - .send(Message::Text( - completed_websocket_response("resp_global_lru_second", 2, 1).into(), - )) - .await - .unwrap(); - - let close = timeout(Duration::from_secs(2), first.next()) - .await - .expect("global LRU eviction should close the oldest idle socket") - .expect("oldest idle socket should emit a close frame") - .expect("oldest idle socket close should be valid"); - std::assert_matches!(close, Message::Close(_)); - - let (third_stream, _) = listener.accept().await.unwrap(); - let mut third = accept_codex_test_websocket(third_stream).await; - let _ = third.next().await.unwrap().unwrap(); - third - .send(Message::Text( - completed_websocket_response("resp_global_lru_third", 2, 1).into(), - )) - .await - .unwrap(); }); let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { - max_per_account: 8, - max_total: 2, - max_connecting: 2, + max_connecting: 1, ..websocket_pool_config_for_tests(None, None, None) })); let backend = CodexBackendClient::new( @@ -983,36 +875,52 @@ async fn websocket_pool_should_evict_lru_idle_connection_at_total_cap() { test_wire_profile(), ) .with_websocket_pool(Arc::clone(&pool)); - - let mut responses = Vec::new(); - for index in 0..3 { - let request = pooled_websocket_request(&format!("conversation-global-lru-{index}")); - let account_id = format!("account-global-lru-{index}"); - let request_id = format!("req_global_lru_{index}"); + let mut responses = Vec::with_capacity(REQUEST_COUNT); + for index in 0..REQUEST_COUNT { + let request = pooled_websocket_request(&format!("conversation-same-account-{index}")); responses.push( backend - .create_response(&request, request_context(&request_id, Some(&account_id))) + .create_response( + &request, + request_context( + &format!("req_same_account_{index}"), + Some("chatgpt-account"), + ), + ) .await - .expect("global LRU request should keep using websocket"), + .expect("same-account conversation should use websocket"), ); } server.await.unwrap(); - assert!(responses[0].body.contains("resp_global_lru_first")); - assert!(responses[1].body.contains("resp_global_lru_second")); - assert!(responses[2].body.contains("resp_global_lru_third")); - assert!( - responses - .iter() - .all(|response| response.transport == CodexBackendTransport::WebSocket) + assert_eq!( + ( + responses + .iter() + .filter(|response| response.transport == CodexBackendTransport::WebSocket) + .count(), + accepted_connections.load(Ordering::SeqCst), + ), + (REQUEST_COUNT, REQUEST_COUNT) ); pool.shutdown().await; } -async fn concurrent_pool_transport_counts( - max_total: usize, - max_connecting: usize, -) -> (usize, usize) { +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn websocket_pool_should_bound_concurrent_openings_across_accounts() { + let (websockets, http) = concurrent_pool_transport_counts(2).await; + + assert_eq!((websockets, http), (2, 6)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn websocket_pool_should_allow_all_openings_when_connect_permits_are_available() { + let (websockets, http) = concurrent_pool_transport_counts(8).await; + + assert_eq!((websockets, http), (8, 0)); +} + +async fn concurrent_pool_transport_counts(max_connecting: usize) -> (usize, usize) { const REQUEST_COUNT: usize = 8; let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -1051,7 +959,7 @@ async fn concurrent_pool_transport_counts( websocket .send(Message::Text( completed_websocket_response( - &format!("resp_global_cap_{max_total}_{max_connecting}_{index}"), + &format!("resp_connecting_cap_{max_connecting}_{index}"), 2, 1, ) @@ -1063,8 +971,6 @@ async fn concurrent_pool_transport_counts( (websocket_count, http_count) }); let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { - max_per_account: REQUEST_COUNT, - max_total, max_connecting, ..websocket_pool_config_for_tests(None, None, None) })); @@ -1096,7 +1002,7 @@ async fn concurrent_pool_transport_counts( let response = request .await .unwrap() - .expect("globally capped request should complete"); + .expect("concurrent request should complete"); if response.transport == CodexBackendTransport::WebSocket { websocket_responses += 1; } @@ -1910,7 +1816,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_error_termin .unwrap(); second_websocket.close(None).await.unwrap(); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -1987,7 +1893,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_unknown_resp .unwrap(); second_websocket.close(None).await.unwrap(); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"),