From 83f5d71853f90ea4006a38bee6c916e56844a557 Mon Sep 17 00:00:00 2001 From: Zyy <0x@ainz.cc> Date: Sat, 17 Oct 2026 14:56:30 +0800 Subject: [PATCH] fix(openai): bound WebSocket connections per outbound proxy (#458) Co-authored-by: Zyy Co-authored-by: Codex --- .../crates/gateway-admin/src/model/proxies.rs | 2 + .../crates/gateway-api/src/admin/proxies.rs | 8 +- .../crates/gateway-api/tests/admin/proxies.rs | 58 +- .../crates/gateway-core/src/account/proxy.rs | 28 +- .../src/postgres/account_groups.rs | 2 +- .../provider_accounts/admin_queries.rs | 2 +- .../postgres/provider_accounts/repository.rs | 8 +- .../src/postgres/provider_accounts/rows.rs | 18 +- .../gateway-store/src/postgres/proxies.rs | 45 +- .../tests/postgres/plugins/sources.rs | 1 + .../tests/postgres/provider_accounts/mod.rs | 1 + .../gateway-store/tests/postgres/proxies.rs | 96 ++- .../openai/src/transport/client_sse.rs | 3 +- .../src/transport/websocket/coordinator.rs | 39 +- .../src/transport/websocket/handshake.rs | 6 +- .../openai/src/transport/websocket/model.rs | 11 +- .../src/transport/websocket/pool/capacity.rs | 142 ++++ .../src/transport/websocket/pool/lease.rs | 7 + .../src/transport/websocket/pool/mod.rs | 238 ++++--- .../src/transport/websocket/pool/state.rs | 15 +- .../transport/websocket/pool/supervisor.rs | 4 + .../openai/src/transport/websocket/pump.rs | 34 +- .../providers/openai/tests/transport/mod.rs | 1 + .../openai/tests/transport/proxy_capacity.rs | 625 ++++++++++++++++++ backend/migrations/.frozen-sha256 | 1 + .../0033_proxy_websocket_capacity.sql | 3 + deploy/README.md | 12 + docs/api.md | 10 +- frontend/src/api/modules/proxies.ts | 5 +- .../proxies/components/ProxyFormModal.vue | 4 + .../views/proxies/composables/useProxies.ts | 9 + frontend/src/views/proxies/index.vue | 1 + 32 files changed, 1277 insertions(+), 162 deletions(-) create mode 100644 backend/crates/providers/openai/src/transport/websocket/pool/capacity.rs create mode 100644 backend/crates/providers/openai/tests/transport/proxy_capacity.rs create mode 100644 backend/migrations/0033_proxy_websocket_capacity.sql diff --git a/backend/crates/gateway-admin/src/model/proxies.rs b/backend/crates/gateway-admin/src/model/proxies.rs index 3082aac3c..1f2f299fb 100644 --- a/backend/crates/gateway-admin/src/model/proxies.rs +++ b/backend/crates/gateway-admin/src/model/proxies.rs @@ -104,6 +104,8 @@ pub struct NewProxy { #[derive(Debug, Clone)] pub struct UpdateProxy { + /// 省略保留配置，零值取消限制 + pub max_websocket_connections: Option, pub enabled: Option, pub auto_location: Option, pub test: Option, diff --git a/backend/crates/gateway-api/src/admin/proxies.rs b/backend/crates/gateway-api/src/admin/proxies.rs index 08e2e1ad0..0766907ae 100644 --- a/backend/crates/gateway-api/src/admin/proxies.rs +++ b/backend/crates/gateway-api/src/admin/proxies.rs @@ -50,6 +50,8 @@ struct RemoveAccountRequest { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] struct CreateRequest { + #[serde(default)] + max_websocket_connections: u32, #[serde(default)] auto_location: bool, location: Option, @@ -60,6 +62,7 @@ struct CreateRequest { #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] struct UpdateRequest { + max_websocket_connections: Option, enabled: Option, auto_location: Option, #[serde(default, deserialize_with = "deserialize_location_update")] @@ -146,6 +149,7 @@ struct ProxyAccountView { #[derive(Serialize)] #[serde(rename_all = "camelCase")] struct ProxyView { + max_websocket_connections: u32, enabled: bool, auto_location: bool, detected_location: Option, @@ -170,6 +174,7 @@ impl From<(ProxyRecord, crate::time::TimePresenter)> for ProxyView { let endpoint = record.proxy.endpoint(); Self { auto_location: record.auto_location, + max_websocket_connections: record.proxy.max_websocket_connections(), enabled: record.enabled, detected_location: record.detected_location, location: record.location, @@ -411,7 +416,7 @@ where test: None, location: request.location, name: request.name, - proxy, + proxy: proxy.with_max_websocket_connections(request.max_websocket_connections), }) }) .collect(); @@ -472,6 +477,7 @@ where .proxies() .update( UpdateProxy { + max_websocket_connections: request.max_websocket_connections, enabled: request.enabled, auto_location: request.auto_location, test: None, diff --git a/backend/crates/gateway-api/tests/admin/proxies.rs b/backend/crates/gateway-api/tests/admin/proxies.rs index fb8648684..a686ae2d8 100644 --- a/backend/crates/gateway-api/tests/admin/proxies.rs +++ b/backend/crates/gateway-api/tests/admin/proxies.rs @@ -213,7 +213,10 @@ impl ProxyStore for MemoryProxies { _: &MutationContext, ) -> AdminStoreResult { let mut stored = self.0.lock().unwrap(); - if stored.iter().any(|record| record.proxy == command.proxy) { + if stored + .iter() + .any(|record| record.proxy.expose_url() == command.proxy.expose_url()) + { return Err(AdminStoreError::new( AdminStoreErrorKind::Conflict, "proxy", @@ -274,16 +277,20 @@ impl ProxyStore for MemoryProxies { record.location = location; } if let Some(proxy) = command.proxy { - if record.proxy != proxy { + if record.proxy.expose_url() != proxy.expose_url() { record.detected_location = None; } - record.proxy = proxy; + record.proxy = + proxy.with_max_websocket_connections(record.proxy.max_websocket_connections()); } if let Some(result) = command.test { record.detected_location = record.detected_location_after_test(&result); record.last_test = Some(result); record.last_test_at = Some(Utc::now()); } + if let Some(limit) = command.max_websocket_connections { + record.proxy = record.proxy.clone().with_max_websocket_connections(limit); + } record.revision = Revision::new(record.revision.get() + 1).unwrap(); Ok(ProxyMutation { config_revision: record.revision, @@ -933,3 +940,48 @@ async fn proxy_arrays_reject_invalid_shape_size_and_duplicate_test_ids_before_ex let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await; assert_eq!(listed["data"]["page"]["total"], 0); } + +#[tokio::test] +async fn proxy_websocket_capacity_round_trips_and_preserves_omitted_updates() { + let fixture = AdminTestFixture::new().await; + fixture.auth.insert_session("valid-session"); + let (status, created) = request(&fixture, "/api/admin/proxies/create", Some(json!([ + {"name":"Limited", "proxyUrl":"http://proxy.example:8080", "maxWebsocketConnections":24}, + {"name":"Unlimited", "proxyUrl":"http://other.example:8080"} + ])), true).await; + assert_eq!(status, StatusCode::OK); + assert_eq!( + created["data"][0]["data"]["record"]["maxWebsocketConnections"], + 24 + ); + assert_eq!( + created["data"][1]["data"]["record"]["maxWebsocketConnections"], + 0 + ); + for (revision, change, expected) in [(1, None, 24), (2, Some(12), 12), (3, Some(0), 0)] { + let mut body = json!({"id":"proxy_test", "revision":revision, "name":"Limited"}); + if let Some(limit) = change { + body["maxWebsocketConnections"] = json!(limit); + } + let (status, updated) = + request(&fixture, "/api/admin/proxies/update", Some(body), true).await; + assert_eq!(status, StatusCode::OK); + assert_eq!( + updated["data"]["record"]["maxWebsocketConnections"], + expected + ); + let (_, listed) = request(&fixture, "/api/admin/proxies", None, true).await; + assert_eq!( + listed["data"]["items"][0]["maxWebsocketConnections"], + expected + ); + assert_eq!(listed["data"]["items"][1]["maxWebsocketConnections"], 0); + } + for invalid in [json!(-1), json!(1.5), json!(4294967296u64)] { + let (status, result) = request(&fixture, "/api/admin/proxies/create", Some(json!([ + {"name":"Invalid", "proxyUrl":"http://invalid.example:8080", "maxWebsocketConnections":invalid} + ])), true).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(result["data"][0]["httpStatus"], 400); + } +} diff --git a/backend/crates/gateway-core/src/account/proxy.rs b/backend/crates/gateway-core/src/account/proxy.rs index 1011fc12c..a6fa94e43 100644 --- a/backend/crates/gateway-core/src/account/proxy.rs +++ b/backend/crates/gateway-core/src/account/proxy.rs @@ -1,10 +1,13 @@ -//! 账号出站代理端点与认证值，诊断及普通管理投影不暴露凭据 +//! 账号出站代理连接与容量配置，诊断及普通管理投影不暴露凭据 use std::fmt; use url::Url; #[derive(Clone, PartialEq, Eq, Hash)] -pub struct OutboundProxy(Url); +pub struct OutboundProxy { + url: Url, + max_websocket_connections: u32, +} #[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] #[error("invalid outbound proxy; expected http, https, socks5 or socks5h URL with a host and port")] @@ -25,17 +28,32 @@ impl OutboundProxy { { return Err(InvalidOutboundProxy); } - Ok(Self(url)) + Ok(Self { + url, + max_websocket_connections: 0, + }) + } + + /// 此出口的上游 WebSocket 容量，零值表示不限制 + #[must_use] + pub fn with_max_websocket_connections(mut self, limit: u32) -> Self { + self.max_websocket_connections = limit; + self + } + + #[must_use] + pub const fn max_websocket_connections(&self) -> u32 { + self.max_websocket_connections } #[must_use] pub fn expose_url(&self) -> &str { - self.0.as_str() + self.url.as_str() } #[must_use] pub fn endpoint(&self) -> String { - let mut url = self.0.clone(); + let mut url = self.url.clone(); let _ = url.set_username(""); let _ = url.set_password(None); url.to_string() diff --git a/backend/crates/gateway-store/src/postgres/account_groups.rs b/backend/crates/gateway-store/src/postgres/account_groups.rs index 6e4442d84..2e4d3015d 100644 --- a/backend/crates/gateway-store/src/postgres/account_groups.rs +++ b/backend/crates/gateway-store/src/postgres/account_groups.rs @@ -226,7 +226,7 @@ impl AccountGroupStore for PgAccountGroupRepository { .map(|group_id| group_id.as_str().to_owned()) .collect::>(); let rows = sqlx::query( - "select coalesce(proxy.enabled, true) as outbound_proxy_enabled, proxy.auto_location, proxy.detected_location_json, proxy.location_country, proxy.location_region, proxy.location_city, proxy.location_timezone, membership.account_group_id, + "select coalesce(proxy.enabled, true) as outbound_proxy_enabled, coalesce(proxy.max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, proxy.auto_location, proxy.detected_location_json, proxy.location_country, proxy.location_region, proxy.location_city, proxy.location_timezone, membership.account_group_id, account.id, account.provider_kind, account.name, account.notes, account.email, account.upstream_user_id, account.upstream_account_id, account.plan_type, account.authentication_kind, account.credential_revision, account.outbound_proxy_url, diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/admin_queries.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/admin_queries.rs index efafa92ea..b618d2d54 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/admin_queries.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/admin_queries.rs @@ -106,7 +106,7 @@ pub(crate) async fn load_admin_account_page( from runtime_settings where id = 1 ) - select coalesce(p.enabled, true) as outbound_proxy_enabled, p.auto_location, p.detected_location_json, p.location_country, p.location_region, p.location_city, p.location_timezone, a.outbound_proxy_url, a.id, a.provider_kind, a.name, a.notes, a.email, a.upstream_user_id, + select coalesce(p.enabled, true) as outbound_proxy_enabled, coalesce(p.max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, p.auto_location, p.detected_location_json, p.location_country, p.location_region, p.location_city, p.location_timezone, a.outbound_proxy_url, a.id, a.provider_kind, a.name, a.notes, a.email, a.upstream_user_id, a.upstream_account_id, a.plan_type, a.authentication_kind, a.credential_revision, a.has_refresh_token, a.access_token_expires_at, a.next_refresh_at, a.enabled, a.concurrency_limit, a.weight, a.model_access_json, diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs index 6e741b9fd..a99000d09 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs @@ -122,14 +122,14 @@ impl ProviderAccountRepository for PgProviderAccountRepository { return Err(invalid("plugin account page limit must be positive")); } let rows = sqlx::query( - "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, + "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, coalesce(max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, upstream_account_id, plan_type, authentication_kind, credential_revision, has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state, credential_observed_at, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at, quota_observed_at, last_error_reason, last_error_message, created_at, updated_at from provider_accounts - left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location + left join (select id as location_proxy_id, enabled as proxy_enabled, max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location on outbound_proxy_id = location_proxy_id where ($1::text is null or provider_kind = $1) and ($2::text is null or id > $2) @@ -151,14 +151,14 @@ impl ProviderAccountRepository for PgProviderAccountRepository { include_disabled: bool, ) -> StoreResult> { let rows = sqlx::query( - "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, + "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, coalesce(max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, upstream_account_id, plan_type, authentication_kind, credential_revision, has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state, credential_observed_at, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at, quota_observed_at, last_error_reason, last_error_message, created_at, updated_at from provider_accounts - left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location + left join (select id as location_proxy_id, enabled as proxy_enabled, max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location on outbound_proxy_id = location_proxy_id where ($1::text is null or provider_kind = $1) and ($2 or enabled) order by provider_kind, name, id", diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs index 1bc2a1cb4..792175226 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs @@ -324,37 +324,37 @@ impl ProviderAccountStateUpdate { } } -pub(crate) const ACCOUNT_SELECT: &str = "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, +pub(crate) const ACCOUNT_SELECT: &str = "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, coalesce(max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision, has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state, provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at, last_error_reason, last_error_message, credential_observed_at, quota_observed_at, created_at, updated_at from provider_accounts - left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location + left join (select id as location_proxy_id, enabled as proxy_enabled, max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location on outbound_proxy_id = location_proxy_id where id = $1"; -pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, +pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, coalesce(max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision, has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state, provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at, last_error_reason, last_error_message, credential_observed_at, quota_observed_at, created_at, updated_at from provider_accounts - left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location + left join (select id as location_proxy_id, enabled as proxy_enabled, max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location on outbound_proxy_id = location_proxy_id where id = any($1::text[]) and provider_kind = $2 order by id"; -pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, +pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select coalesce(proxy_enabled, true) as outbound_proxy_enabled, coalesce(max_websocket_connections, 0) as outbound_proxy_max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone, outbound_proxy_url, id, provider_kind, name, notes, email, upstream_user_id, upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision, has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, model_access_json, credential_state, provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at, last_error_reason, last_error_message, credential_observed_at, quota_observed_at, created_at, updated_at from provider_accounts - left join (select id as location_proxy_id, enabled as proxy_enabled, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location + left join (select id as location_proxy_id, enabled as proxy_enabled, max_websocket_connections, auto_location, detected_location_json, location_country, location_region, location_city, location_timezone from outbound_proxies) proxy_location on outbound_proxy_id = location_proxy_id where provider_kind = $1 and has_refresh_token @@ -469,7 +469,13 @@ pub(crate) fn account_summary_from_row( request_location: super::super::proxies::effective_location_from_row(&row)?, outbound_proxy: get::>(&row, "outbound_proxy_url")? .map(|url| { + let limit = u32::try_from(get::( + &row, + "outbound_proxy_max_websocket_connections", + )?) + .map_err(|source| invalid("invalid proxy websocket limit").with_source(source))?; gateway_core::account::OutboundProxy::parse(&url) + .map(|proxy| proxy.with_max_websocket_connections(limit)) .map_err(|source| invalid("invalid outbound proxy").with_source(source)) }) .transpose()?, diff --git a/backend/crates/gateway-store/src/postgres/proxies.rs b/backend/crates/gateway-store/src/postgres/proxies.rs index 0ae6d3d8c..eb9c6910c 100644 --- a/backend/crates/gateway-store/src/postgres/proxies.rs +++ b/backend/crates/gateway-store/src/postgres/proxies.rs @@ -131,7 +131,14 @@ fn record(row: PgRow) -> StoreResult { &row.try_get::("proxy_url") .map_err(|source| invalid().with_source(source))?, ) - .map_err(|source| invalid().with_source(source))?, + .map_err(|source| invalid().with_source(source))? + .with_max_websocket_connections( + u32::try_from( + row.try_get::("max_websocket_connections") + .map_err(|source| invalid().with_source(source))?, + ) + .map_err(|source| invalid().with_source(source))?, + ), revision: AdminRevision::new( u64::try_from( row.try_get::("revision") @@ -348,8 +355,8 @@ pub(crate) async fn resolve_proxy_selection( Ok((Some(id), Some(proxy.clone()))) } AccountProxySelection::Saved(id) => { - let (value, enabled): (String, bool) = sqlx::query_as( - "select proxy_url, enabled from outbound_proxies where id = $1 for share", + let (value, enabled, limit): (String, bool, i64) = sqlx::query_as( + "select proxy_url, enabled, max_websocket_connections from outbound_proxies where id = $1 for share", ) .bind(id) .fetch_optional(&mut **transaction) @@ -365,7 +372,13 @@ pub(crate) async fn resolve_proxy_selection( } Ok(( Some(id.clone()), - Some(OutboundProxy::parse(&value).map_err(|source| invalid().with_source(source))?), + Some( + OutboundProxy::parse(&value) + .map_err(|source| invalid().with_source(source))? + .with_max_websocket_connections( + u32::try_from(limit).map_err(|source| invalid().with_source(source))?, + ), + ), )) } } @@ -674,9 +687,10 @@ impl ProxyStore for PgProxyRepository { save_location(&mut transaction, &id, command.location.as_ref()) .await .map_err(store_error)?; - sqlx::query("update outbound_proxies set auto_location = $2 where id = $1") + sqlx::query("update outbound_proxies set auto_location = $2, max_websocket_connections = $3 where id = $1") .bind(&id) .bind(command.auto_location) + .bind(i64::from(command.proxy.max_websocket_connections())) .execute(&mut *transaction) .await .map_err(|source| store_error(unavailable(source)))?; @@ -693,7 +707,13 @@ impl ProxyStore for PgProxyRepository { context, MutationAuditOperation::OutboundProxyCreate, &id, - &["name", "proxy_url", "location", "auto_location"], + &[ + "name", + "proxy_url", + "location", + "auto_location", + "max_websocket_connections", + ], revision, ) .await @@ -742,7 +762,7 @@ impl ProxyStore for PgProxyRepository { } let changed = sqlx::query( "update outbound_proxies set name = $3, proxy_url = coalesce($4, proxy_url), revision = revision + 1, updated_at = now(), - auto_location = coalesce($5, auto_location), enabled = coalesce($6, enabled), + auto_location = coalesce($5, auto_location), enabled = coalesce($6, enabled), max_websocket_connections = coalesce($7, max_websocket_connections), detected_location_json = case when ($4 is not null and $4 <> proxy_url) or ($5 is not null and $5 <> auto_location) then null else detected_location_json end, last_location_detection_json = case when ($4 is not null and $4 <> proxy_url) or ($5 is not null and $5 <> auto_location) then '{\"status\":\"notRequested\"}'::jsonb else last_location_detection_json end, last_test_at = case when $4 is not null and $4 <> proxy_url then null else last_test_at end, @@ -755,7 +775,7 @@ impl ProxyStore for PgProxyRepository { where id = $1 and revision = $2") .bind(&command.id).bind(i64::try_from(command.revision.get()).map_err(|_| store_error(invalid()))?) .bind(&command.name).bind(command.proxy.as_ref().map(OutboundProxy::expose_url)).bind(command.auto_location) - .bind(command.enabled) + .bind(command.enabled).bind(command.max_websocket_connections.map(i64::from)) .execute(&mut *transaction).await.map_err(|source| store_error(unavailable(source)))?; if changed.rows_affected() != 1 { return Err(store_error(conflict(&command.id))); @@ -780,7 +800,14 @@ impl ProxyStore for PgProxyRepository { context, MutationAuditOperation::OutboundProxyUpdate, &command.id, - &["name", "proxy_url", "location", "auto_location", "enabled"], + &[ + "name", + "proxy_url", + "location", + "auto_location", + "enabled", + "max_websocket_connections", + ], revision, ) .await diff --git a/backend/crates/gateway-store/tests/postgres/plugins/sources.rs b/backend/crates/gateway-store/tests/postgres/plugins/sources.rs index 64daf52f7..f8bfaf9e6 100644 --- a/backend/crates/gateway-store/tests/postgres/plugins/sources.rs +++ b/backend/crates/gateway-store/tests/postgres/plugins/sources.rs @@ -291,6 +291,7 @@ async fn source_proxy_is_revision_fenced_and_referenced_until_the_artifact_is_de let updated = proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: None, auto_location: None, test: None, diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs index 15b369262..2c5b29690 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs @@ -2083,6 +2083,7 @@ async fn authorization_import_rejects_a_saved_proxy_changed_during_oauth() { let edited = proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: None, auto_location: None, test: None, diff --git a/backend/crates/gateway-store/tests/postgres/proxies.rs b/backend/crates/gateway-store/tests/postgres/proxies.rs index 75ce36449..c4e748ed7 100644 --- a/backend/crates/gateway-store/tests/postgres/proxies.rs +++ b/backend/crates/gateway-store/tests/postgres/proxies.rs @@ -68,6 +68,7 @@ async fn proxy_location_is_shared_preserved_cleared_and_removed_with_binding() { saved = proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: None, auto_location: None, test: None, @@ -88,6 +89,7 @@ async fn proxy_location_is_shared_preserved_cleared_and_removed_with_binding() { saved = proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: None, auto_location: None, test: None, @@ -117,6 +119,7 @@ async fn proxy_location_is_shared_preserved_cleared_and_removed_with_binding() { proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: None, auto_location: None, test: None, @@ -135,6 +138,7 @@ async fn proxy_location_is_shared_preserved_cleared_and_removed_with_binding() { proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: None, auto_location: None, test: None, @@ -234,6 +238,7 @@ async fn disabled_proxy_projects_direct_while_preserving_bindings_and_credential saved = proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: Some(enabled), auto_location: None, test: None, @@ -305,6 +310,7 @@ async fn disabled_proxy_projects_direct_while_preserving_bindings_and_credential proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: Some(!enabled), auto_location: None, test: None, @@ -774,6 +780,7 @@ async fn stale_proxy_test_releases_lock_before_next_update() { let update = store .update( UpdateProxy { + max_websocket_connections: None, enabled: None, id: saved.id.clone(), revision: saved.revision, @@ -811,7 +818,9 @@ async fn import_reservation_blocks_proxy_mutations_until_rotated_credentials_are test: None, location: None, name: "导入出口".to_owned(), - proxy: OutboundProxy::parse("http://127.0.0.1:8080").unwrap(), + proxy: OutboundProxy::parse("http://127.0.0.1:8080") + .unwrap() + .with_max_websocket_connections(24), }, &context, ) @@ -834,6 +843,7 @@ async fn import_reservation_blocks_proxy_mutations_until_rotated_credentials_are .unwrap(); let reservation = store.reserve_import(&saved.id).await.unwrap(); let replacement = UpdateProxy { + max_websocket_connections: None, enabled: None, auto_location: None, test: None, @@ -1015,6 +1025,7 @@ async fn managed_proxies_persist_bind_update_all_accounts_and_protect_stale_test let renamed = store .update( UpdateProxy { + max_websocket_connections: None, enabled: None, auto_location: None, test: None, @@ -1036,6 +1047,7 @@ async fn managed_proxies_persist_bind_update_all_accounts_and_protect_stale_test let edited = store .update( UpdateProxy { + max_websocket_connections: None, enabled: None, auto_location: None, test: None, @@ -1358,6 +1370,7 @@ async fn automatic_proxy_location_projects_to_accounts_and_guards_refresh_result saved = proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: None, id: saved.id.clone(), revision: saved.revision, @@ -1377,6 +1390,7 @@ async fn automatic_proxy_location_projects_to_accounts_and_guards_refresh_result saved = proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: None, id: saved.id.clone(), revision: saved.revision, @@ -1404,6 +1418,7 @@ async fn automatic_proxy_location_projects_to_accounts_and_guards_refresh_result saved = proxies .update( UpdateProxy { + max_websocket_connections: None, enabled: None, id: saved.id.clone(), revision: saved.revision, @@ -1430,3 +1445,82 @@ async fn automatic_proxy_location_projects_to_accounts_and_guards_refresh_result ); database.close().await; } + +#[tokio::test] +async fn proxy_websocket_capacity_is_shared_by_bound_accounts_and_updates_without_rebinding() { + use gateway_core::account::ProviderAccountStore; + let Some(database) = TestDatabase::create("proxy_ws_capacity").await else { + return; + }; + let proxies = PgProxyRepository::new(database.pool.clone()); + let accounts = PgProviderAccountRepository::new(database.pool.clone()); + let mut saved = proxies + .create( + NewProxy { + auto_location: false, + test: None, + location: None, + name: "Limited".to_owned(), + proxy: OutboundProxy::parse("http://127.0.0.1:18080") + .unwrap() + .with_max_websocket_connections(24), + }, + &context(), + ) + .await + .unwrap() + .record; + assert_eq!(saved.proxy.max_websocket_connections(), 24); + for id in ["acct_proxy_budget_a", "acct_proxy_budget_b"] { + let mut input = account(id, id); + input.outbound_proxy = Some(saved.proxy.clone()); + accounts.insert_provider_account(input).await.unwrap(); + } + let mut direct = account("acct_other_proxy", "Other"); + direct.outbound_proxy = Some(OutboundProxy::parse("http://127.0.0.1:18081").unwrap()); + accounts.insert_provider_account(direct).await.unwrap(); + for (change, expected) in [(None, 24), (Some(8), 8), (Some(0), 0)] { + saved = proxies + .update( + UpdateProxy { + max_websocket_connections: change, + enabled: None, + auto_location: None, + test: None, + location: None, + id: saved.id.clone(), + revision: saved.revision, + name: saved.name.clone(), + proxy: None, + }, + &context(), + ) + .await + .unwrap() + .record; + assert_eq!(saved.proxy.max_websocket_connections(), expected); + for id in ["acct_proxy_budget_a", "acct_proxy_budget_b"] { + let projected = accounts + .get_account(&ProviderAccountId::new(id).unwrap()) + .await + .unwrap() + .unwrap(); + assert_eq!( + projected + .outbound_proxy() + .unwrap() + .max_websocket_connections(), + expected + ); + } + let other = accounts + .get_account(&ProviderAccountId::new("acct_other_proxy").unwrap()) + .await + .unwrap() + .unwrap(); + assert_eq!( + other.outbound_proxy().unwrap().max_websocket_connections(), + 0 + ); + } +} diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index 7f4c594b3..63cf30a40 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -341,7 +341,7 @@ impl CodexBackendClient { let pool_key = self.websocket_pool_key(request, context, pool_account_id, &connection_profile); let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key); - let pool = self.websocket_pool.as_deref().zip(pool_key); + let pool = self.websocket_pool.as_deref().map(|pool| (pool, pool_key)); let fast_path_budget = match requirement { TransportRequirement::PersistedContinuation | TransportRequirement::NewChain => { Some(WEBSOCKET_FAST_PATH_BUDGET) @@ -552,6 +552,7 @@ impl CodexBackendClient { .or(request.previous_response_id())?; let mut key = CodexWebSocketPoolKey::new(&self.base_url, account_id, conversation_id) .with_egress_key(&self.egress_key) + .with_outbound_proxy(self.outbound_proxy.as_ref()) .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); diff --git a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs index 693a4b1e1..0b3197ea9 100644 --- a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs +++ b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs @@ -113,7 +113,7 @@ impl PreparedWebSocket { /// 只建立或租用 WebSocket，不发送 payload pub(crate) async fn prepare_response_create_request_with_pool( request: &CodexWebSocketRequest, - pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>, + pool: Option<(&CodexWebSocketPool, Option)>, breaker: &WebSocketOriginBreaker, origin_key: &str, fast_path_budget: Option, @@ -121,7 +121,7 @@ pub(crate) async fn prepare_response_create_request_with_pool( fallback_stream_idle_timeout: Option, ) -> Result, CodexWebSocketExchangeError> { let decision_started_at = Instant::now(); - let Some((pool, key)) = pool else { + let Some((pool, Some(key))) = pool else { if require_pool || !request.continuation().permits_fresh_connection() { return Err(continuation_unavailable( PreviousResponseUnavailableReason::PoolUnavailable, @@ -134,6 +134,7 @@ pub(crate) async fn prepare_response_create_request_with_pool( fast_path_budget, fallback_stream_idle_timeout, decision_started_at, + pool.map(|(pool, _)| pool), ) .await; }; @@ -146,7 +147,14 @@ pub(crate) async fn prepare_response_create_request_with_pool( | WebSocketContinuationRequirement::Persisted { .. } | WebSocketContinuationRequirement::ExternalUnknown { .. } => None, }; - match pool.acquire(&key, required_response_id).await { + match pool + .acquire( + &key, + required_response_id, + request.connection().outbound_proxy.as_ref(), + ) + .await + { WebSocketPoolAcquire::Reused { connection, lease } => { if let WebSocketContinuationRequirement::ConnectionLocal { response_id } = request.continuation() @@ -218,7 +226,14 @@ async fn prepare_unpooled_websocket( fast_path_budget: Option, stream_idle_timeout: Option, decision_started_at: Instant, + capacity_pool: Option<&CodexWebSocketPool>, ) -> Result, CodexWebSocketExchangeError> { + let proxy_permit = if let Some(pool) = capacity_pool { + pool.acquire_proxy_capacity(request.connection().outbound_proxy.as_ref(), None) + .await? + } else { + None + }; let permit = acquire_breaker_permit(breaker, origin_key, fast_path_budget.is_some())?; let context = pump_log_context_from_connection(request.connection()); let connected = connect_with_budget( @@ -226,6 +241,7 @@ async fn prepare_unpooled_websocket( PumpKeepalive::disabled(), fast_path_budget, context, + proxy_permit, ) .await; let connected = match connected { @@ -323,10 +339,11 @@ impl PooledWebSocketConnectWaiter { fn start_pooled_websocket_connect( connection: CodexWebSocketConnection, pool: CodexWebSocketPool, - connect_lease: WebSocketPoolConnectLease, + mut connect_lease: WebSocketPoolConnectLease, permit: Option, fast_path: bool, ) -> PooledWebSocketConnectWaiter { + let proxy_permit = connect_lease.take_proxy_permit(); let task_key = connect_lease.key().clone(); let started_at = connect_lease.started_at(); let cancellation = connect_lease.cancellation_token(); @@ -355,7 +372,7 @@ fn start_pooled_websocket_connect( ); return; } - result = connect_websocket_connection(&connection, keepalive, context, fast_path) => result, + result = connect_websocket_connection(&connection, keepalive, context, fast_path, proxy_permit) => result, }; match finish_breaker_attempt(permit, connected) { Ok((connection, connect_elapsed)) => { @@ -459,11 +476,18 @@ async fn connect_with_budget( keepalive: PumpKeepalive, fast_path_budget: Option, context: PumpLogContext, + proxy_permit: Option, ) -> Result, CodexWebSocketExchangeError> { match wait_for_fast_path( TokioInstant::now(), fast_path_budget, - connect_websocket_connection(connection, keepalive, context, fast_path_budget.is_some()), + connect_websocket_connection( + connection, + keepalive, + context, + fast_path_budget.is_some(), + proxy_permit, + ), ) .await { @@ -477,10 +501,11 @@ async fn connect_websocket_connection( keepalive: PumpKeepalive, context: PumpLogContext, fast_path: bool, + proxy_permit: Option, ) -> Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError> { let started_at = Instant::now(); let (websocket, response) = - connect_pumped_websocket(connection, keepalive, context, fast_path).await?; + connect_pumped_websocket(connection, keepalive, context, fast_path, proxy_permit).await?; Ok(( PooledWebSocketConnection { websocket, diff --git a/backend/crates/providers/openai/src/transport/websocket/handshake.rs b/backend/crates/providers/openai/src/transport/websocket/handshake.rs index bb8bc55a3..91d5704cd 100644 --- a/backend/crates/providers/openai/src/transport/websocket/handshake.rs +++ b/backend/crates/providers/openai/src/transport/websocket/handshake.rs @@ -129,9 +129,13 @@ pub(super) async fn connect_pumped_websocket( keepalive: PumpKeepalive, context: PumpLogContext, fast_path: bool, + proxy_permit: Option, ) -> Result<(PumpedWebSocket, WsResponse>>), CodexWebSocketExchangeError> { let (raw, response) = connect_websocket(connection, fast_path).await?; - Ok((PumpedWebSocket::new(raw, keepalive, context), response)) + Ok(( + PumpedWebSocket::new(raw, keepalive, context, proxy_permit), + response, + )) } pub(super) async fn send_websocket_request( diff --git a/backend/crates/providers/openai/src/transport/websocket/model.rs b/backend/crates/providers/openai/src/transport/websocket/model.rs index 00299a011..e23f174c4 100644 --- a/backend/crates/providers/openai/src/transport/websocket/model.rs +++ b/backend/crates/providers/openai/src/transport/websocket/model.rs @@ -15,12 +15,19 @@ pub struct CodexWebSocketConnection { pub(super) headers: Vec<(String, String)>, } -// 请求预算不是连接身份；连接池复用不能由请求局部状态决定 +// 请求预算和代理额度不是连接身份，额度调整不能改变已有连接的复用合同 impl PartialEq for CodexWebSocketConnection { fn eq(&self, other: &Self) -> bool { self.endpoint == other.endpoint && self.headers == other.headers - && self.outbound_proxy == other.outbound_proxy + && self + .outbound_proxy + .as_ref() + .map(gateway_core::account::OutboundProxy::expose_url) + == other + .outbound_proxy + .as_ref() + .map(gateway_core::account::OutboundProxy::expose_url) } } impl Eq for CodexWebSocketConnection {} diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/capacity.rs b/backend/crates/providers/openai/src/transport/websocket/pool/capacity.rs new file mode 100644 index 000000000..9008a5ffc --- /dev/null +++ b/backend/crates/providers/openai/src/transport/websocket/pool/capacity.rs @@ -0,0 +1,142 @@ +//! 同一出站代理配置共享物理 WebSocket 预算，仅在容量不足时回收普通空闲连接 + +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use gateway_core::account::OutboundProxy; +use sha2::{Digest, Sha256}; +use tokio::sync::Notify; + +use super::state::{WebSocketPoolSlot, close_pooled_connection}; +use super::{CodexWebSocketPool, CodexWebSocketPoolKey}; +use crate::transport::websocket::CodexWebSocketExchangeError; + +/// 代理容量按完整连接配置合计，关联的多个上游账号共享名额，不混合不同代理认证 +pub(crate) fn proxy_capacity_key(proxy: Option<&OutboundProxy>) -> String { + proxy.map_or_else(String::new, |proxy| { + hex::encode(Sha256::digest(proxy.expose_url().as_bytes())) + }) +} + +pub(super) struct ProxyCapacity { + connections: AtomicUsize, + changed: Arc, +} + +/// 由建连任务移交给实际 socket，取消、断开和关闭均在析构时归还容量 +pub(crate) struct ProxyConnectionPermit(Arc); + +impl Drop for ProxyConnectionPermit { + fn drop(&mut self) { + self.0.connections.fetch_sub(1, Ordering::AcqRel); + self.0.changed.notify_waiters(); + } +} + +impl CodexWebSocketPool { + fn proxy_capacity(&self, proxy_key: &str) -> Arc { + let mut capacities = self + .proxy_capacities + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(capacity) = capacities.get(proxy_key).and_then(std::sync::Weak::upgrade) { + return capacity; + } + // socket 和等待者共同持有计数，无存活连接的出口无需永久保留注册项 + capacities.retain(|_, capacity| capacity.strong_count() > 0); + let capacity = Arc::new(ProxyCapacity { + connections: AtomicUsize::new(0), + changed: Arc::clone(&self.capacity_changed), + }); + capacities.insert(proxy_key.to_owned(), Arc::downgrade(&capacity)); + capacity + } + + pub(crate) async fn acquire_proxy_capacity( + &self, + proxy: Option<&OutboundProxy>, + retry_key: Option<&CodexWebSocketPoolKey>, + ) -> Result, CodexWebSocketExchangeError> { + let Some(proxy) = proxy else { + return Ok(None); + }; + let proxy_key = proxy_capacity_key(Some(proxy)); + let capacity = self.proxy_capacity(&proxy_key); + let limit = proxy.max_websocket_connections() as usize; + loop { + let changed = self.capacity_changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + if self.shutdown.is_cancelled() { + return Err(CodexWebSocketExchangeError::SharedConnectFailed); + } + // 等待期间同一会话可能已完成建连或归还，先回到现有复用路径 + if retry_key.is_some_and(|key| self.lock_state().slots.contains_key(key)) { + return Ok(None); + } + // 不限额时也跟踪 socket，后续调低或启用预算不会丢失已有连接的占用 + if capacity + .connections + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |used| { + if limit != 0 && used >= limit { + None + } else { + used.checked_add(1) + } + }) + .is_ok() + { + return Ok(Some(ProxyConnectionPermit(Arc::clone(&capacity)))); + } + let evicted = { + let mut state = self.lock_state(); + let oldest = state + .slots + .iter() + .filter_map(|(key, slot)| match slot { + WebSocketPoolSlot::Idle { idle_since, .. } + if key.proxy_capacity_key == proxy_key + && key.downstream_connection_id.is_empty() => + { + Some((key, *idle_since)) + } + _ => None, + }) + .min_by_key(|(_, idle_since)| *idle_since) + .map(|(key, _)| key.clone()); + oldest.and_then(|key| { + let WebSocketPoolSlot::Idle { connection, .. } = state.slots.remove(&key)? + else { + return None; + }; + state.remember_continuation_loss( + &key, + connection.continuation.latest_response_id(), + connection + .websocket + .observation() + .with_exit_reason("proxy_capacity_evicted"), + tokio::time::Instant::now(), + ); + tracing::info!( + proxy_scope_hash = proxy_key, + websocket_connection_id = %connection.websocket.connection_id(), + "Evicted idle WebSocket for proxy capacity" + ); + Some(*connection) + }) + }; + if let Some(connection) = evicted { + // 等待旧 socket 真正关闭后再占用它释放的名额，避免 CONNECT 瞬时超额 + close_pooled_connection(connection).await; + continue; + } + tokio::select! { + _ = &mut changed => {} + _ = self.shutdown.cancelled() => return Err(CodexWebSocketExchangeError::SharedConnectFailed), + } + } + } +} diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/lease.rs b/backend/crates/providers/openai/src/transport/websocket/pool/lease.rs index b0f7a9de4..0a5475218 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/lease.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/lease.rs @@ -65,6 +65,7 @@ pub(crate) struct WebSocketPoolConnectLease { // slot 分配时即注册，封闭 acquire 与后台 task spawn 之间的 shutdown 竞态 _task_registration: TaskTrackerToken, _connect_permit: OwnedSemaphorePermit, + proxy_permit: Option, armed: bool, } @@ -73,6 +74,7 @@ impl WebSocketPoolConnectLease { pool: CodexWebSocketPool, key: CodexWebSocketPoolKey, connect_permit: OwnedSemaphorePermit, + proxy_permit: Option, ) -> Self { let (outcome, _) = watch::channel(WebSocketPoolConnectOutcome::Pending); let cancellation = pool.shutdown.child_token(); @@ -86,6 +88,7 @@ impl WebSocketPoolConnectLease { cancellation, _task_registration: task_registration, _connect_permit: connect_permit, + proxy_permit, armed: true, } } @@ -98,6 +101,10 @@ impl WebSocketPoolConnectLease { &self.key } + pub(crate) fn take_proxy_permit(&mut self) -> Option { + self.proxy_permit.take() + } + pub(crate) fn cancellation_token(&self) -> CancellationToken { self.cancellation.clone() } 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 10cd0f3bc..d2471e504 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs @@ -1,16 +1,18 @@ //! Codex WebSocket 连接池 +mod capacity; mod lease; mod state; mod supervisor; use std::{ + collections::HashMap, future::Future, - sync::{Arc, Mutex, MutexGuard, atomic::AtomicBool}, + sync::{Arc, Mutex, MutexGuard, Weak, atomic::AtomicBool}, time::Duration, }; -use tokio::sync::Semaphore; +use tokio::sync::{Notify, Semaphore}; use tokio_util::{sync::CancellationToken, task::TaskTracker}; use uuid::Uuid; @@ -21,6 +23,8 @@ use self::state::{ use super::pump::PumpKeepalive; use super::pump::WebSocketConnectionObservation; +use self::capacity::ProxyCapacity; +pub(crate) use self::capacity::{ProxyConnectionPermit, proxy_capacity_key}; pub use self::state::CodexWebSocketPoolKey; pub(crate) use self::{ lease::{ @@ -49,6 +53,8 @@ pub struct CodexWebSocketPool { shutdown: CancellationToken, connect_semaphore: Arc, maintenance_started: Arc, + proxy_capacities: Arc>>>, + capacity_changed: Arc, } impl Default for CodexWebSocketPool { @@ -107,7 +113,7 @@ impl CodexWebSocketPoolConfig { } impl CodexWebSocketPool { - /// 构造不限制累计 slot 数的连接池策略和状态 + /// 使用指定连接总寿命构造连接池，不启动定时保活与清理 pub fn new(max_age: Duration) -> Self { Self::with_config(CodexWebSocketPoolConfig { max_age, @@ -127,6 +133,8 @@ impl CodexWebSocketPool { shutdown: CancellationToken::new(), connect_semaphore: Arc::new(Semaphore::new(config.max_connecting)), maintenance_started: Arc::new(AtomicBool::new(false)), + proxy_capacities: Arc::default(), + capacity_changed: Arc::default(), }; pool.spawn_maintenance_task(); pool @@ -151,121 +159,137 @@ impl CodexWebSocketPool { &self, key: &CodexWebSocketPoolKey, required_response_id: Option<&str>, + proxy: Option<&gateway_core::account::OutboundProxy>, ) -> WebSocketPoolAcquire { self.spawn_maintenance_task(); - let mut connections_to_close = Vec::new(); - let acquire = { - let mut state = self.lock_state(); - let mut continuation_loss = None; - if !self.config.enabled || state.shutting_down { - return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Disabled); - } - let key = if let Some(response_id) = required_response_id { - let Some(key) = state.slots.iter().find_map(|(candidate, slot)| { - (candidate.same_logical_connection(key) - && slot.latest_response_id() == Some(response_id)) - .then(|| candidate.clone()) - }) else { - return state - .continuation_loss(key, response_id, tokio::time::Instant::now()) - .map_or( - WebSocketPoolAcquire::Bypass( - WebSocketPoolBypassReason::ContinuationNotFound, - ), - WebSocketPoolAcquire::ContinuationLost, - ); + let mut proxy_permit = None; + loop { + let mut connections_to_close = Vec::new(); + let acquire = { + let mut state = self.lock_state(); + let mut continuation_loss = None; + if !self.config.enabled || state.shutting_down { + return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Disabled); + } + let key = if let Some(response_id) = required_response_id { + let Some(key) = state.slots.iter().find_map(|(candidate, slot)| { + (candidate.same_logical_connection(key) + && slot.latest_response_id() == Some(response_id)) + .then(|| candidate.clone()) + }) else { + return state + .continuation_loss(key, response_id, tokio::time::Instant::now()) + .map_or( + WebSocketPoolAcquire::Bypass( + WebSocketPoolBypassReason::ContinuationNotFound, + ), + WebSocketPoolAcquire::ContinuationLost, + ); + }; + key + } else { + key.clone() }; - key - } else { - key.clone() - }; - match state.slots.get(&key) { - Some(WebSocketPoolSlot::Busy(_)) => { - return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Busy); - } - Some(WebSocketPoolSlot::Connecting(connecting)) => { - return WebSocketPoolAcquire::Wait(WebSocketPoolConnectWaiter { - receiver: connecting.outcome.subscribe(), - started_at: connecting.started_at, - }); - } - Some(WebSocketPoolSlot::Idle { .. }) => { - let Some(WebSocketPoolSlot::Idle { connection, .. }) = state.slots.remove(&key) - else { + match state.slots.get(&key) { + Some(WebSocketPoolSlot::Busy(_)) => { return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Busy); - }; - // 零成本探活：后台 pump 已实时感知连接死亡（RST/Close/EOF/失活）， - // 复用前只需读取 is_closed 标志，避免复用到静默死连接后卡到超时 - let expired = connection.created_at.elapsed() >= self.config.max_age; - let closed = connection.websocket.is_closed(); - if !expired && !closed { - let lease = WebSocketPoolLease::reserve( - self.clone(), - key.clone(), + } + Some(WebSocketPoolSlot::Connecting(connecting)) => { + return WebSocketPoolAcquire::Wait(WebSocketPoolConnectWaiter { + receiver: connecting.outcome.subscribe(), + started_at: connecting.started_at, + }); + } + Some(WebSocketPoolSlot::Idle { .. }) => { + let Some(WebSocketPoolSlot::Idle { connection, .. }) = + state.slots.remove(&key) + else { + return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Busy); + }; + // 零成本探活：后台 pump 已实时感知连接死亡（RST/Close/EOF/失活）， + // 复用前只需读取 is_closed 标志，避免复用到静默死连接后卡到超时 + let expired = connection.created_at.elapsed() >= self.config.max_age; + let closed = connection.websocket.is_closed(); + if !expired && !closed { + let lease = WebSocketPoolLease::reserve( + self.clone(), + key.clone(), + connection.continuation.latest_response_id(), + ); + state.slots.insert( + key.clone(), + WebSocketPoolSlot::Busy(lease.reservation.clone()), + ); + return WebSocketPoolAcquire::Reused { connection, lease }; + } + let observation = if expired && !closed { + connection + .websocket + .observation() + .with_exit_reason("max_age_expired") + } else { + connection.websocket.observation() + }; + state.remember_continuation_loss( + &key, connection.continuation.latest_response_id(), + observation.clone(), + tokio::time::Instant::now(), ); - state.slots.insert( - key.clone(), - WebSocketPoolSlot::Busy(lease.reservation.clone()), - ); - return WebSocketPoolAcquire::Reused { connection, lease }; - } - let observation = if expired && !closed { - connection - .websocket - .observation() - .with_exit_reason("max_age_expired") - } else { - connection.websocket.observation() - }; - state.remember_continuation_loss( - &key, - connection.continuation.latest_response_id(), - observation.clone(), - tokio::time::Instant::now(), - ); - connections_to_close.push(*connection); - if required_response_id.is_some() { - continuation_loss = Some(observation); + connections_to_close.push(*connection); + if required_response_id.is_some() { + continuation_loss = Some(observation); + } } + None => {} } - None => {} - } - let acquire = if let Some(observation) = continuation_loss { - WebSocketPoolAcquire::ContinuationLost(observation) - } else if required_response_id.is_some() { - WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::ContinuationNotFound) - } else { - let connect_permit = self.connect_semaphore.clone().try_acquire_owned().ok(); - match connect_permit { - Some(connect_permit) => { - 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), - } + let acquire = if let Some(observation) = continuation_loss { + Some(WebSocketPoolAcquire::ContinuationLost(observation)) + } else if required_response_id.is_some() { + Some(WebSocketPoolAcquire::Bypass( + WebSocketPoolBypassReason::ContinuationNotFound, + )) + } else if proxy.is_some() && proxy_permit.is_none() { + None + } else { + let connect_permit = self.connect_semaphore.clone().try_acquire_owned().ok(); + Some(match connect_permit { + Some(connect_permit) => { + let lease = WebSocketPoolConnectLease::reserve( + self.clone(), + key.clone(), + connect_permit, + proxy_permit.take(), + ); + 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), + }) + }; + drop(state); + acquire }; - drop(state); - acquire - }; - close_pooled_connections(connections_to_close).await; + close_pooled_connections(connections_to_close).await; - acquire + if let Some(acquire) = acquire { + return acquire; + } + match self.acquire_proxy_capacity(proxy, Some(key)).await { + Ok(permit) => proxy_permit = permit, + Err(_) => return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Disabled), + } + } } async fn put_reserved( @@ -303,6 +327,7 @@ impl CodexWebSocketPool { key.clone(), WebSocketPoolSlot::Idle { connection: Box::new(connection), + idle_since: tokio::time::Instant::now(), }, ); } @@ -310,6 +335,7 @@ impl CodexWebSocketPool { if let Some(connection) = connection { close_pooled_connection(connection).await; } + self.capacity_changed.notify_waiters(); } fn discard_reserved_now(&self, key: &CodexWebSocketPoolKey, reservation_id: Uuid) { 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 5930970ec..b5e439250 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/state.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/state.rs @@ -23,8 +23,9 @@ pub struct CodexWebSocketPoolKey { account_id: String, conversation_id: String, connection_profile: String, - downstream_connection_id: String, + pub(super) downstream_connection_id: String, egress_key: String, + pub(super) proxy_capacity_key: String, } impl CodexWebSocketPoolKey { @@ -41,6 +42,7 @@ impl CodexWebSocketPoolKey { connection_profile: String::new(), downstream_connection_id: String::new(), egress_key: String::new(), + proxy_capacity_key: String::new(), } } @@ -68,6 +70,14 @@ impl CodexWebSocketPoolKey { self } + pub(crate) fn with_outbound_proxy( + mut self, + proxy: Option<&gateway_core::account::OutboundProxy>, + ) -> Self { + self.proxy_capacity_key = super::proxy_capacity_key(proxy); + self + } + pub(crate) fn conversation_id_hash(&self) -> String { short_sha256([self.conversation_id.as_str()]) } @@ -192,6 +202,7 @@ impl WebSocketContinuationState { pub(super) enum WebSocketPoolSlot { Idle { connection: Box, + idle_since: Instant, }, Busy(WebSocketPoolReservation), Connecting(WebSocketPoolConnecting), @@ -222,7 +233,7 @@ pub(super) struct WebSocketPoolReservation { } pub(super) async fn close_pooled_connection(connection: PooledWebSocketConnection) { - connection.websocket.close().await; + connection.websocket.close_and_wait().await; } pub(super) async fn close_pooled_connections(connections: Vec) { diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/supervisor.rs b/backend/crates/providers/openai/src/transport/websocket/pool/supervisor.rs index 25fed2b2a..0b53b189c 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/supervisor.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/supervisor.rs @@ -110,6 +110,8 @@ impl CodexWebSocketPool { let shutdown = self.shutdown.clone(); let connect_semaphore = Arc::clone(&self.connect_semaphore); let maintenance_started = Arc::clone(&self.maintenance_started); + let proxy_capacities = Arc::clone(&self.proxy_capacities); + let capacity_changed = Arc::clone(&self.capacity_changed); drop(self.tasks.spawn(async move { let mut interval = tokio::time::interval(interval_duration); loop { @@ -127,6 +129,8 @@ impl CodexWebSocketPool { shutdown: shutdown.clone(), connect_semaphore: Arc::clone(&connect_semaphore), maintenance_started: Arc::clone(&maintenance_started), + proxy_capacities: Arc::clone(&proxy_capacities), + capacity_changed: Arc::clone(&capacity_changed), }; if pool.is_shutdown().await { break; diff --git a/backend/crates/providers/openai/src/transport/websocket/pump.rs b/backend/crates/providers/openai/src/transport/websocket/pump.rs index eaff47e5b..db4e2f816 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pump.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pump.rs @@ -298,12 +298,19 @@ impl std::fmt::Debug for PumpedWebSocket { } } +// 字段按声明顺序析构，先释放物理 socket，再归还代理名额 +struct PumpSocket { + stream: RawWsStream, + _proxy_permit: Option, +} + impl PumpedWebSocket { /// 用底层流启动一个 pump 任务并返回句柄 pub(crate) fn new( inner: RawWsStream, keepalive: PumpKeepalive, context: PumpLogContext, + proxy_permit: Option, ) -> Self { let (tx_command, rx_command) = mpsc::channel::(PUMP_COMMAND_BUFFER); let (tx_message, rx_message) = mpsc::channel(PUMP_MESSAGE_BUFFER); @@ -318,6 +325,10 @@ impl PumpedWebSocket { exit_reason: None, })); let lifecycle_for_task = Arc::clone(&lifecycle); + let inner = PumpSocket { + stream: inner, + _proxy_permit: proxy_permit, + }; let pump = tokio::spawn(async move { pump_loop( connection_id, @@ -394,6 +405,15 @@ impl PumpedWebSocket { .observation(self.connection_id, Instant::now()) } + pub(crate) async fn close_and_wait(mut self) { + self.close().await; + if let Some(pump) = self.pump.as_mut() { + // 沿用 Drop 的停止语义，并等待 socket 析构完成后再让容量等待者继续 + pump.abort(); + let _ = pump.await; + } + } + /// 主动关闭连接（best-effort 发送 Close 帧） pub(crate) async fn close(&self) { if !self.is_closed() { @@ -416,7 +436,7 @@ impl Drop for PumpedWebSocket { )] async fn pump_loop( connection_id: Uuid, - mut inner: RawWsStream, + mut inner: PumpSocket, mut rx_command: mpsc::Receiver, tx_message: mpsc::Sender>, keepalive: PumpKeepalive, @@ -465,7 +485,7 @@ async fn pump_loop( let Some(command) = command else { break 'pump PumpExitReason::CommandChannelClosed; }; - if let Some(reason) = handle_command(&mut inner, command, closed.load(Ordering::Acquire)).await { + if let Some(reason) = handle_command(&mut inner.stream, command, closed.load(Ordering::Acquire)).await { break 'pump reason; } } @@ -480,11 +500,11 @@ async fn pump_loop( let Some(command) = command else { break 'pump PumpExitReason::CommandChannelClosed; }; - if let Some(reason) = handle_command(&mut inner, command, closed.load(Ordering::Acquire)).await { + if let Some(reason) = handle_command(&mut inner.stream, command, closed.load(Ordering::Acquire)).await { break 'pump reason; } } - message = inner.next() => { + message = inner.stream.next() => { match message { None => break 'pump PumpExitReason::UpstreamEof, Some(Ok(message)) => { @@ -496,7 +516,7 @@ async fn pump_loop( } match message { Message::Ping(payload) => { - if let Err(error) = inner.send(Message::Pong(payload)).await { + if let Err(error) = inner.stream.send(Message::Pong(payload)).await { break 'pump PumpExitReason::KeepaliveTransportError { message: error.to_string(), metric_reason: transport_metric_reason(&error), @@ -558,7 +578,7 @@ async fn pump_loop( _ = tick(&mut ping_ticker) => { if pending_ping.is_none() { match send_keepalive_ping( - &mut inner, + &mut inner.stream, &mut ping_sequence, ping_timeout, ).await { @@ -601,7 +621,7 @@ async fn pump_loop( ); if reason.should_send_close() { - let _ = inner.send(Message::Close(None)).await; + let _ = inner.stream.send(Message::Close(None)).await; } } diff --git a/backend/crates/providers/openai/tests/transport/mod.rs b/backend/crates/providers/openai/tests/transport/mod.rs index b28de67cc..511d37528 100644 --- a/backend/crates/providers/openai/tests/transport/mod.rs +++ b/backend/crates/providers/openai/tests/transport/mod.rs @@ -63,6 +63,7 @@ mod profile; mod profile_avatar; mod profile_contract; mod protocol; +mod proxy_capacity; mod request; mod reset_credits; mod subscription; diff --git a/backend/crates/providers/openai/tests/transport/proxy_capacity.rs b/backend/crates/providers/openai/tests/transport/proxy_capacity.rs new file mode 100644 index 000000000..2dbf222d6 --- /dev/null +++ b/backend/crates/providers/openai/tests/transport/proxy_capacity.rs @@ -0,0 +1,625 @@ +//! 使用受限 CONNECT 代理验证跨账号 WebSocket 容量、回收与续接恢复 + +use super::*; +use gateway_core::{ + account::{CredentialRevision, OutboundProxy, ProviderAccount, ProviderAccountId}, + routing::ProviderKind, +}; +use provider_openai::transport::websocket::PreviousResponseUnavailableReason; + +fn limited_pool() -> Arc { + Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { + max_age: Duration::from_mins(55), + ..websocket_pool_config_for_tests(None, None, None) + })) +} + +async fn wait_until(mut condition: impl FnMut() -> bool) { + timeout(Duration::from_secs(3), async { + while !condition() { + tokio::task::yield_now().await; + } + }) + .await + .expect("proxy state should progress"); +} + +async fn respond(client: &CodexBackendClient, conversation: &str) { + client + .create_response( + &pooled_websocket_request(conversation), + request_context("capacity", Some("test-upstream")), + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn proxy_capacity_shares_budget_across_accounts() { + for budget in [8, 16, 24, 28] { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let first = proxy.client(pool.clone(), "acct_a", budget); + let second = proxy.client(pool.clone(), "acct_b", budget); + for index in 0..64 { + respond( + if index % 2 == 0 { &first } else { &second }, + &format!("shared-{index}"), + ) + .await; + } + wait_until(|| proxy.active.load(Ordering::SeqCst) == budget as usize).await; + assert_eq!(proxy.opened.load(Ordering::SeqCst), 64); + assert!( + proxy.peak.load(Ordering::SeqCst) <= budget as usize, + "physical connections must stay within budget {budget}" + ); + pool.shutdown().await; + } +} + +#[tokio::test] +async fn proxy_capacity_reuses_and_evicts_lru_with_original_replay_contract() { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let client = proxy.client(pool.clone(), "acct_a", 2); + respond(&client, "a").await; + respond(&client, "b").await; + respond(&client, "a").await; + assert_eq!( + proxy.opened.load(Ordering::SeqCst), + 2, + "reuse while capacity is available" + ); + respond(&client, "c").await; + respond(&client, "a").await; + assert_eq!( + proxy.opened.load(Ordering::SeqCst), + 3, + "recently used a must survive" + ); + let mut continuation = pooled_websocket_request("b"); + continuation.set_previous_response_id(Some("resp_proxy_1".to_owned())); + continuation.previous_response_scope = Some(PreviousResponseScope::ConnectionLocal); + let error = client + .create_response( + &continuation, + request_context("continuation", Some("test-upstream")), + ) + .await + .unwrap_err(); + let CodexClientError::WebSocket(error) = error else { + panic!("expected existing replay error"); + }; + assert_eq!( + error.continuation_unavailable_reason(), + Some(PreviousResponseUnavailableReason::ReusedConnectionLost) + ); + assert_eq!( + error.connection_observation().unwrap().exit_reason(), + "proxy_capacity_evicted" + ); + assert_eq!( + proxy.opened.load(Ordering::SeqCst), + 3, + "do not retry the evicted socket" + ); + let mut replay = codex_request( + "gpt-test", + "be brief", + vec![json!({"role":"user","content":"full history"})], + ); + replay.local_conversation_id = Some("b".to_owned()); + client + .create_response(&replay, request_context("replay", Some("test-upstream"))) + .await + .unwrap(); + { + let requests = proxy.requests.lock().unwrap(); + let replayed = requests.last().unwrap(); + assert!(replayed.get("previous_response_id").is_none()); + assert_eq!(replayed["input"][0]["content"], "full history"); + } + assert_eq!(proxy.opened.load(Ordering::SeqCst), 4); + pool.shutdown().await; +} + +#[tokio::test] +async fn proxy_capacity_preserves_idle_reuse_beyond_two_minutes_without_pressure() { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let client = proxy.client(pool.clone(), "acct_a", 1); + respond(&client, "idle").await; + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(121)).await; + tokio::time::resume(); + respond(&client, "idle").await; + assert_eq!(proxy.opened.load(Ordering::SeqCst), 1); + pool.shutdown().await; +} + +#[tokio::test] +async fn proxy_capacity_does_not_limit_direct_connections() { + let pool = limited_pool(); + let proxy = LimitedProxy::start(32).await; + respond(&proxy.client(pool.clone(), "acct_proxied", 1), "proxied").await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let mut sockets = tokio::task::JoinSet::new(); + for _ in 0..3 { + let (stream, _) = listener.accept().await.unwrap(); + sockets.spawn(async move { + let mut ws = accept_codex_test_websocket(stream).await; + ws.next().await.unwrap().unwrap(); + ws.send(Message::Text( + completed_websocket_response("resp_direct", 1, 1).into(), + )) + .await + .unwrap(); + std::assert_matches!(ws.next().await.unwrap().unwrap(), Message::Close(_)); + }); + } + while let Some(result) = sockets.join_next().await { + result.unwrap(); + } + }); + let direct = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{address}"), + test_wire_profile(), + ) + .with_websocket_pool(pool.clone()); + for index in 0..3 { + timeout( + Duration::from_secs(3), + respond(&direct, &format!("direct-{index}")), + ) + .await + .unwrap(); + } + assert_eq!(proxy.active.load(Ordering::SeqCst), 1); + pool.shutdown().await; + server.await.unwrap(); +} + +#[tokio::test] +async fn proxy_capacity_releases_failed_handshake_before_the_next_connection() { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let client = proxy.client(pool.clone(), "acct_a", 1); + proxy.reject_connects.store(1, Ordering::SeqCst); + let mut request = pooled_websocket_request("rejected"); + request.downstream_websocket_connection_id = Some("downstream-rejected".to_owned()); + request.set_previous_response_id(None); + request.previous_response_scope = None; + assert!( + client + .create_response(&request, request_context("rejected", Some("test-upstream"))) + .await + .is_err() + ); + timeout(Duration::from_secs(3), respond(&client, "after-failure")) + .await + .unwrap(); + assert_eq!(proxy.opened.load(Ordering::SeqCst), 1); + pool.shutdown().await; +} + +#[tokio::test] +async fn proxy_capacity_protects_downstream_bindings_and_blocks_unpooled_bypass() { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let client = proxy.client(pool.clone(), "acct_a", 1); + let mut pinned = pooled_websocket_request("pinned"); + pinned.downstream_websocket_connection_id = Some("downstream-a".to_owned()); + client + .create_response(&pinned, request_context("pinned", Some("test-upstream"))) + .await + .unwrap(); + let unpooled = codex_request("gpt-test", "be brief", Vec::new()); + assert!( + timeout( + Duration::from_millis(1100), + client.create_response( + &unpooled, + request_context("unpooled", Some("test-upstream")) + ) + ) + .await + .is_err() + ); + assert_eq!(proxy.opened.load(Ordering::SeqCst), 1); + client + .create_response( + &pinned, + request_context("pinned-reuse", Some("test-upstream")), + ) + .await + .unwrap(); + assert_eq!( + proxy.opened.load(Ordering::SeqCst), + 1, + "capacity cancellation must not remove pinned idle connection" + ); + pool.shutdown().await; +} + +#[tokio::test] +async fn proxy_capacity_waits_for_busy_connections_then_reclaims_returned_idle() { + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let proxy = LimitedProxy::with_gates(32, None, Some(gate.clone())).await; + let pool = limited_pool(); + let client = proxy.client(pool.clone(), "acct_a", 24); + let mut requests = Vec::new(); + for index in 0..25 { + let client = client.clone(); + requests.push(tokio::spawn(async move { + respond(&client, &format!("busy-{index}")).await + })); + if index < 24 { + wait_until(|| proxy.requests.lock().unwrap().len() == index + 1).await; + } + } + wait_until(|| proxy.requests.lock().unwrap().len() == 24).await; + tokio::time::sleep(Duration::from_millis(1100)).await; + assert_eq!( + proxy.opened.load(Ordering::SeqCst), + 24, + "all busy must wait beyond the HTTP fast-path budget" + ); + gate.add_permits(1); + wait_until(|| proxy.requests.lock().unwrap().len() == 25).await; + assert_eq!(proxy.opened.load(Ordering::SeqCst), 25); + gate.add_permits(24); + for task in requests { + timeout(Duration::from_secs(3), task) + .await + .unwrap() + .unwrap(); + } + pool.shutdown().await; +} + +#[tokio::test] +async fn proxy_capacity_counts_connecting_and_shutdown_cancels_waiters() { + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let proxy = LimitedProxy::with_gates(32, Some(gate), None).await; + let pool = limited_pool(); + let client = proxy.client(pool.clone(), "acct_a", 2); + let mut requests = Vec::new(); + for index in 0..3 { + let client = client.clone(); + requests.push(tokio::spawn(async move { + let mut request = pooled_websocket_request(&format!("connecting-{index}")); + request.downstream_websocket_connection_id = Some(format!("downstream-{index}")); + request.set_previous_response_id(None); + request.previous_response_scope = None; + client + .create_response( + &request, + request_context("connecting", Some("test-upstream")), + ) + .await + })); + } + wait_until(|| proxy.active.load(Ordering::SeqCst) == 2).await; + tokio::time::sleep(Duration::from_millis(1100)).await; + assert_eq!(proxy.active.load(Ordering::SeqCst), 2); + assert_eq!(proxy.opened.load(Ordering::SeqCst), 0); + timeout(Duration::from_secs(3), pool.shutdown()) + .await + .unwrap(); + for task in requests { + assert!( + timeout(Duration::from_secs(3), task) + .await + .unwrap() + .unwrap() + .is_err() + ); + } +} + +#[tokio::test] +async fn proxy_capacity_isolates_different_proxy_endpoints() { + let first = LimitedProxy::start(32).await; + let second = LimitedProxy::start(32).await; + let pool = limited_pool(); + respond(&first.client(pool.clone(), "acct_a", 1), "first").await; + respond(&second.client(pool.clone(), "acct_a", 2), "second").await; + respond(&second.client(pool.clone(), "acct_a", 2), "second-extra").await; + assert_eq!(first.active.load(Ordering::SeqCst), 1); + assert_eq!(second.active.load(Ordering::SeqCst), 2); + pool.shutdown().await; +} + +#[tokio::test] +async fn proxy_capacity_configured_24_leaves_eight_http_tunnels_on_a_32_slot_proxy() { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let client = proxy.client(pool.clone(), "acct_a", 24); + for index in 0..24 { + respond(&client, &format!("headroom-{index}")).await; + } + let mut http_tunnels = Vec::new(); + for _ in 0..8 { + let mut stream = TcpStream::connect(proxy.url.trim_start_matches("http://")) + .await + .unwrap(); + stream.write_all(b"CONNECT upstream.invalid:80 HTTP/1.1\r\nHost: upstream.invalid:80\r\nX-Capacity-Probe: health\r\n\r\n").await.unwrap(); + let response = read_http_request(&mut stream).await; + assert!(response.starts_with("HTTP/1.1 200")); + stream + .write_all(b"GET /health HTTP/1.1\r\nHost: upstream.invalid\r\n\r\n") + .await + .unwrap(); + let response = read_http_request(&mut stream).await; + assert!(response.starts_with("HTTP/1.1 200")); + http_tunnels.push(stream); + } + assert_eq!(proxy.active.load(Ordering::SeqCst), 32); + timeout(Duration::from_secs(3), respond(&client, "headroom-new")) + .await + .unwrap(); + assert_eq!(proxy.opened.load(Ordering::SeqCst), 25); + assert_eq!(proxy.peak.load(Ordering::SeqCst), 32); + drop(http_tunnels); + pool.shutdown().await; +} + +struct LimitedProxy { + url: String, + active: Arc, + peak: Arc, + opened: Arc, + requests: Arc>>, + reject_connects: Arc, + task: tokio::task::JoinHandle<()>, +} + +struct ActiveTunnel(Arc); + +impl Drop for ActiveTunnel { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } +} + +impl Drop for LimitedProxy { + fn drop(&mut self) { + self.task.abort(); + } +} + +impl LimitedProxy { + async fn start(limit: usize) -> Self { + Self::with_gates(limit, None, None).await + } + + async fn with_gates( + limit: usize, + connect_gate: Option>, + response_gate: Option>, + ) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + let active = Arc::new(AtomicUsize::new(0)); + let peak = Arc::new(AtomicUsize::new(0)); + let opened = Arc::new(AtomicUsize::new(0)); + let requests = Arc::new(Mutex::new(Vec::new())); + let reject_connects = Arc::new(AtomicUsize::new(0)); + let rejects = reject_connects.clone(); + let recorded = requests.clone(); + let (a, p, o) = (active.clone(), peak.clone(), opened.clone()); + let task = tokio::spawn(async move { + let slots = Arc::new(tokio::sync::Semaphore::new(limit)); + let mut tunnels = tokio::task::JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let (mut stream, _) = accepted.unwrap(); + let (slots, active, peak, opened) = (slots.clone(), a.clone(), p.clone(), o.clone()); + let (connect_gate, response_gate, requests) = (connect_gate.clone(), response_gate.clone(), recorded.clone()); + let reject_connects = rejects.clone(); + tunnels.spawn(async move { + let _slot = slots.acquire_owned().await.unwrap(); + let count = active.fetch_add(1, Ordering::SeqCst) + 1; + let _active = ActiveTunnel(active); + peak.fetch_max(count, Ordering::SeqCst); + let request = read_http_request(&mut stream).await; + assert!(request.starts_with("CONNECT upstream.invalid:80 ")); + if reject_connects.fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| remaining.checked_sub(1)).is_ok() { + stream.write_all(b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\n\r\n").await.unwrap(); + return; + } + if let Some(gate) = connect_gate { gate.acquire().await.unwrap().forget(); } + stream.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n").await.unwrap(); + if request.contains("X-Capacity-Probe: health") { + let health = read_http_request(&mut stream).await; + assert!(health.starts_with("GET /health HTTP/1.1")); + stream.write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok").await.unwrap(); + let mut byte = [0]; + let _ = stream.read(&mut byte).await; + return; + } + let Ok(mut ws) = accept_hdr_async_with_config(stream, TestWebSocketCallback(|_request: &WebSocketRequest, response: &mut WebSocketResponse| { response.headers_mut().insert("sec-websocket-extensions", "permessage-deflate".parse().unwrap()); }), Some(websocket_accept_config())).await else { return; }; + let id = opened.fetch_add(1, Ordering::SeqCst); + while let Some(Ok(message)) = ws.next().await { + match message { + Message::Text(payload) => { + requests.lock().unwrap().push(serde_json::from_str(&payload).unwrap()); + if let Some(gate) = &response_gate { gate.acquire().await.unwrap().forget(); } + if ws.send(Message::Text(completed_websocket_response(&format!("resp_proxy_{id}"), 1, 1).into())).await.is_err() { break; } + } + Message::Ping(payload) => { if ws.send(Message::Pong(payload)).await.is_err() { break; } } + Message::Close(_) => { let _ = ws.flush().await; break; } + _ => {} + } + } + }); + } + _ = tunnels.join_next(), if !tunnels.is_empty() => {} + } + } + }); + Self { + url, + active, + peak, + opened, + requests, + reject_connects, + task, + } + } + + fn client( + &self, + pool: Arc, + account_id: &str, + limit: u32, + ) -> CodexBackendClient { + self.client_with_proxy_user(pool, account_id, limit, "proxy-user") + } + + fn client_with_proxy_user( + &self, + pool: Arc, + account_id: &str, + limit: u32, + proxy_user: &str, + ) -> CodexBackendClient { + let account = ProviderAccount::new( + ProviderAccountId::new(account_id).unwrap(), + ProviderKind::new("openai").unwrap(), + account_id.to_owned(), + Some(account_id.to_owned()), + "oauth".to_owned(), + CredentialRevision::new(1).unwrap(), + None, + ) + .with_outbound_proxy(Some( + OutboundProxy::parse( + &self + .url + .replace("http://", &format!("http://{proxy_user}:test@")), + ) + .unwrap() + .with_max_websocket_connections(limit), + )); + CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + "http://upstream.invalid", + test_wire_profile(), + ) + .with_websocket_pool(pool) + .for_account(&account) + .unwrap() + } +} + +#[tokio::test] +async fn proxy_capacity_unconfigured_proxy_keeps_existing_unbounded_behavior() { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let client = proxy.client(pool.clone(), "acct_a", 0); + for index in 0..32 { + client + .create_response( + &pooled_websocket_request(&format!("baseline-{index}")), + request_context("baseline", Some("acct_a")), + ) + .await + .unwrap(); + } + let next = pooled_websocket_request("baseline-33"); + assert!( + timeout( + Duration::from_millis(300), + client.create_response(&next, request_context("baseline-33", Some("acct_a"))) + ) + .await + .is_err() + ); + assert_eq!(proxy.active.load(Ordering::SeqCst), 32); + assert_eq!(proxy.opened.load(Ordering::SeqCst), 32); + assert_eq!(proxy.peak.load(Ordering::SeqCst), 32); + pool.shutdown().await; +} + +#[tokio::test] +async fn proxy_capacity_enabling_and_lowering_limits_keeps_existing_socket_count() { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let unlimited = proxy.client(pool.clone(), "acct_a", 0); + for index in 0..3 { + respond(&unlimited, &format!("unlimited-{index}")).await; + } + assert_eq!(proxy.active.load(Ordering::SeqCst), 3); + let limited = proxy.client(pool.clone(), "acct_a", 1); + respond(&limited, "unlimited-2").await; + assert_eq!(proxy.opened.load(Ordering::SeqCst), 3); + respond(&limited, "enabled-limit").await; + assert_eq!(proxy.active.load(Ordering::SeqCst), 1); + assert_eq!(proxy.peak.load(Ordering::SeqCst), 3); + let raised = proxy.client(pool.clone(), "acct_a", 2); + respond(&raised, "raised-limit").await; + assert_eq!(proxy.active.load(Ordering::SeqCst), 2); + respond(&limited, "lowered-again").await; + assert_eq!(proxy.active.load(Ordering::SeqCst), 1); + pool.shutdown().await; +} + +#[tokio::test] +async fn proxy_capacity_enabling_limit_preserves_existing_downstream_bindings() { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let unlimited = proxy.client(pool.clone(), "acct_a", 0); + for index in 0..2 { + let mut request = pooled_websocket_request(&format!("pinned-{index}")); + request.downstream_websocket_connection_id = Some(format!("downstream-{index}")); + unlimited + .create_response(&request, request_context("pinned", Some("test-upstream"))) + .await + .unwrap(); + } + let limited = proxy.client(pool.clone(), "acct_a", 1); + let mut retained = pooled_websocket_request("pinned-0"); + retained.downstream_websocket_connection_id = Some("downstream-0".to_owned()); + limited + .create_response( + &retained, + request_context("pinned-reused", Some("test-upstream")), + ) + .await + .unwrap(); + assert_eq!(proxy.opened.load(Ordering::SeqCst), 2); + assert!( + timeout(Duration::from_millis(100), respond(&limited, "must-wait")) + .await + .is_err() + ); + assert_eq!(proxy.opened.load(Ordering::SeqCst), 2); + let raised = proxy.client(pool.clone(), "acct_a", 3); + respond(&raised, "raised-limit").await; + assert_eq!(proxy.active.load(Ordering::SeqCst), 3); + pool.shutdown().await; +} + +#[tokio::test] +async fn proxy_capacity_keeps_separate_proxy_credentials_independent() { + let proxy = LimitedProxy::start(32).await; + let pool = limited_pool(); + let first = proxy.client(pool.clone(), "acct_a", 1); + let second = proxy.client_with_proxy_user(pool.clone(), "acct_b", 2, "other-proxy-user"); + respond(&first, "first").await; + respond(&second, "second").await; + respond(&second, "second-extra").await; + assert_eq!(proxy.active.load(Ordering::SeqCst), 3); + respond(&first, "first-new").await; + assert_eq!(proxy.active.load(Ordering::SeqCst), 3); + respond(&second, "second").await; + assert_eq!(proxy.opened.load(Ordering::SeqCst), 4); + pool.shutdown().await; +} diff --git a/backend/migrations/.frozen-sha256 b/backend/migrations/.frozen-sha256 index c6f4fbb3b..131a12e66 100644 --- a/backend/migrations/.frozen-sha256 +++ b/backend/migrations/.frozen-sha256 @@ -30,3 +30,4 @@ d4b599066d7382a1c9cb6b3c003200046579e4fce3531f232cbd665a9d92c542 0026_codex_pri f10058e2b4c974b6d2ffa11be78b085504a233205253948c03357ef6fe234bca 0030_analytics_and_scheduling.sql b49832e205405de7ef078f3b65c0c19f2c030b8984d42411598738b6b6c146bf 0031_proxy_and_privacy.sql a0e409402611dbc76bb71393b8eda7c717756c18974dbabca305fdcb4b115f31 0032_usage_and_scheduling.sql +6acd13f237cd163fdf9b52ec9c003ca75b723980e5db5c52c9776413cfe85c4b 0033_proxy_websocket_capacity.sql diff --git a/backend/migrations/0033_proxy_websocket_capacity.sql b/backend/migrations/0033_proxy_websocket_capacity.sql new file mode 100644 index 000000000..b21868423 --- /dev/null +++ b/backend/migrations/0033_proxy_websocket_capacity.sql @@ -0,0 +1,3 @@ +alter table outbound_proxies + add column max_websocket_connections bigint not null default 0 + check (max_websocket_connections between 0 and 4294967295); diff --git a/deploy/README.md b/deploy/README.md index db2ef405b..b1be072c2 100644 --- a/deploy/README.md +++ b/deploy/README.md @@ -302,6 +302,18 @@ OpenAI 上游池化 WebSocket 默认每 25 秒发送一次 Ping，发出后允 收到 Pong 或其他入站帧即解除本次心跳截止，持续无响应则以 `pong_timeout` 关闭连接。 此策略也覆盖正在生成的请求，与等待下一条上游消息的 `stream_idle_timeout_ms` 分别计时 +每条已保存代理可在管理端配置“上游 WS 连接上限”，默认留空、不限制连接数。 +预算限制单个 CPR 进程使用该代理的 OpenAI Responses 上游 WS 总数，关联账号共享，包含建连中、在途、空闲及非池化连接；直连不受影响。 +上限应按代理实际容量、其他流量与 CPR 实例数设置。例如 32 槽代理可配置为 24，为 HTTP 等流量留出 8 槽，但这不是统一默认值。 +不同代理 URL（含认证）分别计数；不同 URL 若共用服务商额度，需自行分配预算，各 CPR 实例之间也不协调容量。 +修改上限对后续请求生效，已存在的连接继续计数；调低上限不会强制中断在途或绑定下游 WS 的连接。 +预算仅覆盖宿主原生 OpenAI Responses 连接，插件自建连接及其他 HTTP/WS 流量仍需预留容量 + +有容量时继续复用，连接总寿命仍由 `max_age_ms` 控制，默认 55 分钟，不设置固定空闲回收时间。 +满额时优先关闭同代理下最久未使用的普通空闲 WS；在途、建连中和绑定下游 WS 的连接不参与容量淘汰。 +没有可回收连接时等待容量，受原请求期限、客户端取消和连接池关闭约束；无显式请求期限时不会额外设置排队超时。 +被回收连接上的精确续接沿用既有客户端完整回放合同，可能多一次恢复往返 + 若 Codex 在压缩或长时间生成时出现 `error decoding response body`，这表示 客户端读取 HTTP 响应体失败。请结合网关请求诊断中的 `upstream.read.failed`、 `downstream.body.closed` 和反向代理日志判断断开位置，不能仅凭此消息认定是 JSON 格式错误。 diff --git a/docs/api.md b/docs/api.md index a5e6ec48b..13297bab4 100644 --- a/docs/api.md +++ b/docs/api.md @@ -805,8 +805,8 @@ Images 按请求正文的图片模型 ID 检查，不依赖文本模型目录； | `GET` | `/api/admin/proxies` | `page`、`pageSize`（1-200）、`search`（名称） | `{ items, page }` | | `GET` | `/api/admin/proxies/accounts` | `proxyId`、`page`、`pageSize`（1-200）、`search`（账号名称或邮箱） | `{ items, page }` | | `POST` | `/api/admin/proxies/accounts/remove` | `{ proxyId, accountId }` | `{ configRevision }` | -| `POST` | `/api/admin/proxies/create` | `[{ name, proxyUrl, location?, autoLocation? }]` | 逐项结果数组，成功项的 `data` 为 `{ record, configRevision }` | -| `POST` | `/api/admin/proxies/update` | `{ id, revision, name, enabled?, proxyUrl?, location?, autoLocation? }` | `{ record, configRevision }` | +| `POST` | `/api/admin/proxies/create` | `[{ name, proxyUrl, location?, autoLocation?, maxWebsocketConnections? }]` | 逐项结果数组，成功项的 `data` 为 `{ record, configRevision }` | +| `POST` | `/api/admin/proxies/update` | `{ id, revision, name, enabled?, proxyUrl?, location?, autoLocation?, maxWebsocketConnections? }` | `{ record, configRevision }` | | `POST` | `/api/admin/proxies/probe` | `{ proxyUrl, detectLocation? }` | 测试未保存的地址，返回连通性及可选位置结果，不创建代理 | | `POST` | `/api/admin/proxies/test` | `[{ id, revision, detectLocation? }]` | 逐项结果数组，成功项的 `data` 为最新代理记录 | | `POST` | `/api/admin/proxies/delete` | `{ id, revision }` | `{ configRevision }` | @@ -820,7 +820,7 @@ Images 按请求正文的图片模型 ID 检查，不依赖文本模型目录； 调用方应按超时预算分组发送，请求中断不会回滚已提交项 `record` 包含 `id`、`name`、`enabled`、`endpoint`、`hasAuthentication`、`revision`、`accountCount`、`location`、 -`autoLocation`、`detectedLocation`、`lastTestAt`、 +`autoLocation`、`detectedLocation`、`maxWebsocketConnections`、`lastTestAt`、 `lastTest: { success, latencyMs, exitIp, exitIpv4, exitIpv6, message, location }`、`createdAt`、`updatedAt`。 未测试时 `lastTestAt` / `lastTest` 为 `null`。连通性失败返回 HTTP 200 和 `lastTest.success=false`； 记录版本过期、重复 URL 或代理正在用于账号导入时对应 409，并发测试满载对应 429； @@ -836,6 +836,10 @@ Images 按请求正文的图片模型 ID 检查，不依赖文本模型目录； 删除代理会自动解除全部关联账号的代理绑定与连接地址，账号保留其余设置并改为直连。 代理仍被插件来源或插件包引用时返回 409，账号绑定保持不变 +`maxWebsocketConnections` 为此代理的 OpenAI Responses 上游 WS 连接上限，范围 0–4294967295，0 表示不限制。 +创建时省略默认为 0，更新时省略或 `null` 保留原值。关联账号共享预算，修改后新请求使用新值，已存在的连接仍计数。 +容量回收、等待及多实例边界见[部署说明](../deploy/README.md#公网访问) + 更新省略 `proxyUrl` 保留认证；连接配置改变时清除测试结果并更新所有绑定账号。 测试结果只在请求中的版本仍匹配时保存 diff --git a/frontend/src/api/modules/proxies.ts b/frontend/src/api/modules/proxies.ts index 71749dcd0..7854b2902 100644 --- a/frontend/src/api/modules/proxies.ts +++ b/frontend/src/api/modules/proxies.ts @@ -27,6 +27,7 @@ export interface OutboundProxyTest { } export interface OutboundProxyRecord { + maxWebsocketConnections: number enabled: boolean autoLocation: boolean detectedLocation: DetectedProxyLocation | null @@ -104,7 +105,7 @@ export function getProxies(data: { page: number, pageSize: number, search?: stri }) } -export function createProxy(data: { name: string, proxyUrl: string, autoLocation?: boolean, location?: RequestLocation | null }[], options: RequestOptions = {}) { +export function createProxy(data: { name: string, proxyUrl: string, maxWebsocketConnections?: number, autoLocation?: boolean, location?: RequestLocation | null }[], options: RequestOptions = {}) { return request[]>({ url: '/api/admin/proxies/create', timeout: 25000 * Math.max(1, data.length), @@ -114,7 +115,7 @@ export function createProxy(data: { name: string, proxyUrl: string, autoLocation }) } -export function updateProxy(data: { id: string, revision: number, name: string, enabled?: boolean, proxyUrl?: string, autoLocation?: boolean, location?: RequestLocation | null }) { +export function updateProxy(data: { id: string, revision: number, name: string, enabled?: boolean, maxWebsocketConnections?: number, proxyUrl?: string, autoLocation?: boolean, location?: RequestLocation | null }) { return request({ url: '/api/admin/proxies/update', timeout: 25000, diff --git a/frontend/src/views/proxies/components/ProxyFormModal.vue b/frontend/src/views/proxies/components/ProxyFormModal.vue index 8b7123c80..86461c8ae 100644 --- a/frontend/src/views/proxies/components/ProxyFormModal.vue +++ b/frontend/src/views/proxies/components/ProxyFormModal.vue @@ -28,6 +28,7 @@ const emit = defineEmits<{ const open = defineModel({ required: true }) const name = defineModel('name', { required: true }) const proxyUrl = defineModel('proxyUrl', { required: true }) +const maxWebsocketConnections = defineModel('maxWebsocketConnections', { required: true }) const customLocation = defineModel('customLocation', { required: true }) const location = defineModel('location', { required: true }) const mode = shallowRef('single') @@ -95,6 +96,9 @@ function clearForm() { placeholder="请输入代理地址" /> + + +
diff --git a/frontend/src/views/proxies/composables/useProxies.ts b/frontend/src/views/proxies/composables/useProxies.ts index 801b5140d..00e5e72dc 100644 --- a/frontend/src/views/proxies/composables/useProxies.ts +++ b/frontend/src/views/proxies/composables/useProxies.ts @@ -24,6 +24,7 @@ export function useProxies() { const form = reactive({ name: '', proxyUrl: '', + maxWebsocketConnections: '', customLocation: false, location: { country: '', region: '', city: '', timezone: '' }, }) @@ -50,6 +51,7 @@ export function useProxies() { editing.value = proxy form.name = proxy?.name ?? '' form.proxyUrl = '' + form.maxWebsocketConnections = proxy?.maxWebsocketConnections ? String(proxy.maxWebsocketConnections) : '' formTestResult.value = null const currentLocation = proxy?.autoLocation ? proxy.detectedLocation?.location ?? proxy.location : proxy?.location form.customLocation = currentLocation != null @@ -186,6 +188,11 @@ export function useProxies() { ZNotification.warning({ message: '请填写代理名称和连接地址' }) return } + const maxWebsocketConnections = Number(form.maxWebsocketConnections) + if (!Number.isInteger(maxWebsocketConnections) || maxWebsocketConnections < 0 || maxWebsocketConnections > 4294967295) { + ZNotification.warning({ message: '连接上限须为 0 到 4294967295 之间的整数' }) + return + } const location = form.customLocation ? normalizeRequestLocation(form.location) : null const locationError = location ? requestLocationError(location) : '' if (locationError) { @@ -200,12 +207,14 @@ export function useProxies() { revision: editing.value.revision, name, autoLocation: false, + maxWebsocketConnections, proxyUrl: proxyUrl || undefined, location, }) : createProxy([{ name, autoLocation: false, + maxWebsocketConnections, proxyUrl, location, }]).then(results => proxyResultData(results[0]))) diff --git a/frontend/src/views/proxies/index.vue b/frontend/src/views/proxies/index.vue index 047a66928..2cf927d7c 100644 --- a/frontend/src/views/proxies/index.vue +++ b/frontend/src/views/proxies/index.vue @@ -164,6 +164,7 @@ watch(search, () => { v-model="showForm" v-model:name="form.name" v-model:proxy-url="form.proxyUrl" + v-model:max-websocket-connections="form.maxWebsocketConnections" v-model:custom-location="form.customLocation" v-model:location="form.location" :test-result="formTestResult"