mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
feat: add account egress and client key USD budgets
Import sub2api proxy bindings, enforce durable daily and seven-day budgets alongside existing key concurrency, and integrate matching admin controls and audited reconciliation.
This commit is contained in:
@@ -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 数据库备份、同大版本在线更新与回滚。
|
||||
|
||||
Generated
+3
@@ -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",
|
||||
|
||||
@@ -89,6 +89,7 @@ pub struct AccountRecord {
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
pub credential_state: CredentialState,
|
||||
pub credential_observed_at: DateTime<Utc>,
|
||||
pub quota: QuotaState,
|
||||
@@ -211,6 +212,7 @@ pub struct UpdateAccount {
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
|
||||
pub outbound_proxy: Option<Option<gateway_core::account::OutboundProxy>>,
|
||||
}
|
||||
|
||||
/// 账号更新结果。
|
||||
@@ -228,6 +230,7 @@ pub struct BatchUpdateAccounts {
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
|
||||
pub outbound_proxy: Option<Option<gateway_core::account::OutboundProxy>>,
|
||||
}
|
||||
|
||||
/// 批量账号更新结果。
|
||||
|
||||
@@ -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<DateTime<Utc>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
@@ -146,6 +148,7 @@ pub struct CreateClientKey {
|
||||
pub label: Option<String>,
|
||||
pub group_ids: Vec<AccountGroupId>,
|
||||
pub limits: RateLimits,
|
||||
pub budget: ClientBudgetLimits,
|
||||
}
|
||||
|
||||
/// 管理用例生成 ID 与明文后的持久化命令。
|
||||
@@ -156,6 +159,7 @@ pub struct NewClientKey {
|
||||
pub label: Option<String>,
|
||||
pub group_ids: Vec<AccountGroupId>,
|
||||
pub limits: RateLimits,
|
||||
pub budget: ClientBudgetLimits,
|
||||
pub plaintext: String,
|
||||
}
|
||||
|
||||
@@ -179,6 +183,8 @@ pub struct UpdateClientKey {
|
||||
pub label: Option<String>,
|
||||
pub group_ids: Vec<AccountGroupId>,
|
||||
pub limits: RateLimits,
|
||||
pub daily_limit_usd: Option<gateway_core::metering::Decimal>,
|
||||
pub weekly_limit_usd: Option<gateway_core::metering::Decimal>,
|
||||
}
|
||||
|
||||
/// 修改 Client Key 启用状态。
|
||||
@@ -217,3 +223,19 @@ pub struct ClientKeyMutation {
|
||||
pub record: Option<ClientKeyRecord>,
|
||||
pub id: ClientApiKeyId,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UnresolvedClientCharge {
|
||||
pub request_id: String,
|
||||
pub started_at: DateTime<Utc>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
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,
|
||||
}
|
||||
|
||||
@@ -144,6 +144,7 @@ impl fmt::Debug for PrepareCredentialImport {
|
||||
/// Provider 已验证、可由 Store 原子创建的一份 credential。
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct PreparedCredentialCreate {
|
||||
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
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<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
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<gateway_core::account::OutboundProxy>,
|
||||
) -> 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<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -382,6 +414,7 @@ pub struct StartAuthorization {
|
||||
pub context: MutationContext,
|
||||
pub name: String,
|
||||
pub reauthorization: Option<ProviderAccountId>,
|
||||
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
/// OAuth 流程启动结果。
|
||||
|
||||
@@ -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<Vec<crate::model::client_keys::UnresolvedClientCharge>> {
|
||||
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<ClientKeyPage>;
|
||||
|
||||
async fn reveal_client_key(
|
||||
|
||||
@@ -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<Vec<crate::model::client_keys::UnresolvedClientCharge>, 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<ClientKeyPage, AdminError>;
|
||||
async fn reveal(&self, id: &ClientApiKeyId) -> Result<ClientKeySecret, AdminError>;
|
||||
async fn create(
|
||||
@@ -64,6 +77,34 @@ impl DefaultClientKeyService {
|
||||
|
||||
#[async_trait]
|
||||
impl ClientKeyService for DefaultClientKeyService {
|
||||
async fn unresolved_charges(
|
||||
&self,
|
||||
id: &ClientApiKeyId,
|
||||
) -> Result<Vec<crate::model::client_keys::UnresolvedClientCharge>, 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<ClientKeyPage, AdminError> {
|
||||
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,
|
||||
|
||||
@@ -115,9 +115,13 @@ async fn pending_authorization(
|
||||
command: &StartAuthorization,
|
||||
resource: &'static str,
|
||||
) -> Result<PendingAuthorizationMutation, AdminError> {
|
||||
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(
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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")),
|
||||
|
||||
@@ -60,6 +60,7 @@ pub struct StartAccountAuthorizationRequest {
|
||||
pub provider: String,
|
||||
pub name: String,
|
||||
pub account_id: Option<String>,
|
||||
pub outbound_proxy_url: Option<super::wire::AccountProxyUpdate>,
|
||||
}
|
||||
|
||||
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<super::wire::AccountProxyUpdate>,
|
||||
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<UpdateAccount, WireValidationError> {
|
||||
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)?,
|
||||
|
||||
@@ -99,6 +99,10 @@ pub(super) fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> 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
|
||||
|
||||
@@ -2,6 +2,21 @@
|
||||
|
||||
use super::*;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AccountProxyUpdate(pub Option<gateway_core::account::OutboundProxy>);
|
||||
|
||||
impl<'de> Deserialize<'de> for AccountProxyUpdate {
|
||||
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
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<AccountProxyUpdate>,
|
||||
pub account_ids: Vec<String>,
|
||||
pub enabled: bool,
|
||||
#[serde(deserialize_with = "deserialize_required_nullable")]
|
||||
@@ -68,6 +84,7 @@ impl BatchUpdateAccountsRequest {
|
||||
pub(super) fn into_command(self) -> Result<BatchUpdateAccounts, WireValidationError> {
|
||||
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<String>,
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub provider: String,
|
||||
|
||||
@@ -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<String>,
|
||||
field: &'static str,
|
||||
) -> Result<Option<Decimal>, 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<String>,
|
||||
max_concurrency: u64,
|
||||
requests_per_minute: u64,
|
||||
daily_limit_usd: Option<String>,
|
||||
weekly_limit_usd: Option<String>,
|
||||
}
|
||||
|
||||
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<String>,
|
||||
max_concurrency: u64,
|
||||
requests_per_minute: u64,
|
||||
daily_limit_usd: Option<String>,
|
||||
weekly_limit_usd: Option<String>,
|
||||
}
|
||||
|
||||
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<DateTime<Utc>>,
|
||||
weekly_resets_at: Option<DateTime<Utc>>,
|
||||
unresolved_requests: u64,
|
||||
created_at: DateTime<Utc>,
|
||||
updated_at: DateTime<Utc>,
|
||||
last_used_at: Option<DateTime<Utc>>,
|
||||
@@ -292,6 +321,13 @@ impl From<ClientKeyRecord> 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::<S>))
|
||||
.route(
|
||||
"/api/admin/client-keys/unresolved-charges",
|
||||
get(unresolved_client_charges::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/client-keys/reconcile-charge",
|
||||
post(reconcile_client_charge::<S>),
|
||||
)
|
||||
.route(
|
||||
"/api/admin/client-keys/create",
|
||||
post(create_client_key::<S>),
|
||||
@@ -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<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminQuery(query): AdminQuery<ClientKeyIdQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
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::<Vec<_>>();
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(serde_json::json!({"items": items})),
|
||||
))
|
||||
}
|
||||
|
||||
async fn reconcile_client_charge<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
AdminJson(payload): AdminJson<ReconcileClientChargeRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
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::<Decimal>()
|
||||
.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<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
|
||||
@@ -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 失败响应。
|
||||
|
||||
@@ -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::<CreateClientKeyRequest>(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::<CreateClientKeyRequest>(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::<UpdateClientKeyRequest>(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::<AdminTestState>().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::<ListClientKeysQuery>(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"]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
mod error;
|
||||
mod model;
|
||||
mod proxy;
|
||||
pub use proxy::{InvalidOutboundProxy, OutboundProxy};
|
||||
pub mod scope;
|
||||
mod selection;
|
||||
mod store;
|
||||
|
||||
@@ -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<SystemTime>,
|
||||
next_refresh_at: Option<SystemTime>,
|
||||
has_refresh_token: bool,
|
||||
outbound_proxy: Option<super::OutboundProxy>,
|
||||
}
|
||||
|
||||
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<super::OutboundProxy>) -> 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,
|
||||
|
||||
@@ -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<Self, InvalidOutboundProxy> {
|
||||
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(<redacted>)")
|
||||
}
|
||||
}
|
||||
@@ -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<SystemTime>,
|
||||
pub weekly_resets_at: Option<SystemTime>,
|
||||
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<Decimal>,
|
||||
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>>;
|
||||
}
|
||||
@@ -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<S: ?Sized> {
|
||||
request_id: ModelRequestId,
|
||||
client_api_key_ref: crate::policy::ClientApiKeyId,
|
||||
observation: ResponseObservation,
|
||||
budget_prior_attempts_usd: Option<Decimal>,
|
||||
budget_attempt_send_state: UpstreamSendState,
|
||||
budget_attempt_already_counted: bool,
|
||||
trace: TraceContext,
|
||||
deadline: SystemTime,
|
||||
/// 会话级 deadline 计时器;deadline 固定,帧循环内复用而非逐事件新建。
|
||||
@@ -288,6 +296,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
delivery_pending: bool,
|
||||
upstream_complete: bool,
|
||||
finalized: bool,
|
||||
finalized_at: Option<SystemTime>,
|
||||
image_generation_requested: bool,
|
||||
/// 最近一次为无感恢复而被丢弃的原始上游错误;只在后续空选路时成为终态。
|
||||
last_retryable_failure: Option<ProviderError>,
|
||||
@@ -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<Decimal> {
|
||||
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
|
||||
|
||||
@@ -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<dyn ProviderCircuitPort>,
|
||||
continuation: Arc<dyn NativeContinuationPort>,
|
||||
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
|
||||
budget: Option<Arc<dyn ClientBudgetPort>>,
|
||||
}
|
||||
|
||||
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<dyn ClientBudgetPort>) -> Self {
|
||||
self.budget = Some(budget);
|
||||
self
|
||||
}
|
||||
|
||||
async fn start_inner(&self, request: StartExecution) -> Result<StartedExecution, GatewayError> {
|
||||
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<dyn NativeContinuationPort>,
|
||||
observed_provider_outcomes: usize,
|
||||
continuation_recorded: bool,
|
||||
budget: Option<Arc<dyn ClientBudgetPort>>,
|
||||
}
|
||||
|
||||
impl DefaultExecutionSession {
|
||||
@@ -949,6 +989,7 @@ impl DefaultExecutionSession {
|
||||
admission: AdmissionLease,
|
||||
circuits: Arc<dyn ProviderCircuitPort>,
|
||||
continuation: Arc<dyn NativeContinuationPort>,
|
||||
budget: Option<Arc<dyn ClientBudgetPort>>,
|
||||
) -> 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;
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! 模型请求生命周期、单行持久化 port 与 commit/send/cancellation 边界。
|
||||
|
||||
pub mod admission;
|
||||
pub mod budget;
|
||||
pub mod continuation;
|
||||
pub mod coordinator;
|
||||
pub mod execution;
|
||||
|
||||
@@ -975,6 +975,13 @@ pub struct GatewayError {
|
||||
message: &'static str,
|
||||
diagnostic: Option<Box<ProviderDiagnostic>>,
|
||||
client_visible_upstream_error: Option<ClientVisibleUpstreamError>,
|
||||
client_details: Option<Box<GatewayClientErrorDetails>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
struct GatewayClientErrorDetails {
|
||||
code: Option<&'static str>,
|
||||
retry_after: Option<Duration>,
|
||||
}
|
||||
|
||||
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<Duration> {
|
||||
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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -52,6 +52,7 @@ pub struct CoreStorePorts {
|
||||
snapshots: Arc<dyn SnapshotStorePort>,
|
||||
snapshot_subscriptions: Arc<dyn SnapshotSubscriptionPort>,
|
||||
client_api_key_usage: Arc<dyn ClientApiKeyUsageSink>,
|
||||
budget: Option<Arc<dyn engine::budget::ClientBudgetPort>>,
|
||||
}
|
||||
|
||||
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<dyn engine::budget::ClientBudgetPort>) -> 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<dyn ExecutionService> = service.clone();
|
||||
let account_probe: Arc<dyn AccountProbe> = service;
|
||||
let snapshot_control: Arc<dyn SnapshotControl> = publisher;
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
mod proxy;
|
||||
mod selection;
|
||||
mod selection_diagnostics;
|
||||
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
|
||||
@@ -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<AtomicBool>,
|
||||
}
|
||||
|
||||
impl ClientAdmissionPort for Admissions {
|
||||
fn admit(
|
||||
&self,
|
||||
_: ClientAdmissionRequest,
|
||||
) -> BoxFuture<'_, Result<ClientAdmissionDecision, ClientAdmissionError>> {
|
||||
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<bool, ClientAdmissionError>> {
|
||||
Box::pin(async { Ok(self.active.swap(false, Ordering::SeqCst)) })
|
||||
}
|
||||
fn restore(
|
||||
&self,
|
||||
_: ClientAdmissionRecovery,
|
||||
) -> BoxFuture<'_, Result<ClientAdmissionRestoreResult, ClientAdmissionError>> {
|
||||
Box::pin(async { Ok(ClientAdmissionRestoreResult::default()) })
|
||||
}
|
||||
}
|
||||
|
||||
struct Budget {
|
||||
reject: bool,
|
||||
active: Arc<AtomicBool>,
|
||||
charges: Mutex<Vec<ClientBudgetCharge>>,
|
||||
}
|
||||
|
||||
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<Admissions>, budget: Arc<Budget>) -> 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};
|
||||
|
||||
@@ -171,7 +171,8 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
)?),
|
||||
),
|
||||
Arc::new(client_key_usage),
|
||||
);
|
||||
)
|
||||
.with_budget(Arc::new(postgres::PgClientBudgetStore::new(pool.clone())));
|
||||
|
||||
let provider_ports = ProviderStorePorts::new(
|
||||
account_store,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<BTreeMap<String, ClientBudgetCharge>>,
|
||||
}
|
||||
|
||||
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::<Vec<_>>();
|
||||
if !retries.is_empty() {
|
||||
let ids = retries
|
||||
.iter()
|
||||
.map(|charge| charge.request_id.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
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::<bool, _>("enabled") {
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::PolicyDenied,
|
||||
"client API key is disabled",
|
||||
));
|
||||
}
|
||||
let limits = ClientBudgetLimits {
|
||||
daily_usd: row
|
||||
.get::<String, _>("daily_limit_usd")
|
||||
.parse()
|
||||
.map_err(|_| unavailable())?,
|
||||
weekly_usd: row
|
||||
.get::<String, _>("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::<String, _>("daily_used_usd")
|
||||
.parse()
|
||||
.map_err(|_| unavailable())?;
|
||||
let weekly: Decimal = window
|
||||
.get::<String, _>("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<Utc> = window.get("daily_end");
|
||||
let weekly_end: DateTime<Utc> = 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::<Utc>::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::<Utc>::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::<Utc>::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<Utc>,
|
||||
) -> 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::<Vec<_>>();
|
||||
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<Decimal> {
|
||||
row.get::<String, _>(field)
|
||||
.parse()
|
||||
.map_err(|_| postgres_unavailable("decode client budget"))
|
||||
};
|
||||
budgets.insert(
|
||||
row.get::<String, _>("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::<Option<DateTime<Utc>>, _>("daily_end")
|
||||
.map(Into::into),
|
||||
weekly_resets_at: row
|
||||
.get::<Option<DateTime<Utc>>, _>("weekly_end")
|
||||
.map(Into::into),
|
||||
unresolved_requests: row.get::<i64, _>("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<gateway_admin::model::client_keys::UnresolvedClientCharge>,
|
||||
> {
|
||||
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::<String, _>("amount_usd").parse::<Decimal>().ok()
|
||||
== Some(command.amount_usd)
|
||||
{
|
||||
Ok(())
|
||||
} else {
|
||||
Err(conflict())
|
||||
};
|
||||
}
|
||||
if state == "pending" && !row.get::<bool, _>("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::<DateTime<Utc>, _>("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",
|
||||
)
|
||||
}
|
||||
@@ -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<DateTime<Utc>>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
@@ -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<String>,
|
||||
pub max_concurrency: u64,
|
||||
pub requests_per_minute: u64,
|
||||
pub daily_limit_usd: Option<Decimal>,
|
||||
pub weekly_limit_usd: Option<Decimal>,
|
||||
}
|
||||
|
||||
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<Option<ClientApiKeyRecord>> {
|
||||
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<Vec<gateway_admin::model::client_keys::UnresolvedClientCharge>> {
|
||||
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<Admin
|
||||
max_concurrency: record.max_concurrency,
|
||||
requests_per_minute: record.requests_per_minute,
|
||||
},
|
||||
budget: record.budget,
|
||||
last_used_at: record.last_used_at,
|
||||
created_at: record.created_at,
|
||||
updated_at: record.updated_at,
|
||||
@@ -956,8 +996,8 @@ pub(crate) async fn insert_client_api_key_in_transaction(
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, label, key, enabled, max_concurrency, requests_per_minute,
|
||||
last_used_at, created_at, updated_at
|
||||
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now())",
|
||||
last_used_at, created_at, updated_at, daily_limit_usd, weekly_limit_usd
|
||||
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now(), $7::text::numeric, $8::text::numeric)",
|
||||
)
|
||||
.bind(&key.id)
|
||||
.bind(&key.name)
|
||||
@@ -965,6 +1005,8 @@ pub(crate) async fn insert_client_api_key_in_transaction(
|
||||
.bind(&key.key)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
.bind(to_i64(key.requests_per_minute)?)
|
||||
.bind(key.budget.daily_usd.canonical())
|
||||
.bind(key.budget.weekly_usd.canonical())
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("insert client API key in transaction"))?;
|
||||
@@ -980,7 +1022,9 @@ pub(crate) async fn update_client_api_key_in_transaction(
|
||||
let result = sqlx::query(
|
||||
"update client_api_keys
|
||||
set name = $2, label = $3, max_concurrency = $4,
|
||||
requests_per_minute = $5, updated_at = now()
|
||||
requests_per_minute = $5, updated_at = now(),
|
||||
daily_limit_usd = coalesce($6::text::numeric, daily_limit_usd),
|
||||
weekly_limit_usd = coalesce($7::text::numeric, weekly_limit_usd)
|
||||
where id = $1",
|
||||
)
|
||||
.bind(&key.id)
|
||||
@@ -988,6 +1032,8 @@ pub(crate) async fn update_client_api_key_in_transaction(
|
||||
.bind(&key.label)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
.bind(to_i64(key.requests_per_minute)?)
|
||||
.bind(key.daily_limit_usd.map(|amount| amount.canonical()))
|
||||
.bind(key.weekly_limit_usd.map(|amount| amount.canonical()))
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("update client API key in transaction"))?;
|
||||
@@ -1142,6 +1188,7 @@ fn client_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<ClientApiK
|
||||
row.try_get("requests_per_minute")
|
||||
.map_err(|_| invalid("invalid requests per minute"))?,
|
||||
)?,
|
||||
budget: ClientBudgetStatus::default(),
|
||||
last_used_at: row
|
||||
.try_get("last_used_at")
|
||||
.map_err(|_| invalid("invalid last used at"))?,
|
||||
|
||||
@@ -15,6 +15,7 @@ mod account_groups;
|
||||
mod admin_security_audit;
|
||||
mod admission_recovery;
|
||||
mod backup;
|
||||
mod client_budgets;
|
||||
mod client_keys;
|
||||
mod execution;
|
||||
mod execution_buffer;
|
||||
@@ -30,6 +31,7 @@ pub use account_groups::*;
|
||||
pub use admin_security_audit::*;
|
||||
pub use admission_recovery::*;
|
||||
pub use backup::*;
|
||||
pub use client_budgets::PgClientBudgetStore;
|
||||
pub use client_keys::*;
|
||||
pub use execution::*;
|
||||
pub use execution_buffer::*;
|
||||
|
||||
@@ -702,6 +702,15 @@ impl AccountStore for PgAdminAccountStore {
|
||||
"invalid provider account ID",
|
||||
)
|
||||
})?;
|
||||
let mut changed_fields = vec![
|
||||
"enabled".to_owned(),
|
||||
"concurrency_limit".to_owned(),
|
||||
"weight".to_owned(),
|
||||
"groups".to_owned(),
|
||||
];
|
||||
if command.outbound_proxy.is_some() {
|
||||
changed_fields.push("outbound_proxy".to_owned());
|
||||
}
|
||||
let config_revision = self
|
||||
.accounts
|
||||
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
|
||||
@@ -710,17 +719,13 @@ impl AccountStore for PgAdminAccountStore {
|
||||
concurrency_limit: command.concurrency_limit,
|
||||
weight: command.weight,
|
||||
group_ids: command.group_ids,
|
||||
outbound_proxy: command.outbound_proxy,
|
||||
audit: mutation_audit(
|
||||
context,
|
||||
"update",
|
||||
"provider_account",
|
||||
&command.account_id,
|
||||
vec![
|
||||
"enabled".to_owned(),
|
||||
"concurrency_limit".to_owned(),
|
||||
"weight".to_owned(),
|
||||
"groups".to_owned(),
|
||||
],
|
||||
changed_fields,
|
||||
),
|
||||
})
|
||||
.await
|
||||
@@ -790,6 +795,15 @@ impl AccountStore for PgAdminAccountStore {
|
||||
} else {
|
||||
"provider_accounts".to_owned()
|
||||
};
|
||||
let mut changed_fields = vec![
|
||||
"enabled".to_owned(),
|
||||
"concurrency_limit".to_owned(),
|
||||
"weight".to_owned(),
|
||||
"groups".to_owned(),
|
||||
];
|
||||
if command.outbound_proxy.is_some() {
|
||||
changed_fields.push("outbound_proxy".to_owned());
|
||||
}
|
||||
let config_revision = self
|
||||
.accounts
|
||||
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
|
||||
@@ -798,17 +812,13 @@ impl AccountStore for PgAdminAccountStore {
|
||||
concurrency_limit: command.concurrency_limit,
|
||||
weight: command.weight,
|
||||
group_ids: command.group_ids,
|
||||
outbound_proxy: command.outbound_proxy,
|
||||
audit: mutation_audit(
|
||||
context,
|
||||
"batch_update",
|
||||
"provider_account",
|
||||
&audit_target,
|
||||
vec![
|
||||
"enabled".to_owned(),
|
||||
"concurrency_limit".to_owned(),
|
||||
"weight".to_owned(),
|
||||
"groups".to_owned(),
|
||||
],
|
||||
changed_fields,
|
||||
),
|
||||
})
|
||||
.await
|
||||
|
||||
@@ -116,7 +116,7 @@ pub(crate) async fn load_admin_account_page(
|
||||
from runtime_settings
|
||||
where id = 1
|
||||
)
|
||||
select a.id, a.provider_kind, a.name, a.email, a.upstream_user_id,
|
||||
select a.outbound_proxy_url, a.id, a.provider_kind, a.name, a.email, a.upstream_user_id,
|
||||
a.upstream_account_id, a.plan_type, a.authentication_kind,
|
||||
a.credential_revision, a.has_refresh_token, a.access_token_expires_at,
|
||||
a.next_refresh_at, a.enabled, a.concurrency_limit, a.weight,
|
||||
|
||||
@@ -24,6 +24,7 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
)
|
||||
.map_err(core_store_error)?;
|
||||
self.insert_provider_account(NewProviderAccount {
|
||||
outbound_proxy: account.account.outbound_proxy().cloned(),
|
||||
id: account.account.id().as_str().to_owned(),
|
||||
provider_kind: account.account.provider().as_str().to_owned(),
|
||||
name: account.account.name().to_owned(),
|
||||
|
||||
@@ -6,6 +6,7 @@ pub(crate) fn admin_account_record(
|
||||
summary: ProviderAccountSummary,
|
||||
) -> AdminStoreResult<AccountRecord> {
|
||||
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<NewProviderAccount> {
|
||||
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,
|
||||
|
||||
@@ -108,7 +108,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
include_disabled: bool,
|
||||
) -> StoreResult<Vec<ProviderAccountSummary>> {
|
||||
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<AccountConcurrencyLimit>,
|
||||
weight: AccountWeight,
|
||||
outbound_proxy: Option<&Option<gateway_core::account::OutboundProxy>>,
|
||||
) -> 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"))?
|
||||
|
||||
@@ -120,6 +120,7 @@ pub(crate) fn parse_error_reason(value: Option<String>) -> StoreResult<Option<Ac
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ProviderAccountSummary {
|
||||
pub outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
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<gateway_core::account::OutboundProxy>,
|
||||
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<Option<gateway_core::account::OutboundProxy>>,
|
||||
pub account_ids: Vec<String>,
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
@@ -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::<Option<String>>(&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")?,
|
||||
|
||||
@@ -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<AccountGroupId>) -> 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,
|
||||
|
||||
@@ -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<Utc>, DateTime<Utc>, DateTime<Utc>) =
|
||||
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;
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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<String>, 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<serde_json::Value> =
|
||||
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(),
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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<ProviderAccount, Provi
|
||||
revision,
|
||||
account.access_token_expires_at.map(SystemTime::from),
|
||||
)
|
||||
.with_outbound_proxy(account.outbound_proxy.clone())
|
||||
.with_profile(
|
||||
account.email.clone(),
|
||||
account.upstream_account_id.clone(),
|
||||
|
||||
@@ -122,6 +122,7 @@ struct ParsedCodexImportAccount {
|
||||
name: Option<String>,
|
||||
email: Option<String>,
|
||||
authentication: ParsedCodexAuthentication,
|
||||
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
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<String>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
refresh_token: Option<String>,
|
||||
id_token: Option<String>,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<(CodexOAuthSecret, Option<DateTime<Utc>>), 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<Vec<ParsedCodexImportAccount>, 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<Option<gateway_core::account::OutboundProxy>, 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::<std::net::IpAddr>() {
|
||||
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<Vec<&Value>, CodexCredentialAdminError> {
|
||||
if let Some(accounts) = payload.get("accounts") {
|
||||
return accounts
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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<SecretString>,
|
||||
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;
|
||||
|
||||
@@ -133,13 +133,16 @@ impl CodexCredentialProfileService {
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
) -> Result<CodexProfileStatistics, CodexProfileStatisticsError> {
|
||||
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<String>), CodexProfileStatisticsError> {
|
||||
) -> Result<
|
||||
(
|
||||
SecretString,
|
||||
Option<String>,
|
||||
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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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<TokenPair, RefreshFailure>;
|
||||
async fn refresh_with_proxy(
|
||||
&self,
|
||||
refresh_token: &str,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<TokenPair, RefreshFailure> {
|
||||
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<AuthorizationTokenSet, AuthorizationCodeExchangeError>;
|
||||
async fn exchange_with_proxy(
|
||||
&self,
|
||||
grant: AuthorizationCodeGrant,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<AuthorizationTokenSet, AuthorizationCodeExchangeError> {
|
||||
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<Self, TokenClientBuildError> {
|
||||
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<TokenPair, RefreshFailure> {
|
||||
self.with_proxy(proxy)
|
||||
.map_err(|_| proxy_refresh_failure())?
|
||||
.refresh(refresh_token)
|
||||
.await
|
||||
}
|
||||
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
|
||||
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<AuthorizationTokenSet, AuthorizationCodeExchangeError> {
|
||||
self.with_proxy(proxy)
|
||||
.map_err(|_| AuthorizationCodeExchangeError::Unavailable)?
|
||||
.exchange_authorization_code(grant)
|
||||
.await
|
||||
}
|
||||
async fn exchange_authorization_code(
|
||||
&self,
|
||||
grant: AuthorizationCodeGrant,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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<String>;
|
||||
type ReqwestClientCacheKey = (Option<String>, String);
|
||||
type ReqwestClientCache = Mutex<HashMap<ReqwestClientCacheKey, Client>>;
|
||||
|
||||
/// 构建带缓存、自动协商 HTTP/2 的 reqwest Client。
|
||||
pub fn build_reqwest_client() -> Result<Client, CustomCaError> {
|
||||
build_account_http_client("", None)
|
||||
}
|
||||
|
||||
pub fn build_account_http_client(
|
||||
account_id: &str,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<Client, CustomCaError> {
|
||||
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<ReqwestClientCache> = 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<Client, CustomCaError> {
|
||||
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<Client, CustomCaError> {
|
||||
.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<Arc<CodexWebSocketPool>>,
|
||||
pub(super) websocket_origin_breaker: WebSocketOriginBreaker,
|
||||
pub(super) websocket_origin_key: String,
|
||||
pub(super) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
pub(super) egress_key: String,
|
||||
}
|
||||
|
||||
impl CodexBackendClient {
|
||||
pub fn for_account(
|
||||
&self,
|
||||
account: &gateway_core::account::ProviderAccount,
|
||||
) -> Result<Self, CodexClientError> {
|
||||
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。
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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<Optio
|
||||
maybe_build_rustls_client_config_with_env(&ProcessEnv)
|
||||
}
|
||||
|
||||
pub(crate) fn account_proxy_tls_config() -> CustomCaResult<Arc<ClientConfig>> {
|
||||
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,
|
||||
|
||||
@@ -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<MaybeTlsStream<tokio::net::TcpStream>, 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<tokio::net::TcpStream, tungstenite::Error> {
|
||||
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::<std::net::IpAddr>().map_or_else(
|
||||
|_| format!("{host}:{port}"),
|
||||
|ip| std::net::SocketAddr::new(ip, port).to_string(),
|
||||
);
|
||||
let uri = format!("http://{authority}")
|
||||
.parse::<hyper::Uri>()
|
||||
.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<WsRequest, tungstenite::http::Error> {
|
||||
|
||||
@@ -9,6 +9,7 @@ use crate::transport::protocol::responses::{
|
||||
/// WebSocket opening 描述。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CodexWebSocketConnection {
|
||||
pub(crate) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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<MaybeTlsStream<TcpStream>>;
|
||||
pub(crate) trait WebSocketIo:
|
||||
Stream<Item = Result<Message, tungstenite::Error>>
|
||||
+ Sink<Message, Error = tungstenite::Error>
|
||||
+ Send
|
||||
+ Unpin
|
||||
{
|
||||
}
|
||||
|
||||
impl<T> WebSocketIo for T where
|
||||
T: Stream<Item = Result<Message, tungstenite::Error>>
|
||||
+ Sink<Message, Error = tungstenite::Error>
|
||||
+ Send
|
||||
+ Unpin
|
||||
{
|
||||
}
|
||||
|
||||
pub(crate) type RawWsStream = Box<dyn WebSocketIo>;
|
||||
|
||||
const PUMP_COMMAND_BUFFER: usize = 32;
|
||||
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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<TokenPair, RefreshFailure> {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<String>) {
|
||||
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<String>) {
|
||||
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::<std::net::IpAddr>().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()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -45,6 +45,7 @@ use tokio_tungstenite::{
|
||||
},
|
||||
};
|
||||
|
||||
mod account_proxy;
|
||||
mod canonical;
|
||||
mod catalog;
|
||||
mod client;
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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::<Utc>::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<ProviderAccount, Provi
|
||||
revision,
|
||||
account.access_token_expires_at.map(SystemTime::from),
|
||||
)
|
||||
.with_outbound_proxy(account.outbound_proxy.clone())
|
||||
.with_profile(
|
||||
account.email.clone(),
|
||||
account.upstream_account_id.clone(),
|
||||
|
||||
@@ -1162,7 +1162,8 @@ async fn fetch_candidate_catalog(
|
||||
.map(|value| SecretValue::new(value.to_owned())),
|
||||
wire_profile,
|
||||
)
|
||||
.map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?;
|
||||
.map_err(|_| GrokCredentialCatalogError::InvalidCredentialData)?
|
||||
.with_outbound_proxy(candidate.account.outbound_proxy().cloned());
|
||||
let snapshot = client
|
||||
.fetch(&session)
|
||||
.await
|
||||
@@ -1214,6 +1215,7 @@ fn billing_session(
|
||||
.map(|value| SecretValue::new(value.to_owned())),
|
||||
wire_profile.clone(),
|
||||
)
|
||||
.map(|session| session.with_outbound_proxy(loaded.account.outbound_proxy().cloned()))
|
||||
.map_err(|_| GrokQuotaError::InvalidData)
|
||||
}
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ pub struct GrokOAuthClient {
|
||||
wire_profile: crate::XaiWireProfileState,
|
||||
transport: Arc<dyn OAuthHttpTransport>,
|
||||
verifier: Arc<dyn TokenVerifier>,
|
||||
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
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<gateway_core::account::OutboundProxy>) -> 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<OAuthHttpResponse, OAuthError> {
|
||||
request.outbound_proxy = self.outbound_proxy.clone();
|
||||
self.transport
|
||||
.execute(request)
|
||||
.await
|
||||
|
||||
@@ -121,6 +121,7 @@ pub struct OAuthHttpRequest {
|
||||
url: Url,
|
||||
headers: Vec<HttpHeader>,
|
||||
form: Vec<FormField>,
|
||||
pub(crate) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -114,10 +114,16 @@ impl GrokOAuthImportMetadata {
|
||||
pub struct GrokOAuthImportEntry {
|
||||
name: String,
|
||||
email: Option<String>,
|
||||
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
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<Self, GrokOAuthImportError> {
|
||||
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),
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -92,6 +92,7 @@ pub struct DueGrokCredential {
|
||||
upstream_account_id: Option<String>,
|
||||
plan_type: Option<String>,
|
||||
refresh_token_expires_at: Option<DateTime<Utc>>,
|
||||
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
struct DueGrokRefreshBatch {
|
||||
@@ -171,6 +172,16 @@ pub trait GrokCredentialRefresher: Send + Sync {
|
||||
&self,
|
||||
refresh_token: &SecretValue,
|
||||
) -> Result<GrokRefreshTokens, GrokRefreshFailure>;
|
||||
async fn refresh_with_proxy(
|
||||
&self,
|
||||
token: &SecretValue,
|
||||
proxy: Option<&gateway_core::account::OutboundProxy>,
|
||||
) -> Result<GrokRefreshTokens, GrokRefreshFailure> {
|
||||
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<GrokRefreshTokens, GrokRefreshFailure> {
|
||||
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::<Vec<_>>();
|
||||
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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -44,6 +44,7 @@ pub struct GrokModelCatalogSession {
|
||||
user_id: SecretValue,
|
||||
email: Option<SecretValue>,
|
||||
wire_profile: XaiWireProfileState,
|
||||
outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
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<gateway_core::account::OutboundProxy>,
|
||||
) -> 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<GrokHeader>,
|
||||
pub(crate) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
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<GrokHeader>,
|
||||
pub(crate) outbound_proxy: Option<gateway_core::account::OutboundProxy>,
|
||||
}
|
||||
|
||||
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。
|
||||
|
||||
@@ -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<Duration>,
|
||||
) -> Result<Client, GrokReqwestTransportBuildError> {
|
||||
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)
|
||||
|
||||
@@ -17,7 +17,7 @@ use crate::SecretValue;
|
||||
|
||||
/// 注入的推理 transport 可识别的不透明假名化 egress/session 键。
|
||||
#[derive(Clone, PartialEq, Eq, Hash)]
|
||||
pub struct GrokSessionBinding(String);
|
||||
pub struct GrokSessionBinding(String, Option<gateway_core::account::OutboundProxy>);
|
||||
|
||||
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<gateway_core::account::OutboundProxy>,
|
||||
) -> 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 {
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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<String>,
|
||||
) {
|
||||
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;
|
||||
|
||||
@@ -1 +1,3 @@
|
||||
b7d3df7b48d835a00a3b74cbefc958a039d38fea5529cdad04bfbc703dca72aa 0001_initial.sql
|
||||
78c427c17d8b5759405f1629a80ceb624012cb6b31371b762078b000266d86c6 0002_client_key_budgets.sql
|
||||
46f493ff410cf77550114542e0418475876ee2b2c975436290f0f72440e098c4 0003_account_outbound_proxy.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';
|
||||
@@ -0,0 +1 @@
|
||||
alter table provider_accounts add column outbound_proxy_url text;
|
||||
@@ -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)。
|
||||
+38
-6
@@ -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. 运行设置
|
||||
|
||||
| 方法 | 路由 | 说明 |
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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<boolean>({ required: true })
|
||||
const enabled = defineModel<boolean>('enabled', { required: true })
|
||||
const concurrencyLimit = defineModel<string>('concurrencyLimit', { required: true })
|
||||
const weight = defineModel<string>('weight', { required: true })
|
||||
const proxyMode = defineModel<string>('proxyMode', { required: true })
|
||||
const proxyUrl = defineModel<string>('proxyUrl', { required: true })
|
||||
const selectedGroupIds = defineModel<string[]>('selectedGroupIds', { required: true })
|
||||
</script>
|
||||
|
||||
@@ -69,6 +72,7 @@ const selectedGroupIds = defineModel<string[]>('selectedGroupIds', { required: t
|
||||
</BaseFormItem>
|
||||
</div>
|
||||
|
||||
<AccountProxyField v-model:mode="proxyMode" v-model:url="proxyUrl" :disabled="saving" />
|
||||
<BaseFormItem label="所属分组">
|
||||
<AccountGroupCheckboxGrid
|
||||
v-model="selectedGroupIds"
|
||||
|
||||
@@ -16,6 +16,7 @@ import BaseScrollbar from '@/components/base/BaseScrollbar.vue'
|
||||
import BaseSegmented from '@/components/base/BaseSegmented.vue'
|
||||
import BaseTextarea from '@/components/base/BaseTextarea.vue'
|
||||
import { useCopyText } from '@/composables/useCopyText'
|
||||
import AccountProxyField from '../AccountProxyField.vue'
|
||||
import AccountProviderChooser from './AccountProviderChooser.vue'
|
||||
import { resolveAccountCreatePresentation } from './presenter'
|
||||
|
||||
@@ -148,6 +149,7 @@ async function copyText(value: string, successText: string) {
|
||||
/>
|
||||
|
||||
<div v-if="mode === 'oauth'" class="flex flex-col gap-4">
|
||||
<AccountProxyField v-if="!reauthorizing" v-model:mode="form.proxyMode" v-model:url="form.proxyUrl" :preserve="false" :disabled="saving || oauthLoading || Boolean(form.oauthFlowId)" />
|
||||
<div class="rounded-cp bg-cp-fill-quaternary px-4 py-3">
|
||||
<div class="flex items-start gap-3">
|
||||
<div
|
||||
|
||||
@@ -6,6 +6,8 @@ export interface AccountCreateForm {
|
||||
oauthFlowId: string
|
||||
oauthAuthUrl: string
|
||||
oauthCallback: string
|
||||
proxyMode: string
|
||||
proxyUrl: string
|
||||
}
|
||||
|
||||
export function emptyAccountCreateForm(): AccountCreateForm {
|
||||
@@ -17,5 +19,7 @@ export function emptyAccountCreateForm(): AccountCreateForm {
|
||||
oauthFlowId: '',
|
||||
oauthAuthUrl: '',
|
||||
oauthCallback: '',
|
||||
proxyMode: 'direct',
|
||||
proxyUrl: '',
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ import BaseSwitch from '@/components/base/BaseSwitch.vue'
|
||||
import ProviderIconGroup from '@/components/ProviderIconGroup.vue'
|
||||
import AccountIdentityCell from './AccountIdentityCell.vue'
|
||||
import AccountPlanBadge from './AccountPlanBadge.vue'
|
||||
import AccountProxyField from './AccountProxyField.vue'
|
||||
|
||||
defineProps<{
|
||||
account: AccountRow | null
|
||||
@@ -27,6 +28,8 @@ const open = defineModel<boolean>({ required: true })
|
||||
const enabled = defineModel<boolean>('enabled', { required: true })
|
||||
const concurrencyLimit = defineModel<string>('concurrencyLimit', { required: true })
|
||||
const weight = defineModel<string>('weight', { required: true })
|
||||
const proxyMode = defineModel<string>('proxyMode', { required: true })
|
||||
const proxyUrl = defineModel<string>('proxyUrl', { required: true })
|
||||
const selectedGroupIds = defineModel<string[]>('selectedGroupIds', { required: true })
|
||||
</script>
|
||||
|
||||
@@ -98,6 +101,7 @@ const selectedGroupIds = defineModel<string[]>('selectedGroupIds', { required: t
|
||||
:disabled="saving"
|
||||
/>
|
||||
</BaseFormItem>
|
||||
<AccountProxyField v-model:mode="proxyMode" v-model:url="proxyUrl" :endpoint="account.outboundProxyEndpoint" :disabled="saving" />
|
||||
</div>
|
||||
|
||||
<template #footer>
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
<script setup lang="ts">
|
||||
import { computed } from 'vue'
|
||||
import BaseFormItem from '@/components/base/BaseForm/FormItem.vue'
|
||||
import BaseInput from '@/components/base/BaseInput.vue'
|
||||
import BaseSelect from '@/components/base/BaseSelect.vue'
|
||||
|
||||
const props = withDefaults(defineProps<{ endpoint?: string | null, disabled?: boolean, preserve?: boolean }>(), { preserve: true })
|
||||
const mode = defineModel<string>('mode', { required: true })
|
||||
const url = defineModel<string>('url', { required: true })
|
||||
const options = computed(() => [
|
||||
...(props.preserve ? [{ label: '保持当前', value: 'preserve' }] : []),
|
||||
{ label: '直连', value: 'direct' },
|
||||
{ label: '指定代理', value: 'proxy' },
|
||||
])
|
||||
</script>
|
||||
|
||||
<template>
|
||||
<div class="grid gap-3">
|
||||
<BaseFormItem label="出站隧道">
|
||||
<BaseSelect v-model="mode" :options="options" :disabled="disabled" aria-label="出站隧道" />
|
||||
</BaseFormItem>
|
||||
<p v-if="endpoint && mode === 'preserve'" class="m-0 font-mono text-xs break-all text-cp-text-tertiary">
|
||||
{{ endpoint }}
|
||||
</p>
|
||||
<BaseFormItem v-if="mode === 'proxy'" label="代理 URL" required>
|
||||
<BaseInput v-model="url" type="password" autocomplete="new-password" aria-label="代理 URL" :disabled="disabled" />
|
||||
</BaseFormItem>
|
||||
</div>
|
||||
</template>
|
||||
@@ -20,6 +20,8 @@ export function useAccountBatchEditor(options: {
|
||||
const schedulingEnabled = shallowRef(true)
|
||||
const concurrencyLimit = shallowRef('')
|
||||
const weight = shallowRef('1')
|
||||
const proxyMode = shallowRef('preserve')
|
||||
const proxyUrl = shallowRef('')
|
||||
const selectedGroupIds = ref<string[]>([])
|
||||
const saveAction = useAsyncAction()
|
||||
const saving = saveAction.loading
|
||||
@@ -30,6 +32,8 @@ export function useAccountBatchEditor(options: {
|
||||
return
|
||||
|
||||
schedulingEnabled.value = accounts.every(account => account.enabled)
|
||||
proxyMode.value = 'preserve'
|
||||
proxyUrl.value = ''
|
||||
concurrencyLimit.value = sharedConcurrencyLimit(accounts)
|
||||
weight.value = sharedWeight(accounts)
|
||||
selectedGroupIds.value = sharedGroupIds(accounts)
|
||||
@@ -40,6 +44,10 @@ export function useAccountBatchEditor(options: {
|
||||
if (saving.value || options.selectedIds.value.size === 0)
|
||||
return
|
||||
const scheduling = parseAccountSchedulingForm(concurrencyLimit.value, weight.value)
|
||||
if (proxyMode.value === 'proxy' && !proxyUrl.value.trim()) {
|
||||
toast.warning('请输入代理 URL')
|
||||
return
|
||||
}
|
||||
if (!scheduling.valid) {
|
||||
toast.warning(scheduling.message)
|
||||
return
|
||||
@@ -49,6 +57,7 @@ export function useAccountBatchEditor(options: {
|
||||
const accountIds = selectedAccounts().map(account => account.id)
|
||||
await batchUpdateAccounts({
|
||||
accountIds,
|
||||
outboundProxyUrl: proxyMode.value === 'preserve' ? undefined : proxyMode.value === 'direct' ? '' : proxyUrl.value.trim(),
|
||||
enabled: schedulingEnabled.value,
|
||||
concurrencyLimit: scheduling.values.concurrencyLimit,
|
||||
weight: scheduling.values.weight,
|
||||
@@ -89,6 +98,8 @@ export function useAccountBatchEditor(options: {
|
||||
if (open || isSaving)
|
||||
return
|
||||
schedulingEnabled.value = true
|
||||
proxyMode.value = 'preserve'
|
||||
proxyUrl.value = ''
|
||||
concurrencyLimit.value = ''
|
||||
weight.value = '1'
|
||||
selectedGroupIds.value = []
|
||||
@@ -99,6 +110,8 @@ export function useAccountBatchEditor(options: {
|
||||
schedulingEnabled,
|
||||
concurrencyLimit,
|
||||
weight,
|
||||
proxyMode,
|
||||
proxyUrl,
|
||||
selectedGroupIds,
|
||||
saving,
|
||||
open,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user