diff --git a/README.md b/README.md index fbbd90f8a..cfaed35a0 100644 --- a/README.md +++ b/README.md @@ -21,8 +21,8 @@ ## 功能 -- **账号管理**：OAuth 授权、账号导入、凭据刷新与额度查看。 -- **访问控制**：通过账号分组分配密钥权限，设置并发和请求频率。 +- **账号管理**：OAuth 授权、账号导入、凭据刷新与额度查看；支持独立出站代理及 sub2api OpenAI 账号代理导入。 +- **访问控制**：通过账号分组分配密钥权限，设置日/周 USD 限额、并发和请求频率。详见 [账号代理与 Key 限额](docs/account-egress-key-limits.md)。 - **客户端接入**：支持 Codex CLI、桌面端与 Responses API 客户端，提供 HTTP 流式响应和 WebSocket。 - **请求诊断**：查看请求记录、用量、费用、延迟和错误详情。 - **部署运维**：Docker Compose 部署、S3/R2 数据库备份、同大版本在线更新与回滚。 diff --git a/backend/Cargo.lock b/backend/Cargo.lock index e959a653d..f22106884 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -1488,6 +1488,7 @@ dependencies = [ "sha2 0.11.0", "thiserror", "tracing", + "url", "uuid", ] @@ -2610,6 +2611,7 @@ dependencies = [ "hex", "hyper", "hyper-rustls", + "hyper-util", "reqwest", "roxmltree", "rustls", @@ -2627,6 +2629,7 @@ dependencies = [ "tokio-rustls", "tokio-tungstenite 0.28.0", "tokio-util", + "tower-service", "tracing", "tracing-subscriber", "tungstenite 0.27.0", diff --git a/backend/crates/gateway-admin/src/model/accounts.rs b/backend/crates/gateway-admin/src/model/accounts.rs index 7db36bc33..fe3e7601d 100644 --- a/backend/crates/gateway-admin/src/model/accounts.rs +++ b/backend/crates/gateway-admin/src/model/accounts.rs @@ -89,6 +89,7 @@ pub struct AccountRecord { pub enabled: bool, pub concurrency_limit: Option, pub weight: AccountWeight, + pub outbound_proxy: Option, pub credential_state: CredentialState, pub credential_observed_at: DateTime, pub quota: QuotaState, @@ -211,6 +212,7 @@ pub struct UpdateAccount { pub concurrency_limit: Option, pub weight: AccountWeight, pub group_ids: Vec, + pub outbound_proxy: Option>, } /// 账号更新结果。 @@ -228,6 +230,7 @@ pub struct BatchUpdateAccounts { pub concurrency_limit: Option, pub weight: AccountWeight, pub group_ids: Vec, + pub outbound_proxy: Option>, } /// 批量账号更新结果。 diff --git a/backend/crates/gateway-admin/src/model/client_keys.rs b/backend/crates/gateway-admin/src/model/client_keys.rs index 3176833f9..59e5eb131 100644 --- a/backend/crates/gateway-admin/src/model/client_keys.rs +++ b/backend/crates/gateway-admin/src/model/client_keys.rs @@ -5,6 +5,7 @@ use std::{fmt, num::NonZeroU16}; use chrono::{DateTime, Utc}; use gateway_core::{ + engine::budget::{ClientBudgetLimits, ClientBudgetStatus}, policy::{ClientApiKeyId, RateLimits}, routing::{AccountGroupId, ProviderKind}, }; @@ -93,6 +94,7 @@ pub struct ClientKeyRecord { pub prefix: String, pub enabled: bool, pub limits: RateLimits, + pub budget: ClientBudgetStatus, pub last_used_at: Option>, pub created_at: DateTime, pub updated_at: DateTime, @@ -146,6 +148,7 @@ pub struct CreateClientKey { pub label: Option, pub group_ids: Vec, pub limits: RateLimits, + pub budget: ClientBudgetLimits, } /// 管理用例生成 ID 与明文后的持久化命令。 @@ -156,6 +159,7 @@ pub struct NewClientKey { pub label: Option, pub group_ids: Vec, pub limits: RateLimits, + pub budget: ClientBudgetLimits, pub plaintext: String, } @@ -179,6 +183,8 @@ pub struct UpdateClientKey { pub label: Option, pub group_ids: Vec, pub limits: RateLimits, + pub daily_limit_usd: Option, + pub weekly_limit_usd: Option, } /// 修改 Client Key 启用状态。 @@ -217,3 +223,19 @@ pub struct ClientKeyMutation { pub record: Option, pub id: ClientApiKeyId, } + +#[derive(Debug, Clone)] +pub struct UnresolvedClientCharge { + pub request_id: String, + pub started_at: DateTime, + pub completed_at: Option>, + pub state: String, +} + +#[derive(Debug, Clone)] +pub struct ReconcileClientCharge { + pub key_id: ClientApiKeyId, + pub request_id: String, + pub amount_usd: gateway_core::metering::Decimal, + pub reason: String, +} diff --git a/backend/crates/gateway-admin/src/model/provider_credentials.rs b/backend/crates/gateway-admin/src/model/provider_credentials.rs index a4f0f218f..6b7e247fb 100644 --- a/backend/crates/gateway-admin/src/model/provider_credentials.rs +++ b/backend/crates/gateway-admin/src/model/provider_credentials.rs @@ -144,6 +144,7 @@ impl fmt::Debug for PrepareCredentialImport { /// Provider 已验证、可由 Store 原子创建的一份 credential。 #[derive(Debug, Clone, PartialEq)] pub struct PreparedCredentialCreate { + pub outbound_proxy: Option, pub account_id: ProviderAccountId, pub provider_kind: ProviderKind, pub name: String, @@ -249,6 +250,7 @@ pub struct PendingAuthorizationMutation { provider_kind: ProviderKind, target: AuthorizationMutationTarget, owner_binding: AuthorizationOwnerBinding, + outbound_proxy: Option, } impl PendingAuthorizationMutation { @@ -262,6 +264,7 @@ impl PendingAuthorizationMutation { provider_kind, target, owner_binding, + outbound_proxy: None, } } @@ -279,6 +282,19 @@ impl PendingAuthorizationMutation { pub const fn owner_binding(&self) -> &AuthorizationOwnerBinding { &self.owner_binding } + + #[must_use] + pub fn with_outbound_proxy( + mut self, + proxy: Option, + ) -> Self { + self.outbound_proxy = proxy; + self + } + + pub fn outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> { + self.outbound_proxy.as_ref() + } } impl PendingAuthorizationMutation { @@ -300,7 +316,7 @@ impl PendingAuthorizationMutation { AuthorizationOwner::AdminApiKey => serde_json::json!({"kind": "admin_api_key"}), AuthorizationOwner::System => serde_json::json!({"kind": "system"}), }; - Map::from_iter([ + let mut document = Map::from_iter([ ( "provider_kind".to_owned(), Value::String(self.provider_kind.as_str().to_owned()), @@ -311,7 +327,14 @@ impl PendingAuthorizationMutation { "started_request_id".to_owned(), Value::String(self.owner_binding.started_request_id().to_owned()), ), - ]) + ]); + if let Some(proxy) = &self.outbound_proxy { + document.insert( + "outbound_proxy_url".to_owned(), + Value::String(proxy.expose_url().to_owned()), + ); + } + document } /// 恢复 adapter 已确认版本为 v1 的中立事务字段。 @@ -341,6 +364,12 @@ impl PendingAuthorizationMutation { StoredAuthorizationOwnerV1::AdminApiKey => AuthorizationOwner::AdminApiKey, StoredAuthorizationOwnerV1::System => AuthorizationOwner::System, }; + let proxy = document + .outbound_proxy_url + .as_deref() + .map(gateway_core::account::OutboundProxy::parse) + .transpose() + .map_err(|_| invalid())?; Ok(Self::new( provider_kind, target, @@ -348,7 +377,8 @@ impl PendingAuthorizationMutation { owner, started_request_id: document.started_request_id, }, - )) + ) + .with_outbound_proxy(proxy)) } } @@ -359,6 +389,8 @@ struct StoredAuthorizationMutationV1 { target: StoredAuthorizationTargetV1, owner: StoredAuthorizationOwnerV1, started_request_id: String, + #[serde(default)] + outbound_proxy_url: Option, } #[derive(Deserialize)] @@ -382,6 +414,7 @@ pub struct StartAuthorization { pub context: MutationContext, pub name: String, pub reauthorization: Option, + pub outbound_proxy: Option, } /// OAuth 流程启动结果。 diff --git a/backend/crates/gateway-admin/src/ports/store.rs b/backend/crates/gateway-admin/src/ports/store.rs index 181cfe95f..9cc4cc02a 100644 --- a/backend/crates/gateway-admin/src/ports/store.rs +++ b/backend/crates/gateway-admin/src/ports/store.rs @@ -220,6 +220,29 @@ pub trait AuthStore: Send + Sync { /// Client API Key 管理写入。 #[async_trait] pub trait ClientKeyStore: Send + Sync { + async fn unresolved_charges( + &self, + _id: &gateway_core::policy::ClientApiKeyId, + ) -> AdminStoreResult> { + Err(AdminStoreError::new( + AdminStoreErrorKind::Unavailable, + "client budget", + "unsupported", + )) + } + + async fn reconcile_charge( + &self, + _command: crate::model::client_keys::ReconcileClientCharge, + _context: &MutationContext, + ) -> AdminStoreResult<()> { + Err(AdminStoreError::new( + AdminStoreErrorKind::Unavailable, + "client budget", + "unsupported", + )) + } + async fn list_client_keys(&self, query: ClientKeyListQuery) -> AdminStoreResult; async fn reveal_client_key( diff --git a/backend/crates/gateway-admin/src/use_case/client_keys.rs b/backend/crates/gateway-admin/src/use_case/client_keys.rs index 7517ef717..5a8f5b3fe 100644 --- a/backend/crates/gateway-admin/src/use_case/client_keys.rs +++ b/backend/crates/gateway-admin/src/use_case/client_keys.rs @@ -26,6 +26,19 @@ use super::{map_store_error, publish_committed}; /// API 消费的 Client Key 管理服务。 #[async_trait] pub trait ClientKeyService: Send + Sync { + async fn unresolved_charges( + &self, + _id: &ClientApiKeyId, + ) -> Result, AdminError> { + Err(AdminError::unavailable("client budget unavailable")) + } + async fn reconcile_charge( + &self, + _context: &MutationContext, + _command: crate::model::client_keys::ReconcileClientCharge, + ) -> Result<(), AdminError> { + Err(AdminError::unavailable("client budget unavailable")) + } async fn list(&self, query: ClientKeyListQuery) -> Result; async fn reveal(&self, id: &ClientApiKeyId) -> Result; async fn create( @@ -64,6 +77,34 @@ impl DefaultClientKeyService { #[async_trait] impl ClientKeyService for DefaultClientKeyService { + async fn unresolved_charges( + &self, + id: &ClientApiKeyId, + ) -> Result, AdminError> { + self.store + .unresolved_charges(id) + .await + .map_err(|error| map_store_error(error, "client budget")) + } + + async fn reconcile_charge( + &self, + context: &MutationContext, + command: crate::model::client_keys::ReconcileClientCharge, + ) -> Result<(), AdminError> { + if command.request_id.is_empty() + || command.request_id.len() > 128 + || command.reason.trim().is_empty() + || command.reason.len() > 1024 + || command.reason.chars().any(char::is_control) + { + return Err(AdminError::invalid("invalid charge reconciliation")); + } + self.store + .reconcile_charge(command, context) + .await + .map_err(|error| map_store_error(error, "client budget")) + } async fn list(&self, query: ClientKeyListQuery) -> Result { validate_cursor(&query)?; self.store @@ -99,6 +140,7 @@ impl ClientKeyService for DefaultClientKeyService { label: command.label, group_ids: command.group_ids, limits: command.limits, + budget: command.budget, plaintext: plaintext.clone(), }, context, diff --git a/backend/crates/gateway-admin/src/use_case/mod.rs b/backend/crates/gateway-admin/src/use_case/mod.rs index 4223529d2..e871bc3dc 100644 --- a/backend/crates/gateway-admin/src/use_case/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/mod.rs @@ -115,9 +115,13 @@ async fn pending_authorization( command: &StartAuthorization, resource: &'static str, ) -> Result { + let mut proxy = command.outbound_proxy.clone(); let target = match &command.reauthorization { Some(account_id) => { - required_credential(accounts, provider_kind, account_id, resource).await?; + proxy = required_credential(accounts, provider_kind, account_id, resource) + .await? + .credential + .outbound_proxy; AuthorizationMutationTarget::Reauthorize { account_id: account_id.clone(), } @@ -132,7 +136,8 @@ async fn pending_authorization( crate::model::provider_credentials::AuthorizationOwnerBinding::from_context( &command.context, ), - )) + ) + .with_outbound_proxy(proxy)) } fn validate_prepared_import( diff --git a/backend/crates/gateway-admin/tests/use_case/accounts.rs b/backend/crates/gateway-admin/tests/use_case/accounts.rs index 6f5f35e9c..28b1a8647 100644 --- a/backend/crates/gateway-admin/tests/use_case/accounts.rs +++ b/backend/crates/gateway-admin/tests/use_case/accounts.rs @@ -1213,6 +1213,7 @@ async fn accounts_update_should_commit_then_release_disabled_account_and_publish .update( &context("update-request"), UpdateAccount { + outbound_proxy: None, account_id: "acct_test".to_owned(), enabled: false, concurrency_limit: None, @@ -1249,6 +1250,7 @@ async fn accounts_update_should_not_notify_provider_when_store_commit_fails() { .update( &context("update-failure"), UpdateAccount { + outbound_proxy: None, account_id: "acct_test".to_owned(), enabled: false, concurrency_limit: None, @@ -1290,6 +1292,7 @@ async fn accounts_batch_update_should_commit_once_and_notify_each_provider() { .batch_update( &context("batch-update-request"), BatchUpdateAccounts { + outbound_proxy: None, account_ids: vec!["acct_openai".to_owned(), "acct_xai".to_owned()], enabled: false, concurrency_limit: None, @@ -1931,6 +1934,7 @@ fn quota_local_usage(account_id: &str, total_tokens: u64) -> AccountUsage { pub(super) fn account_record(kind: &str) -> AccountRecord { let now = Utc::now(); AccountRecord { + outbound_proxy: None, id: "acct_test".to_owned(), provider_kind: ProviderKind::new(kind).expect("provider kind"), groups: Vec::new(), @@ -1972,6 +1976,7 @@ fn prepared_create_with_id( ) -> PreparedCredentialCreate { let now = Utc::now(); PreparedCredentialCreate { + outbound_proxy: None, account_id: ProviderAccountId::new(account_id).expect("prepared account ID"), provider_kind, name: name.to_owned(), diff --git a/backend/crates/gateway-admin/tests/use_case/openai.rs b/backend/crates/gateway-admin/tests/use_case/openai.rs index 3d83fd2d7..836d8ff01 100644 --- a/backend/crates/gateway-admin/tests/use_case/openai.rs +++ b/backend/crates/gateway-admin/tests/use_case/openai.rs @@ -157,6 +157,7 @@ async fn openai_authorization_create_should_observe_initial_quota() { services .openai() .start_authorization(StartAuthorization { + outbound_proxy: None, context: context("oauth-start-openai-create"), name: "new OpenAI credential".to_owned(), reauthorization: None, @@ -207,6 +208,7 @@ async fn openai_authorization_create_should_remain_successful_when_initial_quota services .openai() .start_authorization(StartAuthorization { + outbound_proxy: None, context: context("oauth-start-openai-quota-failure"), name: "new OpenAI credential".to_owned(), reauthorization: None, @@ -246,6 +248,7 @@ async fn openai_authorization_store_failure_should_release_claim_for_retry() { services .openai() .start_authorization(StartAuthorization { + outbound_proxy: None, context: context("oauth-store-retry"), name: "new OpenAI credential".to_owned(), reauthorization: None, @@ -311,6 +314,7 @@ async fn openai_reauthorization_should_commit_after_credential_revision_advances services .openai() .start_authorization(StartAuthorization { + outbound_proxy: None, context: context("oauth-start-openai"), name: "reauthorize".to_owned(), reauthorization: Some(ProviderAccountId::new("acct_test").expect("account ID")), diff --git a/backend/crates/gateway-api/src/admin/accounts/credentials.rs b/backend/crates/gateway-api/src/admin/accounts/credentials.rs index 0895f6279..bd9efc9f5 100644 --- a/backend/crates/gateway-api/src/admin/accounts/credentials.rs +++ b/backend/crates/gateway-api/src/admin/accounts/credentials.rs @@ -60,6 +60,7 @@ pub struct StartAccountAuthorizationRequest { pub provider: String, pub name: String, pub account_id: Option, + pub outbound_proxy_url: Option, } impl StartAccountAuthorizationRequest { @@ -89,6 +90,7 @@ impl StartAccountAuthorizationRequest { context, name: self.name, reauthorization, + outbound_proxy: self.outbound_proxy_url.and_then(|value| value.0), }, )) } @@ -139,6 +141,7 @@ impl CompleteAccountAuthorizationRequest { #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct UpdateAccountRequest { + pub outbound_proxy_url: Option, pub account_id: String, pub enabled: bool, #[serde(deserialize_with = "deserialize_required_nullable")] @@ -159,6 +162,7 @@ impl UpdateAccountRequest { pub(super) fn into_command(self) -> Result { self.validate()?; Ok(UpdateAccount { + outbound_proxy: self.outbound_proxy_url.map(|value| value.0), account_id: self.account_id, enabled: self.enabled, concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?, diff --git a/backend/crates/gateway-api/src/admin/accounts/presenter.rs b/backend/crates/gateway-api/src/admin/accounts/presenter.rs index d15c7fbfa..3be2ca21b 100644 --- a/backend/crates/gateway-api/src/admin/accounts/presenter.rs +++ b/backend/crates/gateway-api/src/admin/accounts/presenter.rs @@ -99,6 +99,10 @@ pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime) -> Ac enabled: account.enabled, concurrency_limit: account.concurrency_limit.map(|limit| limit.get()), weight: account.weight.get(), + outbound_proxy_endpoint: account + .outbound_proxy + .as_ref() + .map(|proxy| proxy.endpoint()), access_token_expires_at: expires_at, access_token_expires_at_display: account .access_token_expires_at diff --git a/backend/crates/gateway-api/src/admin/accounts/wire.rs b/backend/crates/gateway-api/src/admin/accounts/wire.rs index bcc94cfd6..6eca72108 100644 --- a/backend/crates/gateway-api/src/admin/accounts/wire.rs +++ b/backend/crates/gateway-api/src/admin/accounts/wire.rs @@ -2,6 +2,21 @@ use super::*; +#[derive(Debug, Clone)] +pub struct AccountProxyUpdate(pub Option); + +impl<'de> Deserialize<'de> for AccountProxyUpdate { + fn deserialize>(deserializer: D) -> Result { + let value = String::deserialize(deserializer)?; + if value.is_empty() { + return Ok(Self(None)); + } + gateway_core::account::OutboundProxy::parse(&value) + .map(|proxy| Self(Some(proxy))) + .map_err(serde::de::Error::custom) + } +} + /// 账号列表查询参数。 #[derive(Debug, Clone, Default, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -19,6 +34,7 @@ pub struct ListQuery { #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct BatchUpdateAccountsRequest { + pub outbound_proxy_url: Option, pub account_ids: Vec, pub enabled: bool, #[serde(deserialize_with = "deserialize_required_nullable")] @@ -68,6 +84,7 @@ impl BatchUpdateAccountsRequest { pub(super) fn into_command(self) -> Result { self.validate()?; Ok(BatchUpdateAccounts { + outbound_proxy: self.outbound_proxy_url.map(|value| value.0), account_ids: self.account_ids, enabled: self.enabled, concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?, @@ -192,6 +209,7 @@ pub struct AccountSummaryView { #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct AccountView { + pub outbound_proxy_endpoint: Option, pub id: String, pub name: String, pub provider: String, diff --git a/backend/crates/gateway-api/src/admin/client_keys.rs b/backend/crates/gateway-api/src/admin/client_keys.rs index c72d23ec6..9e0952949 100644 --- a/backend/crates/gateway-api/src/admin/client_keys.rs +++ b/backend/crates/gateway-api/src/admin/client_keys.rs @@ -11,6 +11,8 @@ use gateway_admin::model::client_keys::{ CreatedClientKey, DeleteClientKey, SetClientKeyEnabled, SortDirection, UpdateClientKey, }; use gateway_core::{ + engine::budget::ClientBudgetLimits, + metering::Decimal, policy::{ClientApiKeyId, RateLimits}, routing::AccountGroupId, }; @@ -33,6 +35,15 @@ const MAX_CURSOR_BYTES: usize = 512; const MAX_SEARCH_BYTES: usize = 256; const DEFAULT_PAGE_SIZE: u16 = 50; +fn parse_budget( + value: Option, + field: &'static str, +) -> Result, WireValidationError> { + value + .map(|value| value.parse().map_err(|_| WireValidationError::new(field))) + .transpose() +} + /// Client Key 列表查询。 #[derive(Debug, Clone, PartialEq, Eq, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] @@ -140,6 +151,8 @@ pub struct CreateClientKeyRequest { group_ids: Vec, max_concurrency: u64, requests_per_minute: u64, + daily_limit_usd: Option, + weekly_limit_usd: Option, } impl CreateClientKeyRequest { @@ -154,6 +167,11 @@ impl CreateClientKeyRequest { name: self.name, label: self.label, group_ids, + budget: ClientBudgetLimits { + daily_usd: parse_budget(self.daily_limit_usd, "dailyLimitUsd")?.unwrap_or_default(), + weekly_usd: parse_budget(self.weekly_limit_usd, "weeklyLimitUsd")? + .unwrap_or_default(), + }, limits: RateLimits { max_concurrency: self.max_concurrency, requests_per_minute: self.requests_per_minute, @@ -172,6 +190,8 @@ pub struct UpdateClientKeyRequest { group_ids: Vec, max_concurrency: u64, requests_per_minute: u64, + daily_limit_usd: Option, + weekly_limit_usd: Option, } impl UpdateClientKeyRequest { @@ -188,6 +208,8 @@ impl UpdateClientKeyRequest { name: self.name, label: self.label, group_ids, + daily_limit_usd: parse_budget(self.daily_limit_usd, "dailyLimitUsd")?, + weekly_limit_usd: parse_budget(self.weekly_limit_usd, "weeklyLimitUsd")?, limits: RateLimits { max_concurrency: self.max_concurrency, requests_per_minute: self.requests_per_minute, @@ -247,6 +269,13 @@ pub struct ClientKeyView { enabled: bool, max_concurrency: u64, requests_per_minute: u64, + daily_limit_usd: String, + weekly_limit_usd: String, + daily_used_usd: String, + weekly_used_usd: String, + daily_resets_at: Option>, + weekly_resets_at: Option>, + unresolved_requests: u64, created_at: DateTime, updated_at: DateTime, last_used_at: Option>, @@ -292,6 +321,13 @@ impl From for ClientKeyView { enabled: record.enabled, max_concurrency: record.limits.max_concurrency, requests_per_minute: record.limits.requests_per_minute, + daily_limit_usd: record.budget.limits.daily_usd.canonical(), + weekly_limit_usd: record.budget.limits.weekly_usd.canonical(), + daily_used_usd: record.budget.daily_used_usd.canonical(), + weekly_used_usd: record.budget.weekly_used_usd.canonical(), + daily_resets_at: record.budget.daily_resets_at.map(DateTime::from), + weekly_resets_at: record.budget.weekly_resets_at.map(DateTime::from), + unresolved_requests: record.budget.unresolved_requests, created_at: record.created_at, updated_at: record.updated_at, last_used_at: record.last_used_at, @@ -643,6 +679,14 @@ where { Router::new() .route("/api/admin/client-keys", get(list_client_keys::)) + .route( + "/api/admin/client-keys/unresolved-charges", + get(unresolved_client_charges::), + ) + .route( + "/api/admin/client-keys/reconcile-charge", + post(reconcile_client_charge::), + ) .route( "/api/admin/client-keys/create", post(create_client_key::), @@ -666,6 +710,66 @@ where ) } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct ReconcileClientChargeRequest { + id: String, + request_id: String, + amount_usd: String, + reason: String, +} + +async fn unresolved_client_charges( + _auth: AdminAuth, + State(state): State, + AdminQuery(query): AdminQuery, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let id = query.into_domain_id().map_err(map_wire_error)?; + let items = state.admin_services().client_keys().unresolved_charges(&id).await.map_err(map_service_error)? + .into_iter().map(|event| serde_json::json!({ + "requestId": event.request_id, "startedAt": event.started_at.to_rfc3339(), + "completedAt": event.completed_at.map(|time| time.to_rfc3339()), "state": event.state, + })).collect::>(); + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(serde_json::json!({"items": items})), + )) +} + +async fn reconcile_client_charge( + auth: AdminAuth, + State(state): State, + AdminJson(payload): AdminJson, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + let key_id = client_key_id(payload.id, "id").map_err(map_wire_error)?; + let amount_usd = payload + .amount_usd + .parse::() + .map_err(|_| map_wire_error(WireValidationError::new("amountUsd")))?; + let command = gateway_admin::model::client_keys::ReconcileClientCharge { + key_id, + request_id: payload.request_id, + amount_usd, + reason: payload.reason, + }; + state + .admin_services() + .client_keys() + .reconcile_charge(&auth.context().mutation_context(), command) + .await + .map_err(map_service_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(serde_json::json!({"reconciled": true})), + )) +} + async fn list_client_keys( _auth: AdminAuth, State(state): State, diff --git a/backend/crates/gateway-api/src/openai/error.rs b/backend/crates/gateway-api/src/openai/error.rs index 565efeda8..ee53f8003 100644 --- a/backend/crates/gateway-api/src/openai/error.rs +++ b/backend/crates/gateway-api/src/openai/error.rs @@ -143,13 +143,27 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError { /// Gateway 错误的 OpenAI HTTP 表达。 pub fn gateway_error_response(error: &GatewayError) -> Response { let (status, default_type, default_code) = gateway_error_contract(error.kind()); - openai_error_response( + let mut response = openai_error_response( status, error.client_message(), error.client_error_type().unwrap_or(default_type), error.client_error_code().unwrap_or(default_code), ) - .into_response() + .into_response(); + if let Some(delay) = error.retry_after() + && let Ok(value) = HeaderValue::from_str( + &delay + .as_secs() + .saturating_add(u64::from(delay.subsec_nanos() > 0)) + .max(1) + .to_string(), + ) + { + response + .headers_mut() + .insert(axum::http::header::RETRY_AFTER, value); + } + response } /// 在下游尚未提交时优先交付 Provider 的原始 HTTP 失败响应。 diff --git a/backend/crates/gateway-api/tests/admin/client_keys.rs b/backend/crates/gateway-api/tests/admin/client_keys.rs index 97c9b5234..f53cf6014 100644 --- a/backend/crates/gateway-api/tests/admin/client_keys.rs +++ b/backend/crates/gateway-api/tests/admin/client_keys.rs @@ -10,6 +10,84 @@ use serde_json::json; use super::{AdminTestFixture, AdminTestState}; +#[test] +fn budget_inputs_preserve_decimal_precision_and_omitted_updates() { + let payload = json!({"name": "budget", "groupIds": [], "maxConcurrency": 3, + "requestsPerMinute": 0, "dailyLimitUsd": "0.1234567891", "weeklyLimitUsd": "15"}); + let command = serde_json::from_value::(payload.clone()) + .unwrap() + .into_command() + .unwrap(); + assert_eq!(command.budget.daily_usd.canonical(), "0.1234567891"); + assert_eq!(command.budget.weekly_usd.canonical(), "15"); + assert_eq!(command.limits.max_concurrency, 3); + for field in ["dailyLimitUsd", "weeklyLimitUsd"] { + for invalid in ["-1", "NaN", "1e3", "10000000000", "0.00000000001", ""] { + let mut payload = payload.clone(); + payload[field] = json!(invalid); + assert_eq!( + serde_json::from_value::(payload) + .unwrap() + .into_command() + .unwrap_err() + .field(), + field + ); + } + } + let mut update = payload; + update["id"] = json!("key_budget"); + update.as_object_mut().unwrap().remove("dailyLimitUsd"); + update["weeklyLimitUsd"] = json!("0"); + let command = serde_json::from_value::(update) + .unwrap() + .into_command() + .unwrap(); + assert_eq!(command.daily_limit_usd, None); + assert_eq!( + command.weekly_limit_usd, + Some(gateway_core::metering::Decimal::ZERO) + ); +} + +#[tokio::test] +async fn budget_reconciliation_routes_require_an_admin_session() { + use axum::{ + body::Body, + http::{Request, StatusCode}, + }; + use tower::ServiceExt as _; + let fixture = AdminTestFixture::new().await; + let router = client_keys::router::().with_state(fixture.state()); + for (method, path, body) in [ + ( + "GET", + "/api/admin/client-keys/unresolved-charges?id=key_1", + "", + ), + ( + "POST", + "/api/admin/client-keys/reconcile-charge", + r#"{"id":"key_1","requestId":"req_charge","amountUsd":"0","reason":"verified"}"#, + ), + ] { + let response = router + .clone() + .oneshot( + Request::builder() + .method(method) + .uri(path) + .header("content-type", "application/json") + .header("x-request-id", "req_auth_budget") + .body(Body::from(body)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + } +} + #[test] fn client_key_queries_should_reject_unknown_zero_and_oversized_fields() { let unknown = serde_json::from_value::(json!({ "other": true })); @@ -207,6 +285,7 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() { .single() .expect("valid time"); let view = ClientKeyView::from(gateway_admin::model::client_keys::ClientKeyRecord { + budget: Default::default(), id: gateway_core::policy::ClientApiKeyId::new("key_visible").expect("Client Key ID"), name: "visible".to_owned(), label: None, @@ -238,6 +317,10 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() { assert_eq!(list["items"][0]["requestsPerMinute"], 60); assert!(list["items"][0].get("tokensPerMinute").is_none()); assert!(list["items"][0].get("policyJson").is_none()); + assert_eq!(list["items"][0]["dailyLimitUsd"], "0"); + assert_eq!(list["items"][0]["weeklyLimitUsd"], "0"); + assert_eq!(list["items"][0]["dailyUsedUsd"], "0"); + assert_eq!(list["items"][0]["unresolvedRequests"], 0); assert_eq!( DateTime::parse_from_rfc3339( list["items"][0]["lastUsedAt"] diff --git a/backend/crates/gateway-api/tests/admin/mod.rs b/backend/crates/gateway-api/tests/admin/mod.rs index b5c8d54f3..cdc9e339f 100644 --- a/backend/crates/gateway-api/tests/admin/mod.rs +++ b/backend/crates/gateway-api/tests/admin/mod.rs @@ -640,6 +640,7 @@ impl ClientKeyStore for MemoryClientKeyStore { let now = Utc::now(); Ok(Some(ClientKeySecret::new( ClientKeyRecord { + budget: Default::default(), id: id.clone(), name: "revealed".to_owned(), label: None, diff --git a/backend/crates/gateway-api/tests/openai/error.rs b/backend/crates/gateway-api/tests/openai/error.rs index 2b54af1a0..fceae885e 100644 --- a/backend/crates/gateway-api/tests/openai/error.rs +++ b/backend/crates/gateway-api/tests/openai/error.rs @@ -11,6 +11,21 @@ use gateway_api::openai::error::{ openai_error_response, }; +#[tokio::test] +async fn key_budget_errors_preserve_limit_code_and_retry_after() { + for code in ["key_daily_budget_exceeded", "key_weekly_budget_exceeded"] { + let error = GatewayError::new(GatewayErrorKind::RateLimited, "key budget exhausted") + .with_client_code(code) + .with_retry_after(std::time::Duration::from_millis(1501)); + let response = gateway_error_response(&error); + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(response.headers()["retry-after"], "2"); + let body: serde_json::Value = + serde_json::from_slice(&to_bytes(response.into_body(), 4096).await.unwrap()).unwrap(); + assert_eq!(body["error"]["code"], code); + } +} + #[test] fn invalid_request_error_should_map_to_openai_bad_request() { assert_eq!( diff --git a/backend/crates/gateway-core/Cargo.toml b/backend/crates/gateway-core/Cargo.toml index cb8c5ccbf..ce72ee456 100644 --- a/backend/crates/gateway-core/Cargo.toml +++ b/backend/crates/gateway-core/Cargo.toml @@ -24,6 +24,7 @@ semver = "1.0.27" thiserror = "2.0.18" tracing = "0.1.44" uuid = { version = "1.23.3", features = ["v7"] } +url = "2.5.8" [lints] workspace = true diff --git a/backend/crates/gateway-core/src/account/mod.rs b/backend/crates/gateway-core/src/account/mod.rs index 348306157..c6fb4472e 100644 --- a/backend/crates/gateway-core/src/account/mod.rs +++ b/backend/crates/gateway-core/src/account/mod.rs @@ -2,6 +2,8 @@ mod error; mod model; +mod proxy; +pub use proxy::{InvalidOutboundProxy, OutboundProxy}; pub mod scope; mod selection; mod store; diff --git a/backend/crates/gateway-core/src/account/model.rs b/backend/crates/gateway-core/src/account/model.rs index ee97e6832..5ce2762c3 100644 --- a/backend/crates/gateway-core/src/account/model.rs +++ b/backend/crates/gateway-core/src/account/model.rs @@ -654,7 +654,7 @@ const fn status_projection(status: AccountStatus) -> AccountStatusProjection { } } -/// 不含 secret 的账号持久事实。 +/// 账号持久事实；代理认证信息只通过显式 secret accessor 读取。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct ProviderAccount { id: ProviderAccountId, @@ -676,6 +676,7 @@ pub struct ProviderAccount { access_token_expires_at: Option, next_refresh_at: Option, has_refresh_token: bool, + outbound_proxy: Option, } impl ProviderAccount { @@ -710,6 +711,7 @@ impl ProviderAccount { access_token_expires_at, next_refresh_at: None, has_refresh_token: false, + outbound_proxy: None, } } @@ -726,6 +728,17 @@ impl ProviderAccount { self } + #[must_use] + pub fn with_outbound_proxy(mut self, proxy: Option) -> Self { + self.outbound_proxy = proxy; + self + } + + #[must_use] + pub const fn outbound_proxy(&self) -> Option<&super::OutboundProxy> { + self.outbound_proxy.as_ref() + } + #[must_use] pub fn with_account_facts( mut self, diff --git a/backend/crates/gateway-core/src/account/proxy.rs b/backend/crates/gateway-core/src/account/proxy.rs new file mode 100644 index 000000000..bbb02edf0 --- /dev/null +++ b/backend/crates/gateway-core/src/account/proxy.rs @@ -0,0 +1,49 @@ +//! Explicit account egress. Credentials never appear in Debug or ordinary admin projections. + +use std::fmt; +use url::Url; + +#[derive(Clone, PartialEq, Eq, Hash)] +pub struct OutboundProxy(Url); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +#[error("invalid outbound proxy; expected http, https, socks5 or socks5h URL with a host and port")] +pub struct InvalidOutboundProxy; + +impl OutboundProxy { + pub fn parse(value: &str) -> Result { + if value.len() > 4096 || value.chars().any(char::is_control) { + return Err(InvalidOutboundProxy); + } + let url = Url::parse(value).map_err(|_| InvalidOutboundProxy)?; + if !matches!(url.scheme(), "http" | "https" | "socks5" | "socks5h") + || url.host_str().is_none() + || url.port_or_known_default().is_none_or(|port| port == 0) + || !matches!(url.path(), "" | "/") + || url.query().is_some() + || url.fragment().is_some() + { + return Err(InvalidOutboundProxy); + } + Ok(Self(url)) + } + + #[must_use] + pub fn expose_url(&self) -> &str { + self.0.as_str() + } + + #[must_use] + pub fn endpoint(&self) -> String { + let mut url = self.0.clone(); + let _ = url.set_username(""); + let _ = url.set_password(None); + url.to_string() + } +} + +impl fmt::Debug for OutboundProxy { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("OutboundProxy()") + } +} diff --git a/backend/crates/gateway-core/src/engine/budget.rs b/backend/crates/gateway-core/src/engine/budget.rs new file mode 100644 index 000000000..92126d6cb --- /dev/null +++ b/backend/crates/gateway-core/src/engine/budget.rs @@ -0,0 +1,59 @@ +//! Durable per-key USD budgets, independent of best-effort request observations. + +use std::time::SystemTime; + +use futures::future::BoxFuture; + +use crate::{error::GatewayError, metering::Decimal, policy::ClientApiKeyId}; + +use super::ModelRequestId; + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ClientBudgetLimits { + pub daily_usd: Decimal, + pub weekly_usd: Decimal, +} + +impl ClientBudgetLimits { + #[must_use] + pub fn is_limited(self) -> bool { + self.daily_usd != Decimal::ZERO || self.weekly_usd != Decimal::ZERO + } +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ClientBudgetStatus { + pub limits: ClientBudgetLimits, + pub daily_used_usd: Decimal, + pub weekly_used_usd: Decimal, + pub daily_resets_at: Option, + pub weekly_resets_at: Option, + pub unresolved_requests: u64, +} + +#[derive(Debug, Clone)] +pub struct ClientBudgetAdmission { + pub key_id: ClientApiKeyId, + pub request_id: ModelRequestId, + pub deadline_at: SystemTime, +} + +#[derive(Debug, Clone)] +pub struct ClientBudgetCharge { + pub request_id: ModelRequestId, + /// None means the request may have incurred an unknown charge, never zero. + pub amount_usd: Option, + pub completed_at: SystemTime, +} + +#[derive(Debug, Clone, Copy, thiserror::Error)] +#[error("client budget store is unavailable")] +pub struct ClientBudgetError; + +pub trait ClientBudgetPort: Send + Sync { + /// Persist a pending charge before any upstream work; read the current policy atomically. + fn admit(&self, request: ClientBudgetAdmission) -> BoxFuture<'_, Result<(), GatewayError>>; + + /// Idempotent by gateway request ID. Failed writes leave the durable pending charge intact. + fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>>; +} diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index f9366dc6e..bdd9250bf 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -30,6 +30,7 @@ use crate::event::{ GatewayEvent, ProviderEvent, ProviderResponseHeader, ProviderResponseObservation, }; use crate::lifecycle::CancellationToken; +use crate::metering::Decimal; use crate::operation::{Operation, ProviderSessionState}; use crate::routing::RoutingPlan; use futures::future::Fuse; @@ -166,6 +167,9 @@ where request_id, client_api_key_ref, observation: ResponseObservation::new(timing_started_at), + budget_prior_attempts_usd: Some(Decimal::ZERO), + budget_attempt_send_state: UpstreamSendState::NotSent, + budget_attempt_already_counted: false, trace, deadline, deadline_timer: Delay::new( @@ -197,6 +201,7 @@ where delivery_pending: false, upstream_complete: false, finalized: false, + finalized_at: None, image_generation_requested, last_retryable_failure: None, last_retryable_failure_events: Vec::new(), @@ -253,6 +258,9 @@ pub struct ResponseExecutionSession { request_id: ModelRequestId, client_api_key_ref: crate::policy::ClientApiKeyId, observation: ResponseObservation, + budget_prior_attempts_usd: Option, + budget_attempt_send_state: UpstreamSendState, + budget_attempt_already_counted: bool, trace: TraceContext, deadline: SystemTime, /// 会话级 deadline 计时器；deadline 固定，帧循环内复用而非逐事件新建。 @@ -288,6 +296,7 @@ pub struct ResponseExecutionSession { delivery_pending: bool, upstream_complete: bool, finalized: bool, + finalized_at: Option, image_generation_requested: bool, /// 最近一次为无感恢复而被丢弃的原始上游错误；只在后续空选路时成为终态。 last_retryable_failure: Option, @@ -475,6 +484,34 @@ where self.finalized } + #[must_use] + pub fn budget_charge(&self) -> super::budget::ClientBudgetCharge { + let amount_usd = self + .budget_prior_attempts_usd + .zip(self.budget_attempt_usd()) + .and_then(|(prior, current)| prior.checked_add(current)); + super::budget::ClientBudgetCharge { + request_id: self.request_id.clone(), + amount_usd, + completed_at: self.finalized_at.unwrap_or_else(SystemTime::now), + } + } + + fn budget_attempt_usd(&self) -> Option { + if self.budget_attempt_already_counted { + return Some(Decimal::ZERO); + } + self.observation + .cost + .total() + .filter(|money| money.currency().as_str() == "USD") + .map(crate::metering::Money::amount) + .or_else(|| { + (self.budget_attempt_send_state == UpstreamSendState::NotSent) + .then_some(Decimal::ZERO) + }) + } + /// 返回截至当前已完成的实际上游调用结果。 /// /// 调用方可保存已消费下标；该切片在会话生命周期内只追加、不改写。 @@ -749,6 +786,8 @@ where }), ); let provider_request = ProviderRequest::new(self.operation.clone(), candidate.clone()); + // Cancellation while a provider is opening cannot prove the request was never sent. + self.budget_attempt_send_state = UpstreamSendState::Ambiguous; let stream = match poll_provider( provider, provider_request, @@ -772,6 +811,7 @@ where ProviderBoundary::Result(result) => match *result { Ok(stream) => stream, Err(error) => { + self.budget_attempt_send_state = error.send_state(); record_trace_error(&attempt_trace, &error); if self.prepare_unavailable_native_continuation_replay(&error) { return Ok(Some(PullOutcome::AttemptDiscarded)); @@ -800,6 +840,7 @@ where let mut events = std::mem::take(&mut self.last_retryable_failure_events); if !events.is_empty() { self.observe_atomic_terminal_events(&mut events); + self.budget_attempt_already_counted = true; return Ok(Some(PullOutcome::TerminalFailure { events, error: last_failure, @@ -1027,6 +1068,7 @@ where .await; } current.send_observed = true; + self.budget_attempt_send_state = UpstreamSendState::Sent; self.send_state_watermark = UpstreamSendState::Sent; } } @@ -1066,6 +1108,7 @@ where error.send_state() }; let send_state = self.raise_send_watermark(attempt_send_state); + self.budget_attempt_send_state = attempt_send_state; if self.request_persisted { best_effort_store_write( "mark_send_state", @@ -1143,6 +1186,11 @@ where "sendState": format!("{attempt_send_state:?}"), })); if retryable { + for event in &atomic_client_events { + for fact in event.canonical_facts() { + self.observation.observe_event(fact); + } + } // 普通 clone 只保留稳定事实；原始 wire/HTTP response 由 request-local // 所有权保留到下一次 attempt 成功，或最终空选路时返回客户端。 let persistence_error = self.request_persisted.then(|| error.clone()); @@ -1317,6 +1365,13 @@ where } fn reset_uncommitted_observations(&mut self) { + // Response facts are attempt-local; the key must also pay for discarded attempts. + self.budget_prior_attempts_usd = self + .budget_prior_attempts_usd + .zip(self.budget_attempt_usd()) + .and_then(|(prior, current)| prior.checked_add(current)); + self.budget_attempt_send_state = UpstreamSendState::NotSent; + self.budget_attempt_already_counted = false; self.observation.reset_for_attempt(); self.upstream_complete = false; } @@ -1326,6 +1381,7 @@ where return Ok(()); } let completed_at = SystemTime::now(); + self.finalized_at = Some(completed_at); self.observation.finish(); let upstream_request_id = self.current.as_ref().and_then(|current| { current @@ -1505,11 +1561,12 @@ where "sendState": format!("{:?}", finalization.send_state), "attempts": self.attempts, }), ); + let completed_at = SystemTime::now(); + self.finalized_at = Some(completed_at); if !self.request_persisted { self.finalized = true; return Ok(()); } - let completed_at = SystemTime::now(); self.observation.finish(); let upstream_request_id = self.current.as_ref().and_then(|current| { current diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs index 0325eb6a6..7d134d42b 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -14,6 +14,7 @@ use uuid::Uuid; use crate::engine::admission::{ ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest, }; +use crate::engine::budget::{ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetPort}; use crate::engine::continuation::{ ContinuationBinding, NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreErrorKind, PreviousResponseId, @@ -256,6 +257,7 @@ pub struct DefaultExecutionService { circuits: Arc, continuation: Arc, client_api_key_usage: Arc, + budget: Option>, } impl DefaultExecutionService { @@ -283,9 +285,16 @@ impl DefaultExecutionService { circuits, continuation, client_api_key_usage, + budget: None, } } + #[must_use] + pub fn with_budget(mut self, budget: Arc) -> Self { + self.budget = Some(budget); + self + } + async fn start_inner(&self, request: StartExecution) -> Result { let StartExecution { client, @@ -504,6 +513,18 @@ impl DefaultExecutionService { client_api_key_id: client.policy.key_id().clone(), model_request_id: request_id.clone(), }; + if let Some(budget) = &self.budget + && let Err(error) = budget + .admit(ClientBudgetAdmission { + key_id: client.policy.key_id().clone(), + request_id: request_id.clone(), + deadline_at, + }) + .await + { + admission.release().await; + return Err(error); + } let observation = plan .candidates() .first() @@ -552,6 +573,17 @@ impl DefaultExecutionService { { Ok(core) => core, Err(error) => { + if let Some(budget) = &self.budget { + settle_budget( + budget.as_ref(), + ClientBudgetCharge { + request_id: request_id.clone(), + amount_usd: Some(crate::metering::Decimal::ZERO), + completed_at: SystemTime::now(), + }, + ) + .await; + } admission.release().await; return Err(gateway_error_from_engine(&error)); } @@ -565,6 +597,7 @@ impl DefaultExecutionService { admission, Arc::clone(&self.circuits), Arc::clone(&self.continuation), + self.budget.clone(), )), }) } @@ -922,6 +955,12 @@ struct AdmissionLease { model_request_id: ModelRequestId, } +async fn settle_budget(port: &dyn ClientBudgetPort, charge: ClientBudgetCharge) { + if let Err(error) = port.settle(charge).await { + tracing::error!(%error, "Client budget settlement pending; durable admission requires reconciliation"); + } +} + impl AdmissionLease { async fn release(self) { if let Err(error) = self @@ -941,6 +980,7 @@ struct DefaultExecutionSession { continuation: Arc, observed_provider_outcomes: usize, continuation_recorded: bool, + budget: Option>, } impl DefaultExecutionSession { @@ -949,6 +989,7 @@ impl DefaultExecutionSession { admission: AdmissionLease, circuits: Arc, continuation: Arc, + budget: Option>, ) -> Self { Self { core: Some(core), @@ -957,6 +998,7 @@ impl DefaultExecutionSession { continuation, observed_provider_outcomes: 0, continuation_recorded: false, + budget, } } @@ -973,6 +1015,11 @@ impl DefaultExecutionSession { .is_some_and(ResponseExecutionSession::is_finalized) && let Some(admission) = self.admission.take() { + if let Some(budget) = self.budget.take() + && let Some(core) = &self.core + { + settle_budget(budget.as_ref(), core.budget_charge()).await; + } admission.release().await; } } @@ -1023,6 +1070,9 @@ impl DefaultExecutionSession { .get(self.observed_provider_outcomes..) .unwrap_or_default(); publish_provider_attempt_outcomes(self.circuits.as_ref(), pending).await; + if let Some(budget) = self.budget.take() { + settle_budget(budget.as_ref(), core.budget_charge()).await; + } if let Some(admission) = self.admission.take() { admission.release().await; } diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index 6cd4eb580..b3ebb0d30 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -1,6 +1,7 @@ //! 模型请求生命周期、单行持久化 port 与 commit/send/cancellation 边界。 pub mod admission; +pub mod budget; pub mod continuation; pub mod coordinator; pub mod execution; diff --git a/backend/crates/gateway-core/src/error.rs b/backend/crates/gateway-core/src/error.rs index 4530f132c..8903886e1 100644 --- a/backend/crates/gateway-core/src/error.rs +++ b/backend/crates/gateway-core/src/error.rs @@ -975,6 +975,13 @@ pub struct GatewayError { message: &'static str, diagnostic: Option>, client_visible_upstream_error: Option, + client_details: Option>, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +struct GatewayClientErrorDetails { + code: Option<&'static str>, + retry_after: Option, } impl GatewayError { @@ -986,9 +993,33 @@ impl GatewayError { message, diagnostic: None, client_visible_upstream_error: None, + client_details: None, } } + #[must_use] + pub fn with_client_code(mut self, code: &'static str) -> Self { + self.client_details + .get_or_insert_with(Default::default) + .code = Some(code); + self + } + + #[must_use] + pub fn with_retry_after(mut self, delay: Duration) -> Self { + self.client_details + .get_or_insert_with(Default::default) + .retry_after = Some(delay); + self + } + + #[must_use] + pub fn retry_after(&self) -> Option { + self.client_details + .as_ref() + .and_then(|details| details.retry_after) + } + /// 将 Provider 错误归一为客户端无关错误。 #[must_use] pub fn from_provider(error: &ProviderError) -> Self { @@ -1103,6 +1134,11 @@ impl GatewayError { self.client_visible_upstream_error .as_ref() .and_then(ClientVisibleUpstreamError::code) + .or_else(|| { + self.client_details + .as_ref() + .and_then(|details| details.code) + }) } } diff --git a/backend/crates/gateway-core/src/lib.rs b/backend/crates/gateway-core/src/lib.rs index 89e6c6196..ca9fe8723 100644 --- a/backend/crates/gateway-core/src/lib.rs +++ b/backend/crates/gateway-core/src/lib.rs @@ -52,6 +52,7 @@ pub struct CoreStorePorts { snapshots: Arc, snapshot_subscriptions: Arc, client_api_key_usage: Arc, + budget: Option>, } impl CoreStorePorts { @@ -79,8 +80,15 @@ impl CoreStorePorts { snapshots, snapshot_subscriptions, client_api_key_usage, + budget: None, } } + + #[must_use] + pub fn with_budget(mut self, budget: Arc) -> Self { + self.budget = Some(budget); + self + } } pub struct CoreBundle { @@ -145,7 +153,7 @@ pub async fn initialize( Arc::clone(&ports.snapshot_subscriptions), )); let worker_contributions = publisher.worker_contributions()?; - let service = Arc::new(DefaultExecutionService::new( + let mut service = DefaultExecutionService::new( snapshots.clone(), ports.execution, providers, @@ -153,7 +161,11 @@ pub async fn initialize( ports.circuits, ports.continuation, ports.client_api_key_usage, - )); + ); + if let Some(budget) = ports.budget { + service = service.with_budget(budget); + } + let service = Arc::new(service); let execution: Arc = service.clone(); let account_probe: Arc = service; let snapshot_control: Arc = publisher; diff --git a/backend/crates/gateway-core/tests/account/mod.rs b/backend/crates/gateway-core/tests/account/mod.rs index c18cfd4e0..f1b775a5e 100644 --- a/backend/crates/gateway-core/tests/account/mod.rs +++ b/backend/crates/gateway-core/tests/account/mod.rs @@ -1,3 +1,4 @@ +mod proxy; mod selection; mod selection_diagnostics; diff --git a/backend/crates/gateway-core/tests/account/proxy.rs b/backend/crates/gateway-core/tests/account/proxy.rs new file mode 100644 index 000000000..21b6d92ce --- /dev/null +++ b/backend/crates/gateway-core/tests/account/proxy.rs @@ -0,0 +1,33 @@ +use gateway_core::account::OutboundProxy; + +#[test] +fn proxy_endpoints_support_explicit_schemes_and_redact_credentials() { + for scheme in ["http", "https", "socks5", "socks5h"] { + let proxy = + OutboundProxy::parse(&format!("{scheme}://user:p%40ss%3Aword@[::1]:1080")).unwrap(); + assert!(proxy.expose_url().contains("p%40ss%3Aword")); + assert!(!proxy.endpoint().contains("user")); + assert!(!proxy.endpoint().contains("p%40ss")); + assert!(!format!("{proxy:?}").contains("user")); + assert!(proxy.endpoint().contains("[::1]:1080")); + } +} + +#[test] +fn invalid_proxy_never_silently_becomes_direct() { + for url in [ + "", + "localhost:1080", + "file:///etc/passwd", + "ftp://host:21", + "http://host:0", + "socks5://host", + "http://host:80/path", + "http://host:80?secret=1", + "http://host:80#fragment", + "http://user:secret@host:80\n", + ] { + let error = OutboundProxy::parse(url).unwrap_err(); + assert!(!error.to_string().contains("secret")); + } +} diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index d4ea8825b..579ccbaa3 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -2158,6 +2158,80 @@ fn discarded_attempt_cost_never_leaks_into_retry_result() { let state = store.state.lock().expect("store lock"); assert_eq!(state.finalizations[0].cost_source, CostSource::Unavailable); assert_eq!(state.finalizations[0].cost_ticks, None); + assert_eq!( + session.budget_charge().amount_usd, + None, + "unknown final usage cannot erase a charged retry" + ); +} + +#[test] +fn key_budget_accumulates_retries_and_preserves_unknown_attempts() { + for (first_cost, first_send, expected) in [ + (Some(999), UpstreamSendState::Sent, Some(1122)), + (None, UpstreamSendState::Sent, None), + (None, UpstreamSendState::NotSent, Some(123)), + ] { + let operation = generate_operation(); + let route_plan = plan(&operation); + let mut first = Vec::new(); + if first_send != UpstreamSendState::NotSent { + first.push(Ok(GatewayEvent::Started(ResponseMeta::new( + "resp_first", + "gpt-5", + )))); + } + if let Some(cost) = first_cost { + first.push(Ok(GatewayEvent::ProviderCost( + ProviderReportedCost::from_usd_ticks(cost).unwrap(), + ))); + } + first.push(Err(ProviderError::new( + ProviderErrorKind::Unavailable, + first_send, + ) + .with_replay_safe())); + let (coordinator, _, _) = coordinator(vec![ + Script::Stream { + account_id: "acct_first", + items: first, + }, + Script::Stream { + account_id: "acct_second", + items: vec![ + Ok(GatewayEvent::Started(ResponseMeta::new( + "resp_done", + "gpt-5", + ))), + Ok(GatewayEvent::ProviderCost( + ProviderReportedCost::from_usd_ticks(123).unwrap(), + )), + Ok(GatewayEvent::Completed(ResponseMeta::new( + "resp_done", + "gpt-5", + ))), + ], + }, + ]); + let mut session = block_on(coordinator.start( + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + None, + CancellationToken::new(), + )) + .unwrap(); + block_on(session.collect_uncommitted()).unwrap(); + block_on(session.commit_downstream(Some(200))).unwrap(); + let charge = session.budget_charge(); + assert_eq!(charge.amount_usd.map(|amount| amount.scaled()), expected); + assert_eq!( + session.budget_charge().completed_at, + charge.completed_at, + "final settlement time stays stable" + ); + } } #[test] @@ -2688,10 +2762,16 @@ fn exhausted_atomic_response_failed_should_deliver_only_the_last_failure_once() fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire() { let operation = generate_operation(); let route_plan = plan(&operation); + let mut failure = bare_atomic_response_failed("response-only-failed"); + let mut events = failure.take_atomic_client_events(); + events.push(ProviderEvent::canonical(GatewayEvent::ProviderCost( + ProviderReportedCost::from_usd_ticks(500).unwrap(), + ))); + let failure = failure.with_atomic_client_events(events); let (coordinator, store, provider) = coordinator(vec![ Script::ObservedStream { account_id: "acct_only", - items: vec![Err(bare_atomic_response_failed("response-only-failed"))], + items: vec![Err(failure)], }, Script::Error(ProviderError::new( ProviderErrorKind::NoEligibleAccount, @@ -2731,6 +2811,14 @@ fn empty_selection_after_atomic_failure_should_deliver_the_last_upstream_wire() )); assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2); + assert_eq!( + session + .budget_charge() + .amount_usd + .map(|amount| amount.scaled()), + Some(500), + "restoring a discarded failure must not charge its usage twice" + ); let state = store.state.lock().expect("store lock"); assert_eq!(state.attempts.len(), 1); assert_eq!(state.intermediate_failures, 1); diff --git a/backend/crates/gateway-core/tests/engine/execution.rs b/backend/crates/gateway-core/tests/engine/execution.rs index c31b70603..fe7e67570 100644 --- a/backend/crates/gateway-core/tests/engine/execution.rs +++ b/backend/crates/gateway-core/tests/engine/execution.rs @@ -7,6 +7,165 @@ use std::{ time::{Duration, Instant, SystemTime}, }; +use gateway_core::engine::budget::{ + ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetError, ClientBudgetPort, +}; +use gateway_core::error::GatewayError; + +#[derive(Default)] +struct Admissions { + active: Arc, +} + +impl ClientAdmissionPort for Admissions { + fn admit( + &self, + _: ClientAdmissionRequest, + ) -> BoxFuture<'_, Result> { + Box::pin(async { + assert!(!self.active.swap(true, Ordering::SeqCst)); + Ok(ClientAdmissionDecision::Granted) + }) + } + fn release<'a>( + &'a self, + _: &'a ClientApiKeyId, + _: &'a ModelRequestId, + ) -> BoxFuture<'a, Result> { + Box::pin(async { Ok(self.active.swap(false, Ordering::SeqCst)) }) + } + fn restore( + &self, + _: ClientAdmissionRecovery, + ) -> BoxFuture<'_, Result> { + Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) }) + } +} + +struct Budget { + reject: bool, + active: Arc, + charges: Mutex>, +} + +impl ClientBudgetPort for Budget { + fn admit(&self, _: ClientBudgetAdmission) -> BoxFuture<'_, Result<(), GatewayError>> { + Box::pin(async { + assert!(self.active.load(Ordering::SeqCst)); + if self.reject { + Err( + GatewayError::new(GatewayErrorKind::RateLimited, "budget exhausted") + .with_client_code("key_daily_budget_exceeded"), + ) + } else { + Ok(()) + } + }) + } + fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> { + Box::pin(async move { + assert!( + self.active.load(Ordering::SeqCst), + "settle before releasing the concurrent request slot" + ); + self.charges.lock().unwrap().push(charge); + Ok(()) + }) + } +} + +fn service(admissions: Arc, budget: Arc) -> DefaultExecutionService { + DefaultExecutionService::new( + RuntimeSnapshotHandle::new(start_snapshot()), + Arc::new(TrackingExecutionStore::default()), + ProviderRegistry::default(), + admissions, + Arc::new(UnusedCircuits), + Arc::new(UnusedContinuation), + Arc::new(RecordingClientApiKeyUsage::default()), + ) + .with_budget(budget) +} + +fn request(service: &DefaultExecutionService, transport: ClientTransport) -> StartExecution { + StartExecution { + client: service.authenticate("sk_start_test").unwrap(), + public_model: PublicModelId::new("gpt-start").unwrap(), + operation: start_operation(), + metadata: ExecutionRequestMetadata { + protocol: "openai".to_owned(), + endpoint: "/v1/responses".to_owned(), + transport, + stream: transport != ClientTransport::HttpJson, + client_ip: None, + user_agent: None, + previous_response_id: None, + }, + } +} + +#[test] +fn budget_rejection_releases_client_concurrency_without_creating_a_charge() { + let admissions = Arc::new(Admissions::default()); + let budget = Arc::new(Budget { + reject: true, + active: admissions.active.clone(), + charges: Mutex::default(), + }); + let service = service(admissions.clone(), budget.clone()); + for transport in [ + ClientTransport::HttpJson, + ClientTransport::HttpSse, + ClientTransport::WebSocket, + ] { + let result = block_on(service.start(request(&service, transport))); + assert!( + matches!(result, Err(error) if error.client_error_code() == Some("key_daily_budget_exceeded")) + ); + assert!(!admissions.active.load(Ordering::SeqCst)); + assert!(budget.charges.lock().unwrap().is_empty()); + } +} + +#[test] +fn cancellation_and_pre_send_failure_settle_zero_and_release_concurrency_for_all_transports() { + let admissions = Arc::new(Admissions::default()); + let budget = Arc::new(Budget { + reject: false, + active: admissions.active.clone(), + charges: Mutex::default(), + }); + let service = service(admissions.clone(), budget.clone()); + for transport in [ + ClientTransport::HttpJson, + ClientTransport::HttpSse, + ClientTransport::WebSocket, + ] { + for detach in [true, false] { + let mut started = block_on(service.start(request(&service, transport))).unwrap(); + assert!(admissions.active.load(Ordering::SeqCst)); + if detach { + block_on(started.session.detach_finalize()); + } else { + assert!(block_on(started.session.next_event()).is_err()); + block_on(started.session.detach_finalize()); + } + assert!(!admissions.active.load(Ordering::SeqCst)); + } + } + let charges = budget.charges.lock().unwrap(); + assert_eq!( + charges.len(), + 6, + "one settlement per request, including detached finalized sessions" + ); + assert!( + charges + .iter() + .all(|charge| charge.amount_usd == Some(gateway_core::metering::Decimal::ZERO)) + ); +} + use async_trait::async_trait; use bytes::Bytes; use futures::{executor::block_on, future::BoxFuture}; diff --git a/backend/crates/gateway-store/src/bundle.rs b/backend/crates/gateway-store/src/bundle.rs index adbce7fd4..1a954f502 100644 --- a/backend/crates/gateway-store/src/bundle.rs +++ b/backend/crates/gateway-store/src/bundle.rs @@ -171,7 +171,8 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult { )?), ), Arc::new(client_key_usage), - ); + ) + .with_budget(Arc::new(postgres::PgClientBudgetStore::new(pool.clone()))); let provider_ports = ProviderStorePorts::new( account_store, diff --git a/backend/crates/gateway-store/src/postgres/account_groups.rs b/backend/crates/gateway-store/src/postgres/account_groups.rs index 329b0c1ba..56dd24a34 100644 --- a/backend/crates/gateway-store/src/postgres/account_groups.rs +++ b/backend/crates/gateway-store/src/postgres/account_groups.rs @@ -167,7 +167,7 @@ impl AccountGroupStore for PgAccountGroupRepository { "select membership.account_group_id, account.id, account.provider_kind, account.name, account.email, account.upstream_user_id, account.upstream_account_id, account.plan_type, - account.authentication_kind, account.credential_revision, + account.authentication_kind, account.credential_revision, account.outbound_proxy_url, account.has_refresh_token, account.access_token_expires_at, account.next_refresh_at, account.enabled, account.concurrency_limit, account.weight, account.credential_state, account.quota_access_state, diff --git a/backend/crates/gateway-store/src/postgres/client_budgets.rs b/backend/crates/gateway-store/src/postgres/client_budgets.rs new file mode 100644 index 000000000..8db149bdf --- /dev/null +++ b/backend/crates/gateway-store/src/postgres/client_budgets.rs @@ -0,0 +1,412 @@ +//! Durable admission and idempotent USD settlement, serialized per client key. + +use std::{collections::BTreeMap, sync::Mutex, time::Duration}; + +use chrono::{DateTime, Utc}; +use futures::future::BoxFuture; +use gateway_core::{ + engine::budget::{ + ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetError, ClientBudgetLimits, + ClientBudgetPort, ClientBudgetStatus, + }, + error::{GatewayError, GatewayErrorKind}, + metering::Decimal, +}; +use sqlx::{PgPool, Postgres, Row, Transaction}; + +use crate::{StoreResult, postgres_unavailable}; + +pub struct PgClientBudgetStore { + pool: PgPool, + retry: Mutex>, +} + +impl PgClientBudgetStore { + #[must_use] + pub fn new(pool: PgPool) -> Self { + Self { + pool, + retry: Mutex::new(BTreeMap::new()), + } + } + + async fn admit_inner(&self, request: ClientBudgetAdmission) -> Result<(), GatewayError> { + // Preserve exact charges across a transient outage. On process loss the durable + // pending event still blocks a limited key after its deadline, requiring reconciliation. + let retries = self + .retry + .lock() + .map_err(|_| unavailable())? + .values() + .cloned() + .collect::>(); + if !retries.is_empty() { + let ids = retries + .iter() + .map(|charge| charge.request_id.as_str()) + .collect::>(); + let owned = sqlx::query_scalar::<_, String>( + "select request_id from client_key_charge_events + where client_api_key_id = $1 and request_id = any($2)", + ) + .bind(request.key_id.as_str()) + .bind(ids) + .fetch_all(&self.pool) + .await + .map_err(|_| unavailable())?; + for charge in retries { + if owned.iter().any(|id| id == charge.request_id.as_str()) { + self.settle(charge).await.map_err(|_| unavailable())?; + } + } + } + let mut tx = self.pool.begin().await.map_err(|_| unavailable())?; + let row = sqlx::query( + "select daily_limit_usd::text, weekly_limit_usd::text, enabled + from client_api_keys where id = $1 for update", + ) + .bind(request.key_id.as_str()) + .fetch_optional(&mut *tx) + .await + .map_err(|_| unavailable())? + .ok_or_else(|| { + GatewayError::new( + GatewayErrorKind::Unauthorized, + "client API key no longer exists", + ) + })?; + if !row.get::("enabled") { + return Err(GatewayError::new( + GatewayErrorKind::PolicyDenied, + "client API key is disabled", + )); + } + let limits = ClientBudgetLimits { + daily_usd: row + .get::("daily_limit_usd") + .parse() + .map_err(|_| unavailable())?, + weekly_usd: row + .get::("weekly_limit_usd") + .parse() + .map_err(|_| unavailable())?, + }; + let now = Utc::now(); + advance_windows(&mut tx, request.key_id.as_str(), now) + .await + .map_err(|_| unavailable())?; + if limits.is_limited() { + let unresolved = sqlx::query_scalar::<_, bool>("select exists(select 1 from client_key_charge_events + where client_api_key_id = $1 and (state = 'unknown' or (state = 'pending' and deadline_at <= $2)))") + .bind(request.key_id.as_str()).bind(now).fetch_one(&mut *tx).await.map_err(|_| unavailable())?; + if unresolved { + return Err(GatewayError::new( + GatewayErrorKind::RateLimited, + "client API key has unresolved usage; administrator reconciliation is required", + ) + .with_client_code("key_budget_unresolved")); + } + let window = sqlx::query( + "select daily_used_usd::text, weekly_used_usd::text, daily_end, weekly_end + from client_key_budget_windows where client_api_key_id = $1", + ) + .bind(request.key_id.as_str()) + .fetch_one(&mut *tx) + .await + .map_err(|_| unavailable())?; + let daily: Decimal = window + .get::("daily_used_usd") + .parse() + .map_err(|_| unavailable())?; + let weekly: Decimal = window + .get::("weekly_used_usd") + .parse() + .map_err(|_| unavailable())?; + let daily_exceeded = limits.daily_usd != Decimal::ZERO && daily >= limits.daily_usd; + let weekly_exceeded = limits.weekly_usd != Decimal::ZERO && weekly >= limits.weekly_usd; + if daily_exceeded || weekly_exceeded { + let daily_end: DateTime = window.get("daily_end"); + let weekly_end: DateTime = window.get("weekly_end"); + let reset = if weekly_exceeded { + weekly_end + } else { + daily_end + }; + let retry = (reset - now).to_std().unwrap_or(Duration::from_secs(1)); + return Err(GatewayError::new( + GatewayErrorKind::RateLimited, + "client API key USD budget is exhausted", + ) + .with_client_code(if weekly_exceeded { + "key_weekly_budget_exceeded" + } else { + "key_daily_budget_exceeded" + }) + .with_retry_after(retry)); + } + } + sqlx::query( + "insert into client_key_charge_events (request_id, client_api_key_id, deadline_at) + values ($1, $2, $3)", + ) + .bind(request.request_id.as_str()) + .bind(request.key_id.as_str()) + .bind(DateTime::::from(request.deadline_at)) + .execute(&mut *tx) + .await + .map_err(|_| unavailable())?; + tx.commit().await.map_err(|_| unavailable()) + } + + async fn settle_inner(&self, charge: &ClientBudgetCharge) -> Result<(), ClientBudgetError> { + let mut tx = self.pool.begin().await.map_err(|_| ClientBudgetError)?; + // Match admission's key -> window -> event lock order, including concurrent settlements. + let key = sqlx::query_scalar::<_, String>( + "select k.id from client_api_keys k + join client_key_charge_events e on e.client_api_key_id = k.id + where e.request_id = $1 for update of k", + ) + .bind(charge.request_id.as_str()) + .fetch_optional(&mut *tx) + .await + .map_err(|_| ClientBudgetError)?; + let Some(key) = key else { return Ok(()) }; // Deleting a key also deletes its ledger. + settle_in_transaction(&mut tx, &key, charge) + .await + .map_err(|_| ClientBudgetError)?; + tx.commit().await.map_err(|_| ClientBudgetError) + } +} + +async fn settle_in_transaction( + tx: &mut Transaction<'_, Postgres>, + key: &str, + charge: &ClientBudgetCharge, +) -> Result<(), sqlx::Error> { + advance_windows(tx, key, Utc::now()).await?; + let changed = sqlx::query( + "update client_key_charge_events + set state = $2, amount_usd = $3::text::numeric, completed_at = $4 + where request_id = $1 and state <> 'settled'", + ) + .bind(charge.request_id.as_str()) + .bind(if charge.amount_usd.is_some() { + "settled" + } else { + "unknown" + }) + .bind(charge.amount_usd.map(|amount| amount.canonical())) + .bind(DateTime::::from(charge.completed_at)) + .execute(&mut **tx) + .await? + .rows_affected(); + if changed == 1 + && let Some(amount) = charge.amount_usd + { + sqlx::query("update client_key_budget_windows set + daily_used_usd = daily_used_usd + case when $3 >= daily_start and $3 < daily_end then $2::text::numeric else 0 end, + weekly_used_usd = weekly_used_usd + case when $3 >= weekly_start and $3 < weekly_end then $2::text::numeric else 0 end + where client_api_key_id = $1") + .bind(key).bind(amount.canonical()).bind(DateTime::::from(charge.completed_at)) + .execute(&mut **tx).await?; + } + Ok(()) +} + +impl ClientBudgetPort for PgClientBudgetStore { + fn admit(&self, request: ClientBudgetAdmission) -> BoxFuture<'_, Result<(), GatewayError>> { + Box::pin(async move { self.admit_inner(request).await }) + } + + fn settle(&self, charge: ClientBudgetCharge) -> BoxFuture<'_, Result<(), ClientBudgetError>> { + Box::pin(async move { + let result = self.settle_inner(&charge).await; + let mut retry = self.retry.lock().map_err(|_| ClientBudgetError)?; + if result.is_err() { + retry.insert(charge.request_id.as_str().to_owned(), charge); + } else { + retry.remove(charge.request_id.as_str()); + } + result + }) + } +} + +async fn advance_windows( + tx: &mut Transaction<'_, Postgres>, + key: &str, + now: DateTime, +) -> Result<(), sqlx::Error> { + sqlx::query("insert into client_key_budget_windows + (client_api_key_id, daily_start, daily_end, weekly_start, weekly_end) + select $1, day, day + interval '24 hours', day, day + interval '168 hours' + from (select date_trunc('day', $2::timestamptz at time zone 'Asia/Shanghai') at time zone 'Asia/Shanghai' as day) d + on conflict (client_api_key_id) do update set + daily_start = case when client_key_budget_windows.daily_end <= $2 then excluded.daily_start else client_key_budget_windows.daily_start end, + daily_end = case when client_key_budget_windows.daily_end <= $2 then excluded.daily_end else client_key_budget_windows.daily_end end, + daily_used_usd = case when client_key_budget_windows.daily_end <= $2 then 0 else client_key_budget_windows.daily_used_usd end, + weekly_start = case when client_key_budget_windows.weekly_end <= $2 then excluded.weekly_start else client_key_budget_windows.weekly_start end, + weekly_end = case when client_key_budget_windows.weekly_end <= $2 then excluded.weekly_end else client_key_budget_windows.weekly_end end, + weekly_used_usd = case when client_key_budget_windows.weekly_end <= $2 then 0 else client_key_budget_windows.weekly_used_usd end") + .bind(key).bind(now).execute(&mut **tx).await?; + Ok(()) +} + +pub(super) async fn load_client_key_budgets( + pool: &PgPool, + records: &mut [super::ClientApiKeyRecord], +) -> StoreResult<()> { + if records.is_empty() { + return Ok(()); + } + let ids = records + .iter() + .map(|record| record.id.as_str()) + .collect::>(); + let rows = sqlx::query("select k.id, k.daily_limit_usd::text, k.weekly_limit_usd::text, + (case when w.daily_end > now() then w.daily_used_usd else 0 end)::text as daily_used, + (case when w.weekly_end > now() then w.weekly_used_usd else 0 end)::text as weekly_used, + case when w.daily_end > now() then w.daily_end end as daily_end, + case when w.weekly_end > now() then w.weekly_end end as weekly_end, + (select count(*) from client_key_charge_events e where e.client_api_key_id = k.id + and (e.state = 'unknown' or (e.state = 'pending' and e.deadline_at <= now()))) as unresolved + from client_api_keys k left join client_key_budget_windows w on w.client_api_key_id = k.id + where k.id = any($1)") + .bind(ids).fetch_all(pool).await.map_err(|_| postgres_unavailable("load client budgets"))?; + let mut budgets = BTreeMap::new(); + for row in rows { + let parse = |field| -> StoreResult { + row.get::(field) + .parse() + .map_err(|_| postgres_unavailable("decode client budget")) + }; + budgets.insert( + row.get::("id"), + ClientBudgetStatus { + limits: ClientBudgetLimits { + daily_usd: parse("daily_limit_usd")?, + weekly_usd: parse("weekly_limit_usd")?, + }, + daily_used_usd: parse("daily_used")?, + weekly_used_usd: parse("weekly_used")?, + daily_resets_at: row + .get::>, _>("daily_end") + .map(Into::into), + weekly_resets_at: row + .get::>, _>("weekly_end") + .map(Into::into), + unresolved_requests: row.get::("unresolved").unsigned_abs(), + }, + ); + } + for record in records { + record.budget = budgets + .remove(&record.id) + .ok_or_else(|| postgres_unavailable("load client budget policy"))?; + } + Ok(()) +} + +fn unavailable() -> GatewayError { + GatewayError::new( + GatewayErrorKind::ProviderInfrastructureUnavailable, + "client budget service is temporarily unavailable", + ) + .with_client_code("key_budget_unavailable") +} + +pub(super) async fn unresolved_charges( + pool: &PgPool, + key: &gateway_core::policy::ClientApiKeyId, +) -> gateway_admin::ports::store::AdminStoreResult< + Vec, +> { + let rows = sqlx::query("select request_id, started_at, completed_at, state from client_key_charge_events + where client_api_key_id = $1 and (state = 'unknown' or (state = 'pending' and deadline_at <= now())) + order by started_at, request_id limit 200") + .bind(key.as_str()).fetch_all(pool).await.map_err(|_| budget_admin_unavailable())?; + Ok(rows + .into_iter() + .map( + |row| gateway_admin::model::client_keys::UnresolvedClientCharge { + request_id: row.get("request_id"), + started_at: row.get("started_at"), + completed_at: row.get("completed_at"), + state: row.get("state"), + }, + ) + .collect()) +} + +pub(super) async fn reconcile_charge( + pool: &PgPool, + command: gateway_admin::model::client_keys::ReconcileClientCharge, + context: &gateway_admin::model::MutationContext, + revision: gateway_admin::model::Revision, +) -> gateway_admin::ports::store::AdminStoreResult<()> { + use gateway_admin::ports::store::{AdminStoreError, AdminStoreErrorKind}; + let conflict = || { + AdminStoreError::new( + AdminStoreErrorKind::Conflict, + "client budget", + "charge is missing, active, or already settled with another amount", + ) + }; + let mut tx = pool.begin().await.map_err(|_| budget_admin_unavailable())?; + sqlx::query("select id from client_api_keys where id = $1 for update") + .bind(command.key_id.as_str()) + .fetch_optional(&mut *tx) + .await + .map_err(|_| budget_admin_unavailable())? + .ok_or_else(conflict)?; + let row = sqlx::query("select state, amount_usd::text, coalesce(completed_at, deadline_at) as completed_at, deadline_at <= now() as expired + from client_key_charge_events where request_id = $1 and client_api_key_id = $2 for update") + .bind(&command.request_id).bind(command.key_id.as_str()).fetch_optional(&mut *tx).await.map_err(|_| budget_admin_unavailable())?.ok_or_else(conflict)?; + let state: &str = row.get("state"); + if state == "settled" { + return if row.get::("amount_usd").parse::().ok() + == Some(command.amount_usd) + { + Ok(()) + } else { + Err(conflict()) + }; + } + if state == "pending" && !row.get::("expired") { + return Err(conflict()); + } + let charge = ClientBudgetCharge { + request_id: gateway_core::engine::ModelRequestId::new(command.request_id.clone()) + .map_err(|_| conflict())?, + amount_usd: Some(command.amount_usd), + completed_at: row.get::, _>("completed_at").into(), + }; + settle_in_transaction(&mut tx, command.key_id.as_str(), &charge) + .await + .map_err(|_| budget_admin_unavailable())?; + sqlx::query("update client_key_charge_events set reconciled_at = now(), reconciliation_reason = $2 where request_id = $1") + .bind(&command.request_id).bind(command.reason).execute(&mut *tx).await.map_err(|_| budget_admin_unavailable())?; + let audit = crate::mutation_audit( + context, + "reconcile", + "client_key_charge", + &command.request_id, + vec!["amount_usd".to_owned(), "reconciliation_reason".to_owned()], + ); + super::admin_security_audit::append_admin_audit_event_in_transaction( + &mut tx, + audit, + crate::store_revision(revision)?, + ) + .await + .map_err(|_| budget_admin_unavailable())?; + tx.commit().await.map_err(|_| budget_admin_unavailable()) +} + +fn budget_admin_unavailable() -> gateway_admin::ports::store::AdminStoreError { + gateway_admin::ports::store::AdminStoreError::new( + gateway_admin::ports::store::AdminStoreErrorKind::Unavailable, + "client budget", + "storage unavailable", + ) +} diff --git a/backend/crates/gateway-store/src/postgres/client_keys.rs b/backend/crates/gateway-store/src/postgres/client_keys.rs index 6805479c4..26f7ec873 100644 --- a/backend/crates/gateway-store/src/postgres/client_keys.rs +++ b/backend/crates/gateway-store/src/postgres/client_keys.rs @@ -27,8 +27,10 @@ use gateway_admin::{ ports::store::{AdminStoreResult, ClientKeyStore}, }; use gateway_core::{ + engine::budget::{ClientBudgetLimits, ClientBudgetStatus}, engine::execution::ClientApiKeyUsageSink, lifecycle::CancellationToken, + metering::Decimal, policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits}, routing::{AccountGroupId, ProviderKind}, task::{DaemonTask, WorkerTaskError}, @@ -114,6 +116,7 @@ pub struct ClientApiKeyRecord { pub enabled: bool, pub max_concurrency: u64, pub requests_per_minute: u64, + pub budget: ClientBudgetStatus, pub last_used_at: Option>, pub created_at: DateTime, pub updated_at: DateTime, @@ -271,6 +274,7 @@ pub struct NewClientApiKey { pub key: String, pub max_concurrency: u64, pub requests_per_minute: u64, + pub budget: ClientBudgetLimits, } impl fmt::Debug for NewClientApiKey { @@ -301,6 +305,8 @@ pub struct UpdateClientApiKeyDetails { pub group_ids: Vec, pub max_concurrency: u64, pub requests_per_minute: u64, + pub daily_limit_usd: Option, + pub weekly_limit_usd: Option, } impl UpdateClientApiKeyDetails { @@ -383,6 +389,7 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository { None }; load_client_key_memberships(&self.pool, &mut items).await?; + super::client_budgets::load_client_key_budgets(&self.pool, &mut items).await?; Ok(ClientApiKeyPage { items, total, @@ -406,7 +413,7 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository { async fn get_client_api_key(&self, id: &str) -> StoreResult> { require_nonempty(ENTITY, "id", id)?; - sqlx::query( + let record = sqlx::query( "select k.id, k.name, k.label, left(k.key, 10) as prefix, k.enabled, k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at, k.updated_at, coalesce(groups.groups, '[]'::jsonb) as groups, @@ -445,7 +452,16 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository { .map_err(|_| postgres_unavailable("get client API key"))? .as_ref() .map(client_record_from_row) - .transpose() + .transpose()?; + let Some(mut record) = record else { + return Ok(None); + }; + super::client_budgets::load_client_key_budgets( + &self.pool, + std::slice::from_mut(&mut record), + ) + .await?; + Ok(Some(record)) } async fn touch_client_api_keys( @@ -662,6 +678,22 @@ impl PgAdminClientKeyStore { #[async_trait] impl ClientKeyStore for PgAdminClientKeyStore { + async fn unresolved_charges( + &self, + id: &ClientApiKeyId, + ) -> AdminStoreResult> { + self.required_record(id).await?; + super::client_budgets::unresolved_charges(&self.keys.pool, id).await + } + + async fn reconcile_charge( + &self, + command: gateway_admin::model::client_keys::ReconcileClientCharge, + context: &MutationContext, + ) -> AdminStoreResult<()> { + let revision = self.revision().await?; + super::client_budgets::reconcile_charge(&self.keys.pool, command, context, revision).await + } async fn list_client_keys( &self, query: AdminClientKeyListQuery, @@ -719,6 +751,7 @@ impl ClientKeyStore for PgAdminClientKeyStore { .map(|id| id.as_str().to_owned()) .collect(), key: command.plaintext, + budget: command.budget, max_concurrency: command.limits.max_concurrency, requests_per_minute: command.limits.requests_per_minute, }, @@ -735,6 +768,8 @@ impl ClientKeyStore for PgAdminClientKeyStore { "enabled", "max_concurrency", "requests_per_minute", + "daily_limit_usd", + "weekly_limit_usd", ] .into_iter() .map(str::to_owned) @@ -766,6 +801,8 @@ impl ClientKeyStore for PgAdminClientKeyStore { .collect(), max_concurrency: command.limits.max_concurrency, requests_per_minute: command.limits.requests_per_minute, + daily_limit_usd: command.daily_limit_usd, + weekly_limit_usd: command.weekly_limit_usd, }, mutation_audit( context, @@ -778,6 +815,8 @@ impl ClientKeyStore for PgAdminClientKeyStore { "group_ids", "max_concurrency", "requests_per_minute", + "daily_limit_usd", + "weekly_limit_usd", ] .into_iter() .map(str::to_owned) @@ -942,6 +981,7 @@ fn admin_client_key_record(record: ClientApiKeyRecord) -> AdminStoreResult StoreResult AdminStoreResult { Ok(AccountRecord { + outbound_proxy: summary.outbound_proxy, id: summary.id, provider_kind: ProviderKind::new(summary.provider_kind).map_err(|_| { AdminStoreError::new( @@ -42,6 +43,7 @@ pub(crate) fn prepared_account( credential: PreparedCredentialCreate, ) -> StoreResult { Ok(NewProviderAccount { + outbound_proxy: credential.outbound_proxy, id: credential.account_id.as_str().to_owned(), provider_kind: credential.provider_kind.as_str().to_owned(), name: credential.name, 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 4e51f2ada..1ac6c51f8 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs @@ -108,7 +108,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository { include_disabled: bool, ) -> StoreResult> { let rows = sqlx::query( - "select id, provider_kind, name, email, upstream_user_id, + "select outbound_proxy_url, id, provider_kind, name, 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, credential_state, credential_observed_at, quota_access_state, quota_evidence, @@ -135,13 +135,13 @@ impl ProviderAccountRepository for PgProviderAccountRepository { }; sqlx::query( "insert into provider_accounts ( - id, provider_kind, name, email, upstream_user_id, + outbound_proxy_url, id, provider_kind, name, 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, credential_state, provider_quota_json, credential_observed_at, quota_access_observed_at, quota_observed_at, created_at, updated_at ) values ( - $1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13, + $18, $1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13, $14, $15, $16, null, $17, null, null, now(), greatest(now(), $17) )", ) @@ -162,6 +162,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository { .bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?) .bind(credential_state.as_str()) .bind(account.credential_observed_at) + .bind(account.outbound_proxy.as_ref().map(|proxy| proxy.expose_url())) .execute(&self.pool) .await .map_err(|_| postgres_unavailable("insert provider account"))?; @@ -524,6 +525,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository { command.enabled, command.concurrency_limit, command.weight, + command.outbound_proxy.as_ref(), ) .await?; replace_account_group_assignments_in_transaction( @@ -678,13 +680,13 @@ pub(crate) async fn upsert_provider_account_in_transaction( }; let imported_id = sqlx::query_scalar::<_, String>( "insert into provider_accounts ( - id, provider_kind, name, email, upstream_user_id, + outbound_proxy_url, id, provider_kind, name, 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, credential_state, provider_quota_json, credential_observed_at, quota_access_observed_at, quota_observed_at, created_at, updated_at ) values ( - $1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13, + $18, $1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13, $14, $15, $16, null, $17, null, null, now(), greatest(now(), $17) ) on conflict ( @@ -697,6 +699,7 @@ pub(crate) async fn upsert_provider_account_in_transaction( plan_type = excluded.plan_type, authentication_kind = excluded.authentication_kind, provider_credentials_json = excluded.provider_credentials_json, + outbound_proxy_url = coalesce(excluded.outbound_proxy_url, provider_accounts.outbound_proxy_url), credential_revision = provider_accounts.credential_revision + 1, has_refresh_token = excluded.has_refresh_token, access_token_expires_at = excluded.access_token_expires_at, @@ -732,6 +735,7 @@ pub(crate) async fn upsert_provider_account_in_transaction( .bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?) .bind(credential_state.as_str()) .bind(account.credential_observed_at) + .bind(account.outbound_proxy.as_ref().map(|proxy| proxy.expose_url())) .fetch_optional(&mut **transaction) .await .map_err(|error| { @@ -838,10 +842,13 @@ pub(crate) async fn update_provider_accounts_scheduling_in_transaction( enabled: bool, concurrency_limit: Option, weight: AccountWeight, + outbound_proxy: Option<&Option>, ) -> StoreResult<()> { let updated = sqlx::query_scalar::<_, String>( "update provider_accounts - set enabled = $2, concurrency_limit = $3, weight = $4, updated_at = now() + set enabled = $2, concurrency_limit = $3, weight = $4, updated_at = now(), + outbound_proxy_url = case when $5 then $6 else outbound_proxy_url end, + credential_revision = credential_revision + case when $5 and outbound_proxy_url is distinct from $6 then 1 else 0 end where id = any($1::text[]) returning id", ) @@ -849,6 +856,8 @@ pub(crate) async fn update_provider_accounts_scheduling_in_transaction( .bind(enabled) .bind(concurrency_limit.map(|limit| i64::from(limit.get()))) .bind(i16::try_from(weight.get()).map_err(|_| invalid("invalid weight"))?) + .bind(outbound_proxy.is_some()) + .bind(outbound_proxy.and_then(Option::as_ref).map(|proxy| proxy.expose_url())) .fetch_all(&mut **transaction) .await .map_err(|_| postgres_unavailable("set provider accounts state in admin transaction"))? 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 79db91f88..69ade944d 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs @@ -120,6 +120,7 @@ pub(crate) fn parse_error_reason(value: Option) -> StoreResult, pub id: String, pub provider_kind: String, pub name: String, @@ -167,6 +168,7 @@ impl fmt::Debug for ProviderAccountRecord { #[derive(Clone)] pub struct NewProviderAccount { + pub outbound_proxy: Option, pub id: String, pub provider_kind: String, pub name: String, @@ -307,6 +309,7 @@ impl fmt::Debug for RotateProviderAccount { #[derive(Debug, Clone)] pub struct BatchUpdateProviderAccountsAdmin { + pub outbound_proxy: Option>, pub account_ids: Vec, pub enabled: bool, pub concurrency_limit: Option, @@ -372,7 +375,7 @@ impl ProviderAccountStateUpdate { } } -pub(crate) const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email, upstream_user_id, +pub(crate) const ACCOUNT_SELECT: &str = "select outbound_proxy_url, id, provider_kind, name, 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, credential_state, provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at, @@ -380,7 +383,7 @@ pub(crate) const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email, credential_observed_at, quota_observed_at, created_at, updated_at from provider_accounts where id = $1"; -pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name, email, upstream_user_id, +pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select outbound_proxy_url, id, provider_kind, name, 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, credential_state, provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at, @@ -390,7 +393,7 @@ pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name, where id = any($1::text[]) and provider_kind = $2 order by id"; -pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select id, provider_kind, name, email, upstream_user_id, +pub(crate) const REFRESH_CANDIDATES_SELECT: &str = "select outbound_proxy_url, id, provider_kind, name, 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, credential_state, provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at, @@ -465,6 +468,7 @@ pub(crate) fn core_account_from_summary( summary.last_error_message, ) .with_scheduling(summary.concurrency_limit, summary.weight) + .with_outbound_proxy(summary.outbound_proxy) .with_refresh_schedule( summary.has_refresh_token, summary.next_refresh_at.map(Into::into), @@ -523,6 +527,12 @@ pub(crate) fn account_summary_from_row( .and_then(AccountWeight::new) .ok_or_else(|| invalid("invalid weight"))?; Ok(ProviderAccountSummary { + outbound_proxy: get::>(&row, "outbound_proxy_url")? + .map(|url| { + gateway_core::account::OutboundProxy::parse(&url) + .map_err(|_| invalid("invalid outbound proxy")) + }) + .transpose()?, id: get(&row, "id")?, provider_kind: get(&row, "provider_kind")?, name: get(&row, "name")?, diff --git a/backend/crates/gateway-store/tests/postgres/account_groups.rs b/backend/crates/gateway-store/tests/postgres/account_groups.rs index 25e494cd7..03df67fa7 100644 --- a/backend/crates/gateway-store/tests/postgres/account_groups.rs +++ b/backend/crates/gateway-store/tests/postgres/account_groups.rs @@ -168,6 +168,8 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip let (scope_revision, widened) = keys .update_client_key( UpdateClientKey { + daily_limit_usd: None, + weekly_limit_usd: None, id: client_key_id("key_group_one"), name: "key_group_one".to_owned(), label: None, @@ -191,6 +193,8 @@ async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multip let (restricted_revision, restricted) = keys .update_client_key( UpdateClientKey { + daily_limit_usd: None, + weekly_limit_usd: None, id: client_key_id("key_group_one"), name: "key_group_one".to_owned(), label: None, @@ -311,6 +315,7 @@ async fn group_costs_should_include_statusless_websocket_but_reject_statusless_h fn new_key(id: &str, group_ids: Vec) -> NewClientKey { let marker = char::from(id.as_bytes().last().copied().unwrap_or(b'k')); NewClientKey { + budget: Default::default(), id: client_key_id(id), name: id.to_owned(), label: None, diff --git a/backend/crates/gateway-store/tests/postgres/client_budgets.rs b/backend/crates/gateway-store/tests/postgres/client_budgets.rs new file mode 100644 index 000000000..82fdc8a17 --- /dev/null +++ b/backend/crates/gateway-store/tests/postgres/client_budgets.rs @@ -0,0 +1,383 @@ +use std::time::{Duration, SystemTime}; + +use chrono::{DateTime, Utc}; +use gateway_admin::{ + model::{ + MutationActor, MutationContext, + client_keys::{ReconcileClientCharge, UpdateClientKey}, + }, + ports::store::ClientKeyStore as _, +}; +use gateway_core::{ + engine::{ + ModelRequestId, + budget::{ClientBudgetAdmission, ClientBudgetCharge, ClientBudgetPort, ClientBudgetStatus}, + }, + error::GatewayErrorKind, + policy::{ClientApiKeyId, RateLimits}, +}; +use gateway_store::postgres::{ + ClientApiKeyRepository as _, PgAdminClientKeyStore, PgClientApiKeyRepository, + PgClientBudgetStore, +}; + +use super::TestDatabase; + +fn admission(key: &str, request: &str) -> ClientBudgetAdmission { + ClientBudgetAdmission { + key_id: ClientApiKeyId::new(key).unwrap(), + request_id: ModelRequestId::new(format!("req_{request}")).unwrap(), + deadline_at: SystemTime::now() + Duration::from_secs(600), + } +} + +fn charge(request: &str, amount: Option<&str>) -> ClientBudgetCharge { + ClientBudgetCharge { + request_id: ModelRequestId::new(format!("req_{request}")).unwrap(), + amount_usd: amount.map(|value| value.parse().unwrap()), + completed_at: SystemTime::now(), + } +} + +async fn seed(database: &TestDatabase, key: &str, daily: &str, weekly: &str) { + sqlx::query( + "insert into client_api_keys (id, name, key, daily_limit_usd, weekly_limit_usd, created_at, updated_at) + values ($1, $1, $2, $3::text::numeric, $4::text::numeric, now(), now())", + ) + .bind(key) + .bind(format!("sk_{key:a<43}")) + .bind(daily) + .bind(weekly) + .execute(&database.pool) + .await + .unwrap(); +} + +async fn status(database: &TestDatabase, key: &str) -> ClientBudgetStatus { + PgClientApiKeyRepository::new(database.pool.clone()) + .get_client_api_key(key) + .await + .unwrap() + .unwrap() + .budget +} + +fn context() -> MutationContext { + MutationContext { + actor: MutationActor::System, + request_id: "reconcile-test".to_owned(), + } +} + +#[tokio::test] +async fn budgets_settle_exactly_once_and_enforce_each_threshold_across_store_instances() { + let Some(database) = TestDatabase::create("budgets_exact").await else { + return; + }; + seed(&database, "day", "0.3", "2").await; + seed(&database, "week", "2", "0.2").await; + let first = PgClientBudgetStore::new(database.pool.clone()); + let second = PgClientBudgetStore::new(database.pool.clone()); + for (key, prefix, amount) in [("day", "d", "0.1"), ("week", "w", "0.1")] { + for index in 0..3 { + let id = format!("{prefix}-{index}"); + // Admitted work can finish above the threshold; no estimated cost reservation. + first.admit(admission(key, &id)).await.unwrap(); + } + for index in 0..3 { + let id = format!("{prefix}-{index}"); + let (a, b) = tokio::join!( + first.settle(charge(&id, Some(amount))), + second.settle(charge(&id, Some(amount))) + ); + a.unwrap(); + b.unwrap(); + } + } + let day = status(&database, "day").await; + assert_eq!(day.daily_used_usd.canonical(), "0.3"); + assert_eq!(day.weekly_used_usd.canonical(), "0.3"); + let day_error = second + .admit(admission("day", "d-rejected")) + .await + .unwrap_err(); + assert_eq!(day_error.kind(), GatewayErrorKind::RateLimited); + assert_eq!( + day_error.client_error_code(), + Some("key_daily_budget_exceeded") + ); + assert!(day_error.retry_after().is_some()); + let week_error = first + .admit(admission("week", "w-rejected")) + .await + .unwrap_err(); + assert_eq!( + week_error.client_error_code(), + Some("key_weekly_budget_exceeded") + ); + let rejected_count: i64 = sqlx::query_scalar( + "select count(*) from client_key_charge_events where request_id like '%rejected'", + ) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(rejected_count, 0); + database.close().await; +} + +#[tokio::test] +async fn window_rollover_is_shanghai_midnight_and_seven_days_with_late_settlement() { + let Some(database) = TestDatabase::create("budgets_windows").await else { + return; + }; + seed(&database, "key", "1", "2").await; + let store = PgClientBudgetStore::new(database.pool.clone()); + store.admit(admission("key", "first")).await.unwrap(); + let (day_start, day_end, week_end): (DateTime, DateTime, DateTime) = + sqlx::query_as("select daily_start, daily_end, weekly_end from client_key_budget_windows where client_api_key_id = 'key'") + .fetch_one(&database.pool).await.unwrap(); + assert_eq!(day_start.timestamp().rem_euclid(86400), 16 * 3600); + assert_eq!((day_end - day_start).num_hours(), 24); + assert_eq!((week_end - day_start).num_hours(), 168); + store.settle(charge("first", Some("1"))).await.unwrap(); + sqlx::query("update client_key_budget_windows set daily_start = daily_start - interval '1 day', daily_end = daily_start") + .execute(&database.pool).await.unwrap(); + store.admit(admission("key", "second")).await.unwrap(); + assert_eq!( + status(&database, "key").await.daily_used_usd.canonical(), + "0" + ); + assert_eq!( + status(&database, "key").await.weekly_used_usd.canonical(), + "1" + ); + store.settle(charge("second", Some("1"))).await.unwrap(); + sqlx::query("update client_key_budget_windows set daily_end = now() - interval '1 second', weekly_end = now() - interval '1 second'") + .execute(&database.pool).await.unwrap(); + let virtual_reset = status(&database, "key").await; + assert_eq!(virtual_reset.daily_used_usd.canonical(), "0"); + assert_eq!(virtual_reset.weekly_used_usd.canonical(), "0"); + store.admit(admission("key", "after-reset")).await.unwrap(); + let old = ClientBudgetCharge { + completed_at: (day_start - chrono::Duration::seconds(1)).into(), + ..charge("after-reset", Some("0.9")) + }; + store.settle(old).await.unwrap(); + let reset = status(&database, "key").await; + assert_eq!(reset.daily_used_usd.canonical(), "0"); + assert_eq!(reset.weekly_used_usd.canonical(), "0"); + assert!(reset.weekly_resets_at.is_some()); + database.close().await; +} + +#[tokio::test] +async fn unknown_and_orphan_charges_block_limited_keys_until_audited_reconciliation() { + let Some(database) = TestDatabase::create("budgets_unknown").await else { + return; + }; + seed(&database, "key", "1", "5").await; + seed(&database, "other", "0", "0").await; + let store = PgClientBudgetStore::new(database.pool.clone()); + let admin = PgAdminClientKeyStore::new(database.pool.clone()); + store.admit(admission("key", "unknown")).await.unwrap(); + store.settle(charge("unknown", None)).await.unwrap(); + assert_eq!( + store + .admit(admission("key", "blocked")) + .await + .unwrap_err() + .client_error_code(), + Some("key_budget_unresolved") + ); + let key_id = ClientApiKeyId::new("key").unwrap(); + assert_eq!(admin.unresolved_charges(&key_id).await.unwrap().len(), 1); + assert_eq!(status(&database, "key").await.unresolved_requests, 1); + let reconcile = |key: &str, request: &str, amount: &str| ReconcileClientCharge { + key_id: ClientApiKeyId::new(key).unwrap(), + request_id: format!("req_{request}"), + amount_usd: amount.parse().unwrap(), + reason: "verified against upstream usage".to_owned(), + }; + assert!( + admin + .reconcile_charge(reconcile("other", "unknown", "0.4"), &context()) + .await + .is_err() + ); + admin + .reconcile_charge(reconcile("key", "unknown", "0.4"), &context()) + .await + .unwrap(); + admin + .reconcile_charge(reconcile("key", "unknown", "0.4"), &context()) + .await + .unwrap(); + assert!( + admin + .reconcile_charge(reconcile("key", "unknown", "0.5"), &context()) + .await + .is_err() + ); + assert_eq!( + status(&database, "key").await.daily_used_usd.canonical(), + "0.4" + ); + let audits: i64 = sqlx::query_scalar( + "select count(*) from admin_audit_events where entity_kind = 'client_key_charge'", + ) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(audits, 1); + store.admit(admission("key", "orphan")).await.unwrap(); + assert!( + admin + .reconcile_charge(reconcile("key", "orphan", "0"), &context()) + .await + .is_err() + ); + sqlx::query("update client_key_charge_events set deadline_at = now() - interval '1 second' where request_id = 'req_orphan'") + .execute(&database.pool).await.unwrap(); + let restarted = PgClientBudgetStore::new(database.pool.clone()); + assert_eq!( + restarted + .admit(admission("key", "restart-blocked")) + .await + .unwrap_err() + .client_error_code(), + Some("key_budget_unresolved") + ); + admin + .reconcile_charge(reconcile("key", "orphan", "0"), &context()) + .await + .unwrap(); + restarted + .admit(admission("key", "recovered")) + .await + .unwrap(); + assert_eq!(status(&database, "key").await.unresolved_requests, 0); + store + .admit(admission("other", "unlimited-unknown")) + .await + .unwrap(); + store + .settle(charge("unlimited-unknown", None)) + .await + .unwrap(); + store + .admit(admission("other", "unlimited-next")) + .await + .unwrap(); + database.close().await; +} + +#[tokio::test] +async fn budget_updates_preserve_omitted_limits_and_do_not_clear_usage() { + let Some(database) = TestDatabase::create("budgets_policy").await else { + return; + }; + seed(&database, "key", "0", "0").await; + let store = PgClientBudgetStore::new(database.pool.clone()); + let admin = PgAdminClientKeyStore::new(database.pool.clone()); + store.admit(admission("key", "unlimited")).await.unwrap(); + store + .settle(charge("unlimited", Some("2.75"))) + .await + .unwrap(); + let update = UpdateClientKey { + id: ClientApiKeyId::new("key").unwrap(), + name: "key".to_owned(), + label: None, + group_ids: vec![], + limits: RateLimits { + max_concurrency: 3, + requests_per_minute: 0, + }, + daily_limit_usd: Some("2".parse().unwrap()), + weekly_limit_usd: Some("10".parse().unwrap()), + }; + admin + .update_client_key(update.clone(), &context()) + .await + .unwrap(); + admin + .update_client_key( + UpdateClientKey { + daily_limit_usd: None, + weekly_limit_usd: None, + ..update.clone() + }, + &context(), + ) + .await + .unwrap(); + let current = status(&database, "key").await; + assert_eq!(current.limits.daily_usd.canonical(), "2"); + assert_eq!(current.limits.weekly_usd.canonical(), "10"); + assert_eq!(current.daily_used_usd.canonical(), "2.75"); + assert_eq!( + store + .admit(admission("key", "blocked")) + .await + .unwrap_err() + .client_error_code(), + Some("key_daily_budget_exceeded") + ); + admin + .update_client_key( + UpdateClientKey { + daily_limit_usd: Some("0".parse().unwrap()), + weekly_limit_usd: None, + ..update + }, + &context(), + ) + .await + .unwrap(); + store.admit(admission("key", "allowed")).await.unwrap(); + sqlx::query("update client_api_keys set enabled = false where id = 'key'") + .execute(&database.pool) + .await + .unwrap(); + assert_eq!( + store + .admit(admission("key", "disabled")) + .await + .unwrap_err() + .kind(), + GatewayErrorKind::PolicyDenied + ); + sqlx::query("delete from client_api_keys where id = 'key'") + .execute(&database.pool) + .await + .unwrap(); + assert_eq!( + store + .admit(admission("key", "deleted")) + .await + .unwrap_err() + .kind(), + GatewayErrorKind::Unauthorized + ); + store.settle(charge("allowed", Some("1"))).await.unwrap(); + database.close().await; +} + +#[tokio::test] +async fn budget_database_outage_fails_closed() { + let Some(database) = TestDatabase::create("budgets_outage").await else { + return; + }; + let store = PgClientBudgetStore::new(database.pool.clone()); + database.pool.close().await; + assert_eq!( + store + .admit(admission("key", "offline")) + .await + .unwrap_err() + .client_error_code(), + Some("key_budget_unavailable") + ); + assert!(store.settle(charge("offline", Some("1"))).await.is_err()); + database.close().await; +} diff --git a/backend/crates/gateway-store/tests/postgres/client_keys.rs b/backend/crates/gateway-store/tests/postgres/client_keys.rs index 88ec35b25..38c01d876 100644 --- a/backend/crates/gateway-store/tests/postgres/client_keys.rs +++ b/backend/crates/gateway-store/tests/postgres/client_keys.rs @@ -24,6 +24,7 @@ use super::TestDatabase; #[test] fn client_key_requires_the_frozen_plaintext_format() { let key = NewClientApiKey { + budget: Default::default(), id: "key-1".to_owned(), name: "default".to_owned(), label: None, @@ -382,6 +383,7 @@ async fn dedicated_reveal_returns_plaintext_without_debug_exposure() { fn client_key_debug_redacts_plaintext() { let secret = format!("sk_{}", "s".repeat(43)); let key = NewClientApiKey { + budget: Default::default(), id: "key-1".to_owned(), name: "default".to_owned(), label: None, diff --git a/backend/crates/gateway-store/tests/postgres/mod.rs b/backend/crates/gateway-store/tests/postgres/mod.rs index e274651d2..14484a076 100644 --- a/backend/crates/gateway-store/tests/postgres/mod.rs +++ b/backend/crates/gateway-store/tests/postgres/mod.rs @@ -14,6 +14,7 @@ mod account_groups; mod admin_security_audit; mod admission_recovery; mod backup; +mod client_budgets; mod client_keys; mod execution; mod execution_buffer; @@ -230,6 +231,8 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly "backup_settings", "client_api_key_groups", "client_api_keys", + "client_key_budget_windows", + "client_key_charge_events", "model_requests", "ops_events", "provider_accounts", 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 a82c93fe5..fcfc1cc3e 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs @@ -1124,6 +1124,7 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() { let result = store .update_account( UpdateAccount { + outbound_proxy: None, account_id: "acct_terminal_mutation".to_owned(), enabled: false, concurrency_limit: None, @@ -1184,6 +1185,78 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() { database.close().await; } +#[tokio::test] +async fn account_proxy_edits_preserve_clear_and_fence_credential_caches_without_audit_secrets() { + let Some(database) = TestDatabase::create("account_proxy").await else { + return; + }; + let repository = PgProviderAccountRepository::new(database.pool.clone()); + let mut seed = account("acct_proxy", "proxy-user"); + seed.outbound_proxy = Some( + gateway_core::account::OutboundProxy::parse("http://initial:secret@127.0.0.1:18080") + .unwrap(), + ); + repository.insert_provider_account(seed).await.unwrap(); + let store = admin_account_store(&database.pool); + let context = MutationContext { + actor: MutationActor::System, + request_id: "proxy-edit".to_owned(), + }; + let command = UpdateAccount { + account_id: "acct_proxy".to_owned(), + enabled: true, + concurrency_limit: None, + weight: gateway_core::account::AccountWeight::DEFAULT, + group_ids: vec![], + outbound_proxy: None, + }; + store + .update_account(command.clone(), &context) + .await + .unwrap(); + let read = || async { + sqlx::query_as::<_, (Option, i64)>("select outbound_proxy_url, credential_revision from provider_accounts where id = 'acct_proxy'").fetch_one(&database.pool).await.unwrap() + }; + assert_eq!( + read().await, + (Some("http://initial:secret@127.0.0.1:18080/".to_owned()), 1) + ); + store + .update_account( + UpdateAccount { + outbound_proxy: Some(Some( + gateway_core::account::OutboundProxy::parse( + "socks5h://next:new-secret@127.0.0.1:1080", + ) + .unwrap(), + )), + ..command.clone() + }, + &context, + ) + .await + .unwrap(); + assert_eq!(read().await.1, 2); + store + .update_account( + UpdateAccount { + outbound_proxy: Some(None), + ..command + }, + &context, + ) + .await + .unwrap(); + assert_eq!(read().await, (None, 3)); + let audits: Vec = + sqlx::query_scalar("select to_jsonb(a) from admin_audit_events a") + .fetch_all(&database.pool) + .await + .unwrap(); + assert!(!serde_json::to_string(&audits).unwrap().contains("secret")); + database.close().await; +} + #[tokio::test] async fn account_recovery_resets_status_facts_without_changing_credentials_or_scheduling() { const GROUP_ID: &str = "grp_00000000000000000000000000000070"; @@ -1342,6 +1415,7 @@ async fn terminal_batch_update_replaces_state_and_groups_once_or_rolls_back_ever let result = store .batch_update_accounts( BatchUpdateAccounts { + outbound_proxy: None, account_ids: account_ids.clone(), enabled: false, concurrency_limit: gateway_core::account::AccountConcurrencyLimit::new(7), @@ -1377,6 +1451,7 @@ async fn terminal_batch_update_replaces_state_and_groups_once_or_rolls_back_ever store .batch_update_accounts( BatchUpdateAccounts { + outbound_proxy: None, account_ids: account_ids.clone(), enabled: true, concurrency_limit: None, @@ -1571,6 +1646,7 @@ async fn authorization_create_returns_existing_account_id_when_identity_is_upser AuthorizationOwnerBinding::from_context(&context), ), credential: AuthorizationCredentialCommit::Create(PreparedCredentialCreate { + outbound_proxy: None, account_id: ProviderAccountId::new("acct_authorization_candidate") .expect("candidate account ID"), provider_kind, @@ -1614,6 +1690,7 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() { let repository = PgProviderAccountRepository::new(database.pool.clone()); repository .insert_provider_account(NewProviderAccount { + outbound_proxy: None, id: "acct_core_refresh".to_owned(), provider_kind: "xai".to_owned(), name: "before refresh".to_owned(), @@ -1947,6 +2024,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() { let revision = repository .batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin { + outbound_proxy: None, account_ids: vec!["acct_admin_a".to_owned()], enabled: false, concurrency_limit: None, @@ -2305,6 +2383,7 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() { fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount { NewProviderAccount { + outbound_proxy: None, id: id.to_owned(), provider_kind: "openai".to_owned(), name: id.to_owned(), diff --git a/backend/crates/providers/openai/Cargo.toml b/backend/crates/providers/openai/Cargo.toml index a2b167d06..907ec6a0c 100644 --- a/backend/crates/providers/openai/Cargo.toml +++ b/backend/crates/providers/openai/Cargo.toml @@ -26,8 +26,9 @@ hex = "0.4.3" aws-lc-rs = "=1.18.1" hyper = { version = "=1.11.1", default-features = false } hyper-rustls = "=0.27.7" +hyper-util = { version = "0.1.20", features = ["client-legacy", "http1", "tokio"] } # 与官方 Codex Desktop 一致：reqwest 不启用压缩特性，上游请求不带 accept-encoding。 -reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "http2"] } +reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "http2", "socks"] } roxmltree = "0.21.1" semver = "1.0.27" # 与官方 rustls 路径一致：aws-lc 支持 P-521，优先提供 X25519MLKEM768。 @@ -44,6 +45,7 @@ thiserror = "2.0.18" tokio = { version = "=1.53.1", default-features = false, features = ["fs", "macros", "rt", "sync", "time"] } tokio-rustls = { version = "=0.26.4", default-features = false } tokio-util = { version = "0.7.18", features = ["rt"] } +tower-service = "0.3.3" # 官方 Codex WebSocket 需要上游尚未发布的 proxy 与 permessage-deflate 实现，固定到已审计提交。 tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungstenite", rev = "0e5b2d73aa18dd9f0a50ee9ff199d5aef7594186", features = ["proxy", "rustls-tls-native-roots"] } tracing = "0.1.44" diff --git a/backend/crates/providers/openai/src/admin.rs b/backend/crates/providers/openai/src/admin.rs index 92b209e73..39f38c754 100644 --- a/backend/crates/providers/openai/src/admin.rs +++ b/backend/crates/providers/openai/src/admin.rs @@ -595,6 +595,7 @@ fn prepared_create( credential, } = prepared; Ok(PreparedCredentialCreate { + outbound_proxy: account.outbound_proxy().cloned(), account_id: account.id().clone(), provider_kind: account.provider().clone(), name: account.name().to_owned(), @@ -698,6 +699,7 @@ fn account_from_record(account: &AccountRecord) -> Result, email: Option, authentication: ParsedCodexAuthentication, + outbound_proxy: Option, } struct ParsedCodexAuthentication { @@ -226,6 +227,8 @@ struct CodexCprExportCommon { status: &'static str, added_at: String, updated_at: String, + #[serde(skip_serializing_if = "Option::is_none")] + outbound_proxy_url: Option, } #[derive(Serialize)] @@ -502,6 +505,9 @@ impl CodexCredentialAdmin { status: cpr_status(&account), added_at: china_rfc3339(item.added_at), updated_at: china_rfc3339(item.updated_at), + outbound_proxy_url: account + .outbound_proxy() + .map(|proxy| proxy.expose_url().to_owned()), }; let exported = match data { CodexCredentialData::OAuth(data) => { @@ -743,7 +749,10 @@ impl CodexCredentialAdminService { }; let tokens = self .refresher - .refresh(refresh_token.expose_secret()) + .refresh_with_proxy( + refresh_token.expose_secret(), + current.account.outbound_proxy(), + ) .await .inspect_err(|error| log_manual_refresh_failure(&account_id, error)) .map_err(map_refresh_failure)?; @@ -819,6 +828,7 @@ impl CodexCredentialAdminService { candidate.authentication.access_token.clone(), candidate.authentication.refresh_token.clone(), candidate.authentication.id_token.clone(), + candidate.outbound_proxy.as_ref(), ) .await?; // 与官方 `parse_chatgpt_jwt_claims` 一致：先读 ID token，缺失字段 @@ -859,7 +869,12 @@ impl CodexCredentialAdminService { next_refresh_at: None, enabled: true, })?; - accounts.push(prepared); + accounts.push(NewProviderAccount { + account: prepared + .account + .with_outbound_proxy(candidate.outbound_proxy), + credential: prepared.credential, + }); } Ok(PreparedCodexAccountImport { accounts }) } @@ -870,6 +885,7 @@ impl CodexCredentialAdminService { access_token: Option, refresh_token: Option, id_token: Option, + proxy: Option<&gateway_core::account::OutboundProxy>, ) -> Result<(CodexOAuthSecret, Option>), CodexCredentialAdminError> { let id_token = id_token.map(SecretString::from); if let Some(access_token) = access_token { @@ -889,7 +905,7 @@ impl CodexCredentialAdminService { let refresh_token = refresh_token.ok_or(CodexCredentialAdminError::InvalidCredential)?; let tokens = self .refresher - .refresh(&refresh_token) + .refresh_with_proxy(&refresh_token, proxy) .await .map_err(map_refresh_failure)?; let access_token = tokens @@ -986,12 +1002,18 @@ fn log_manual_refresh_failure(account_id: &ProviderAccountId, error: &RefreshFai fn parse_import_document( payload: &Value, ) -> Result, CodexCredentialAdminError> { + let payload = payload + .get("data") + .filter(|data| data.get("accounts").is_some()) + .unwrap_or(payload); let mut accounts = Vec::new(); for value in import_account_values(payload)? { if !is_openai_oauth_candidate(value) { continue; } - accounts.push(parse_oauth_import_account(value)?); + let mut account = parse_oauth_import_account(value)?; + account.outbound_proxy = import_proxy(payload, value)?; + accounts.push(account); } Ok(accounts) } @@ -1004,9 +1026,89 @@ fn parse_oauth_import_account( name: first_string(value, &["name", "label"]), email: first_string(value, &["email"]).or_else(|| first_string(credentials, &["email"])), authentication: parse_oauth_import_tokens(value)?, + outbound_proxy: None, }) } +fn import_proxy( + payload: &Value, + account: &Value, +) -> Result, CodexCredentialAdminError> { + let invalid = || CodexCredentialAdminError::InvalidInput; + if let Some(value) = account + .get("outbound_proxy_url") + .or_else(|| account.get("outboundProxyUrl")) + .filter(|value| !value.is_null()) + { + let text = value.as_str().ok_or_else(invalid)?; + if text.is_empty() { + return Ok(None); + } + return gateway_core::account::OutboundProxy::parse(text) + .map(Some) + .map_err(|_| invalid()); + } + let Some(key) = account.get("proxy_key").filter(|value| !value.is_null()) else { + return Ok(None); + }; + let key = key + .as_str() + .filter(|key| !key.is_empty()) + .ok_or_else(invalid)?; + let proxies = payload + .get("proxies") + .and_then(Value::as_array) + .ok_or_else(invalid)?; + let mut matches = proxies + .iter() + .filter(|proxy| proxy.get("proxy_key").and_then(Value::as_str) == Some(key)); + let proxy = matches.next().ok_or_else(invalid)?; + if matches.next().is_some() + || proxy + .get("status") + .and_then(Value::as_str) + .is_some_and(|value| value != "active") + || proxy + .get("expires_at") + .is_some_and(|value| !value.is_null()) + || proxy + .get("fallback_mode") + .and_then(Value::as_str) + .is_some_and(|value| !value.is_empty() && value != "none") + { + return Err(invalid()); + } + let text = |field| proxy.get(field).and_then(Value::as_str).ok_or_else(invalid); + let scheme = text("protocol")?; + // sub2api resolves SOCKS targets at the proxy. + let scheme = if scheme == "socks5" { + "socks5h" + } else { + scheme + }; + let mut url = url::Url::parse(&format!("{scheme}://localhost")).map_err(|_| invalid())?; + let host = text("host")?; + if let Ok(address) = host.parse::() { + url.set_ip_host(address).map_err(|_| invalid())?; + } else { + url.set_host(Some(host)).map_err(|_| invalid())?; + } + let port = proxy + .get("port") + .and_then(Value::as_u64) + .and_then(|port| u16::try_from(port).ok()) + .filter(|port| *port != 0) + .ok_or_else(invalid)?; + url.set_port(Some(port)).map_err(|_| invalid())?; + url.set_username(proxy.get("username").and_then(Value::as_str).unwrap_or("")) + .map_err(|_| invalid())?; + url.set_password(proxy.get("password").and_then(Value::as_str)) + .map_err(|_| invalid())?; + gateway_core::account::OutboundProxy::parse(url.as_str()) + .map(Some) + .map_err(|_| invalid()) +} + fn import_account_values(payload: &Value) -> Result, CodexCredentialAdminError> { if let Some(accounts) = payload.get("accounts") { return accounts diff --git a/backend/crates/providers/openai/src/credential/catalog.rs b/backend/crates/providers/openai/src/credential/catalog.rs index 7cdb5f650..555cba692 100644 --- a/backend/crates/providers/openai/src/credential/catalog.rs +++ b/backend/crates/providers/openai/src/credential/catalog.rs @@ -436,6 +436,10 @@ impl CodexCredentialCatalogService { .authorization_header() .map_err(|_| CodexCredentialCatalogError::InvalidCredentialData)?; let result = client + .for_account(account) + .map_err(|error| CodexCredentialCatalogError::Upstream { + detail: error.to_string(), + })? .fetch_models_with_context(CodexRequestContext::auxiliary( authorization.expose_secret(), account.upstream_account_id(), diff --git a/backend/crates/providers/openai/src/credential/oauth.rs b/backend/crates/providers/openai/src/credential/oauth.rs index 32d834f2b..dec168eae 100644 --- a/backend/crates/providers/openai/src/credential/oauth.rs +++ b/backend/crates/providers/openai/src/credential/oauth.rs @@ -524,7 +524,16 @@ impl CodexOAuthAdminService { .transpose()? .flatten(); let (mut secret, id_token) = self - .exchange_pending(&pending, callback_url, fallback_refresh_token) + .exchange_pending( + &pending, + callback_url, + fallback_refresh_token, + current + .as_ref() + .map_or(pending.mutation.outbound_proxy(), |current| { + current.account.outbound_proxy() + }), + ) .await?; let mutation = pending.mutation.clone(); let access_token_expires_at = @@ -545,20 +554,25 @@ impl CodexOAuthAdminService { ) } else { let account_id = format!("acct_{}", Uuid::now_v7().simple()); - CompletedCodexOAuthCredential::Create( - self.credentials - .prepare_unresolved_oauth(UnresolvedCodexOAuthCredential { - account_id, - name: pending.name, - installation_id: pending.installation_id, - secret, - metadata, - access_token_expires_at, - next_refresh_at: None, - enabled: true, - }) - .map_err(map_admin_error)?, - ) + let prepared = self + .credentials + .prepare_unresolved_oauth(UnresolvedCodexOAuthCredential { + account_id, + name: pending.name, + installation_id: pending.installation_id, + secret, + metadata, + access_token_expires_at, + next_refresh_at: None, + enabled: true, + }) + .map_err(map_admin_error)?; + CompletedCodexOAuthCredential::Create(gateway_core::account::NewProviderAccount { + account: prepared + .account + .with_outbound_proxy(mutation.outbound_proxy().cloned()), + credential: prepared.credential, + }) }; Ok((mutation, credential)) } @@ -666,6 +680,7 @@ impl CodexOAuthAdminService { pending: &CodexPendingAuthorization, callback_url: &str, fallback_refresh_token: Option, + proxy: Option<&gateway_core::account::OutboundProxy>, ) -> Result<(super::types::CodexOAuthSecret, SecretString), CodexOAuthAdminError> { let (code, callback_state) = callback_parts(callback_url)?; if !constant_time_equal( @@ -676,10 +691,13 @@ impl CodexOAuthAdminService { } let tokens = self .exchanger - .exchange_authorization_code(AuthorizationCodeGrant { - code, - code_verifier: pending.code_verifier.clone(), - }) + .exchange_with_proxy( + AuthorizationCodeGrant { + code, + code_verifier: pending.code_verifier.clone(), + }, + proxy, + ) .await .map_err(map_exchange_error)?; let mut secret = tokens.secret; diff --git a/backend/crates/providers/openai/src/credential/profile_statistics.rs b/backend/crates/providers/openai/src/credential/profile_statistics.rs index a5175e702..32008fe59 100644 --- a/backend/crates/providers/openai/src/credential/profile_statistics.rs +++ b/backend/crates/providers/openai/src/credential/profile_statistics.rs @@ -133,13 +133,16 @@ impl CodexCredentialProfileService { &self, account_id: &ProviderAccountId, ) -> Result { - let (authorization, upstream_account_id) = self.account_authentication(account_id).await?; + let (authorization, upstream_account_id, account) = + self.account_authentication(account_id).await?; let request_id = format!("profile_statistics_{}", Uuid::now_v7().simple()); let statistics = CodexBackendClient::new( self.http.clone(), self.base_url.clone(), self.profile.clone(), ) + .for_account(&account) + .map_err(map_client_error)? .fetch_profile_statistics(CodexRequestContext::auxiliary( authorization.expose_secret(), upstream_account_id.as_deref(), @@ -156,7 +159,14 @@ impl CodexCredentialProfileService { async fn account_authentication( &self, account_id: &ProviderAccountId, - ) -> Result<(SecretString, Option), CodexProfileStatisticsError> { + ) -> Result< + ( + SecretString, + Option, + gateway_core::account::ProviderAccount, + ), + CodexProfileStatisticsError, + > { let account = self .repository .store() @@ -187,6 +197,7 @@ impl CodexCredentialProfileService { Ok(( authorization, account.upstream_account_id().map(str::to_owned), + account, )) } @@ -202,9 +213,19 @@ impl CodexCredentialProfileService { .image_url .ok_or(CodexProfileAvatarError::Missing)?, }; - let (authorization, upstream_account_id) = self.account_authentication(account_id).await?; + let (authorization, upstream_account_id, account) = + self.account_authentication(account_id).await?; + let http = if account.outbound_proxy().is_some() { + crate::transport::client::build_account_http_client( + account.id().as_str(), + account.outbound_proxy(), + ) + .map_err(|_| CodexProfileAvatarError::TransportUnavailable)? + } else { + self.http.clone() + }; fetch_profile_avatar( - &self.http, + &http, &self.base_url, &self.profile.snapshot(), &source, diff --git a/backend/crates/providers/openai/src/credential/quota/mod.rs b/backend/crates/providers/openai/src/credential/quota/mod.rs index a31c19485..da291690a 100644 --- a/backend/crates/providers/openai/src/credential/quota/mod.rs +++ b/backend/crates/providers/openai/src/credential/quota/mod.rs @@ -1225,6 +1225,8 @@ async fn list_reset_credits_once( .authorization_header() .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; client + .for_account(&prepared.account) + .map_err(ResetCreditAttemptError::Upstream)? .list_rate_limit_reset_credits(CodexRequestContext::auxiliary( authorization.expose_secret(), prepared.account.upstream_account_id(), @@ -1248,6 +1250,8 @@ async fn consume_reset_credit_once( .authorization_header() .map_err(|_| ResetCreditAttemptError::InvalidCredential)?; client + .for_account(&prepared.account) + .map_err(ResetCreditAttemptError::Upstream)? .consume_rate_limit_reset_credit( CodexRequestContext::auxiliary( authorization.expose_secret(), @@ -1329,6 +1333,8 @@ async fn fetch_usage_once( .map_err(|_| CodexQuotaFetchAttemptError::InvalidCredential)?; let request_id = format!("quota_{}", Uuid::now_v7().simple()); client + .for_account(&prepared.account) + .map_err(CodexQuotaFetchAttemptError::Upstream)? .fetch_usage(CodexRequestContext::auxiliary( authorization.expose_secret(), prepared.account.upstream_account_id(), diff --git a/backend/crates/providers/openai/src/credential/refresh.rs b/backend/crates/providers/openai/src/credential/refresh.rs index 1a0205026..2168b3d2a 100644 --- a/backend/crates/providers/openai/src/credential/refresh.rs +++ b/backend/crates/providers/openai/src/credential/refresh.rs @@ -288,7 +288,10 @@ impl CodexCredentialRefreshService { .refresh_token .as_ref() .ok_or(CodexCredentialRefreshError::InvalidRefreshResponse)?; - let refresh_result = self.refresher.refresh(refresh_token.expose_secret()).await; + let refresh_result = self + .refresher + .refresh_with_proxy(refresh_token.expose_secret(), due.account.outbound_proxy()) + .await; if recovery_window_exhausted && let Err(failure) = &refresh_result { let message = failure.message().map(str::to_owned); let upstream = failure.upstream(); diff --git a/backend/crates/providers/openai/src/credential/token_client.rs b/backend/crates/providers/openai/src/credential/token_client.rs index 64010bb41..bc32eac41 100644 --- a/backend/crates/providers/openai/src/credential/token_client.rs +++ b/backend/crates/providers/openai/src/credential/token_client.rs @@ -178,6 +178,22 @@ impl fmt::Debug for RefreshUpstreamFailure { #[async_trait] pub trait TokenRefresher: Send + Sync + 'static { async fn refresh(&self, refresh_token: &str) -> Result; + async fn refresh_with_proxy( + &self, + refresh_token: &str, + proxy: Option<&gateway_core::account::OutboundProxy>, + ) -> Result { + if proxy.is_some() { + return Err(proxy_refresh_failure()); + } + self.refresh(refresh_token).await + } +} + +fn proxy_refresh_failure() -> RefreshFailure { + RefreshFailure::RetryableTransport { + message: "account OAuth egress unavailable".to_owned(), + } } /// Authorization Code + PKCE 的一次性 grant。 @@ -232,6 +248,16 @@ pub trait AuthorizationCodeExchanger: Send + Sync + 'static { &self, grant: AuthorizationCodeGrant, ) -> Result; + async fn exchange_with_proxy( + &self, + grant: AuthorizationCodeGrant, + proxy: Option<&gateway_core::account::OutboundProxy>, + ) -> Result { + if proxy.is_some() { + return Err(AuthorizationCodeExchangeError::Unavailable); + } + self.exchange_authorization_code(grant).await + } } /// OpenAI token 续期客户端配置。 @@ -257,6 +283,24 @@ pub struct OpenAiTokenClient { pub struct TokenClientBuildError; impl OpenAiTokenClient { + fn with_proxy( + &self, + proxy: Option<&gateway_core::account::OutboundProxy>, + ) -> Result { + let Some(proxy) = proxy else { + return Ok(self.clone()); + }; + let builder = Client::builder() + .use_rustls_tls() + .no_proxy() + .proxy(reqwest::Proxy::all(proxy.expose_url()).map_err(|_| TokenClientBuildError)?) + .redirect(Policy::none()) + .connect_timeout(TOKEN_CONNECT_TIMEOUT) + .timeout(TOKEN_REQUEST_TIMEOUT); + let client = + build_reqwest_client_with_custom_ca(builder).map_err(|_| TokenClientBuildError)?; + Ok(Self::new(client, self.config.clone(), self.profile.clone())) + } /// 共享运行时画像；刷新时取快照，授权码交换仍使用 raw auth 请求。 pub fn new(client: Client, config: TokenClientConfig, profile: CodexWireProfileState) -> Self { Self { @@ -357,6 +401,16 @@ struct AuthorizationCodeResponse { #[async_trait] impl TokenRefresher for OpenAiTokenClient { + async fn refresh_with_proxy( + &self, + refresh_token: &str, + proxy: Option<&gateway_core::account::OutboundProxy>, + ) -> Result { + self.with_proxy(proxy) + .map_err(|_| proxy_refresh_failure())? + .refresh(refresh_token) + .await + } async fn refresh(&self, refresh_token: &str) -> Result { let headers = build_codex_profile_headers(&self.profile.snapshot()).map_err(|_| { RefreshFailure::Transport { @@ -389,6 +443,16 @@ impl TokenRefresher for OpenAiTokenClient { #[async_trait] impl AuthorizationCodeExchanger for OpenAiTokenClient { + async fn exchange_with_proxy( + &self, + grant: AuthorizationCodeGrant, + proxy: Option<&gateway_core::account::OutboundProxy>, + ) -> Result { + self.with_proxy(proxy) + .map_err(|_| AuthorizationCodeExchangeError::Unavailable)? + .exchange_authorization_code(grant) + .await + } async fn exchange_authorization_code( &self, grant: AuthorizationCodeGrant, diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index b87d37116..9c32e830d 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -124,7 +124,9 @@ impl CodexProvider { crate::transport::request::scope_turn_metadata(metadata, lease.installation_id(), true) }); let events = cold_json_response_stream(ColdJsonResponse { - client: self.client.clone(), + client: self.client.for_account(lease.account()).map_err(|_| { + provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent) + })?, response_origin: request.response_origin, endpoint_path: request.endpoint_path, body: request.body, diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 67a573d3d..21cb2c457 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -519,7 +519,9 @@ impl Provider for CodexProvider { AttemptTransport::Default | AttemptTransport::Fallback => 0, }; let events = cold_response_stream(ColdResponse { - client: self.client.clone(), + client: self.client.for_account(lease.account()).map_err(|_| { + provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent) + })?, response_origin: self.responses_url.clone(), request: upstream_request, upstream_model: upstream_model.clone(), diff --git a/backend/crates/providers/openai/src/transport/client.rs b/backend/crates/providers/openai/src/transport/client.rs index 0af62aa0d..eacc51a12 100644 --- a/backend/crates/providers/openai/src/transport/client.rs +++ b/backend/crates/providers/openai/src/transport/client.rs @@ -45,13 +45,20 @@ const UPSTREAM_CONNECT_TIMEOUT: Duration = Duration::from_secs(15); pub(super) const UPSTREAM_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60); const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str = "x-codex-ws-stream-request-start-ms"; -type ReqwestClientCacheKey = Option; +type ReqwestClientCacheKey = (Option, String); type ReqwestClientCache = Mutex>; /// 构建带缓存、自动协商 HTTP/2 的 reqwest Client。 pub fn build_reqwest_client() -> Result { + build_account_http_client("", None) +} + +pub fn build_account_http_client( + account_id: &str, + proxy: Option<&gateway_core::account::OutboundProxy>, +) -> Result { super::tls::ensure_rustls_provider(); - let cache_key = custom_ca_env_cache_key(); + let cache_key = (custom_ca_env_cache_key(), egress_key(account_id, proxy)); static CLIENTS: OnceLock = OnceLock::new(); let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new())); if let Some(client) = cache @@ -62,7 +69,7 @@ pub fn build_reqwest_client() -> Result { return Ok(client.clone()); } - let builder = Client::builder() + let mut builder = Client::builder() .use_rustls_tls() .no_proxy() .redirect(reqwest::redirect::Policy::none()) @@ -73,13 +80,35 @@ pub fn build_reqwest_client() -> Result { .http2_keep_alive_interval(Duration::from_secs(30)) .http2_keep_alive_timeout(Duration::from_secs(5)) .http2_keep_alive_while_idle(true); + if let Some(proxy) = proxy { + builder = builder.proxy( + reqwest::Proxy::all(proxy.expose_url()) + .map_err(|_| CustomCaError::ProxyConfiguration)?, + ); + } let client = build_reqwest_client_with_custom_ca(builder)?; let mut clients = cache .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); + if clients.len() >= 256 { + clients.clear(); + } Ok(clients.entry(cache_key).or_insert(client).clone()) } +fn egress_key(account_id: &str, proxy: Option<&gateway_core::account::OutboundProxy>) -> String { + use sha2::{Digest, Sha256}; + let mut hash = Sha256::new(); + hash.update(account_id.as_bytes()); + hash.update([0]); + hash.update( + proxy + .map_or("direct", |proxy| proxy.expose_url()) + .as_bytes(), + ); + hex::encode(hash.finalize()) +} + // --------------------------------------------------------------------------- // 错误类型 // --------------------------------------------------------------------------- @@ -599,11 +628,36 @@ pub struct CodexBackendJsonResponse { #[derive(Clone)] pub struct CodexBackendClient { pub(super) client: Client, + pub(super) direct_client: Client, pub(super) base_url: String, pub(super) profile: CodexWireProfileState, pub(super) websocket_pool: Option>, pub(super) websocket_origin_breaker: WebSocketOriginBreaker, pub(super) websocket_origin_key: String, + pub(super) outbound_proxy: Option, + pub(super) egress_key: String, +} + +impl CodexBackendClient { + pub fn for_account( + &self, + account: &gateway_core::account::ProviderAccount, + ) -> Result { + let mut client = self.clone(); + client.outbound_proxy = account.outbound_proxy().cloned(); + client.egress_key = egress_key(account.id().as_str(), account.outbound_proxy()); + client.websocket_origin_key = format!( + "{}:{}", + websocket_origin_key(&self.base_url), + client.egress_key + ); + client.client = if account.outbound_proxy().is_some() { + build_account_http_client(account.id().as_str(), account.outbound_proxy())? + } else { + self.direct_client.clone() + }; + Ok(client) + } } /// 已完成账号级 opening 准备、但尚未发送 payload 的 transport。 diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index 1e0a27591..1b3143b70 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -52,8 +52,11 @@ impl CodexBackendClient { ) -> Self { let base_url = base_url.into().trim_end_matches('/').to_string(); Self { + direct_client: client.clone(), client, websocket_origin_key: websocket_origin_key(&base_url), + outbound_proxy: None, + egress_key: String::new(), base_url, profile, websocket_pool: None, @@ -234,13 +237,14 @@ impl CodexBackendClient { let websocket_request = websocket_upstream_request(request); let headers = self.request_headers_for_websocket_response(&websocket_request, context)?; - let websocket_create = CodexWebSocketConnection::responses_create_request( + let mut websocket_create = CodexWebSocketConnection::responses_create_request( &self.base_url, &generate_key(), websocket_header_pairs(&headers), &websocket_request, ) .map_err(CodexClientError::WebSocketEncode)?; + websocket_create.connection.outbound_proxy = self.outbound_proxy.clone(); context.trace.cloned().unwrap_or_default().headers( "upstream.request.headers", serde_json::json!({"transport": "websocket", "phase": "prepared_opening"}), @@ -469,6 +473,7 @@ impl CodexBackendClient { .as_deref() .or(request.previous_response_id())?; let mut key = CodexWebSocketPoolKey::new(&self.base_url, account_id, conversation_id) + .with_egress_key(&self.egress_key) .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/tls.rs b/backend/crates/providers/openai/src/transport/tls.rs index 4eadd5305..cdd2d3742 100644 --- a/backend/crates/providers/openai/src/transport/tls.rs +++ b/backend/crates/providers/openai/src/transport/tls.rs @@ -33,6 +33,8 @@ pub const SSL_CERT_FILE_ENV: &str = "SSL_CERT_FILE"; /// 自定义 CA 错误。 #[derive(Debug, Error)] pub enum CustomCaError { + #[error("invalid account proxy configuration")] + ProxyConfiguration, /// 读取 CA 证书文件失败。 #[error( "Failed to read CA certificate file {} selected by {}: {source}. {hint}", @@ -121,6 +123,17 @@ pub fn maybe_build_rustls_client_config_with_custom_ca() -> CustomCaResult CustomCaResult> { + if let Some(config) = maybe_build_rustls_client_config_with_custom_ca()? { + return Ok(config); + } + Ok(Arc::new( + ClientConfig::builder() + .with_root_certificates(native_root_store().map_err(CustomCaError::LoadNativeRoots)?) + .with_no_client_auth(), + )) +} + fn build_reqwest_client_with_env( env_source: &dyn EnvSource, mut builder: reqwest::ClientBuilder, diff --git a/backend/crates/providers/openai/src/transport/websocket/handshake.rs b/backend/crates/providers/openai/src/transport/websocket/handshake.rs index 304673902..4f6d6db31 100644 --- a/backend/crates/providers/openai/src/transport/websocket/handshake.rs +++ b/backend/crates/providers/openai/src/transport/websocket/handshake.rs @@ -5,7 +5,9 @@ use std::time::Duration; use bytes::Bytes; use gateway_protocol::openai::events; use tokio::time::timeout; -use tokio_tungstenite::{Connector, connect_async_tls_with_config}; +use tokio_tungstenite::{ + Connector, MaybeTlsStream, client_async_tls_with_config, connect_async_tls_with_config, +}; use tungstenite::{ self, Message, extensions::{ExtensionsConfig, compression::deflate::DeflateConfig}, @@ -60,7 +62,11 @@ impl CodexWebSocketConnection { "sec-websocket-extensions".to_string(), WEBSOCKET_EXTENSIONS.to_string(), )); - Self { endpoint, headers } + Self { + endpoint, + headers, + outbound_proxy: None, + } } /// 构造 Responses WebSocket opening 与首个 `response.create` 文本帧。 @@ -140,10 +146,39 @@ async fn connect_websocket( ))) })? .map(Connector::Rustls); - let result = timeout( - WEBSOCKET_CONNECT_TIMEOUT, - connect_async_tls_with_config(request, Some(websocket_config()), false, connector), - ) + let result = timeout(WEBSOCKET_CONNECT_TIMEOUT, async { + // Preserve the native direct handshake; explicit egress never inherits a global proxy. + if connection.outbound_proxy.is_none() + && matches!( + tungstenite::proxy::ProxyConfig::from_env(request.uri()), + Ok(None) + ) + { + let (websocket, response) = + connect_async_tls_with_config(request, Some(websocket_config()), false, connector) + .await?; + return Ok((Box::new(websocket) as RawWsStream, response)); + } + let stream = dial_account(connection).await?; + match stream { + MaybeTlsStream::Plain(tcp) => { + let (websocket, response) = + client_async_tls_with_config(request, tcp, Some(websocket_config()), connector) + .await?; + Ok((Box::new(websocket) as RawWsStream, response)) + } + stream => { + let (websocket, response) = client_async_tls_with_config( + request, + stream, + Some(websocket_config()), + connector, + ) + .await?; + Ok((Box::new(websocket) as RawWsStream, response)) + } + } + }) .await .map_err(|_| CodexWebSocketExchangeError::ConnectTimeout { timeout: WEBSOCKET_CONNECT_TIMEOUT, @@ -155,6 +190,89 @@ async fn connect_websocket( } } +async fn dial_account( + connection: &CodexWebSocketConnection, +) -> Result, tungstenite::Error> { + use tokio::net::lookup_host; + use tungstenite::proxy::ProxyConfig; + let invalid = + || tungstenite::Error::Io(std::io::Error::other("account WebSocket egress failed")); + let endpoint = url::Url::parse(connection.endpoint()).map_err(|_| invalid())?; + let host = endpoint.host_str().ok_or_else(invalid)?; + let dns_host = match endpoint.host() { + Some(url::Host::Ipv6(ip)) => ip.to_string(), + _ => host.to_owned(), + }; + let port = endpoint.port_or_known_default().ok_or_else(invalid)?; + let Some(proxy) = connection.outbound_proxy.as_ref() else { + return Ok(MaybeTlsStream::Plain(connect_tcp(host, port).await?)); + }; + let mut proxy_url = url::Url::parse(proxy.expose_url()).map_err(|_| invalid())?; + let tls_proxy = proxy_url.scheme() == "https"; + let proxy_port = proxy_url.port_or_known_default().ok_or_else(invalid)?; + if tls_proxy { + proxy_url.set_scheme("http").map_err(|_| invalid())?; + proxy_url + .set_port(Some(proxy_port)) + .map_err(|_| invalid())?; + } + let config = ProxyConfig::parse(proxy_url.as_str()).map_err(|_| invalid())?; + let proxy_host = match proxy_url.host() { + Some(url::Host::Ipv6(ip)) => ip.to_string(), + _ => config.host.clone(), + }; + let tcp = connect_tcp(&proxy_host, config.port).await?; + let stream = if tls_proxy { + let tls = tls::account_proxy_tls_config().map_err(|_| invalid())?; + let name = rustls_pki_types::ServerName::try_from(proxy_host).map_err(|_| invalid())?; + MaybeTlsStream::Rustls( + tokio_rustls::TlsConnector::from(tls) + .connect(name, tcp) + .await?, + ) + } else { + MaybeTlsStream::Plain(tcp) + }; + // The pinned tungstenite fork sends domains remotely for both SOCKS schemes. + let target = if proxy_url.scheme() == "socks5" { + lookup_host((dns_host.as_str(), port)) + .await? + .next() + .ok_or_else(invalid)? + .ip() + .to_string() + } else if proxy_url.scheme() == "socks5h" { + dns_host + } else { + host.to_owned() + }; + tokio_tungstenite::proxy::connect_via_proxy(stream, &config, &target, port) + .await + .map_err(|_| invalid()) +} + +async fn connect_tcp(host: &str, port: u16) -> Result { + use hyper_util::client::legacy::connect::HttpConnector; + use tower_service::Service; + + // Reuse Hyper's DNS and Happy Eyeballs connector for explicit direct/proxy endpoints. + let authority = host.parse::().map_or_else( + |_| format!("{host}:{port}"), + |ip| std::net::SocketAddr::new(ip, port).to_string(), + ); + let uri = format!("http://{authority}") + .parse::() + .map_err(|_| tungstenite::Error::Io(std::io::Error::other("invalid egress endpoint")))?; + let mut connector = HttpConnector::new(); + connector + .call(uri) + .await + .map(|stream| stream.into_inner()) + .map_err(|_| { + tungstenite::Error::Io(std::io::Error::other("account WebSocket connection failed")) + }) +} + fn websocket_handshake_request( connection: &CodexWebSocketConnection, ) -> Result { diff --git a/backend/crates/providers/openai/src/transport/websocket/model.rs b/backend/crates/providers/openai/src/transport/websocket/model.rs index 2e4231b2e..ecbc0a3bf 100644 --- a/backend/crates/providers/openai/src/transport/websocket/model.rs +++ b/backend/crates/providers/openai/src/transport/websocket/model.rs @@ -9,6 +9,7 @@ use crate::transport::protocol::responses::{ /// WebSocket opening 描述。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct CodexWebSocketConnection { + pub(crate) outbound_proxy: Option, pub(super) endpoint: String, pub(super) headers: Vec<(String, String)>, } @@ -16,7 +17,7 @@ pub struct CodexWebSocketConnection { /// 已构造完成的 Responses WebSocket 请求描述。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct CodexWebSocketRequest { - pub(super) connection: CodexWebSocketConnection, + pub(crate) connection: CodexWebSocketConnection, pub(super) payload_text: String, pub(super) continuation: WebSocketContinuationRequirement, } @@ -113,6 +114,7 @@ impl CodexWebSocketConnection { Self { endpoint: endpoint.into(), headers, + outbound_proxy: None, } } 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 cac0b2533..a4bb5c96f 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/state.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/state.rs @@ -24,6 +24,7 @@ pub struct CodexWebSocketPoolKey { conversation_id: String, connection_profile: String, downstream_connection_id: String, + egress_key: String, } impl CodexWebSocketPoolKey { @@ -39,6 +40,7 @@ impl CodexWebSocketPoolKey { conversation_id: conversation_id.into(), connection_profile: String::new(), downstream_connection_id: String::new(), + egress_key: String::new(), } } @@ -61,6 +63,11 @@ impl CodexWebSocketPoolKey { &self.account_id } + pub(crate) fn with_egress_key(mut self, key: &str) -> Self { + self.egress_key = key.to_owned(); + self + } + pub(crate) fn conversation_id_hash(&self) -> String { short_sha256([self.conversation_id.as_str()]) } @@ -72,6 +79,7 @@ impl CodexWebSocketPoolKey { self.conversation_id.as_str(), self.connection_profile.as_str(), self.downstream_connection_id.as_str(), + self.egress_key.as_str(), ]) } @@ -79,6 +87,7 @@ impl CodexWebSocketPoolKey { self.base_url == other.base_url && self.account_id == other.account_id && self.conversation_id == other.conversation_id + && self.egress_key == other.egress_key } } diff --git a/backend/crates/providers/openai/src/transport/websocket/pump.rs b/backend/crates/providers/openai/src/transport/websocket/pump.rs index 23022ea15..7952a8b1e 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pump.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pump.rs @@ -19,20 +19,35 @@ use std::sync::{ }; use std::time::Duration; -use futures::{SinkExt, StreamExt}; +use futures::{Sink, SinkExt, Stream, StreamExt}; use tokio::{ - net::TcpStream, sync::{mpsc, oneshot}, task::JoinHandle, time::{Instant, MissedTickBehavior}, }; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, tungstenite::Message}; +use tokio_tungstenite::tungstenite::Message; use uuid::Uuid; use super::CodexWebSocketCloseError; /// 底层 tungstenite WebSocket 流。 -pub(crate) type RawWsStream = WebSocketStream>; +pub(crate) trait WebSocketIo: + Stream> + + Sink + + Send + + Unpin +{ +} + +impl WebSocketIo for T where + T: Stream> + + Sink + + Send + + Unpin +{ +} + +pub(crate) type RawWsStream = Box; const PUMP_COMMAND_BUFFER: usize = 32; diff --git a/backend/crates/providers/openai/tests/admin.rs b/backend/crates/providers/openai/tests/admin.rs index f55e08ed8..d9e74f5a7 100644 --- a/backend/crates/providers/openai/tests/admin.rs +++ b/backend/crates/providers/openai/tests/admin.rs @@ -1137,6 +1137,7 @@ fn provider_ports_with_catalog( fn account_record(account: &ProviderAccount) -> AccountRecord { let now = Utc::now(); AccountRecord { + outbound_proxy: None, id: account.id().to_string(), provider_kind: account.provider().clone(), groups: Vec::new(), diff --git a/backend/crates/providers/openai/tests/credential/admin.rs b/backend/crates/providers/openai/tests/credential/admin.rs index 4caffbace..a44331178 100644 --- a/backend/crates/providers/openai/tests/credential/admin.rs +++ b/backend/crates/providers/openai/tests/credential/admin.rs @@ -11,6 +11,73 @@ use crate::support::{TestLeaseCoordinator, runtime_policy}; struct UnusedRefresher; +#[tokio::test] +async fn sub2api_import_resolves_distinct_proxy_bindings_and_encodes_credentials() { + let service = CodexCredentialAdminService::new( + Arc::new(UnusedRefresher), + Arc::new(TestLeaseCoordinator::default()), + runtime_policy(), + ); + let prepared = service.prepare_import_document(serde_json::json!({"data": { + "accounts": [ + {"name": "a", "platform": "openai", "type": "oauth", "proxy_key": "a", "credentials": {"access_token": "at-a"}}, + {"name": "b", "platform": "openai", "type": "oauth", "proxy_key": "b", "credentials": {"access_token": "at-b"}}, + {"name": "direct", "platform": "openai", "type": "oauth", "credentials": {"access_token": "at-c"}} + ], + "proxies": [ + {"proxy_key": "a", "protocol": "http", "host": "127.0.0.1", "port": 18080, "username": "user@a", "password": "p:a/ss", "status": "active"}, + {"proxy_key": "b", "protocol": "socks5", "host": "::1", "port": 1080, "status": "active", "fallback_mode": "none"} + ] + }})).await.unwrap(); + assert_eq!(prepared.accounts().len(), 3); + let first = prepared.accounts()[0].account.outbound_proxy().unwrap(); + assert_eq!(first.endpoint(), "http://127.0.0.1:18080/"); + assert!(first.expose_url().contains("user%40a:p%3Aa%2Fss@")); + assert_eq!( + prepared.accounts()[1] + .account + .outbound_proxy() + .unwrap() + .endpoint(), + "socks5h://[::1]:1080" + ); + assert!(prepared.accounts()[2].account.outbound_proxy().is_none()); + assert!(!format!("{prepared:?}").contains("p:a/ss")); +} + +#[tokio::test] +async fn sub2api_invalid_proxy_bindings_are_rejected_before_any_token_refresh() { + let service = CodexCredentialAdminService::new( + Arc::new(UnusedRefresher), + Arc::new(TestLeaseCoordinator::default()), + runtime_policy(), + ); + let proxy = serde_json::json!({"proxy_key": "bound", "protocol": "http", "host": "127.0.0.1", "port": 18080, "status": "active"}); + let mut cases = vec![ + serde_json::json!([]), + serde_json::json!([proxy.clone(), proxy.clone()]), + ]; + for (field, value) in [ + ("status", serde_json::json!("inactive")), + ("expires_at", serde_json::json!(2000000000)), + ("fallback_mode", serde_json::json!("direct")), + ("port", serde_json::json!(0)), + ] { + let mut invalid = proxy.clone(); + invalid[field] = value; + cases.push(serde_json::json!([invalid])); + } + for proxies in cases { + let result = service.prepare_import_document(serde_json::json!({ + "accounts": [ + {"platform": "openai", "type": "oauth", "credentials": {"refresh_token": "must-not-refresh"}}, + {"platform": "openai", "type": "oauth", "proxy_key": "bound", "credentials": {"access_token": "at"}} + ], "proxies": proxies + })).await; + assert!(result.is_err()); + } +} + #[async_trait] impl TokenRefresher for UnusedRefresher { async fn refresh(&self, _refresh_token: &str) -> Result { diff --git a/backend/crates/providers/openai/tests/credential/token_client.rs b/backend/crates/providers/openai/tests/credential/token_client.rs index b26f6e3d5..673eb4775 100644 --- a/backend/crates/providers/openai/tests/credential/token_client.rs +++ b/backend/crates/providers/openai/tests/credential/token_client.rs @@ -6,6 +6,46 @@ use secrecy::{ExposeSecret, SecretString}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +#[tokio::test] +async fn refresh_and_code_exchange_use_the_selected_account_proxy() { + let direct = MockServer::start().await; + let proxy_a = MockServer::start().await; + let proxy_b = MockServer::start().await; + for (proxy, access) in [(&proxy_a, "exit-a-access"), (&proxy_b, "exit-b-access")] { + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "access_token": access, + "refresh_token": "rotated-refresh", + "id_token": "header.e30.signature" + }))) + .expect(2) + .mount(proxy) + .await; + } + let client = client(&direct); + for (proxy, expected) in [(&proxy_a, "exit-a-access"), (&proxy_b, "exit-b-access")] { + let proxy = gateway_core::account::OutboundProxy::parse(&proxy.uri()).unwrap(); + let refreshed = client + .refresh_with_proxy("initial-refresh", Some(&proxy)) + .await + .unwrap(); + assert_eq!(refreshed.access_token.as_deref(), Some(expected)); + let exchanged = client + .exchange_with_proxy( + AuthorizationCodeGrant { + code: SecretString::from("test-code"), + code_verifier: SecretString::from("test-verifier"), + }, + Some(&proxy), + ) + .await + .unwrap(); + assert_eq!(exchanged.secret.access_token.expose_secret(), expected); + } + assert!(direct.received_requests().await.unwrap().is_empty()); +} + #[tokio::test] async fn refresh_tracks_the_online_profile_without_contaminating_code_exchange() { let server = MockServer::start().await; diff --git a/backend/crates/providers/openai/tests/transport/account_proxy.rs b/backend/crates/providers/openai/tests/transport/account_proxy.rs new file mode 100644 index 000000000..6de5eea89 --- /dev/null +++ b/backend/crates/providers/openai/tests/transport/account_proxy.rs @@ -0,0 +1,330 @@ +use super::*; +use gateway_core::{ + account::{CredentialRevision, OutboundProxy, ProviderAccount, ProviderAccountId}, + routing::ProviderKind, +}; + +fn account(id: &str, proxy: Option<&str>) -> ProviderAccount { + ProviderAccount::new( + ProviderAccountId::new(format!("acct_{id}")).unwrap(), + ProviderKind::new("openai").unwrap(), + id.to_owned(), + Some(id.to_owned()), + "oauth".to_owned(), + CredentialRevision::new(1).unwrap(), + None, + ) + .with_outbound_proxy(proxy.map(|value| OutboundProxy::parse(value).unwrap())) +} + +async fn http_exit(label: &'static str) -> (String, tokio::task::JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let request = read_http_request_with_body(&mut stream).await; + let body = format!( + "event: response.completed\ndata: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_{label}\",\"status\":\"completed\",\"output\":[],\"usage\":{{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}}}}\n\n" + ); + stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); + let header_end = request + .windows(4) + .position(|bytes| bytes == b"\r\n\r\n") + .unwrap(); + String::from_utf8(request[..header_end].to_vec()).unwrap() + }); + (format!("http://{address}"), task) +} + +#[tokio::test] +async fn http_sse_accounts_use_distinct_proxies_and_clearing_restores_direct_client() { + let (exit_a, a) = http_exit("exit_a").await; + let (exit_b, b) = http_exit("exit_b").await; + let (direct_url, direct) = http_exit("direct").await; + let base = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + &direct_url, + test_wire_profile(), + ); + let client_a = base.for_account(&account("a", Some(&exit_a))).unwrap(); + let client_b = base.for_account(&account("b", Some(&exit_b))).unwrap(); + let mut request = codex_request("gpt-5.5", "", Vec::new()); + request.force_http_sse = true; + let (response_a, response_b) = tokio::join!( + client_a.create_response(&request, request_context("proxy-a", Some("a"))), + client_b.create_response(&request, request_context("proxy-b", Some("b"))) + ); + assert!(response_a.unwrap().body.contains("resp_exit_a")); + assert!(response_b.unwrap().body.contains("resp_exit_b")); + let cleared = client_a.for_account(&account("a", None)).unwrap(); + assert!( + cleared + .create_response(&request, request_context("direct", Some("a"))) + .await + .unwrap() + .body + .contains("resp_direct") + ); + for head in [a.await.unwrap(), b.await.unwrap()] { + assert!(head.starts_with(&format!("POST {direct_url}/codex/responses HTTP/1.1"))); + } + assert!( + direct + .await + .unwrap() + .starts_with("POST /codex/responses HTTP/1.1") + ); +} + +async fn websocket_exit(label: &'static str) -> (String, tokio::task::JoinHandle) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let connect = read_http_request(&mut stream).await; + assert!(connect.starts_with("CONNECT upstream.invalid:80 HTTP/1.1")); + stream + .write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n") + .await + .unwrap(); + let mut websocket = accept_codex_test_websocket(stream).await; + assert!(matches!(websocket.next().await, Some(Ok(Message::Text(_))))); + websocket + .send(Message::Text( + completed_websocket_response(label, 2, 1).into(), + )) + .await + .unwrap(); + connect + }); + (format!("http://user:pass@{address}"), task) +} + +#[tokio::test] +async fn websocket_connect_and_pool_are_isolated_when_account_proxy_changes() { + let (exit_a, a) = websocket_exit("resp_exit_a").await; + let (exit_b, b) = websocket_exit("resp_exit_b").await; + let pool = Arc::new(CodexWebSocketPool::new(Duration::from_mins(1))); + let base = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + "http://upstream.invalid", + test_wire_profile(), + ) + .with_websocket_pool(pool); + let mut request = codex_request("gpt-5.5", "", Vec::new()); + request.set_previous_response_id(Some("resp_previous".to_owned())); + request.previous_response_scope = Some(PreviousResponseScope::Persisted); + for (proxy, expected) in [(exit_a, "resp_exit_a"), (exit_b, "resp_exit_b")] { + let client = base + .for_account(&account("same-account", Some(&proxy))) + .unwrap(); + let response = timeout( + Duration::from_secs(5), + client.create_response(&request, request_context("proxy-ws", Some("same-account"))), + ) + .await + .unwrap() + .unwrap(); + assert!(response.body.contains(expected)); + } + for head in [a.await.unwrap(), b.await.unwrap()] { + assert_eq!( + read_header_value(&head, "Proxy-Authorization"), + Some("Basic dXNlcjpwYXNz") + ); + } +} + +#[tokio::test] +async fn unavailable_proxy_fails_without_direct_fallback() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let unreachable = format!("http://{}", listener.local_addr().unwrap()); + drop(listener); + let direct = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{}", direct.local_addr().unwrap()), + test_wire_profile(), + ); + let client = base.for_account(&account("a", Some(&unreachable))).unwrap(); + let mut request = codex_request("gpt-5.5", "", Vec::new()); + request.force_http_sse = true; + assert!( + client + .create_response(&request, request_context("blocked", Some("a"))) + .await + .is_err() + ); + assert!( + timeout(Duration::from_millis(30), direct.accept()) + .await + .is_err() + ); +} + +async fn socks_exit(stream: &mut TcpStream) -> (u8, String, u16) { + assert_eq!(stream.read_u8().await.unwrap(), 5); + let count = stream.read_u8().await.unwrap(); + let mut methods = vec![0; usize::from(count)]; + stream.read_exact(&mut methods).await.unwrap(); + assert!(methods.contains(&2)); + stream.write_all(&[5, 2]).await.unwrap(); + assert_eq!(stream.read_u8().await.unwrap(), 1); + let count = stream.read_u8().await.unwrap(); + let mut user = vec![0; usize::from(count)]; + stream.read_exact(&mut user).await.unwrap(); + let count = stream.read_u8().await.unwrap(); + let mut password = vec![0; usize::from(count)]; + stream.read_exact(&mut password).await.unwrap(); + assert_eq!(user, b"user@exit"); + assert_eq!(password, b"pass:word"); + stream.write_all(&[1, 0]).await.unwrap(); + let mut header = [0; 4]; + stream.read_exact(&mut header).await.unwrap(); + assert_eq!(&header[..3], &[5, 1, 0]); + let host = match header[3] { + 1 => { + let mut ip = [0; 4]; + stream.read_exact(&mut ip).await.unwrap(); + std::net::Ipv4Addr::from(ip).to_string() + } + 4 => { + let mut ip = [0; 16]; + stream.read_exact(&mut ip).await.unwrap(); + std::net::Ipv6Addr::from(ip).to_string() + } + 3 => { + let count = stream.read_u8().await.unwrap(); + let mut host = vec![0; usize::from(count)]; + stream.read_exact(&mut host).await.unwrap(); + String::from_utf8(host).unwrap() + } + _ => panic!("unexpected SOCKS address type"), + }; + let port = stream.read_u16().await.unwrap(); + stream + .write_all(&[5, 0, 0, 1, 127, 0, 0, 1, 0, 80]) + .await + .unwrap(); + (header[3], host, port) +} + +#[tokio::test] +async fn socks_account_egress_preserves_dns_mode_auth_and_ipv6_for_sse_and_websocket() { + for websocket in [false, true] { + let mut targets = vec![("socks5", "localhost"), ("socks5h", "upstream.invalid")]; + // The reqwest SOCKS connector currently encodes bracketed IPv6 as a domain. + if websocket { + targets.push(("socks5h", "[::1]")); + } + for (scheme, host) in targets { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = format!( + "{scheme}://user%40exit:pass%3Aword@{}", + listener.local_addr().unwrap() + ); + let task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let target = socks_exit(&mut stream).await; + if websocket { + let mut stream = accept_codex_test_websocket(stream).await; + assert!(matches!(stream.next().await, Some(Ok(Message::Text(_))))); + stream + .send(Message::Text( + completed_websocket_response("resp_socks", 2, 1).into(), + )) + .await + .unwrap(); + } else { + read_http_request_with_body(&mut stream).await; + let body = format!( + "data: {}\n\n", + completed_websocket_response("resp_socks", 2, 1) + ); + stream.write_all(format!("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.unwrap(); + } + target + }); + let client = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{host}"), + test_wire_profile(), + ) + .for_account(&account("socks", Some(&proxy))) + .unwrap(); + let mut request = codex_request("gpt-5.5", "", Vec::new()); + request.force_http_sse = !websocket; + if websocket { + request.set_previous_response_id(Some("resp_previous".to_owned())); + request.previous_response_scope = Some(PreviousResponseScope::Persisted); + } + let response = timeout( + Duration::from_secs(5), + client.create_response(&request, request_context("socks", Some("socks"))), + ) + .await + .unwrap() + .unwrap(); + assert!(response.body.contains("resp_socks")); + let (address_type, target, port) = task.await.unwrap(); + assert_eq!(port, 80); + match host { + "localhost" => { + assert!(matches!(address_type, 1 | 4)); + assert!(target.parse::().unwrap().is_loopback()); + } + "[::1]" => assert_eq!((address_type, target.as_str()), (4, "::1")), + _ => assert_eq!((address_type, target.as_str()), (3, host)), + } + } + } +} + +#[tokio::test] +async fn https_account_proxy_starts_tls_and_rejection_never_falls_back_to_direct() { + for websocket in [false, true] { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = format!("https://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut header = [0; 5]; + stream.read_exact(&mut header).await.unwrap(); + assert_eq!(&header[..2], &[22, 3]); + let mut hello = vec![0; usize::from(u16::from_be_bytes([header[3], header[4]]))]; + stream.read_exact(&mut hello).await.unwrap(); + stream.write_all(&[21, 3, 3, 0, 2, 2, 40]).await.unwrap(); + }); + let direct = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let client = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{}", direct.local_addr().unwrap()), + test_wire_profile(), + ) + .for_account(&account("tls-proxy", Some(&proxy))) + .unwrap(); + let mut request = codex_request("gpt-5.5", "", Vec::new()); + request.force_http_sse = !websocket; + if websocket { + request.set_previous_response_id(Some("resp_previous".to_owned())); + request.previous_response_scope = Some(PreviousResponseScope::Persisted); + } + assert!( + timeout( + Duration::from_secs(5), + client.create_response(&request, request_context("tls-proxy", Some("tls-proxy"))) + ) + .await + .unwrap() + .is_err() + ); + timeout(Duration::from_secs(5), task) + .await + .unwrap() + .unwrap(); + assert!( + timeout(Duration::from_millis(30), direct.accept()) + .await + .is_err() + ); + } +} diff --git a/backend/crates/providers/openai/tests/transport/mod.rs b/backend/crates/providers/openai/tests/transport/mod.rs index 609b61465..ffcf09faf 100644 --- a/backend/crates/providers/openai/tests/transport/mod.rs +++ b/backend/crates/providers/openai/tests/transport/mod.rs @@ -45,6 +45,7 @@ use tokio_tungstenite::{ }, }; +mod account_proxy; mod canonical; mod catalog; mod client; diff --git a/backend/crates/providers/xai/Cargo.toml b/backend/crates/providers/xai/Cargo.toml index 0923a1159..57e7f2dbf 100644 --- a/backend/crates/providers/xai/Cargo.toml +++ b/backend/crates/providers/xai/Cargo.toml @@ -20,7 +20,7 @@ gateway-protocol = { path = "../../gateway-protocol" } getrandom = "0.4.3" jsonwebtoken = { version = "=11.0.0", default-features = false, features = ["aws_lc_rs"] } # 与官方 Grok CLI 一致：reqwest 不启用压缩特性，上游请求不带 accept-encoding。 -reqwest = { version = "=0.12.28", default-features = false, features = ["stream", "rustls-tls-native-roots", "http2"] } +reqwest = { version = "=0.12.28", default-features = false, features = ["stream", "rustls-tls-native-roots", "http2", "socks"] } serde = { version = "1.0.228", features = ["derive"] } serde_json = { version = "1.0.150", features = ["arbitrary_precision", "preserve_order"] } semver = "1.0.27" diff --git a/backend/crates/providers/xai/src/admin.rs b/backend/crates/providers/xai/src/admin.rs index b92426e61..b44ffac4c 100644 --- a/backend/crates/providers/xai/src/admin.rs +++ b/backend/crates/providers/xai/src/admin.rs @@ -431,7 +431,6 @@ impl ProviderAdmin for XaiAdminProvider { .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; let document = GrokOAuthImportDocument::parse_json(&document) .map_err(|_| provider_error(ProviderAdminErrorKind::Invalid))?; - let discovery = self.oauth.discover().await.map_err(map_oauth_error)?; let mut subjects = BTreeSet::new(); let mut credentials = Vec::new(); // 逐条目独立成败：verify 会真实轮换 RT，后续条目的任何失败都不得 @@ -444,11 +443,16 @@ impl ProviderAdmin for XaiAdminProvider { for entry in document.into_entries() { let name = entry.name().to_owned(); let email = entry.email().map(str::to_owned); - let tokens = match self - .oauth - .verify_imported_credential(&discovery, entry.into_candidate()) - .await - { + let outbound_proxy = entry.outbound_proxy().cloned(); + let oauth = self.oauth.with_outbound_proxy(outbound_proxy.clone()); + let verified = async { + let discovery = oauth.discover().await?; + oauth + .verify_imported_credential(&discovery, entry.into_candidate()) + .await + } + .await; + let tokens = match verified { Ok(tokens) => tokens, Err(error) => { tracing::warn!( @@ -493,6 +497,10 @@ impl ProviderAdmin for XaiAdminProvider { continue; } }; + let prepared = NewProviderAccount { + account: prepared.account.with_outbound_proxy(outbound_proxy), + credential: prepared.credential, + }; match prepared_create(prepared, Utc::now()) { Ok(credential) => credentials.push(credential), Err(error) => { @@ -533,7 +541,12 @@ impl ProviderAdmin for XaiAdminProvider { return Err(provider_error(ProviderAdminErrorKind::NotFound)); } } - let discovery = self.oauth.discover().await.map_err(map_oauth_error)?; + let discovery = self + .oauth + .with_outbound_proxy(pending.outbound_proxy().cloned()) + .discover() + .await + .map_err(map_oauth_error)?; let redirect = RedirectUriAllowlist::new([crate::OFFICIAL_REDIRECT_URI]) .and_then(|allowlist| allowlist.authorize(crate::OFFICIAL_REDIRECT_URI)) .map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?; @@ -623,9 +636,11 @@ impl ProviderAdmin for XaiAdminProvider { rotation.expires_at, ), ); - let discovery = self.oauth.discover().await.map_err(map_oauth_error)?; - let tokens = self + let oauth = self .oauth + .with_outbound_proxy(current.account.outbound_proxy().cloned()); + let discovery = oauth.discover().await.map_err(map_oauth_error)?; + let tokens = oauth .verify_imported_credential(&discovery, candidate) .await .map_err(|error| map_failure_class(error.class()))?; @@ -775,9 +790,28 @@ impl XaiAdminProvider { let grant = authorization .accept_authorization_input(callback_url) .map_err(map_oauth_error)?; - let discovery = self.oauth.discover().await.map_err(map_oauth_error)?; - let tokens = self - .oauth + let current = match stored.mutation.target() { + AuthorizationMutationTarget::Create { .. } => None, + AuthorizationMutationTarget::Reauthorize { account_id } => { + let current = self + .accounts + .load_current_credential(account_id) + .await + .map_err(map_store_error)?; + if current.account.provider() != &self.provider_kind { + return Err(provider_error(ProviderAdminErrorKind::NotFound)); + } + Some(current) + } + }; + let proxy = current + .as_ref() + .map_or(stored.mutation.outbound_proxy(), |current| { + current.account.outbound_proxy() + }); + let oauth = self.oauth.with_outbound_proxy(proxy.cloned()); + let discovery = oauth.discover().await.map_err(map_oauth_error)?; + let tokens = oauth .exchange_authorization_code(&discovery, grant) .await .map_err(map_oauth_error)?; @@ -798,17 +832,17 @@ impl XaiAdminProvider { enabled: true, }) .map_err(map_repository_error)?; + let prepared = NewProviderAccount { + account: prepared + .account + .with_outbound_proxy(stored.mutation.outbound_proxy().cloned()), + credential: prepared.credential, + }; PreparedAuthorizationCredential::Create(prepared_create(prepared, Utc::now())?) } - AuthorizationMutationTarget::Reauthorize { account_id } => { - let current = self - .accounts - .load_current_credential(account_id) - .await - .map_err(map_store_error)?; - if current.account.provider() != &self.provider_kind { - return Err(provider_error(ProviderAdminErrorKind::NotFound)); - } + AuthorizationMutationTarget::Reauthorize { .. } => { + let current = + current.ok_or_else(|| provider_error(ProviderAdminErrorKind::Internal))?; let prepared = verified_rotation(current, tokens)?; PreparedAuthorizationCredential::Reauthorize(prepared_rotation( prepared, @@ -898,6 +932,7 @@ fn prepared_create( access_token_expires_at: account.access_token_expires_at().map(DateTime::::from), next_refresh_at: account.next_refresh_at().map(Into::into), enabled: account.enabled(), + outbound_proxy: account.outbound_proxy().cloned(), credential_state: account.credential_state(), credential_observed_at: observed_at, }) @@ -998,6 +1033,7 @@ fn account_from_record(account: &AccountRecord) -> Result, verifier: Arc, + outbound_proxy: Option, } impl GrokOAuthClient { @@ -32,6 +33,7 @@ impl GrokOAuthClient { wire_profile, transport, verifier, + outbound_proxy: None, } } @@ -41,6 +43,13 @@ impl GrokOAuthClient { &self.config } + #[must_use] + pub fn with_outbound_proxy(&self, proxy: Option) -> Self { + let mut client = self.clone(); + client.outbound_proxy = proxy; + client + } + /// 拉取并校验官方同源 OIDC 发现文档。 /// /// # Errors @@ -129,7 +138,13 @@ impl GrokOAuthClient { tokens.id_token.as_ref(), tokens.expires_in, ); - let evidence = self.verifier.verify(context, candidate).await?; + let evidence = self + .verifier + .verify( + context.with_outbound_proxy(self.outbound_proxy.as_ref()), + candidate, + ) + .await?; if evidence.method() != VerificationMethod::IdToken { return Err(VerificationFailure::WrongEvidence.into()); } @@ -194,7 +209,10 @@ impl GrokOAuthClient { ); let evidence = self .verifier - .verify(context, verification_candidate) + .verify( + context.with_outbound_proxy(self.outbound_proxy.as_ref()), + verification_candidate, + ) .await .map_err(OAuthError::from)?; let expected_method = match flow { @@ -251,8 +269,9 @@ impl GrokOAuthClient { async fn execute( &self, operation: OAuthOperation, - request: OAuthHttpRequest, + mut request: OAuthHttpRequest, ) -> Result { + request.outbound_proxy = self.outbound_proxy.clone(); self.transport .execute(request) .await diff --git a/backend/crates/providers/xai/src/credential/http.rs b/backend/crates/providers/xai/src/credential/http.rs index 6145678f7..582922746 100644 --- a/backend/crates/providers/xai/src/credential/http.rs +++ b/backend/crates/providers/xai/src/credential/http.rs @@ -121,6 +121,7 @@ pub struct OAuthHttpRequest { url: Url, headers: Vec, form: Vec, + pub(crate) outbound_proxy: Option, } impl OAuthHttpRequest { @@ -130,6 +131,7 @@ impl OAuthHttpRequest { url, headers: Vec::new(), form: Vec::new(), + outbound_proxy: None, } } @@ -139,6 +141,7 @@ impl OAuthHttpRequest { url, headers, form, + outbound_proxy: None, } } diff --git a/backend/crates/providers/xai/src/credential/import.rs b/backend/crates/providers/xai/src/credential/import.rs index ded48e350..8794603f0 100644 --- a/backend/crates/providers/xai/src/credential/import.rs +++ b/backend/crates/providers/xai/src/credential/import.rs @@ -114,10 +114,16 @@ impl GrokOAuthImportMetadata { pub struct GrokOAuthImportEntry { name: String, email: Option, + outbound_proxy: Option, candidate: GrokOAuthImportCandidate, } impl GrokOAuthImportEntry { + #[must_use] + pub fn outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> { + self.outbound_proxy.as_ref() + } + #[must_use] pub fn name(&self) -> &str { &self.name @@ -153,7 +159,7 @@ pub struct GrokOAuthImportDocument { impl GrokOAuthImportDocument { /// 从外部 JSON 提取 xAI OAuth 认证字段。 /// - /// 来源包装格式、代理、并发和其他展示 metadata 不参与认证，也不会影响导入。 + /// 独立代理 URL 随账号导入；其他展示 metadata 不参与认证。 /// 实际 token 仍须通过官方 refresh/user-info 验证；API Key 不能混入 OAuth 条目。 pub fn parse_json(document: &[u8]) -> Result { if document.is_empty() || document.len() > MAX_IMPORT_DOCUMENT_BYTES { @@ -290,9 +296,30 @@ fn parse_account_entry( expires_at, ); + let outbound_proxy = match account + .get("outboundProxyUrl") + .or_else(|| account.get("outbound_proxy_url")) + { + Some(Value::String(value)) if value.is_empty() => None, + Some(Value::String(value)) => Some( + gateway_core::account::OutboundProxy::parse(value) + .map_err(|_| GrokOAuthImportError::InvalidField("outboundProxyUrl"))?, + ), + Some(Value::Null) | None => { + if account + .get("proxy_key") + .is_some_and(|value| !value.is_null() && value.as_str() != Some("")) + { + return Err(GrokOAuthImportError::InvalidField("proxy_key")); + } + None + } + Some(_) => return Err(GrokOAuthImportError::InvalidField("outboundProxyUrl")), + }; Ok(Some(GrokOAuthImportEntry { name, email, + outbound_proxy, candidate: GrokOAuthImportCandidate::new(tokens, metadata), })) } diff --git a/backend/crates/providers/xai/src/credential/oidc_verifier.rs b/backend/crates/providers/xai/src/credential/oidc_verifier.rs index 6dc9f0651..2ebe7d9e9 100644 --- a/backend/crates/providers/xai/src/credential/oidc_verifier.rs +++ b/backend/crates/providers/xai/src/credential/oidc_verifier.rs @@ -303,7 +303,25 @@ impl TokenVerifier for ReqwestOidcTokenVerifier { context: TokenVerificationContext<'a>, candidate: TokenCandidate<'a>, ) -> VerificationFuture<'a> { - Box::pin(async move { self.verify_inner(context, candidate).await }) + Box::pin(async move { + if let Some(proxy) = context.outbound_proxy() { + let client = crate::transport::network::account_proxy_client( + &self.client, + Some(proxy), + Some(REQUEST_TIMEOUT), + ) + .map_err(|_| VerificationFailure::Unavailable)?; + let scoped = Self { + client, + endpoint_policy: Arc::clone(&self.endpoint_policy), + jwks_cache: Mutex::new(JwksCache::default()), + cache_ttl: self.cache_ttl, + }; + scoped.verify_inner(context, candidate).await + } else { + self.verify_inner(context, candidate).await + } + }) } } diff --git a/backend/crates/providers/xai/src/credential/refresh.rs b/backend/crates/providers/xai/src/credential/refresh.rs index 44e6a64b3..eb46a4372 100644 --- a/backend/crates/providers/xai/src/credential/refresh.rs +++ b/backend/crates/providers/xai/src/credential/refresh.rs @@ -92,6 +92,7 @@ pub struct DueGrokCredential { upstream_account_id: Option, plan_type: Option, refresh_token_expires_at: Option>, + outbound_proxy: Option, } struct DueGrokRefreshBatch { @@ -171,6 +172,16 @@ pub trait GrokCredentialRefresher: Send + Sync { &self, refresh_token: &SecretValue, ) -> Result; + async fn refresh_with_proxy( + &self, + token: &SecretValue, + proxy: Option<&gateway_core::account::OutboundProxy>, + ) -> Result { + if proxy.is_some() { + return Err(GrokRefreshFailure::Transient); + } + self.refresh(token).await + } } pub struct GrokOAuthRefreshClient { @@ -199,6 +210,25 @@ impl GrokOAuthRefreshClient { #[async_trait] impl GrokCredentialRefresher for GrokOAuthRefreshClient { + async fn refresh_with_proxy( + &self, + token: &SecretValue, + proxy: Option<&gateway_core::account::OutboundProxy>, + ) -> Result { + if proxy.is_none() { + return self.refresh(token).await; + } + let client = self.client.with_outbound_proxy(proxy.cloned()); + let discovery = client + .discover() + .await + .map_err(classify_oauth_refresh_error)?; + let tokens = client + .refresh(&discovery, &RefreshTokenGrant::new(token.clone())) + .await + .map_err(classify_oauth_refresh_error)?; + refreshed_tokens(tokens) + } async fn prepare_cycle(&self) -> Result<(), GrokRefreshFailure> { let mut cached = self.discovery.lock().await; match &*cached { @@ -361,10 +391,12 @@ impl GrokCredentialRefreshService { } Err(error) => return Err(error.into()), }; - self.refresher - .prepare_cycle() - .await - .map_err(|_| GrokCredentialRefreshError::Preparation)?; + if loaded.account.outbound_proxy().is_none() { + self.refresher + .prepare_cycle() + .await + .map_err(|_| GrokCredentialRefreshError::Preparation)?; + } let policy = self.runtime_policy.load_refresh_policy().await?; let subject = loaded .account @@ -383,6 +415,7 @@ impl GrokCredentialRefreshService { upstream_account_id: loaded.account.upstream_account_id().map(str::to_owned), plan_type: loaded.account.plan_type().map(str::to_owned), refresh_token_expires_at: loaded.refresh_token_expires_at, + outbound_proxy: loaded.account.outbound_proxy().cloned(), }; self.refresh_one_with_policy(credential, policy).await } @@ -404,10 +437,12 @@ impl GrokCredentialRefreshService { GrokRefreshFailure::InvalidGrant, )); } - self.refresher - .prepare_cycle() - .await - .map_err(GrokCredentialRefreshError::ManualFailure)?; + if loaded.account.outbound_proxy().is_none() { + self.refresher + .prepare_cycle() + .await + .map_err(GrokCredentialRefreshError::ManualFailure)?; + } let capacity_guard = match self .leases .try_acquire(ProviderLeaseRequest::RefreshCapacity( @@ -434,7 +469,7 @@ impl GrokCredentialRefreshService { }; let tokens = self .refresher - .refresh(&loaded.refresh_token) + .refresh_with_proxy(&loaded.refresh_token, loaded.account.outbound_proxy()) .await .map_err(GrokCredentialRefreshError::ManualFailure)?; if tokens.access_token.is_empty() || tokens.access_token.len() > MAX_SECRET_BYTES { @@ -494,14 +529,32 @@ impl GrokCredentialRefreshService { .into_iter() .map(|account_id| GrokCredentialRefreshOutcome::Failed { account_id }) .collect::>(); - let credentials = batch.credentials; + let mut credentials = batch.credentials; if credentials.is_empty() { return Ok(outcomes); } - self.refresher - .prepare_cycle() - .await - .map_err(|_| GrokCredentialRefreshError::Preparation)?; + if credentials + .iter() + .any(|credential| credential.outbound_proxy.is_none()) + && self.refresher.prepare_cycle().await.is_err() + { + if credentials + .iter() + .all(|credential| credential.outbound_proxy.is_none()) + { + return Err(GrokCredentialRefreshError::Preparation); + } + // A failed direct discovery must not prevent accounts on other exits from refreshing. + credentials.retain(|credential| { + if credential.outbound_proxy.is_some() { + return true; + } + outcomes.push(GrokCredentialRefreshOutcome::Failed { + account_id: credential.account_id.clone(), + }); + false + }); + } outcomes.reserve(credentials.len()); let refreshed = stream::iter(credentials.into_iter().map(|credential| async move { let account_id = credential.account_id.clone(); @@ -587,7 +640,14 @@ impl GrokCredentialRefreshService { } }; - match self.refresher.refresh(&credential.refresh_token).await { + match self + .refresher + .refresh_with_proxy( + &credential.refresh_token, + credential.outbound_proxy.as_ref(), + ) + .await + { Ok(tokens) => self.persist_success(credential, tokens).await, Err(GrokRefreshFailure::InvalidGrant) => { self.persist_terminal_failure( @@ -912,6 +972,7 @@ impl GrokCredentialRepository { upstream_account_id: account.upstream_account_id().map(str::to_owned), plan_type: account.plan_type().map(str::to_owned), refresh_token_expires_at: loaded.refresh_token_expires_at, + outbound_proxy: account.outbound_proxy().cloned(), }); } Ok(DueGrokRefreshBatch { diff --git a/backend/crates/providers/xai/src/credential/repository.rs b/backend/crates/providers/xai/src/credential/repository.rs index afa4ab40d..bef40e7b8 100644 --- a/backend/crates/providers/xai/src/credential/repository.rs +++ b/backend/crates/providers/xai/src/credential/repository.rs @@ -274,7 +274,7 @@ impl GrokCredentialAdmin { if let Some(email) = loaded.account.email() { credentials.insert("email".to_owned(), Value::String(email.to_owned())); } - exported_accounts.push(serde_json::json!({ + let mut exported = serde_json::json!({ "name": loaded.account.name(), "platform": "grok", "type": "oauth", @@ -285,7 +285,11 @@ impl GrokCredentialAdmin { .account .email() .map_or_else(|| serde_json::json!({}), |email| serde_json::json!({"email": email})), - })); + }); + if let Some(proxy) = loaded.account.outbound_proxy() { + exported["outboundProxyUrl"] = Value::String(proxy.expose_url().to_owned()); + } + exported_accounts.push(exported); } Ok(GrokAccountExport(serde_json::json!({ "version": 1, diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index bcf2d2806..453367fb4 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -253,7 +253,8 @@ impl GrokAccountSessionSelector { return Err(GrokSessionSelectorError::InvalidSession); } let binding = GrokSessionBinding::new(selected_id.as_str()) - .map_err(|_| GrokSessionSelectorError::InvalidSession)?; + .map_err(|_| GrokSessionSelectorError::InvalidSession)? + .with_outbound_proxy(loaded.account.outbound_proxy().cloned()); let upstream_user_id = loaded .account .upstream_user_id() diff --git a/backend/crates/providers/xai/src/credential/verification.rs b/backend/crates/providers/xai/src/credential/verification.rs index a941ccab4..8ab1497d8 100644 --- a/backend/crates/providers/xai/src/credential/verification.rs +++ b/backend/crates/providers/xai/src/credential/verification.rs @@ -87,6 +87,7 @@ pub struct TokenVerificationContext<'a> { userinfo_endpoint: &'a Url, signing_algorithms: &'a [String], expected_nonce: Option<&'a SecretValue>, + outbound_proxy: Option<&'a gateway_core::account::OutboundProxy>, } impl<'a> TokenVerificationContext<'a> { @@ -107,9 +108,22 @@ impl<'a> TokenVerificationContext<'a> { userinfo_endpoint, signing_algorithms, expected_nonce, + outbound_proxy: None, } } + pub(crate) fn with_outbound_proxy( + mut self, + proxy: Option<&'a gateway_core::account::OutboundProxy>, + ) -> Self { + self.outbound_proxy = proxy; + self + } + + pub(crate) fn outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> { + self.outbound_proxy + } + /// 返回正在验证的流程。 #[must_use] pub const fn flow(&self) -> VerificationFlow { diff --git a/backend/crates/providers/xai/src/transport/catalog.rs b/backend/crates/providers/xai/src/transport/catalog.rs index febbe70dd..635d6d561 100644 --- a/backend/crates/providers/xai/src/transport/catalog.rs +++ b/backend/crates/providers/xai/src/transport/catalog.rs @@ -44,6 +44,7 @@ pub struct GrokModelCatalogSession { user_id: SecretValue, email: Option, wire_profile: XaiWireProfileState, + outbound_proxy: Option, } impl GrokModelCatalogSession { @@ -71,8 +72,18 @@ impl GrokModelCatalogSession { user_id, email, wire_profile, + outbound_proxy: None, }) } + + #[must_use] + pub fn with_outbound_proxy( + mut self, + proxy: Option, + ) -> Self { + self.outbound_proxy = proxy; + self + } } impl fmt::Debug for GrokModelCatalogSession { @@ -92,6 +103,7 @@ impl fmt::Debug for GrokModelCatalogSession { pub struct GrokModelCatalogRequest { endpoint: Url, headers: Vec, + pub(crate) outbound_proxy: Option, } impl GrokModelCatalogRequest { @@ -115,7 +127,11 @@ impl GrokModelCatalogRequest { if let Some(email) = &session.email { headers.push(GrokHeader::sensitive("x-email", email.clone())); } - Ok(Self { endpoint, headers }) + Ok(Self { + endpoint, + headers, + outbound_proxy: session.outbound_proxy.clone(), + }) } /// 返回固定官方 `/v1/models` URL。 @@ -244,6 +260,7 @@ pub trait GrokModelCatalogTransport: Send + Sync { pub struct GrokBillingRequest { endpoint: Url, headers: Vec, + pub(crate) outbound_proxy: Option, } impl GrokBillingRequest { @@ -267,7 +284,11 @@ impl GrokBillingRequest { if let Some(email) = &session.email { headers.push(GrokHeader::sensitive("x-email", email.clone())); } - Ok(Self { endpoint, headers }) + Ok(Self { + endpoint, + headers, + outbound_proxy: session.outbound_proxy.clone(), + }) } /// 返回固定官方 `/v1/billing?format=credits` URL。 diff --git a/backend/crates/providers/xai/src/transport/network.rs b/backend/crates/providers/xai/src/transport/network.rs index 64648d8a9..5ad040827 100644 --- a/backend/crates/providers/xai/src/transport/network.rs +++ b/backend/crates/providers/xai/src/transport/network.rs @@ -1,4 +1,4 @@ -//! 关闭 redirect、proxy 与业务重试的生产 reqwest transport。 +//! Explicit account egress with redirects and business retries disabled. use gateway_core::diagnostics::{StreamCapture, StreamFormat, TraceContext}; use std::collections::{HashMap, VecDeque}; @@ -268,6 +268,12 @@ impl OAuthHttpTransport for ReqwestOAuthTransport { if !endpoint_policy.validate_oauth(request.url()) { return Err(TransportFailure::new(TransportFailureKind::NotSent)); } + let client = account_proxy_client( + &client, + request.outbound_proxy.as_ref(), + Some(OAUTH_REQUEST_TIMEOUT), + ) + .map_err(|_| TransportFailure::new(TransportFailureKind::NotSent))?; let mut builder = match request.method() { HttpMethod::Get => client.get(request.url().clone()), HttpMethod::Post => client.post(request.url().clone()), @@ -334,7 +340,11 @@ impl ReqwestGrokInferenceTransport { if let Some(client) = clients.get(binding) { return Ok((client, GrokInferenceClientCacheStatus::Hit)); } - clients.take_unbound() + if binding.outbound_proxy().is_none() { + clients.take_unbound() + } else { + None + } }; let client = match unbound_client { @@ -344,6 +354,8 @@ impl ReqwestGrokInferenceTransport { .build_inference_client(None) .map_err(|_| inference_client_pool_unavailable_for_cache_miss())?, }; + let client = account_proxy_client(&client, binding.outbound_proxy(), None) + .map_err(|_| inference_client_pool_unavailable_for_cache_miss())?; let client = self .clients .lock() @@ -592,6 +604,14 @@ impl GrokModelCatalogTransport for ReqwestGrokModelCatalogTransport { GrokModelCatalogTransportErrorKind::Protocol, )); } + let client = account_proxy_client( + &client, + request.outbound_proxy.as_ref(), + Some(OAUTH_REQUEST_TIMEOUT), + ) + .map_err(|_| { + GrokModelCatalogTransportError::new(GrokModelCatalogTransportErrorKind::Protocol) + })?; let mut builder = client.get(request.endpoint().clone()); for header in request.headers() { builder = builder.header(header.name(), header.value().expose()); @@ -643,6 +663,12 @@ impl GrokBillingTransport for ReqwestGrokModelCatalogTransport { .ok_or_else(|| { GrokBillingTransportError::new(GrokBillingTransportErrorKind::Protocol) })?; + let client = account_proxy_client( + &client, + request.outbound_proxy.as_ref(), + Some(OAUTH_REQUEST_TIMEOUT), + ) + .map_err(|_| GrokBillingTransportError::new(GrokBillingTransportErrorKind::Protocol))?; let mut builder = client.get(endpoint); for header in request.headers() { builder = builder.header(header.name(), header.value().expose()); @@ -698,6 +724,34 @@ fn build_official_client( .map_err(|_| GrokReqwestTransportBuildError::ClientInitialization) } +pub(crate) fn account_proxy_client( + base: &Client, + proxy: Option<&gateway_core::account::OutboundProxy>, + timeout: Option, +) -> Result { + let Some(proxy) = proxy else { + return Ok(base.clone()); + }; + let invalid = || GrokReqwestTransportBuildError::ClientInitialization; + let mut builder = Client::builder() + .no_proxy() + .redirect(Policy::none()) + .proxy(reqwest::Proxy::all(proxy.expose_url()).map_err(|_| invalid())?) + .connect_timeout(CONNECT_TIMEOUT) + .pool_idle_timeout(POOL_IDLE_TIMEOUT) + .pool_max_idle_per_host(POOL_MAX_IDLE_PER_HOST) + .http2_adaptive_window(true) + .http2_keep_alive_interval(HTTP2_KEEP_ALIVE_INTERVAL) + .http2_keep_alive_timeout(HTTP2_KEEP_ALIVE_TIMEOUT) + .http2_keep_alive_while_idle(true) + .tcp_nodelay(true) + .https_only(true); + if let Some(timeout) = timeout { + builder = builder.timeout(timeout); + } + builder.build().map_err(|_| invalid()) +} + fn valid_official_url(url: &Url, host: &str, path: Option<&str>) -> bool { url.scheme() == "https" && url.host_str() == Some(host) diff --git a/backend/crates/providers/xai/src/transport/session.rs b/backend/crates/providers/xai/src/transport/session.rs index ef673dd24..e9d126e5b 100644 --- a/backend/crates/providers/xai/src/transport/session.rs +++ b/backend/crates/providers/xai/src/transport/session.rs @@ -17,7 +17,7 @@ use crate::SecretValue; /// 注入的推理 transport 可识别的不透明假名化 egress/session 键。 #[derive(Clone, PartialEq, Eq, Hash)] -pub struct GrokSessionBinding(String); +pub struct GrokSessionBinding(String, Option); impl GrokSessionBinding { /// 创建有界的非敏感绑定引用。 @@ -34,7 +34,7 @@ impl GrokSessionBinding { { return Err(GrokSessionDataError::InvalidBinding); } - Ok(Self(value)) + Ok(Self(value, None)) } /// 返回假名化的 transport 查找键。 @@ -42,6 +42,19 @@ impl GrokSessionBinding { pub fn as_str(&self) -> &str { &self.0 } + + #[must_use] + pub fn with_outbound_proxy( + mut self, + proxy: Option, + ) -> Self { + self.1 = proxy; + self + } + + pub(crate) fn outbound_proxy(&self) -> Option<&gateway_core::account::OutboundProxy> { + self.1.as_ref() + } } impl fmt::Debug for GrokSessionBinding { diff --git a/backend/crates/providers/xai/tests/admin.rs b/backend/crates/providers/xai/tests/admin.rs index d8a743683..42a5cfdaa 100644 --- a/backend/crates/providers/xai/tests/admin.rs +++ b/backend/crates/providers/xai/tests/admin.rs @@ -283,7 +283,11 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() let input = create_input("admin_projection", "subject-admin-projection"); seed_input(&store, &input).await.expect("create account"); let account = store.account(&input.account_id).expect("stored account"); - let record = account_record(&account); + let mut record = account_record(&account); + let proxy = + gateway_core::account::OutboundProxy::parse("socks5h://user:pass@proxy.example:1080") + .unwrap(); + record.outbound_proxy = Some(proxy.clone()); let catalog_cache = Arc::new(TestCatalogCache::default()); catalog_cache.seed("plan:standard", ["grok-4.5"]); let bundle = provider_xai::initialize( @@ -356,6 +360,12 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() document.get("type").and_then(Value::as_str), Some("oauth-account-bundle") ); + let entries = + provider_xai::GrokOAuthImportDocument::parse_json(&serde_json::to_vec(document).unwrap()) + .unwrap() + .into_entries(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].outbound_proxy(), Some(&proxy)); } #[tokio::test] @@ -480,6 +490,7 @@ fn provider_ports_with_catalog( fn account_record(account: &ProviderAccount) -> AccountRecord { let now = Utc::now(); AccountRecord { + outbound_proxy: None, id: account.id().to_string(), provider_kind: account.provider().clone(), groups: Vec::new(), diff --git a/backend/crates/providers/xai/tests/transport/network.rs b/backend/crates/providers/xai/tests/transport/network.rs index 8609eddcf..20bad185a 100644 --- a/backend/crates/providers/xai/tests/transport/network.rs +++ b/backend/crates/providers/xai/tests/transport/network.rs @@ -26,6 +26,86 @@ use provider_xai::{ use crate::support::loopback_endpoint_policy; +async fn rejecting_account_proxy() -> ( + gateway_core::account::OutboundProxy, + tokio::task::JoinHandle, +) { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy = gateway_core::account::OutboundProxy::parse(&format!( + "http://user:pass@{}", + listener.local_addr().unwrap() + )) + .unwrap(); + let task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut header = Vec::new(); + while !header.ends_with(b"\r\n\r\n") { + header.push(stream.read_u8().await.unwrap()); + assert!(header.len() < 8192); + } + stream + .write_all( + b"HTTP/1.1 502 Bad Gateway\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await + .unwrap(); + String::from_utf8(header).unwrap() + }); + (proxy, task) +} + +#[tokio::test] +async fn inference_and_oauth_connect_through_each_accounts_authenticated_proxy() { + let inference = + ReqwestGrokInferenceTransport::new(Arc::new(OfficialGrokEndpointPolicy)).unwrap(); + for _ in 0..2 { + let (proxy, captured) = rejecting_account_proxy().await; + let request = GrokInferenceRequest::new( + Url::parse("https://cli-chat-proxy.grok.com/v1/responses").unwrap(), + Vec::new(), + br#"{"model":"grok-code-test","stream":true}"#.to_vec(), + GrokSessionBinding::new("same-account") + .unwrap() + .with_outbound_proxy(Some(proxy)), + ); + assert!( + tokio::time::timeout(Duration::from_secs(5), inference.execute(request)) + .await + .unwrap() + .is_err() + ); + let connect = tokio::time::timeout(Duration::from_secs(5), captured) + .await + .unwrap() + .unwrap() + .to_ascii_lowercase(); + assert!(connect.starts_with("connect cli-chat-proxy.grok.com:443 http/1.1")); + assert!(connect.contains("proxy-authorization: basic dxnlcjpwyxnz")); + } + let (proxy, captured) = rejecting_account_proxy().await; + let client = GrokOAuthClient::new( + GrokOAuthConfig::official().unwrap(), + crate::support::xai_wire_profile(), + Arc::new(ReqwestOAuthTransport::new(Arc::new(OfficialGrokEndpointPolicy)).unwrap()), + Arc::new(FailClosedTokenVerifier), + ) + .with_outbound_proxy(Some(proxy)); + assert!( + tokio::time::timeout(Duration::from_secs(5), client.discover()) + .await + .unwrap() + .is_err() + ); + let connect = tokio::time::timeout(Duration::from_secs(5), captured) + .await + .unwrap() + .unwrap() + .to_ascii_lowercase(); + assert!(connect.starts_with("connect auth.x.ai:443 http/1.1")); + assert!(connect.contains("proxy-authorization: basic dxnlcjpwyxnz")); +} + #[tokio::test] async fn oauth_transport_should_post_form_once_without_redirect() { let server = MockServer::start().await; diff --git a/backend/migrations/.frozen-sha256 b/backend/migrations/.frozen-sha256 index 319a0eb9f..80e287176 100644 --- a/backend/migrations/.frozen-sha256 +++ b/backend/migrations/.frozen-sha256 @@ -1 +1,3 @@ b7d3df7b48d835a00a3b74cbefc958a039d38fea5529cdad04bfbc703dca72aa 0001_initial.sql +78c427c17d8b5759405f1629a80ceb624012cb6b31371b762078b000266d86c6 0002_client_key_budgets.sql +46f493ff410cf77550114542e0418475876ee2b2c975436290f0f72440e098c4 0003_account_outbound_proxy.sql diff --git a/backend/migrations/0002_client_key_budgets.sql b/backend/migrations/0002_client_key_budgets.sql new file mode 100644 index 000000000..11beaf265 --- /dev/null +++ b/backend/migrations/0002_client_key_budgets.sql @@ -0,0 +1,30 @@ +alter table client_api_keys + add column daily_limit_usd numeric(20,10) not null default 0 check (daily_limit_usd >= 0), + add column weekly_limit_usd numeric(20,10) not null default 0 check (weekly_limit_usd >= 0); + +create table client_key_budget_windows ( + client_api_key_id text primary key references client_api_keys(id) on delete cascade, + daily_start timestamptz not null, + daily_end timestamptz not null, + weekly_start timestamptz not null, + weekly_end timestamptz not null, + daily_used_usd numeric(20,10) not null default 0 check (daily_used_usd >= 0), + weekly_used_usd numeric(20,10) not null default 0 check (weekly_used_usd >= 0) +); + +-- Deliberately independent of model_requests: observations may be dropped or pruned. +create table client_key_charge_events ( + request_id text primary key, + client_api_key_id text not null references client_api_keys(id) on delete cascade, + started_at timestamptz not null default now(), + deadline_at timestamptz not null, + completed_at timestamptz, + amount_usd numeric(20,10) check (amount_usd >= 0), + reconciled_at timestamptz, + reconciliation_reason text, + state text not null default 'pending' check (state in ('pending', 'unknown', 'settled')), + check ((state = 'settled') = (amount_usd is not null)) +); +create index client_key_charge_events_unresolved + on client_key_charge_events(client_api_key_id, deadline_at) + where state <> 'settled'; diff --git a/backend/migrations/0003_account_outbound_proxy.sql b/backend/migrations/0003_account_outbound_proxy.sql new file mode 100644 index 000000000..6c64f7a24 --- /dev/null +++ b/backend/migrations/0003_account_outbound_proxy.sql @@ -0,0 +1 @@ +alter table provider_accounts add column outbound_proxy_url text; diff --git a/docs/account-egress-key-limits.md b/docs/account-egress-key-limits.md new file mode 100644 index 000000000..20136ab82 --- /dev/null +++ b/docs/account-egress-key-limits.md @@ -0,0 +1,100 @@ +# 账号代理与 Key 限额 + +## 关联与范围 + +Client Key 绑定账号分组，网关在已启用分组的账号并集中选择符合模型、状态和账号并发条件的账号。 +代理绑定到账号；Key 的日金额、七天金额、并发和 RPM 在所有账号与 Provider 之间合计。 + +```mermaid +flowchart LR + K[Client Key] --> L[日限额 / 七天限额 / 并发 / RPM] + L --> G[已启用分组的账号并集] + G --> A[账号 A] + G --> B[账号 B] + A --> P[代理 A] + B --> Q[代理 B] +``` + +未绑定分组的 Key 可使用全部账号，包括以后新增的账号。绑定了分组但所有分组停用或为空时， +没有可用账号，不会扩大为全部账号。需要一对一绑定时，为一个账号创建独立分组并只给 Key 绑定该组。 + +## 设置出口 + +账号编辑、批量编辑及新 OAuth 授权复用同一出站隧道控件：保持当前、直连、指定代理。 +支持 `http://`、`https://`、`socks5://`、`socks5h://`，例如 `socks5h://user:password@proxy.example:1080`。 +用户名和密码中的保留字符应使用 URL 编码；SOCKS 必须显式指定端口。 +`socks5` 在网关解析目标 DNS，`socks5h` 交给代理解析。 +HTTP 的 SOCKS 连接沿用 reqwest；其当前版本会把 IPv6 字面量目标编码为域名，配置自定义上游时请使用 +域名或 IPv4 目标。官方 OpenAI/xAI 域名不受影响，WebSocket 路径支持 IPv6 目标。 + +管理 API 的 `outboundProxyUrl` 省略或 `null` 保持当前值，`""` 清除代理。 +列表、详情、Debug 与普通审计不包含代理认证信息；敏感账号导出包含完整代理 URL，应按凭据保存。 +代理 URL 与现有 OAuth 凭据一样持久化到数据库，数据库备份也包含这些秘密。 + +OpenAI 的 HTTP/SSE、WebSocket、图片、搜索、OAuth 交换/刷新、额度、模型目录和资料请求使用账号出口。 +xAI 推理、OAuth、JWKS、userinfo 及辅助请求同样携带出口。 +指定代理不可用时请求失败，不会通过直连重试。账号代理变更推进凭据版本与配置快照， +新选择的账号不会复用旧代理的连接池；已经执行的请求可沿原连接完成。 +未配置代理的生产客户端直连，不继承进程的全局代理环境变量。TLS 继续验证证书。 +第三方授权页由用户浏览器访问，其网络出口不由网关代理设置控制。 + +## 从 sub2api 导入 + +1. 在 sub2api 导出所需 OpenAI OAuth 账号及它们引用的代理。 +2. 在账号管理中选择 OpenAI 的 JSON 文件导入，将导出 object 交给已有导入接口。 +3. 导入器从账号的 `proxy_key` 查找同一文档顶层 `proxies`；`data.accounts` 包装也受支持。 +4. 导入前检查所有 OpenAI 条目的代理引用，再在各自出口上执行必要的 token 刷新。 +5. 导入后检查账号出口和分组，并用账号连接测试验证实际部署网络。 + +sub2api 的 SOCKS5 代理按远端 DNS 语义转换为 `socks5h`。引用不存在、引用重复、代理停用、 +代理配置了到期时间或回退策略时，导入会报错，需要先处理这些不兼容配置。 +代理按 URL 保存在账号上，不创建额外的共享代理资源。 + +CPR 的 OpenAI 与 xAI 导出/导入也支持账号条目的 `outboundProxyUrl`。 +重新导入按已识别的上游账号身份更新，保留既有分组与调度配置。 +导入时未提供代理或提供空值，会保留既有账号的代理；清除已有代理请使用账号编辑中的“直连”。 +首次导入的账号没有分组。未迁移 sub2api 的下游 Key、用户余额、订阅、历史用量、 +分组权限及其他 Provider 账号；账号导出文件不是完整实例备份。 + +## 限额与计费 + +Key 创建/编辑可设置日限额、周限额（USD）、最大并发和 RPM。零表示不限。 +金额以十进制字符串存储，最多十位小数；编辑限额不重置已经结算的费用。 + +- 日窗口：北京时间当天零点到次日零点。 +- 周窗口：首次准入请求当天零点起七天；到期后由下一次使用开启新的七天窗口，不固定为周一。 +- 费用归属：按网关请求完成时间；跨日请求结算到完成日。延迟重试写入保留原完成时间。 +- 计费来源：优先 Provider 上报 USD，否则使用现有模型价格计算。订阅账号的估算 USD 不等于订阅账单。 +- 内部重试：累加各次实际获得的费用；一次未知计费不能被后续成功响应覆盖。 +- 并发：同一个 Key 在所有连接上执行中的请求总数。SSE 持有名额直到终态；空闲 WebSocket 不占名额， + 每个 `response.create` 独立准入。切换账号或内部重试继续使用同一个名额。Key 与账号上限同时生效。 + +金额检查使用已结算费用。达到日或周阈值后拒绝新请求；已经准入的一个或多个请求仍可完成并超出阈值。 +这不提供绝不超支的预占保证。费用受限、并发受限都不会把 Key 改为管理员禁用状态。 +错误字段及接口形状见 [Client Key API](api.md#7-client-key)。 + +## 故障与核账 + +每次计费请求在发送上游前持久化一笔待结算记录；并发名额在预算拒绝或启动失败时释放。 +完成、失败、取消、断连进入同一结算路径，以网关请求 ID 保证重复结算不会重复扣费。 +明确未发送的失败按零结算；已发送或可能已发送而未得到可信费用时保留未知状态。 +没有价格或 Usage 的模型和端点（包括无法提供独立费用的搜索请求）也可能产生未知费用。 + +受限 Key 遇到未知费用，或超过请求期限的未结算记录，会停止接受新的计费请求。 +管理员通过“费用核账”对照上游确认金额并填写原因；核账有审计，零金额也必须显式确认。 +原请求在旧窗口完成时，事后核账不会计入当前窗口，但会解除该笔记录的未知状态。 + +结算失败时，进程内保留精确费用，在同一 Key 下次请求前重试。进程崩溃丢失内存后， +数据库中的待结算记录仍在，超过请求期限后要求核账；不会把遗失费用直接按零放行。 +PostgreSQL 不可用时，所有新的计费请求拒绝准入；Redis 的并发/RPM 租约沿用项目现有恢复策略。 +费用账本不依赖可能丢弃或清理的请求观测日志，也不受 `usageRetentionDays` 控制。 +本版保留所有费用事件直到删除 Key；运维需将其纳入数据库容量与备份规划。 + +## 升级 + +服务启动自动执行 `0002_client_key_budgets.sql` 和 `0003_account_outbound_proxy.sql`。 +已发布的 `0001_initial.sql` 不变。旧账号默认直连，旧 Key 的日/周金额默认零，原并发和 RPM 配置保留。 +已有请求日志不会回填到新账本，使用本版本后的请求开始记录预算费用。 + +升级前备份数据库；先在隔离环境验证账号出口和导入文件。启用新策略后，旧二进制不识别代理与金额限额， +不能仅通过回退镜像来承接要求这些策略的流量。新增迁移的冻结校验见 [数据库迁移](../backend/migrations/README.md)。 diff --git a/docs/api.md b/docs/api.md index 4b29c3050..bd6ea4d40 100644 --- a/docs/api.md +++ b/docs/api.md @@ -175,8 +175,8 @@ Responses wire 之间的协议转换层，转换只在 xAI Provider 内完成。 | `POST` | `/api/admin/accounts/refresh` | `{ accountId }` | 手工刷新 OAuth credential（`idToken` / `accessToken` / `refreshToken`），不刷新额度 | | `POST` | `/api/admin/accounts/recover` | `{ accountId }` | 管理员显式清除该账号的本地错误/额度/cooldown 事实并重新启用，不访问上游 | | `POST` | `/api/admin/accounts/rotate` | OpenAI rotation 字段 | 手工替换 OpenAI OAuth token | -| `POST` | `/api/admin/accounts/update` | `{ accountId, enabled, concurrencyLimit, weight, groupIds }` | 一次更新账号调度状态、并发上限（`null` 表示继承运行参数）、权重（1–100）与所属分组 | -| `POST` | `/api/admin/accounts/batch-update` | `{ accountIds, enabled, concurrencyLimit, weight, groupIds }` | 一次事务统一更新所选账号的全部调度字段与完整分组集合 | +| `POST` | `/api/admin/accounts/update` | `{ accountId, enabled, concurrencyLimit, weight, groupIds, outboundProxyUrl? }` | 一次更新账号调度状态、并发上限（`null` 表示继承运行参数）、权重（1–100）、所属分组与出站代理 | +| `POST` | `/api/admin/accounts/batch-update` | `{ accountIds, enabled, concurrencyLimit, weight, groupIds, outboundProxyUrl? }` | 一次事务统一更新所选账号的调度字段、完整分组集合与可选代理 | | `POST` | `/api/admin/accounts/delete` | `{ provider, accountIds }` | 批量删除 1–200 个账号 | | `GET` | `/api/admin/accounts/quota` | `accountId` | 读取当前额度，不强制访问上游 | | `POST` | `/api/admin/accounts/quota/refresh` | `{ accountId }` | 访问 Provider 并刷新额度，同时同步额度所属状态 | @@ -186,7 +186,7 @@ Responses wire 之间的协议转换层，转换只在 xAI Provider 内完成。 | `GET` | `/api/admin/accounts/models` | `accountId` | 优先读取该 Provider + 套餐的模型 cache，缺失时有限实时拉取 | | `POST` | `/api/admin/accounts/models/refresh` | `{ accountId }` | 强制拉取最新模型并覆盖 cache | | `GET` | `/api/admin/accounts/connection-test` | `accountId`、`modelId` | 通过 SSE 返回实时连接测试事件，不作为业务 Responses 用量记录 | -| `POST` | `/api/admin/accounts/oauth/start` | `{ provider, name, accountId? }` | 创建 OpenAI 或 xAI OAuth flow；`accountId` 表示重新授权 | +| `POST` | `/api/admin/accounts/oauth/start` | `{ provider, name, accountId?, outboundProxyUrl? }` | 创建 OpenAI 或 xAI OAuth flow；`accountId` 表示重新授权 | | `POST` | `/api/admin/accounts/oauth/complete` | `{ provider, flowId, callbackUrl }` | 消费 OAuth callback；新账号保持未分组，重新授权保留所属分组 | 账号列表支持以下稳定值： @@ -197,6 +197,13 @@ Responses wire 之间的协议转换层，转换只在 xAI Provider 内完成。 - `sortBy`: `email`、`status`、`planType`、`usage`、`lastUsedAt`、`expiresAt`； - `sortDirection`: `asc`、`desc`。 +`outboundProxyUrl` 接受 HTTP、HTTPS、SOCKS5、SOCKS5H 代理 URL，可带用户名和密码。 +编辑时省略或 `null` 表示保持原配置，空字符串表示清除代理并直连。列表和详情只返回 +不含认证信息的 `outboundProxyEndpoint`（直连时为 `null`）；只有显式敏感导出包含完整 URL。 +指定代理后，推理、OAuth 服务端交换/刷新及账号辅助请求使用同一出口；代理失败不会退回直连。 +浏览器打开的第三方 OAuth 授权页仍使用浏览器自身网络。 +账号与 Key 的关系及迁移步骤见 [账号代理与 Key 限额](account-egress-key-limits.md)。 + ### 账号连接测试 SSE `GET /api/admin/accounts/connection-test` 固定探测请求指定的账号，不参与普通账号轮换。成功流沿用 @@ -227,12 +234,14 @@ Responses wire 之间的协议转换层，转换只在 xAI Provider 内完成。 导入的 `data` 必须是 JSON object，Admin API 请求上限为 64 MiB；Provider 可以收紧限制， 当前 xAI 导入上限为 16 MiB。内部 schema 由目标 Provider 独占解释： -- OpenAI 接受单账号 OAuth 文档、`accounts` 数组（最多 200 项）和 CPR 账号 bundle； +- OpenAI 接受单账号 OAuth 文档、`accounts` 数组（最多 200 项）、CPR 账号 bundle 和含代理引用的 sub2api 导出； - OpenAI OAuth token 字段接受 `accessToken`、`refreshToken`、`idToken`，以及官方 `auth.json` 中的 `access_token`、`refresh_token`、`id_token`，可以嵌套在 `tokens` 等账号 object 内； 每项至少包含 AT 或 RT。仅含 `OPENAI_API_KEY` 的客户端代理配置不是 OAuth 账号导入材料； RT-only 会在导入时换取 AT，AT-only 不具备自动续期能力； -- xAI 从单账号 object 或 `accounts` 数组中提取 OAuth token；包装中的代理、并发、优先级等字段不参与认证； +- OpenAI 与 xAI 的账号条目接受 `outboundProxyUrl`；OpenAI 还会解析 sub2api 的 `proxy_key` 和顶层 `proxies`。 + 代理在 token 刷新前绑定。缺失、重复、停用、带到期时间或配置回退策略的 sub2api 代理会拒绝导入； +- xAI 从单账号 object 或 `accounts` 数组中提取 OAuth token；并发、优先级等字段不参与认证； - xAI 批量导入逐条独立校验：失败条目跳过并记录日志，不中断其余条目，仅当没有任何条目成功时整个导入才报错； - xAI API Key 不是受支持的账号 credential； - 导入不会只凭文件外形写入账号；目标 Provider 使用认证材料完成必要的 token exchange 或已认证账号资料补全。 @@ -409,13 +418,36 @@ PostgreSQL 或 Redis。管理端只在用户打开弹窗或点击刷新时调用 | `POST` | `/api/admin/client-keys/enable` | `{ id }` | 启用 | | `POST` | `/api/admin/client-keys/disable` | `{ id }` | 禁用 | | `POST` | `/api/admin/client-keys/delete` | `{ id }` | 删除 | +| `GET` | `/api/admin/client-keys/unresolved-charges` | `id` | 按时间读取最早的 200 笔待核账请求 | +| `POST` | `/api/admin/client-keys/reconcile-charge` | `{ id, requestId, amountUsd, reason }` | 管理员确认未知费用并记录审计 | -创建字段为 `name`、可选 `label`、`groupIds`、`maxConcurrency`、`requestsPerMinute`，更新请求再增加 +创建字段为 `name`、可选 `label`、`groupIds`、`maxConcurrency`、`requestsPerMinute`、可选 +`dailyLimitUsd` 和 `weeklyLimitUsd`，更新请求再增加 `id`。`groupIds` 必须显式提交：空数组派生 `routingScope: "all"`，非空数组派生 `routingScope: "groups"`。响应同时返回分组引用 `groups`，以及从当前有效账号池派生、仅供展示的 `providerKinds`；Client Key 不再保存 `providerKind`。创建和 reveal 响应会返回完整明文 Key，调用方 必须立即安全保存。 +金额字段为非负十进制字符串，最多 10 位整数与 10 位小数，`"0"` 表示不限额。 +创建时省略金额字段默认为零；更新时省略或 `null` 保留当前值，修改限额不会清空已用金额。 +`maxConcurrency` 和 `requestsPerMinute` 是非负整数，零表示不限。 + +列表增加 `dailyLimitUsd`、`weeklyLimitUsd`、`dailyUsedUsd`、`weeklyUsedUsd`（均为字符串）、 +`dailyResetsAt`、`weeklyResetsAt`（RFC3339 或 `null`）与 `unresolvedRequests`。 +日窗口按北京时间零点重置；周窗口从首次准入当天零点起持续七天，到期后在下一次使用时重新开启。 +费用按请求完成时间归属窗口。并发按同一 Key 的执行中请求累计，包含 SSE 与每个 WebSocket +`response.create`；空闲连接不占名额，内部重试不重复占用。 + +任一已结算金额达到限额后拒绝新请求，已准入请求可完成并使金额超过阈值。 +HTTP 返回 `429`，`error.code` 为 `key_daily_budget_exceeded` 或 `key_weekly_budget_exceeded`， +并附 `Retry-After`；WebSocket 每次 `response.create` 执行相同检查并返回协议错误事件。 +已发送但费用不明、或超过请求期限仍未结算的请求使受限 Key 返回 `key_budget_unresolved`。 +预算存储不可用时返回 `503`、`key_budget_unavailable`。 + +核账的 `amountUsd` 使用相同金额格式，`reason` 为 1–1024 字节的非空原因。只能处理该 Key 下未知或 +已超过期限的请求；相同金额重复提交不重复计费，修改已结算金额返回 `409`。账本独立于使用统计日志， +记录保留至删除 Key，不受 `usageRetentionDays` 影响。 + ## 8. 运行设置 | 方法 | 路由 | 说明 | diff --git a/frontend/src/api/modules/accounts.ts b/frontend/src/api/modules/accounts.ts index e91d6fbcd..f1cafa5ab 100644 --- a/frontend/src/api/modules/accounts.ts +++ b/frontend/src/api/modules/accounts.ts @@ -110,6 +110,7 @@ export interface AccountUsage { } export interface Account { + outboundProxyEndpoint: string | null id: string name: string provider: string @@ -288,6 +289,7 @@ interface AccountResetCreditConsumeParam extends AccountIdParam { } interface AccountUpdateParam { + outboundProxyUrl?: string accountId: string enabled: boolean concurrencyLimit: number | null @@ -296,6 +298,7 @@ interface AccountUpdateParam { } interface AccountBatchUpdateParam { + outboundProxyUrl?: string accountIds: string[] enabled: boolean concurrencyLimit: number | null @@ -314,6 +317,7 @@ interface AccountImportParam { } interface AccountOAuthStartParam { + outboundProxyUrl?: string provider: string name: string accountId?: string diff --git a/frontend/src/api/modules/api-keys.ts b/frontend/src/api/modules/api-keys.ts index 945bf185f..43c782412 100644 --- a/frontend/src/api/modules/api-keys.ts +++ b/frontend/src/api/modules/api-keys.ts @@ -11,6 +11,13 @@ export interface ApiKey { enabled: boolean maxConcurrency: number requestsPerMinute: number + dailyLimitUsd: string + weeklyLimitUsd: string + dailyUsedUsd: string + weeklyUsedUsd: string + dailyResetsAt: string | null + weeklyResetsAt: string | null + unresolvedRequests: number createdAt: string updatedAt: string lastUsedAt: string | null @@ -40,6 +47,21 @@ export interface ApiKeyMutationResponse { id: string } +export interface UnresolvedClientCharge { + requestId: string + startedAt: string + completedAt: string | null + state: string +} + +export function getUnresolvedClientCharges(id: string) { + return request<{ items: UnresolvedClientCharge[] }>({ url: '/api/admin/client-keys/unresolved-charges', method: 'GET', params: { id } }) +} + +export function reconcileClientCharge(data: { id: string, requestId: string, amountUsd: string, reason: string }) { + return request<{ reconciled: boolean }>({ url: '/api/admin/client-keys/reconcile-charge', method: 'POST', data }) +} + // 请求参数类型：仅定义 API 边界的形状，调用方不依赖显式声明。 interface ApiKeyListParams { cursor?: string @@ -55,6 +77,8 @@ export interface ApiKeyWriteParam { groupIds: string[] maxConcurrency: number requestsPerMinute: number + dailyLimitUsd: string + weeklyLimitUsd: string } interface ApiKeyUpdateParam extends ApiKeyWriteParam { diff --git a/frontend/src/views/accounts/components/AccountBatchEditModal.vue b/frontend/src/views/accounts/components/AccountBatchEditModal.vue index 550ec80a7..f3052990a 100644 --- a/frontend/src/views/accounts/components/AccountBatchEditModal.vue +++ b/frontend/src/views/accounts/components/AccountBatchEditModal.vue @@ -7,6 +7,7 @@ import BaseFormItem from '@/components/base/BaseForm/FormItem.vue' import BaseInput from '@/components/base/BaseInput.vue' import BaseModal from '@/components/base/BaseModal/index.vue' import BaseSwitch from '@/components/base/BaseSwitch.vue' +import AccountProxyField from './AccountProxyField.vue' defineProps<{ selectedCount: number @@ -23,6 +24,8 @@ const open = defineModel({ required: true }) const enabled = defineModel('enabled', { required: true }) const concurrencyLimit = defineModel('concurrencyLimit', { required: true }) const weight = defineModel('weight', { required: true }) +const proxyMode = defineModel('proxyMode', { required: true }) +const proxyUrl = defineModel('proxyUrl', { required: true }) const selectedGroupIds = defineModel('selectedGroupIds', { required: true }) @@ -69,6 +72,7 @@ const selectedGroupIds = defineModel('selectedGroupIds', { required: t +
+
({ required: true }) const enabled = defineModel('enabled', { required: true }) const concurrencyLimit = defineModel('concurrencyLimit', { required: true }) const weight = defineModel('weight', { required: true }) +const proxyMode = defineModel('proxyMode', { required: true }) +const proxyUrl = defineModel('proxyUrl', { required: true }) const selectedGroupIds = defineModel('selectedGroupIds', { required: true }) @@ -98,6 +101,7 @@ const selectedGroupIds = defineModel('selectedGroupIds', { required: t :disabled="saving" /> +