feat(routing): add account groups and weighted scheduling

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-08-13 16:58:36 +08:00
co-authored by Codex
parent 986a128d15
commit 01ad8bba8a
101 changed files with 6927 additions and 950 deletions
+17 -7
View File
@@ -21,9 +21,9 @@ pub mod ports;
mod use_case;
pub use use_case::{
accounts::AccountsService, auth::AuthService, backup::BackupService,
client_keys::ClientKeyService, observability::ObservabilityService, openai::OpenAiService,
settings::SettingsService, system::SystemService, xai::XaiService,
account_groups::AccountGroupService, accounts::AccountsService, auth::AuthService,
backup::BackupService, client_keys::ClientKeyService, observability::ObservabilityService,
openai::OpenAiService, settings::SettingsService, system::SystemService, xai::XaiService,
};
use model::{AdminError, AdminErrorKind};
@@ -33,10 +33,10 @@ use ports::{
system::SystemOperations,
};
use use_case::{
accounts::DefaultAccountsService, auth::DefaultAuthService, backup::DefaultBackupService,
client_keys::DefaultClientKeyService, observability::DefaultObservabilityService,
openai::DefaultOpenAiService, settings::DefaultSettingsService, system::DefaultSystemService,
xai::DefaultXaiService,
account_groups::DefaultAccountGroupService, accounts::DefaultAccountsService,
auth::DefaultAuthService, backup::DefaultBackupService, client_keys::DefaultClientKeyService,
observability::DefaultObservabilityService, openai::DefaultOpenAiService,
settings::DefaultSettingsService, system::DefaultSystemService, xai::DefaultXaiService,
};
const OPENAI_PROVIDER_KIND: &str = "openai";
@@ -147,6 +147,7 @@ pub enum AdminConfigError {
pub struct AdminServices {
auth: Arc<dyn AuthService>,
accounts: Arc<dyn AccountsService>,
account_groups: Arc<dyn AccountGroupService>,
client_keys: Arc<dyn ClientKeyService>,
observability: Arc<dyn ObservabilityService>,
settings: Arc<dyn SettingsService>,
@@ -167,6 +168,11 @@ impl AdminServices {
self.accounts.as_ref()
}
#[must_use]
pub fn account_groups(&self) -> &dyn AccountGroupService {
self.account_groups.as_ref()
}
#[must_use]
pub fn client_keys(&self) -> &dyn ClientKeyService {
self.client_keys.as_ref()
@@ -274,6 +280,10 @@ pub async fn initialize(
let services = AdminServices {
auth,
accounts,
account_groups: Arc::new(DefaultAccountGroupService::new(
store.account_groups(),
snapshot.clone(),
)),
client_keys: Arc::new(DefaultClientKeyService::new(
store.client_keys(),
snapshot.clone(),
@@ -0,0 +1,161 @@
//! Provider-neutral account group commands and projections.
use std::collections::BTreeMap;
use chrono::{DateTime, Utc};
use gateway_core::routing::{AccountGroupId, ProviderKind};
use super::{PageSize, Revision, observability::DecimalAmount};
/// Canonical `#RRGGBBAA` color persisted with an account group.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupColor(String);
impl AccountGroupColor {
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
(value.len() == 9
&& value.starts_with('#')
&& value[1..].bytes().all(|byte| byte.is_ascii_hexdigit()))
.then(|| Self(value.to_ascii_uppercase()))
}
}
/// Group member availability at the current observation time.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccountGroupAccountSummary {
pub available: u64,
pub limited: u64,
pub total: u64,
}
/// Group scheduling slots derived from available accounts and runtime leases.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccountGroupCapacity {
pub used_slots: Option<u64>,
pub total_slots: u64,
}
/// Successful, downstream-committed USD request costs for group accounts.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupUsage {
pub today_usd: DecimalAmount,
pub total_usd: DecimalAmount,
}
/// Lightweight group reference embedded in account and client-key views.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupRef {
pub id: AccountGroupId,
pub name: String,
pub color: AccountGroupColor,
pub enabled: bool,
}
/// Account group list query.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupListQuery {
pub page: u32,
pub page_size: PageSize,
pub search: Option<String>,
pub enabled: Option<bool>,
}
/// Complete account group summary.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupRecord {
pub id: AccountGroupId,
pub name: String,
pub description: Option<String>,
pub color: AccountGroupColor,
pub enabled: bool,
pub member_count: u64,
pub provider_counts: BTreeMap<String, u64>,
pub client_key_count: u64,
pub account_summary: AccountGroupAccountSummary,
pub capacity: AccountGroupCapacity,
pub usage: AccountGroupUsage,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
/// Paginated account groups.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupPage {
pub config_revision: Revision,
pub items: Vec<AccountGroupRecord>,
pub total: u64,
pub page: u32,
pub page_size: u16,
}
/// Lightweight account shown by the membership editor.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupMember {
pub id: String,
pub name: String,
pub provider_kind: ProviderKind,
pub email: Option<String>,
pub enabled: bool,
}
/// Complete membership projection for one group.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupMembers {
pub config_revision: Revision,
pub id: AccountGroupId,
pub items: Vec<AccountGroupMember>,
pub total: u64,
}
/// Create an account group.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreateAccountGroup {
pub name: String,
pub description: Option<String>,
pub color: AccountGroupColor,
}
/// Store-ready create command with a generated stable ID.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NewAccountGroup {
pub id: AccountGroupId,
pub name: String,
pub description: Option<String>,
pub color: AccountGroupColor,
}
/// Update an account group's descriptive fields.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateAccountGroup {
pub id: AccountGroupId,
pub name: String,
pub description: Option<String>,
pub color: AccountGroupColor,
}
/// Enable or disable an account group.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SetAccountGroupEnabled {
pub id: AccountGroupId,
pub enabled: bool,
}
/// Delete an account group that is not referenced by a client key.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeleteAccountGroup {
pub id: AccountGroupId,
}
/// Result of any account-group mutation.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountGroupMutation {
pub config_revision: Revision,
pub id: AccountGroupId,
pub record: Option<AccountGroupRecord>,
}
@@ -7,11 +7,12 @@ use futures::Stream;
use gateway_core::routing::ProviderKind;
use super::{PageSize, Revision, observability::TimeRange};
use super::{PageSize, Revision, account_groups::AccountGroupRef, observability::TimeRange};
pub use gateway_core::engine::credential::{
AccountErrorReason, AccountStatus, AccountStatusFacts, AccountStatusProjection,
CredentialState, QuotaAccessState, QuotaEvidence, QuotaState, resolve_account_status,
AccountConcurrencyLimit, AccountErrorReason, AccountStatus, AccountStatusFacts,
AccountStatusProjection, AccountWeight, CredentialState, QuotaAccessState, QuotaEvidence,
QuotaState, resolve_account_status,
};
/// 账号列表排序字段。
@@ -45,16 +46,25 @@ pub struct AccountListQuery {
pub page: u32,
pub page_size: PageSize,
pub provider_kind: Option<ProviderKind>,
pub group_filter: Option<AccountGroupFilter>,
pub search: Option<String>,
pub status: Option<AccountStatus>,
pub sort: Option<AccountSort>,
}
/// Optional account membership filter.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AccountGroupFilter {
Group(gateway_core::routing::AccountGroupId),
Ungrouped,
}
/// 账号公共存储投影;Provider 专属字段不进入此结构。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountRecord {
pub id: String,
pub provider_kind: ProviderKind,
pub groups: Vec<AccountGroupRef>,
pub name: String,
pub email: Option<String>,
pub upstream_user_id: Option<String>,
@@ -66,6 +76,8 @@ pub struct AccountRecord {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
pub quota: QuotaState,
@@ -180,11 +192,38 @@ pub struct AccountSummary {
pub error: u64,
}
/// 账号启停写入命令。
/// 账号可编辑事实的一次性替换命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SetAccountEnabled {
pub struct UpdateAccount {
pub account_id: String,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
}
/// 账号更新结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountUpdateResult {
pub config_revision: Revision,
pub account_id: gateway_core::engine::credential::ProviderAccountId,
}
/// 一批账号可编辑事实的一次性替换命令。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BatchUpdateAccounts {
pub account_ids: Vec<String>,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
}
/// 批量账号更新结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountsUpdateResult {
pub config_revision: Revision,
pub account_ids: Vec<gateway_core::engine::credential::ProviderAccountId>,
}
/// 账号批量删除命令。
@@ -6,10 +6,10 @@ use chrono::{DateTime, Utc};
use gateway_core::{
policy::{ClientApiKeyId, RateLimits},
routing::ProviderKind,
routing::{AccountGroupId, ProviderKind},
};
use super::{AdminModelError, Revision};
use super::{AdminModelError, Revision, account_groups::AccountGroupRef};
/// Client Key 列表保持旧 HTTP 合同允许的完整非零 `u16` 页大小。
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -88,7 +88,8 @@ pub struct ClientKeyRecord {
pub id: ClientApiKeyId,
pub name: String,
pub label: Option<String>,
pub provider_kind: ProviderKind,
pub groups: Vec<AccountGroupRef>,
pub provider_kinds: Vec<ProviderKind>,
pub prefix: String,
pub enabled: bool,
pub limits: RateLimits,
@@ -143,7 +144,7 @@ impl fmt::Debug for ClientKeySecret {
pub struct CreateClientKey {
pub name: String,
pub label: Option<String>,
pub provider_kind: ProviderKind,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
}
@@ -153,7 +154,7 @@ pub struct NewClientKey {
pub id: ClientApiKeyId,
pub name: String,
pub label: Option<String>,
pub provider_kind: ProviderKind,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
pub plaintext: String,
}
@@ -164,7 +165,7 @@ impl fmt::Debug for NewClientKey {
.debug_struct("NewClientKey")
.field("id", &self.id)
.field("name", &self.name)
.field("provider_kind", &self.provider_kind)
.field("group_ids", &self.group_ids)
.field("plaintext", &"[REDACTED]")
.finish_non_exhaustive()
}
@@ -176,7 +177,7 @@ pub struct UpdateClientKey {
pub id: ClientApiKeyId,
pub name: String,
pub label: Option<String>,
pub provider_kind: ProviderKind,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
}
@@ -2,6 +2,7 @@
use std::num::{NonZeroU16, NonZeroU64};
pub mod account_groups;
pub mod accounts;
pub mod auth;
pub mod backup;
@@ -513,8 +513,10 @@ pub struct DashboardAccountUsage {
/// Dashboard 当前账号池的可重建运行时槽位事实。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DashboardRuntimeSlots {
/// 与持久账号池指标相同谓词下可参与调度的账号数。
pub active_accounts: u64,
/// 继承运行参数默认并发上限的可调度账号数。
pub inherited_accounts: u64,
/// 已配置账号级并发上限的可调度账号槽位总数。
pub overridden_slots: u64,
/// Redis 可用时的实时 in-flight 槽位总数;不可用时为 `None`。
pub used_slots: Option<u64>,
}
@@ -548,6 +550,9 @@ pub struct UsageRecord {
pub id: String,
pub client_api_key_ref: String,
pub config_revision: u64,
pub routing_scope: String,
pub routing_group_refs: Vec<String>,
pub routing_group_names_snapshot: Vec<String>,
pub protocol: String,
pub operation: String,
pub endpoint: String,
@@ -10,9 +10,14 @@ use chrono::{DateTime, Utc};
use super::backup::BackupStorePorts;
use crate::model::{
MutationContext, Revision,
account_groups::{
AccountGroupListQuery, AccountGroupMembers, AccountGroupMutation, AccountGroupPage,
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountPageItem, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountListQuery, AccountPage, AccountPageItem, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -138,11 +143,17 @@ pub trait AccountStore: Send + Sync {
context: &MutationContext,
) -> AdminStoreResult<CredentialMutationResult>;
async fn set_account_enabled(
async fn update_account(
&self,
command: SetAccountEnabled,
command: UpdateAccount,
context: &MutationContext,
) -> AdminStoreResult<Revision>;
) -> AdminStoreResult<AccountUpdateResult>;
async fn batch_update_accounts(
&self,
command: BatchUpdateAccounts,
context: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult>;
async fn delete_accounts(
&self,
@@ -215,6 +226,44 @@ pub trait ClientKeyStore: Send + Sync {
) -> AdminStoreResult<Revision>;
}
/// Provider-neutral account group management and membership transactions.
#[async_trait]
pub trait AccountGroupStore: Send + Sync {
async fn list_account_groups(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage>;
async fn account_group_members(
&self,
id: &gateway_core::routing::AccountGroupId,
) -> AdminStoreResult<AccountGroupMembers>;
async fn create_account_group(
&self,
command: NewAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation>;
async fn update_account_group(
&self,
command: UpdateAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation>;
async fn set_account_group_enabled(
&self,
command: SetAccountGroupEnabled,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation>;
async fn delete_account_group(
&self,
command: DeleteAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation>;
}
/// 用量、趋势、诊断与运维错误的只读能力。
#[async_trait]
pub trait ObservabilityStore: Send + Sync {
@@ -302,6 +351,7 @@ pub trait SettingsStore: Send + Sync {
#[derive(Clone)]
pub struct AdminStorePorts {
accounts: Arc<dyn AccountStore>,
account_groups: Arc<dyn AccountGroupStore>,
auth: Arc<dyn AuthStore>,
client_keys: Arc<dyn ClientKeyStore>,
observability: Arc<dyn ObservabilityStore>,
@@ -313,6 +363,7 @@ impl AdminStorePorts {
#[must_use]
pub fn new(
accounts: Arc<dyn AccountStore>,
account_groups: Arc<dyn AccountGroupStore>,
auth: Arc<dyn AuthStore>,
client_keys: Arc<dyn ClientKeyStore>,
observability: Arc<dyn ObservabilityStore>,
@@ -321,6 +372,7 @@ impl AdminStorePorts {
) -> Self {
Self {
accounts,
account_groups,
auth,
client_keys,
observability,
@@ -334,6 +386,11 @@ impl AdminStorePorts {
self.accounts.clone()
}
#[must_use]
pub fn account_groups(&self) -> Arc<dyn AccountGroupStore> {
self.account_groups.clone()
}
#[must_use]
pub fn auth(&self) -> Arc<dyn AuthStore> {
self.auth.clone()
@@ -0,0 +1,139 @@
//! Account group management use cases.
use std::sync::Arc;
use async_trait::async_trait;
use gateway_core::routing::{AccountGroupId, snapshot::SnapshotControl};
use uuid::Uuid;
use crate::{
model::{
AdminError, MutationContext,
account_groups::{
AccountGroupListQuery, AccountGroupMembers, AccountGroupMutation, AccountGroupPage,
CreateAccountGroup, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
UpdateAccountGroup,
},
},
ports::store::AccountGroupStore,
};
use super::{map_store_error, publish_committed};
/// API-facing account group management service.
#[async_trait]
pub trait AccountGroupService: Send + Sync {
async fn list(&self, query: AccountGroupListQuery) -> Result<AccountGroupPage, AdminError>;
async fn members(&self, id: &AccountGroupId) -> Result<AccountGroupMembers, AdminError>;
async fn create(
&self,
context: &MutationContext,
command: CreateAccountGroup,
) -> Result<AccountGroupMutation, AdminError>;
async fn update(
&self,
context: &MutationContext,
command: UpdateAccountGroup,
) -> Result<AccountGroupMutation, AdminError>;
async fn set_enabled(
&self,
context: &MutationContext,
command: SetAccountGroupEnabled,
) -> Result<AccountGroupMutation, AdminError>;
async fn delete(
&self,
context: &MutationContext,
command: DeleteAccountGroup,
) -> Result<AccountGroupMutation, AdminError>;
}
pub(crate) struct DefaultAccountGroupService {
store: Arc<dyn AccountGroupStore>,
snapshot: Arc<dyn SnapshotControl>,
}
impl DefaultAccountGroupService {
#[must_use]
pub(crate) fn new(
store: Arc<dyn AccountGroupStore>,
snapshot: Arc<dyn SnapshotControl>,
) -> Self {
Self { store, snapshot }
}
async fn publish(
&self,
result: Result<AccountGroupMutation, crate::ports::store::AdminStoreError>,
) -> Result<AccountGroupMutation, AdminError> {
let result = result.map_err(|error| map_store_error(error, "account group"))?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
}
#[async_trait]
impl AccountGroupService for DefaultAccountGroupService {
async fn list(&self, query: AccountGroupListQuery) -> Result<AccountGroupPage, AdminError> {
self.store
.list_account_groups(query)
.await
.map_err(|error| map_store_error(error, "account group"))
}
async fn members(&self, id: &AccountGroupId) -> Result<AccountGroupMembers, AdminError> {
self.store
.account_group_members(id)
.await
.map_err(|error| map_store_error(error, "account group"))
}
async fn create(
&self,
context: &MutationContext,
command: CreateAccountGroup,
) -> Result<AccountGroupMutation, AdminError> {
let id = AccountGroupId::new(format!("grp_{}", Uuid::now_v7().simple()))
.map_err(|_| AdminError::internal("Failed to create account group ID"))?;
self.publish(
self.store
.create_account_group(
NewAccountGroup {
id,
name: command.name,
description: command.description,
color: command.color,
},
context,
)
.await,
)
.await
}
async fn update(
&self,
context: &MutationContext,
command: UpdateAccountGroup,
) -> Result<AccountGroupMutation, AdminError> {
self.publish(self.store.update_account_group(command, context).await)
.await
}
async fn set_enabled(
&self,
context: &MutationContext,
command: SetAccountGroupEnabled,
) -> Result<AccountGroupMutation, AdminError> {
self.publish(self.store.set_account_group_enabled(command, context).await)
.await
}
async fn delete(
&self,
context: &MutationContext,
command: DeleteAccountGroup,
) -> Result<AccountGroupMutation, AdminError> {
self.publish(self.store.delete_account_group(command, context).await)
.await
}
}
@@ -18,7 +18,8 @@ use crate::{
AdminError, MutationContext,
accounts::{
AccountConnectionTestEvent, AccountConnectionTestEventStream, AccountListQuery,
AccountPageItem, AccountUsageWindowQuery,
AccountPageItem, AccountUpdateResult, AccountUsageWindowQuery, AccountsUpdateResult,
BatchUpdateAccounts, UpdateAccount,
},
observability::TimeRange,
provider_credentials::{
@@ -54,6 +55,18 @@ pub trait AccountsService: Send + Sync {
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError>;
async fn update(
&self,
context: &MutationContext,
command: UpdateAccount,
) -> Result<AccountUpdateResult, AdminError>;
async fn batch_update(
&self,
context: &MutationContext,
command: BatchUpdateAccounts,
) -> Result<AccountsUpdateResult, AdminError>;
async fn quota(
&self,
account_id: &ProviderAccountId,
@@ -383,6 +396,76 @@ impl AccountsService for DefaultAccountsService {
})
}
async fn update(
&self,
context: &MutationContext,
command: UpdateAccount,
) -> Result<AccountUpdateResult, AdminError> {
let account_id = ProviderAccountId::new(command.account_id.clone())
.map_err(|_| AdminError::invalid("Invalid provider account ID"))?;
let (_, provider) = self.provider_for_account(&account_id).await?;
let enabled = command.enabled;
let result = self
.accounts
.update_account(command, context)
.await
.map_err(|error| map_store_error(error, "provider account"))?;
if !enabled {
provider.account_unavailable(&account_id).await;
}
provider
.account_facts_changed(std::slice::from_ref(&account_id))
.await;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn batch_update(
&self,
context: &MutationContext,
command: BatchUpdateAccounts,
) -> Result<AccountsUpdateResult, AdminError> {
let account_ids = command
.account_ids
.iter()
.map(|id| {
ProviderAccountId::new(id.clone())
.map_err(|_| AdminError::invalid("Invalid provider account ID"))
})
.collect::<Result<Vec<_>, _>>()?;
let mut providers = BTreeMap::<
ProviderKind,
(
Arc<dyn crate::ports::provider::ProviderAdmin>,
Vec<ProviderAccountId>,
),
>::new();
for account_id in &account_ids {
let (item, provider) = self.provider_for_account(account_id).await?;
providers
.entry(item.account.provider_kind)
.or_insert_with(|| (provider, Vec::new()))
.1
.push(account_id.clone());
}
let enabled = command.enabled;
let result = self
.accounts
.batch_update_accounts(command, context)
.await
.map_err(|error| map_store_error(error, "provider accounts"))?;
for (provider, provider_ids) in providers.values() {
if !enabled {
for account_id in provider_ids {
provider.account_unavailable(account_id).await;
}
}
provider.account_facts_changed(provider_ids).await;
}
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn quota(
&self,
account_id: &ProviderAccountId,
@@ -97,7 +97,7 @@ impl ClientKeyService for DefaultClientKeyService {
id,
name: command.name,
label: command.label,
provider_kind: command.provider_kind,
group_ids: command.group_ids,
limits: command.limits,
plaintext: plaintext.clone(),
},
@@ -1,5 +1,6 @@
//! 管理控制面的用例实现。
pub mod account_groups;
pub mod accounts;
pub mod auth;
pub mod backup;
@@ -13,11 +14,11 @@ pub mod xai;
use crate::{
model::{
AdminError, AdminErrorKind, MutationContext,
accounts::{DeleteAccounts, SetAccountEnabled},
accounts::DeleteAccounts,
provider_credentials::{
AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult,
CredentialDetails, CredentialMutation, CredentialMutationResult,
CredentialRotationCommit, PendingAuthorizationMutation, PreparedAuthorizationCommit,
CredentialDetails, CredentialMutationResult, CredentialRotationCommit,
PendingAuthorizationMutation, PreparedAuthorizationCommit,
PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation,
StartAuthorization,
},
@@ -280,44 +281,6 @@ async fn commit_credential_refresh(
}
}
async fn set_credential_enabled(
accounts: &dyn AccountStore,
provider: &dyn ProviderAdmin,
command: CredentialMutation,
enabled: bool,
resource: &'static str,
) -> Result<CredentialMutationResult, AdminError> {
required_credential(
accounts,
provider.provider_kind(),
&command.account_id,
resource,
)
.await?;
let account_id = command.account_id;
let revision = accounts
.set_account_enabled(
SetAccountEnabled {
account_id: account_id.as_str().to_owned(),
enabled,
},
&command.context,
)
.await
.map_err(|error| map_store_error(error, resource))?;
if !enabled {
provider.account_unavailable(&account_id).await;
}
provider
.account_facts_changed(std::slice::from_ref(&account_id))
.await;
Ok(CredentialMutationResult {
config_revision: revision,
account_id,
credential_revision: None,
})
}
async fn delete_credentials(
accounts: &dyn AccountStore,
provider: &dyn ProviderAdmin,
@@ -126,12 +126,20 @@ impl ObservabilityService for DefaultObservabilityService {
let health_timeline = health_timeline_at(&observation.trend, Utc::now());
let wire_profiles = self.providers.dashboard_wire_profiles();
let max_concurrent_per_account = u64::from(settings.max_concurrent_per_account);
let active_accounts = runtime_slots
.as_ref()
.map_or(observation.provider_accounts.normal, |slots| {
slots.active_accounts
});
let total_slots = active_accounts.saturating_mul(max_concurrent_per_account);
let total_slots = runtime_slots.as_ref().map_or_else(
|| {
observation
.provider_accounts
.normal
.saturating_mul(max_concurrent_per_account)
},
|slots| {
slots
.inherited_accounts
.saturating_mul(max_concurrent_per_account)
.saturating_add(slots.overridden_slots)
},
);
let used_slots = runtime_slots.and_then(|slots| slots.used_slots);
Ok(DashboardResult {
capacity: DashboardCapacity {
@@ -11,9 +11,9 @@ use crate::{
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
CredentialDeletionResult, CredentialDetails, CredentialImportCommit,
CredentialImportResult, CredentialListQuery, CredentialMutation,
CredentialMutationResult, CredentialPage, ImportCredentials, PrepareCredentialImport,
PrepareCredentialRotation, ProviderQuotaRequest, RotateCredential, StartAuthorization,
CredentialImportResult, CredentialListQuery, CredentialMutationResult, CredentialPage,
ImportCredentials, PrepareCredentialImport, PrepareCredentialRotation,
ProviderQuotaRequest, RotateCredential, StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
@@ -22,8 +22,7 @@ use crate::{
use super::{
commit_authorization, commit_credential_rotation, delete_credentials, map_provider_error,
map_store_error, pending_authorization, publish_committed, required_credential,
set_credential_enabled, validate_authorization_commit, validate_prepared_import,
validate_prepared_rotation,
validate_authorization_commit, validate_prepared_import, validate_prepared_rotation,
};
/// OpenAI 固定管理路由消费的服务。
@@ -50,14 +49,6 @@ pub trait OpenAiService: Send + Sync {
&self,
command: RotateCredential,
) -> Result<CredentialMutationResult, AdminError>;
async fn enable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
async fn disable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
async fn delete(
&self,
command: CredentialDeletion,
@@ -252,38 +243,6 @@ impl OpenAiService for DefaultOpenAiService {
Ok(result)
}
async fn enable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = set_credential_enabled(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
true,
"OpenAI credential",
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn disable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = set_credential_enabled(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
false,
"OpenAI credential",
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn delete(
&self,
command: CredentialDeletion,
@@ -11,8 +11,8 @@ use crate::{
provider_credentials::{
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
CredentialDeletionResult, CredentialImportCommit, CredentialImportResult,
CredentialListQuery, CredentialMutation, CredentialMutationResult, CredentialPage,
ImportCredentials, PrepareCredentialImport, ProviderQuotaRequest, StartAuthorization,
CredentialListQuery, CredentialMutationResult, CredentialPage, ImportCredentials,
PrepareCredentialImport, ProviderQuotaRequest, StartAuthorization,
},
},
ports::{provider::ProviderAdmin, store::AccountStore},
@@ -20,8 +20,8 @@ use crate::{
use super::{
commit_authorization, delete_credentials, map_provider_error, map_store_error,
pending_authorization, publish_committed, set_credential_enabled,
validate_authorization_commit, validate_prepared_import,
pending_authorization, publish_committed, validate_authorization_commit,
validate_prepared_import,
};
/// xAI 固定管理路由消费的服务。
@@ -40,14 +40,6 @@ pub trait XaiService: Send + Sync {
&self,
command: CompleteAuthorization,
) -> Result<CredentialMutationResult, AdminError>;
async fn enable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
async fn disable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError>;
async fn delete(
&self,
command: CredentialDeletion,
@@ -166,38 +158,6 @@ impl XaiService for DefaultXaiService {
Ok(result)
}
async fn enable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = set_credential_enabled(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
true,
"xAI credential",
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn disable(
&self,
command: CredentialMutation,
) -> Result<CredentialMutationResult, AdminError> {
let result = set_credential_enabled(
self.accounts.as_ref(),
self.provider.as_ref(),
command,
false,
"xAI credential",
)
.await?;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
Ok(result)
}
async fn delete(
&self,
command: CredentialDeletion,
@@ -22,8 +22,9 @@ use gateway_admin::{
AdminError, MutationContext, Revision,
accounts::{
AccountConnectionTestEvent, AccountListQuery, AccountPage, AccountPageItem,
AccountRecord, AccountSummary, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountRecord, AccountSummary, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
observability::TimeRange,
provider_credentials::{
@@ -661,15 +662,36 @@ impl AccountStore for FakeAccountStore {
Ok(rotation_result(command))
}
async fn set_account_enabled(
async fn update_account(
&self,
_: SetAccountEnabled,
command: UpdateAccount,
context: &MutationContext,
) -> AdminStoreResult<Revision> {
self.record("store.set_enabled");
) -> AdminStoreResult<AccountUpdateResult> {
self.record("store.update_account");
self.record_context(context);
self.require_commit()?;
Ok(revision(2))
Ok(AccountUpdateResult {
config_revision: revision(2),
account_id: ProviderAccountId::new(command.account_id).expect("account ID"),
})
}
async fn batch_update_accounts(
&self,
command: BatchUpdateAccounts,
context: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
self.record("store.batch_update_accounts");
self.record_context(context);
self.require_commit()?;
Ok(AccountsUpdateResult {
config_revision: revision(2),
account_ids: command
.account_ids
.into_iter()
.map(|id| ProviderAccountId::new(id).expect("account ID"))
.collect(),
})
}
async fn delete_accounts(
@@ -1001,6 +1023,122 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
assert_eq!(store.audit_requests(), ["refresh-request"]);
}
#[tokio::test]
async fn accounts_update_should_commit_then_release_disabled_account_and_publish_facts() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = accounts_service(provider, store.clone()).await;
let result = services
.accounts()
.update(
&context("update-request"),
UpdateAccount {
account_id: "acct_test".to_owned(),
enabled: false,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
group_ids: Vec::new(),
},
)
.await
.expect("update account");
assert_eq!(result.config_revision, revision(2));
assert_eq!(
recorded(&events),
[
"store.load_account",
"store.update_account",
"provider.account_unavailable",
"provider.account_facts_changed",
]
);
assert_eq!(store.audit_requests(), ["update-request"]);
}
#[tokio::test]
async fn accounts_update_should_not_notify_provider_when_store_commit_fails() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
store.fail_next_commit();
let services = accounts_service(provider, store).await;
services
.accounts()
.update(
&context("update-failure"),
UpdateAccount {
account_id: "acct_test".to_owned(),
enabled: false,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
group_ids: Vec::new(),
},
)
.await
.expect_err("failed account update");
assert_eq!(
recorded(&events),
["store.load_account", "store.update_account"]
);
}
#[tokio::test]
async fn accounts_batch_update_should_commit_once_and_notify_each_provider() {
let events = events();
let openai = FakeProviderAdmin::new("openai", events.clone());
let xai = FakeProviderAdmin::new("xai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let mut openai_account = account_record("openai");
openai_account.id = "acct_openai".to_owned();
let mut xai_account = account_record("xai");
xai_account.id = "acct_xai".to_owned();
store.set_accounts(vec![openai_account, xai_account]);
let services = super::AdminHarness::new()
.accounts(store.clone())
.settings(Arc::new(StaticSettingsStore))
.provider(openai)
.provider(xai)
.probe(Arc::new(SuccessfulAccountProbe))
.build()
.await;
let result = services
.accounts()
.batch_update(
&context("batch-update-request"),
BatchUpdateAccounts {
account_ids: vec!["acct_openai".to_owned(), "acct_xai".to_owned()],
enabled: false,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
group_ids: Vec::new(),
},
)
.await
.expect("batch update accounts");
assert_eq!(result.config_revision, revision(2));
assert_eq!(result.account_ids.len(), 2);
assert_eq!(store.audit_requests(), ["batch-update-request"]);
assert_eq!(
recorded(&events),
[
"store.load_account",
"store.load_account",
"store.batch_update_accounts",
"provider.account_unavailable",
"provider.account_facts_changed",
"provider.account_unavailable",
"provider.account_facts_changed",
]
);
}
#[tokio::test]
async fn accounts_list_should_return_complete_directory_semantics() {
let provider = FakeProviderAdmin::new("openai", events());
@@ -1012,6 +1150,7 @@ async fn accounts_list_should_return_complete_directory_semantics() {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1073,6 +1212,7 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1115,6 +1255,7 @@ async fn accounts_list_should_prefer_credential_error_over_quota_exhaustion() {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1143,6 +1284,7 @@ async fn accounts_list_should_map_unknown_credential_to_error_not_normal() {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1187,6 +1329,7 @@ async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view()
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1231,6 +1374,7 @@ async fn accounts_list_should_not_derive_exhaustion_from_provider_quota_view() {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1281,6 +1425,7 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1366,6 +1511,7 @@ async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_w
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: None,
@@ -1513,6 +1659,7 @@ pub(super) fn account_record(kind: &str) -> AccountRecord {
AccountRecord {
id: "acct_test".to_owned(),
provider_kind: ProviderKind::new(kind).expect("provider kind"),
groups: Vec::new(),
name: "test account".to_owned(),
email: Some("test@example.invalid".to_owned()),
upstream_user_id: Some("upstream-user".to_owned()),
@@ -1524,6 +1671,8 @@ pub(super) fn account_record(kind: &str) -> AccountRecord {
access_token_expires_at: Some(now + TimeDelta::hours(1)),
next_refresh_at: Some(now + TimeDelta::minutes(30)),
enabled: true,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
credential_state: CredentialState::Ready,
credential_observed_at: now,
quota: QuotaState::allowed(now.into()),
@@ -20,9 +20,14 @@ use gateway_admin::{
AdminConfig, AdminServices, InitialAdminPassword,
model::{
MutationContext, Revision,
account_groups::{
AccountGroupListQuery, AccountGroupMembers, AccountGroupMutation, AccountGroupPage,
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountListQuery, AccountPage, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -51,8 +56,8 @@ use gateway_admin::{
backup::BackupStorePorts,
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
store::{
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult,
AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
AccountGroupStore, AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts,
AdminStoreResult, AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
},
system::{
SystemOperationError, SystemOperationErrorKind, SystemOperations,
@@ -171,6 +176,7 @@ impl AdminHarness {
},
AdminStorePorts::new(
self.accounts,
Arc::new(UnavailableAccountGroupStore),
self.auth,
self.client_keys,
self.observability,
@@ -235,6 +241,57 @@ impl AuthStore for BootstrapAuthStore {
struct UnavailableStore;
struct UnavailableAccountGroupStore;
#[async_trait]
impl AccountGroupStore for UnavailableAccountGroupStore {
async fn list_account_groups(
&self,
_: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage> {
Err(unavailable("account groups"))
}
async fn account_group_members(
&self,
_: &gateway_core::routing::AccountGroupId,
) -> AdminStoreResult<AccountGroupMembers> {
Err(unavailable("account group members"))
}
async fn create_account_group(
&self,
_: NewAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
Err(unavailable("account group create"))
}
async fn update_account_group(
&self,
_: UpdateAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
Err(unavailable("account group update"))
}
async fn set_account_group_enabled(
&self,
_: SetAccountGroupEnabled,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
Err(unavailable("account group state"))
}
async fn delete_account_group(
&self,
_: DeleteAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
Err(unavailable("account group delete"))
}
}
#[async_trait]
impl AccountStore for UnavailableStore {
async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult<AccountPage> {
@@ -319,14 +376,22 @@ impl AccountStore for UnavailableStore {
Err(unavailable("credential refresh"))
}
async fn set_account_enabled(
async fn update_account(
&self,
_: SetAccountEnabled,
_: UpdateAccount,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
) -> AdminStoreResult<AccountUpdateResult> {
Err(unavailable("account enabled"))
}
async fn batch_update_accounts(
&self,
_: BatchUpdateAccounts,
_: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
Err(unavailable("account batch update"))
}
async fn delete_accounts(
&self,
_: DeleteAccounts,
@@ -183,7 +183,8 @@ async fn dashboard_summary_should_project_rebuildable_runtime_slots() {
let now = Utc::now();
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
store.replace_runtime_slots(Some(DashboardRuntimeSlots {
active_accounts: 3,
inherited_accounts: 2,
overridden_slots: 5,
used_slots: Some(2),
}));
let services = observability_services(store).await;
@@ -195,9 +196,9 @@ async fn dashboard_summary_should_project_rebuildable_runtime_slots() {
.expect("dashboard summary")
.capacity;
assert_eq!(capacity.total_slots, 3);
assert_eq!(capacity.total_slots, 7);
assert_eq!(capacity.used_slots, Some(2));
assert_eq!(capacity.available_slots, Some(1));
assert_eq!(capacity.available_slots, Some(5));
}
#[tokio::test]
@@ -654,6 +655,9 @@ fn total_record(
id: id.to_owned(),
client_api_key_ref: "key_billing".to_owned(),
config_revision: 1,
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: Vec::new(),
protocol: "openai".to_owned(),
operation: "generate".to_owned(),
endpoint: "/v1/responses".to_owned(),
@@ -5,8 +5,8 @@ use gateway_core::engine::credential::ProviderAccountId;
use gateway_admin::{
AdminServices,
model::provider_credentials::{
AuthorizationMutationTarget, CompleteAuthorization, CredentialDeletion, CredentialMutation,
ImportCredentials, ProviderQuotaRequest, StartAuthorization,
AuthorizationMutationTarget, CompleteAuthorization, CredentialDeletion, ImportCredentials,
ProviderQuotaRequest, StartAuthorization,
},
ports::provider::ProviderAdminErrorKind,
};
@@ -15,55 +15,6 @@ use super::accounts::{
FakeAccountStore, FakeProviderAdmin, context, document, events, recorded, revision,
};
#[tokio::test]
async fn openai_enable_should_use_scoped_store_without_provider_mutation() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = service(provider, store.clone()).await;
services
.openai()
.enable(mutation("acct_test"))
.await
.expect("enable credential");
assert_eq!(
recorded(&events),
[
"store.credential_details",
"store.set_enabled",
"provider.account_facts_changed",
]
);
assert_eq!(store.audit_requests(), ["request-openai"]);
}
#[tokio::test]
async fn openai_disable_should_commit_then_release_provider_resources() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
let services = service(provider, store.clone()).await;
services
.openai()
.disable(mutation("acct_test"))
.await
.expect("disable credential");
assert_eq!(
recorded(&events),
[
"store.credential_details",
"store.set_enabled",
"provider.account_unavailable",
"provider.account_facts_changed",
]
);
assert_eq!(store.audit_requests(), ["request-openai"]);
}
#[tokio::test]
async fn openai_delete_should_commit_then_release_provider_resources() {
let events = events();
@@ -89,26 +40,6 @@ async fn openai_delete_should_commit_then_release_provider_resources() {
assert_eq!(store.audit_requests(), ["request-openai"]);
}
#[tokio::test]
async fn openai_disable_store_failure_should_not_release_provider_resources() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let store = FakeAccountStore::new("openai", events.clone());
store.fail_next_commit();
let services = service(provider, store).await;
services
.openai()
.disable(mutation("acct_test"))
.await
.expect_err("failed Store transaction");
assert_eq!(
recorded(&events),
["store.credential_details", "store.set_enabled"]
);
}
#[tokio::test]
async fn openai_import_should_prepare_before_atomic_store_commit() {
let events = events();
@@ -450,13 +381,6 @@ async fn await_quota_requests(provider: &FakeProviderAdmin, expected: usize) {
assert_eq!(provider.quota_requests().len(), expected);
}
fn mutation(account_id: &str) -> CredentialMutation {
CredentialMutation {
context: context("request-openai"),
account_id: ProviderAccountId::new(account_id).expect("account ID"),
}
}
fn deletion(account_id: &str) -> CredentialDeletion {
CredentialDeletion {
context: context("request-openai"),
@@ -4,37 +4,12 @@ use gateway_core::engine::credential::ProviderAccountId;
use gateway_admin::{
AdminServices,
model::provider_credentials::{CredentialMutation, ImportCredentials, ProviderQuotaRequest},
model::provider_credentials::{ImportCredentials, ProviderQuotaRequest},
ports::provider::ProviderAdminErrorKind,
};
use super::accounts::{FakeAccountStore, FakeProviderAdmin, context, document, events, recorded};
#[tokio::test]
async fn xai_disable_should_commit_then_notify_the_stateless_provider() {
let events = events();
let provider = FakeProviderAdmin::new("xai", events.clone());
let store = FakeAccountStore::new("xai", events.clone());
let services = service(provider.clone(), store.clone()).await;
services
.xai()
.disable(mutation("acct_test"))
.await
.expect("disable credential");
assert_eq!(
recorded(&events),
[
"store.credential_details",
"store.set_enabled",
"provider.account_unavailable",
"provider.account_facts_changed",
]
);
assert_eq!(store.audit_requests(), ["request-xai"]);
}
#[tokio::test]
async fn xai_import_should_prepare_before_atomic_store_commit() {
let events = events();
@@ -132,10 +107,3 @@ async fn service(provider: Arc<FakeProviderAdmin>, store: Arc<FakeAccountStore>)
.build()
.await
}
fn mutation(account_id: &str) -> CredentialMutation {
CredentialMutation {
context: context("request-xai"),
account_id: ProviderAccountId::new(account_id).expect("account ID"),
}
}
@@ -0,0 +1,488 @@
//! Account group HTTP wire and fixed routes.
use std::collections::BTreeMap;
use axum::{
Json, Router,
extract::{Query, State},
http::StatusCode,
response::IntoResponse,
routing::{get, post},
};
use chrono::{DateTime, Utc};
use gateway_admin::model::{
PageSize,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor, AccountGroupListQuery,
AccountGroupMember, AccountGroupMembers, AccountGroupMutation, AccountGroupPage,
AccountGroupRecord, AccountGroupUsage, CreateAccountGroup, DeleteAccountGroup,
SetAccountGroupEnabled, UpdateAccountGroup,
},
};
use gateway_core::routing::AccountGroupId;
use serde::{Deserialize, Serialize};
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, PageMeta,
WireValidationError, wire::map_admin_service_error,
};
const DEFAULT_PAGE_SIZE: u32 = 50;
const MAX_PAGE_SIZE: u32 = 200;
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ListAccountGroupsQuery {
page: Option<u32>,
page_size: Option<u32>,
search: Option<String>,
enabled: Option<bool>,
}
impl ListAccountGroupsQuery {
fn into_command(self) -> Result<AccountGroupListQuery, WireValidationError> {
let page = self.page.unwrap_or(1);
let page_size = self.page_size.unwrap_or(DEFAULT_PAGE_SIZE);
if page == 0 || page_size == 0 || page_size > MAX_PAGE_SIZE {
return Err(WireValidationError::new("page"));
}
let search = self
.search
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty());
if search
.as_deref()
.is_some_and(|value| value.len() > 256 || value.chars().any(char::is_control))
{
return Err(WireValidationError::new("search"));
}
Ok(AccountGroupListQuery {
page,
page_size: PageSize::new(
u16::try_from(page_size).map_err(|_| WireValidationError::new("pageSize"))?,
)
.map_err(|_| WireValidationError::new("pageSize"))?,
search,
enabled: self.enabled,
})
}
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct CreateAccountGroupRequest {
name: String,
description: Option<String>,
color: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct UpdateAccountGroupRequest {
id: String,
name: String,
description: Option<String>,
color: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct AccountGroupIdRequest {
id: String,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(deny_unknown_fields)]
struct AccountGroupIdQuery {
id: String,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupView {
id: String,
name: String,
description: Option<String>,
color: String,
enabled: bool,
member_count: u64,
provider_counts: BTreeMap<String, u64>,
client_key_count: u64,
account_summary: AccountGroupAccountSummaryView,
capacity: AccountGroupCapacityView,
usage: AccountGroupUsageView,
created_at: DateTime<Utc>,
updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupAccountSummaryView {
available: u64,
limited: u64,
total: u64,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupCapacityView {
used_slots: Option<u64>,
total_slots: u64,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupUsageView {
today_usd: String,
total_usd: String,
}
impl From<AccountGroupRecord> for AccountGroupView {
fn from(record: AccountGroupRecord) -> Self {
Self {
id: record.id.to_string(),
name: record.name,
description: record.description,
color: record.color.as_str().to_owned(),
enabled: record.enabled,
member_count: record.member_count,
provider_counts: record.provider_counts,
client_key_count: record.client_key_count,
account_summary: account_summary_view(record.account_summary),
capacity: capacity_view(record.capacity),
usage: usage_view(record.usage),
created_at: record.created_at,
updated_at: record.updated_at,
}
}
}
fn account_summary_view(summary: AccountGroupAccountSummary) -> AccountGroupAccountSummaryView {
AccountGroupAccountSummaryView {
available: summary.available,
limited: summary.limited,
total: summary.total,
}
}
fn capacity_view(capacity: AccountGroupCapacity) -> AccountGroupCapacityView {
AccountGroupCapacityView {
used_slots: capacity.used_slots,
total_slots: capacity.total_slots,
}
}
fn usage_view(usage: AccountGroupUsage) -> AccountGroupUsageView {
AccountGroupUsageView {
today_usd: usage.today_usd.to_string(),
total_usd: usage.total_usd.to_string(),
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupPageData {
items: Vec<AccountGroupView>,
page: PageMeta,
config_revision: u64,
}
impl From<AccountGroupPage> for AccountGroupPageData {
fn from(page: AccountGroupPage) -> Self {
let total_pages = if page.total == 0 {
0
} else {
page.total.div_ceil(u64::from(page.page_size))
};
Self {
items: page.items.into_iter().map(Into::into).collect(),
page: PageMeta::new(
page.page,
u32::from(page.page_size),
page.total,
u32::try_from(total_pages).unwrap_or(u32::MAX),
),
config_revision: page.config_revision.get(),
}
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupMemberView {
id: String,
name: String,
provider_kind: String,
email: Option<String>,
enabled: bool,
}
impl From<AccountGroupMember> for AccountGroupMemberView {
fn from(member: AccountGroupMember) -> Self {
Self {
id: member.id,
name: member.name,
provider_kind: member.provider_kind.to_string(),
email: member.email,
enabled: member.enabled,
}
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupMembersData {
id: String,
items: Vec<AccountGroupMemberView>,
total: u64,
config_revision: u64,
}
impl From<AccountGroupMembers> for AccountGroupMembersData {
fn from(members: AccountGroupMembers) -> Self {
Self {
id: members.id.to_string(),
items: members.items.into_iter().map(Into::into).collect(),
total: members.total,
config_revision: members.config_revision.get(),
}
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountGroupMutationData {
id: String,
record: Option<AccountGroupView>,
config_revision: u64,
}
impl From<AccountGroupMutation> for AccountGroupMutationData {
fn from(mutation: AccountGroupMutation) -> Self {
Self {
id: mutation.id.to_string(),
record: mutation.record.map(Into::into),
config_revision: mutation.config_revision.get(),
}
}
}
/// Construct all fixed account-group management routes.
pub fn router<S>() -> Router<S>
where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.route("/api/admin/account-groups", get(list::<S>))
.route("/api/admin/account-groups/create", post(create::<S>))
.route("/api/admin/account-groups/update", post(update::<S>))
.route("/api/admin/account-groups/enable", post(enable::<S>))
.route("/api/admin/account-groups/disable", post(disable::<S>))
.route("/api/admin/account-groups/delete", post(delete::<S>))
.route("/api/admin/account-groups/members", get(members::<S>))
}
async fn list<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<ListAccountGroupsQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let result = state
.admin_services()
.account_groups()
.list(query.into_command().map_err(map_wire_error)?)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountGroupPageData::from(result)),
))
}
async fn members<S>(
_auth: AdminAuth,
State(state): State<S>,
Query(query): Query<AccountGroupIdQuery>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let id = group_id(query.id)?;
let result = state
.admin_services()
.account_groups()
.members(&id)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountGroupMembersData::from(result)),
))
}
async fn create<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<CreateAccountGroupRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
validate_group_fields(&request.name, request.description.as_deref())?;
mutation_response(
StatusCode::CREATED,
state
.admin_services()
.account_groups()
.create(
&auth.context().mutation_context(),
CreateAccountGroup {
name: request.name,
description: request.description,
color: group_color(&request.color)?,
},
)
.await,
)
}
async fn update<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<UpdateAccountGroupRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
validate_group_fields(&request.name, request.description.as_deref())?;
mutation_response(
StatusCode::OK,
state
.admin_services()
.account_groups()
.update(
&auth.context().mutation_context(),
UpdateAccountGroup {
id: group_id(request.id)?,
name: request.name,
description: request.description,
color: group_color(&request.color)?,
},
)
.await,
)
}
async fn enable<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountGroupIdRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
set_enabled(auth, state, request, true).await
}
async fn disable<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountGroupIdRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
set_enabled(auth, state, request, false).await
}
async fn set_enabled<S>(
auth: AdminAuth,
state: S,
request: AccountGroupIdRequest,
enabled: bool,
) -> Result<AdminResponse<AdminEnvelope<AccountGroupMutationData>>, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutation_response(
StatusCode::OK,
state
.admin_services()
.account_groups()
.set_enabled(
&auth.context().mutation_context(),
SetAccountGroupEnabled {
id: group_id(request.id)?,
enabled,
},
)
.await,
)
}
async fn delete<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountGroupIdRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutation_response(
StatusCode::OK,
state
.admin_services()
.account_groups()
.delete(
&auth.context().mutation_context(),
DeleteAccountGroup {
id: group_id(request.id)?,
},
)
.await,
)
}
fn mutation_response(
status: StatusCode,
result: Result<AccountGroupMutation, gateway_admin::model::AdminError>,
) -> Result<AdminResponse<AdminEnvelope<AccountGroupMutationData>>, AdminError> {
let result = result.map_err(map_service_error)?;
Ok(AdminResponse::new(
status,
AdminEnvelope::ok(AccountGroupMutationData::from(result)),
))
}
fn group_id(value: String) -> Result<AccountGroupId, AdminError> {
AccountGroupId::new(value).map_err(|_| AdminError::bad_request("Invalid account group ID"))
}
fn group_color(value: &str) -> Result<AccountGroupColor, AdminError> {
AccountGroupColor::parse(value)
.ok_or_else(|| AdminError::bad_request("Invalid account group color"))
}
fn validate_group_fields(name: &str, description: Option<&str>) -> Result<(), AdminError> {
if name.trim() != name
|| name.is_empty()
|| name.chars().count() > 100
|| name.chars().any(char::is_control)
|| description
.is_some_and(|value| value.len() > 4096 || value.chars().any(char::is_control))
{
return Err(AdminError::bad_request("Invalid account group request"));
}
Ok(())
}
fn map_wire_error(_: WireValidationError) -> AdminError {
AdminError::bad_request("Invalid account group query")
}
fn map_service_error(error: gateway_admin::model::AdminError) -> AdminError {
map_admin_service_error(error, "Account group repository unavailable")
}
+217 -78
View File
@@ -18,9 +18,10 @@ use futures::{Stream, StreamExt as _};
use gateway_admin::model::{
AdminError as AdminServiceError, PageSize,
accounts::{
AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost, AccountListQuery,
AccountModelUsage, AccountSort, AccountSortField, AccountStatus as DomainAccountStatus,
AccountUsage, SortDirection,
AccountConcurrencyLimit, AccountConnectionTestEvent as DomainConnectionTestEvent,
AccountCost, AccountGroupFilter, AccountListQuery, AccountModelUsage, AccountSort,
AccountSortField, AccountStatus as DomainAccountStatus, AccountUpdateResult, AccountUsage,
AccountWeight, AccountsUpdateResult, BatchUpdateAccounts, SortDirection, UpdateAccount,
},
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
@@ -32,9 +33,9 @@ use gateway_admin::model::{
};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId},
routing::{ProviderKind, UpstreamModelId},
routing::{AccountGroupId, ProviderKind, UpstreamModelId},
};
use serde::{Deserialize, Serialize};
use serde::{Deserialize, Deserializer, Serialize};
use serde_json::{Map, Value};
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
@@ -54,6 +55,7 @@ const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024;
const MAX_ID_TOKEN_BYTES: usize = 16 * 1024;
const MAX_CALLBACK_URL_BYTES: usize = 64 * 1024;
const MAX_ACCOUNT_DELETE_BATCH: usize = 200;
const MAX_ACCOUNT_GROUP_BATCH: usize = 1000;
/// 账号列表查询参数。
#[derive(Debug, Clone, Default, Deserialize)]
@@ -62,12 +64,74 @@ pub struct ListQuery {
pub page: Option<u32>,
pub page_size: Option<u32>,
pub provider: Option<String>,
pub group_id: Option<String>,
pub search: Option<String>,
pub status: Option<String>,
pub sort_by: Option<String>,
pub sort_direction: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct BatchUpdateAccountsRequest {
pub account_ids: Vec<String>,
pub enabled: bool,
#[serde(deserialize_with = "deserialize_required_nullable")]
pub concurrency_limit: Option<u64>,
pub weight: u64,
pub group_ids: Vec<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct BatchUpdatedAccountsData {
account_ids: Vec<String>,
config_revision: u64,
}
impl From<AccountsUpdateResult> for BatchUpdatedAccountsData {
fn from(result: AccountsUpdateResult) -> Self {
Self {
account_ids: result
.account_ids
.into_iter()
.map(|id| id.to_string())
.collect(),
config_revision: result.config_revision.get(),
}
}
}
impl BatchUpdateAccountsRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
if self.account_ids.is_empty()
|| self.account_ids.len() > MAX_ACCOUNT_GROUP_BATCH
|| self
.account_ids
.iter()
.any(|id| require_account_id(id, "accountIds").is_err())
|| self.account_ids.iter().collect::<BTreeSet<_>>().len() != self.account_ids.len()
{
return Err(WireValidationError::new("accountIds"));
}
validate_wire_group_ids(&self.group_ids)?;
parse_concurrency_limit(self.concurrency_limit)?;
parse_account_weight(self.weight)?;
Ok(())
}
fn into_command(self) -> Result<BatchUpdateAccounts, WireValidationError> {
self.validate()?;
Ok(BatchUpdateAccounts {
account_ids: self.account_ids,
enabled: self.enabled,
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
weight: parse_account_weight(self.weight)?,
group_ids: validate_wire_group_ids(&self.group_ids)?,
})
}
}
impl ListQuery {
/// 解析并校验全部 wire 字段,生成 Admin 查询命令。
pub fn validate(self) -> Result<AccountListQuery, WireValidationError> {
@@ -80,6 +144,14 @@ impl ListQuery {
return Err(WireValidationError::new("pageSize"));
}
let provider_kind = parse_provider(self.provider.as_deref().unwrap_or("all"))?;
let group_filter = match self.group_id.as_deref().map(str::trim) {
None | Some("") => None,
Some("ungrouped") => Some(AccountGroupFilter::Ungrouped),
Some(value) => Some(AccountGroupFilter::Group(
AccountGroupId::new(value.to_owned())
.map_err(|_| WireValidationError::new("groupId"))?,
)),
};
let search = self
.search
.map(|value| value.trim().to_owned())
@@ -112,6 +184,7 @@ impl ListQuery {
)
.map_err(|_| WireValidationError::new("pageSize"))?,
provider_kind,
group_filter,
search,
status,
sort,
@@ -177,6 +250,7 @@ pub struct AccountView {
pub id: String,
pub name: String,
pub provider: String,
pub groups: Vec<AccountGroupRefView>,
pub resource_ref: String,
pub email: Option<String>,
pub account_id: Option<String>,
@@ -191,6 +265,8 @@ pub struct AccountView {
/// 最近一次失败的上游错误描述;仅错误状态存在。
pub error_message: Option<String>,
pub enabled: bool,
pub concurrency_limit: Option<u32>,
pub weight: u16,
pub access_token_expires_at: Option<String>,
pub access_token_expires_at_display: Option<String>,
pub refresh_token_expires_at: Option<String>,
@@ -203,6 +279,15 @@ pub struct AccountView {
pub usage: AccountUsageView,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct AccountGroupRefView {
pub id: String,
pub name: String,
pub color: String,
pub enabled: bool,
}
/// Provider quota 安全视图。
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
@@ -682,31 +767,49 @@ impl CompleteAccountAuthorizationRequest {
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct AccountMutationRequest {
pub provider: String,
pub struct UpdateAccountRequest {
pub account_id: String,
pub enabled: bool,
#[serde(deserialize_with = "deserialize_required_nullable")]
pub concurrency_limit: Option<u64>,
pub weight: u64,
pub group_ids: Vec<String>,
}
impl AccountMutationRequest {
impl UpdateAccountRequest {
pub fn validate(&self) -> Result<(), WireValidationError> {
AccountProvider::parse(&self.provider)?;
require_account_id(&self.account_id, "accountId")
require_account_id(&self.account_id, "accountId")?;
parse_concurrency_limit(self.concurrency_limit)?;
parse_account_weight(self.weight)?;
validate_wire_group_ids(&self.group_ids)?;
Ok(())
}
fn into_command(
self,
context: gateway_admin::model::MutationContext,
) -> Result<(AccountProvider, CredentialMutation), WireValidationError> {
fn into_command(self) -> Result<UpdateAccount, WireValidationError> {
self.validate()?;
let provider = AccountProvider::parse(&self.provider)?;
Ok((
provider,
CredentialMutation {
context,
account_id: ProviderAccountId::new(self.account_id)
.map_err(|_| WireValidationError::new("accountId"))?,
},
))
Ok(UpdateAccount {
account_id: self.account_id,
enabled: self.enabled,
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
weight: parse_account_weight(self.weight)?,
group_ids: validate_wire_group_ids(&self.group_ids)?,
})
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UpdatedAccountData {
pub account_id: String,
pub config_revision: u64,
}
impl From<AccountUpdateResult> for UpdatedAccountData {
fn from(result: AccountUpdateResult) -> Self {
Self {
account_id: result.account_id.to_string(),
config_revision: result.config_revision.get(),
}
}
}
@@ -810,8 +913,8 @@ pub struct AccountImportData {
pub account_ids: Vec<String>,
}
impl From<CredentialImportResult> for AccountImportData {
fn from(result: CredentialImportResult) -> Self {
impl AccountImportData {
pub fn from_result(result: CredentialImportResult) -> Self {
let account_ids = result
.credential_ids
.into_iter()
@@ -888,6 +991,46 @@ fn require_account_id(value: &str, field: &'static str) -> Result<(), WireValida
Ok(())
}
fn parse_concurrency_limit(
value: Option<u64>,
) -> Result<Option<AccountConcurrencyLimit>, WireValidationError> {
value
.map(|value| {
u32::try_from(value)
.ok()
.and_then(AccountConcurrencyLimit::new)
.ok_or_else(|| WireValidationError::new("concurrencyLimit"))
})
.transpose()
}
fn deserialize_required_nullable<'de, D>(deserializer: D) -> Result<Option<u64>, D::Error>
where
D: Deserializer<'de>,
{
Option::<u64>::deserialize(deserializer)
}
fn parse_account_weight(value: u64) -> Result<AccountWeight, WireValidationError> {
u16::try_from(value)
.ok()
.and_then(AccountWeight::new)
.ok_or_else(|| WireValidationError::new("weight"))
}
fn validate_wire_group_ids(values: &[String]) -> Result<Vec<AccountGroupId>, WireValidationError> {
if values.len() > MAX_ACCOUNT_GROUP_BATCH
|| values.iter().collect::<BTreeSet<_>>().len() != values.len()
{
return Err(WireValidationError::new("groupIds"));
}
values
.iter()
.cloned()
.map(|value| AccountGroupId::new(value).map_err(|_| WireValidationError::new("groupIds")))
.collect()
}
/// 构造统一账号管理路由。
pub fn router<S>() -> Router<S>
where
@@ -900,9 +1043,12 @@ where
.route("/api/admin/accounts/import", post(import_accounts::<S>))
.route("/api/admin/accounts/refresh", post(refresh_account::<S>))
.route("/api/admin/accounts/rotate", post(rotate_account::<S>))
.route("/api/admin/accounts/enable", post(enable_account::<S>))
.route("/api/admin/accounts/disable", post(disable_account::<S>))
.route("/api/admin/accounts/update", post(update_account::<S>))
.route("/api/admin/accounts/delete", post(delete_accounts::<S>))
.route(
"/api/admin/accounts/batch-update",
post(batch_update_accounts::<S>),
)
.route("/api/admin/accounts/quota", get(account_quota::<S>))
.route(
"/api/admin/accounts/quota/refresh",
@@ -927,6 +1073,27 @@ where
)
}
async fn batch_update_accounts<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<BatchUpdateAccountsRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let command = request.into_command().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.batch_update(&auth.context().mutation_context(), command)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(BatchUpdatedAccountsData::from(result)),
))
}
async fn list_accounts<S>(
_auth: AdminAuth,
State(state): State<S>,
@@ -1016,7 +1183,7 @@ where
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::CREATED,
AdminEnvelope::ok(AccountImportData::from(result)),
AdminEnvelope::ok(AccountImportData::from_result(result)),
))
}
@@ -1111,64 +1278,24 @@ where
))
}
async fn enable_account<S>(
async fn update_account<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountMutationRequest>,
Json(request): Json<UpdateAccountRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_account(auth, state, request, AccountMutationAction::Enable).await
}
async fn disable_account<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountMutationRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
mutate_account(auth, state, request, AccountMutationAction::Disable).await
}
#[derive(Clone, Copy)]
enum AccountMutationAction {
Enable,
Disable,
}
async fn mutate_account<S>(
auth: AdminAuth,
state: S,
request: AccountMutationRequest,
action: AccountMutationAction,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let (provider, command) = request
.into_command(auth.context().mutation_context())
.map_err(map_wire_error)?;
let result = match (provider, action) {
(AccountProvider::OpenAi, AccountMutationAction::Enable) => {
state.admin_services().openai().enable(command).await
}
(AccountProvider::OpenAi, AccountMutationAction::Disable) => {
state.admin_services().openai().disable(command).await
}
(AccountProvider::Xai, AccountMutationAction::Enable) => {
state.admin_services().xai().enable(command).await
}
(AccountProvider::Xai, AccountMutationAction::Disable) => {
state.admin_services().xai().disable(command).await
}
}
.map_err(map_service_error)?;
let command = request.into_command().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.update(&auth.context().mutation_context(), command)
.await
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(AccountMutationData::from(result)),
AdminEnvelope::ok(UpdatedAccountData::from(result)),
))
}
@@ -1384,6 +1511,16 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
id: account.id.clone(),
name: account.name,
provider: account.provider_kind.to_string(),
groups: account
.groups
.into_iter()
.map(|group| AccountGroupRefView {
id: group.id.to_string(),
name: group.name,
color: group.color.as_str().to_owned(),
enabled: group.enabled,
})
.collect(),
resource_ref: account.id,
email: account.email,
account_id: account.upstream_account_id,
@@ -1398,6 +1535,8 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
.map(|reason| reason.as_str().to_owned()),
error_message: projection.error_message,
enabled: account.enabled,
concurrency_limit: account.concurrency_limit.map(|limit| limit.get()),
weight: account.weight.get(),
access_token_expires_at: expires_at,
access_token_expires_at_display: account
.access_token_expires_at
@@ -12,7 +12,7 @@ use gateway_admin::model::client_keys::{
};
use gateway_core::{
policy::{ClientApiKeyId, RateLimits},
routing::ProviderKind,
routing::AccountGroupId,
};
use serde::{Deserialize, Serialize};
@@ -137,7 +137,7 @@ impl ClientKeySort {
pub struct CreateClientKeyRequest {
name: String,
label: Option<String>,
provider_kind: String,
group_ids: Vec<String>,
max_concurrency: u64,
requests_per_minute: u64,
}
@@ -147,14 +147,13 @@ impl CreateClientKeyRequest {
pub fn into_command(self) -> Result<CreateClientKey, WireValidationError> {
validate_required_text(&self.name, "name")?;
validate_optional_text(self.label.as_deref(), "label")?;
validate_provider_kind(&self.provider_kind)?;
let group_ids = validate_group_ids(self.group_ids)?;
validate_limit(self.max_concurrency, "maxConcurrency")?;
validate_limit(self.requests_per_minute, "requestsPerMinute")?;
Ok(CreateClientKey {
name: self.name,
label: self.label,
provider_kind: ProviderKind::new(self.provider_kind)
.map_err(|_| WireValidationError::new("providerKind"))?,
group_ids,
limits: RateLimits {
max_concurrency: self.max_concurrency,
requests_per_minute: self.requests_per_minute,
@@ -170,7 +169,7 @@ pub struct UpdateClientKeyRequest {
id: String,
name: String,
label: Option<String>,
provider_kind: String,
group_ids: Vec<String>,
max_concurrency: u64,
requests_per_minute: u64,
}
@@ -181,15 +180,14 @@ impl UpdateClientKeyRequest {
validate_required_text(&self.id, "id")?;
validate_required_text(&self.name, "name")?;
validate_optional_text(self.label.as_deref(), "label")?;
validate_provider_kind(&self.provider_kind)?;
let group_ids = validate_group_ids(self.group_ids)?;
validate_limit(self.max_concurrency, "maxConcurrency")?;
validate_limit(self.requests_per_minute, "requestsPerMinute")?;
Ok(UpdateClientKey {
id: client_key_id(self.id, "clientKeyMutationNotFound")?,
name: self.name,
label: self.label,
provider_kind: ProviderKind::new(self.provider_kind)
.map_err(|_| WireValidationError::new("providerKind"))?,
group_ids,
limits: RateLimits {
max_concurrency: self.max_concurrency,
requests_per_minute: self.requests_per_minute,
@@ -242,7 +240,9 @@ pub struct ClientKeyView {
id: String,
name: String,
label: Option<String>,
provider_kind: String,
routing_scope: &'static str,
groups: Vec<ClientKeyGroupView>,
provider_kinds: Vec<String>,
prefix: String,
enabled: bool,
max_concurrency: u64,
@@ -252,13 +252,42 @@ pub struct ClientKeyView {
last_used_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ClientKeyGroupView {
id: String,
name: String,
color: String,
enabled: bool,
}
impl From<ClientKeyRecord> for ClientKeyView {
fn from(record: ClientKeyRecord) -> Self {
let routing_scope = if record.groups.is_empty() {
"all"
} else {
"groups"
};
Self {
id: record.id.to_string(),
name: record.name,
label: record.label,
provider_kind: record.provider_kind.to_string(),
routing_scope,
groups: record
.groups
.into_iter()
.map(|group| ClientKeyGroupView {
id: group.id.to_string(),
name: group.name,
color: group.color.as_str().to_owned(),
enabled: group.enabled,
})
.collect(),
provider_kinds: record
.provider_kinds
.into_iter()
.map(|provider| provider.to_string())
.collect(),
prefix: record.prefix,
enabled: record.enabled,
max_concurrency: record.limits.max_concurrency,
@@ -572,21 +601,20 @@ fn validate_required_text(value: &str, field: &'static str) -> Result<(), WireVa
Ok(())
}
fn validate_provider_kind(value: &str) -> Result<(), WireValidationError> {
let valid = !value.is_empty()
&& value.len() <= 64
&& value.bytes().enumerate().all(|(index, byte)| {
byte.is_ascii_lowercase()
|| byte.is_ascii_digit()
|| (byte == b'-' && index > 0 && index + 1 < value.len())
})
&& !value.starts_with('-')
&& !value.ends_with('-');
if valid {
Ok(())
} else {
Err(WireValidationError::new("providerKind"))
fn validate_group_ids(values: Vec<String>) -> Result<Vec<AccountGroupId>, WireValidationError> {
if values.len() > 1000
|| values
.iter()
.collect::<std::collections::BTreeSet<_>>()
.len()
!= values.len()
{
return Err(WireValidationError::new("groupIds"));
}
values
.into_iter()
.map(|value| AccountGroupId::new(value).map_err(|_| WireValidationError::new("groupIds")))
.collect()
}
fn validate_optional_text(
@@ -7,6 +7,7 @@ use axum::{
response::Response,
};
pub mod account_groups;
pub mod accounts;
pub mod auth;
pub mod backups;
@@ -29,6 +30,7 @@ where
S: AdminSessionState + Clone + Send + Sync + 'static,
{
Router::new()
.merge(account_groups::router::<S>())
.merge(accounts::router::<S>())
.merge(auth::router::<S>())
.merge(backups::router::<S>())
@@ -232,6 +232,9 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
id: record.id.clone(),
request_id: record.id,
client_api_key_id: Some(record.client_api_key_ref),
routing_scope: record.routing_scope,
routing_group_refs: record.routing_group_refs,
routing_group_names_snapshot: record.routing_group_names_snapshot,
kind: record.operation,
provider: record.provider_kind,
authentication_kind: record.provider_account_authentication_kind,
@@ -84,6 +84,9 @@ pub struct UsageRecordView {
pub id: String,
pub request_id: String,
pub client_api_key_id: Option<String>,
pub routing_scope: String,
pub routing_group_refs: Vec<String>,
pub routing_group_names_snapshot: Vec<String>,
pub kind: String,
pub provider: Option<String>,
pub authentication_kind: Option<String>,
@@ -89,6 +89,7 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
| EngineError::ProviderMetadataMismatch
| EngineError::ContinuationPinMismatch
| EngineError::RequiredAccountMismatch
| EngineError::AccountOutsideClientScope
| EngineError::DownstreamCommitRequired
| EngineError::NoActiveAttempt
| EngineError::InvalidDeliveryState => {
@@ -0,0 +1,292 @@
use axum::{
Router,
body::{Body, to_bytes},
http::{Method, Request, StatusCode, header},
};
use gateway_api::admin::account_groups;
use serde_json::{Value, json};
use tower::ServiceExt as _;
use super::{AdminTestFixture, AdminTestState, PRIMARY_GROUP_ID};
const SESSION_COOKIE: &str = "cpr_admin_session=valid-session";
#[tokio::test]
async fn list_route_should_keep_camel_case_group_and_page_wire() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::GET,
"/api/admin/account-groups?page=1&pageSize=1&search=alpha&enabled=true",
None,
true,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let value = response_json(response).await;
assert_eq!(value["data"]["page"]["page"], 1);
assert_eq!(value["data"]["page"]["pageSize"], 1);
assert_eq!(value["data"]["page"]["total"], 1);
assert_eq!(value["data"]["page"]["totalPages"], 1);
assert_eq!(value["data"]["configRevision"], 7);
assert_eq!(value["data"]["items"][0]["id"], PRIMARY_GROUP_ID);
assert_eq!(value["data"]["items"][0]["memberCount"], 2);
assert_eq!(value["data"]["items"][0]["providerCounts"]["openai"], 1);
assert_eq!(value["data"]["items"][0]["clientKeyCount"], 2);
assert_eq!(value["data"]["items"][0]["accountSummary"]["available"], 1);
assert_eq!(value["data"]["items"][0]["accountSummary"]["limited"], 1);
assert_eq!(value["data"]["items"][0]["accountSummary"]["total"], 2);
assert_eq!(value["data"]["items"][0]["capacity"]["usedSlots"], 0);
assert_eq!(value["data"]["items"][0]["capacity"]["totalSlots"], 1);
assert_eq!(value["data"]["items"][0]["usage"]["todayUsd"], "1.25");
assert_eq!(value["data"]["items"][0]["usage"]["totalUsd"], "5.5");
assert!(value["data"]["items"][0].get("member_count").is_none());
}
#[tokio::test]
async fn members_route_should_keep_provider_kind_and_member_wire() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::GET,
&format!("/api/admin/account-groups/members?id={PRIMARY_GROUP_ID}"),
None,
true,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let value = response_json(response).await;
assert_eq!(value["data"]["id"], PRIMARY_GROUP_ID);
assert_eq!(value["data"]["total"], 2);
assert_eq!(value["data"]["items"][0]["providerKind"], "openai");
assert_eq!(value["data"]["items"][0]["email"], "openai@example.invalid");
assert!(value["data"]["items"][0].get("provider_kind").is_none());
}
#[tokio::test]
async fn create_route_should_create_an_empty_group() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/create",
Some(json!({
"name": "Gamma routing",
"description": "New traffic pool",
"color": "#a855f780"
})),
true,
)
.await;
assert_eq!(response.status(), StatusCode::CREATED);
let value = response_json(response).await;
assert!(
value["data"]["id"]
.as_str()
.is_some_and(|id| id.starts_with("grp_") && id.len() == 36)
);
assert_eq!(value["data"]["record"]["name"], "Gamma routing");
assert_eq!(value["data"]["record"]["memberCount"], 0);
assert_eq!(value["data"]["record"]["color"], "#A855F780");
assert_eq!(value["data"]["configRevision"], 8);
}
#[tokio::test]
async fn update_route_should_replace_group_fields() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/update",
Some(json!({
"id": PRIMARY_GROUP_ID,
"name": "Renamed routing",
"description": null,
"color": "#F43F5ECC"
})),
true,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let value = response_json(response).await;
assert_eq!(value["data"]["record"]["name"], "Renamed routing");
assert!(value["data"]["record"]["description"].is_null());
assert_eq!(value["data"]["record"]["color"], "#F43F5ECC");
}
#[tokio::test]
async fn enable_and_disable_routes_should_set_explicit_group_state() {
let fixture = authenticated_fixture().await;
let disabled = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/disable",
Some(json!({ "id": PRIMARY_GROUP_ID })),
true,
)
.await;
let enabled = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/enable",
Some(json!({ "id": PRIMARY_GROUP_ID })),
true,
)
.await;
assert_eq!(disabled.status(), StatusCode::OK);
assert!(
!response_json(disabled).await["data"]["record"]["enabled"]
.as_bool()
.expect("disabled flag")
);
assert_eq!(enabled.status(), StatusCode::OK);
assert!(
response_json(enabled).await["data"]["record"]["enabled"]
.as_bool()
.expect("enabled flag")
);
}
#[tokio::test]
async fn delete_route_should_return_deleted_id_without_a_live_record() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/delete",
Some(json!({ "id": PRIMARY_GROUP_ID })),
true,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let value = response_json(response).await;
assert_eq!(value["data"]["id"], PRIMARY_GROUP_ID);
assert!(value["data"]["record"].is_null());
}
#[tokio::test]
async fn list_route_should_reject_unknown_and_invalid_pagination_fields() {
let fixture = authenticated_fixture().await;
for uri in [
"/api/admin/account-groups?other=true",
"/api/admin/account-groups?page=0",
"/api/admin/account-groups?pageSize=0",
"/api/admin/account-groups?pageSize=201",
] {
let response = request(router(&fixture), Method::GET, uri, None, true).await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{uri}");
}
}
#[tokio::test]
async fn mutation_route_should_reject_unknown_json_fields() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/create",
Some(json!({
"name": "Unknown field",
"description": null,
"color": "#2563EBFF",
"providerKind": "openai"
})),
true,
)
.await;
assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
}
#[tokio::test]
async fn create_route_should_reject_six_digit_color_without_alpha() {
let fixture = authenticated_fixture().await;
let response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/create",
Some(json!({
"name": "Opaque legacy color",
"description": null,
"color": "#2563EB"
})),
true,
)
.await;
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
}
#[tokio::test]
async fn account_group_routes_should_require_admin_authentication() {
let fixture = AdminTestFixture::new().await;
let get_response = request(
router(&fixture),
Method::GET,
"/api/admin/account-groups",
None,
false,
)
.await;
let post_response = request(
router(&fixture),
Method::POST,
"/api/admin/account-groups/create",
Some(json!({ "name": "Unauthorized", "description": null, "color": "#2563EBFF" })),
false,
)
.await;
assert_eq!(get_response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(post_response.status(), StatusCode::UNAUTHORIZED);
}
async fn authenticated_fixture() -> AdminTestFixture {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
fixture
}
fn router(fixture: &AdminTestFixture) -> Router {
account_groups::router::<AdminTestState>().with_state(fixture.state())
}
async fn request(
router: Router,
method: Method,
uri: &str,
body: Option<Value>,
authenticated: bool,
) -> axum::response::Response {
let mut builder = Request::builder()
.method(method)
.uri(uri)
.header("x-request-id", "req_account_groups");
if authenticated {
builder = builder.header(header::COOKIE, SESSION_COOKIE);
}
let body = match body {
Some(value) => {
builder = builder.header(header::CONTENT_TYPE, "application/json");
Body::from(serde_json::to_vec(&value).expect("serialize request body"))
}
None => Body::empty(),
};
router
.oneshot(builder.body(body).expect("account group request"))
.await
.expect("account group response")
}
async fn response_json(response: axum::response::Response) -> Value {
let body = to_bytes(response.into_body(), 64 * 1024)
.await
.expect("account group response body");
serde_json::from_slice(&body).expect("account group response JSON")
}
@@ -70,6 +70,135 @@ mod query {
}
}
mod batch_update {
use gateway_api::admin::accounts::{BatchUpdateAccountsRequest, UpdateAccountRequest};
use serde_json::json;
#[test]
fn batch_update_should_accept_complete_atomic_payload() {
let request: BatchUpdateAccountsRequest = serde_json::from_value(json!({
"accountIds": ["acct_openai", "acct_xai"],
"enabled": false,
"concurrencyLimit": null,
"weight": 1,
"groupIds": ["grp_00000000000000000000000000000001"]
}))
.expect("deserialize batch update");
request.validate().expect("validate batch update");
}
#[test]
fn batch_update_should_reject_duplicate_or_invalid_account_ids() {
for account_ids in [
json!(["acct_same", "acct_same"]),
json!(["not-an-account"]),
json!([]),
] {
let request: BatchUpdateAccountsRequest = serde_json::from_value(json!({
"accountIds": account_ids,
"enabled": true,
"concurrencyLimit": 8,
"weight": 100,
"groupIds": []
}))
.expect("deserialize invalid batch update");
assert_eq!(
request.validate().expect_err("reject account IDs").field(),
"accountIds"
);
}
}
#[test]
fn batch_update_should_require_enabled_and_reject_unknown_fields() {
assert!(
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
"accountIds": ["acct_test"],
"concurrencyLimit": null,
"weight": 1,
"groupIds": []
}))
.is_err()
);
assert!(
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
"accountIds": ["acct_test"],
"enabled": true,
"concurrencyLimit": null,
"weight": 1,
"groupIds": [],
"legacy": true
}))
.is_err()
);
}
#[test]
fn batch_update_should_reject_invalid_scheduling_bounds_and_missing_fields() {
for (concurrency_limit, weight, field) in [
(json!(0), json!(1), "concurrencyLimit"),
(json!(4294967296_u64), json!(1), "concurrencyLimit"),
(json!(null), json!(0), "weight"),
(json!(null), json!(101), "weight"),
] {
let request: BatchUpdateAccountsRequest = serde_json::from_value(json!({
"accountIds": ["acct_test"],
"enabled": true,
"concurrencyLimit": concurrency_limit,
"weight": weight,
"groupIds": []
}))
.expect("deserialize invalid scheduling");
assert_eq!(
request.validate().expect_err("reject scheduling").field(),
field
);
}
assert!(
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
"accountIds": ["acct_test"],
"enabled": true,
"weight": 1,
"groupIds": []
}))
.is_err()
);
assert!(
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
"accountIds": ["acct_test"],
"enabled": true,
"concurrencyLimit": null,
"groupIds": []
}))
.is_err()
);
}
#[test]
fn single_update_should_require_the_same_complete_scheduling_contract() {
let request: UpdateAccountRequest = serde_json::from_value(json!({
"accountId": "acct_test",
"enabled": true,
"concurrencyLimit": 4294967295_u64,
"weight": 100,
"groupIds": []
}))
.expect("deserialize single update");
request.validate().expect("validate single update");
assert!(
serde_json::from_value::<UpdateAccountRequest>(json!({
"accountId": "acct_test",
"enabled": true,
"weight": 1,
"groupIds": []
}))
.is_err()
);
}
}
mod response {
use gateway_api::admin::accounts::AccountUsageView;
use gateway_api::admin::presenter::format_decimal_currency;
@@ -125,6 +254,7 @@ mod response {
}
mod actions {
use base64::Engine as _;
use gateway_admin::model::{
Revision,
accounts::AccountConnectionTestEvent as DomainConnectionTestEvent,
@@ -136,7 +266,8 @@ mod actions {
AccountActionRequest, AccountConnectionTestEvent, AccountDeletionData,
AccountDeletionRequest, AccountExportData, AccountExportQuery, AccountIdQuery,
AccountImportData, AccountImportRequest, AccountMutationData, AccountRefreshRequest,
AccountTestQuery, RotateAccountRequest, StartAccountAuthorizationRequest,
AccountTestQuery, CompleteAccountAuthorizationRequest, RotateAccountRequest,
StartAccountAuthorizationRequest,
};
use gateway_core::engine::credential::ProviderAccountId;
use serde_json::json;
@@ -280,6 +411,35 @@ mod actions {
}))
.is_err()
);
assert!(
serde_json::from_value::<AccountImportRequest>(json!({
"provider": "openai",
"data": {},
"groupIds": []
}))
.is_err()
);
}
#[test]
fn oauth_complete_should_not_accept_group_assignment() {
let flow_id = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([7_u8; 32]);
let request: CompleteAccountAuthorizationRequest = serde_json::from_value(json!({
"provider": "openai",
"flowId": flow_id,
"callbackUrl": "http://localhost/callback?code=ok"
}))
.expect("decode OAuth completion");
assert!(request.validate().is_ok());
assert!(
serde_json::from_value::<CompleteAccountAuthorizationRequest>(json!({
"provider": "xai",
"flowId": "flow_test",
"callbackUrl": "http://localhost/callback?code=ok",
"groupIds": []
}))
.is_err()
);
}
#[test]
@@ -316,7 +476,7 @@ mod actions {
#[test]
fn account_import_response_should_emit_account_ids() {
let response = AccountImportData::from(CredentialImportResult {
let response = AccountImportData::from_result(CredentialImportResult {
config_revision: Revision::new(8).expect("revision"),
credential_ids: vec![ProviderAccountId::new("acct_imported").expect("account ID")],
});
@@ -105,7 +105,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
let valid = serde_json::from_value::<CreateClientKeyRequest>(json!({
"name": "terminal key",
"label": "production",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": 2,
"requestsPerMinute": 60
}))
@@ -118,7 +118,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
(
json!({
"name": " ",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": 0,
"requestsPerMinute": 0
}),
@@ -127,7 +127,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
(
json!({
"name": "key",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": u64::MAX,
"requestsPerMinute": 0
}),
@@ -150,7 +150,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
"id": "key_1",
"expectedConfigRevision": 1,
"name": "key",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": 0,
"requestsPerMinute": 0,
"tokensPerMinute": 0
@@ -161,7 +161,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
serde_json::from_value::<CreateClientKeyRequest>(json!({
"expectedConfigRevision": 7,
"name": "terminal key",
"providerKind": "openai",
"groupIds": [],
"maxConcurrency": 2,
"requestsPerMinute": 60
}))
@@ -210,7 +210,10 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
id: gateway_core::policy::ClientApiKeyId::new("key_visible").expect("Client Key ID"),
name: "visible".to_owned(),
label: None,
provider_kind: gateway_core::routing::ProviderKind::new("openai").expect("Provider kind"),
groups: Vec::new(),
provider_kinds: vec![
gateway_core::routing::ProviderKind::new("openai").expect("Provider kind"),
],
prefix: "sk_visible12".to_owned(),
enabled: true,
limits: gateway_core::policy::RateLimits {
@@ -226,7 +229,11 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
assert!(list.get("configRevision").is_none());
assert_eq!(list["total"], 1);
assert_eq!(list["items"][0]["id"], "key_visible");
assert_eq!(list["items"][0]["providerKind"], "openai");
assert_eq!(
list["items"][0]["providerKinds"],
serde_json::json!(["openai"])
);
assert_eq!(list["items"][0]["routingScope"], "all");
assert_eq!(list["items"][0]["maxConcurrency"], 2);
assert_eq!(list["items"][0]["requestsPerMinute"], 60);
assert!(list["items"][0].get("tokensPerMinute").is_none());
+312 -11
View File
@@ -13,9 +13,16 @@ use gateway_admin::{
AdminConfig, AdminServices, InitialAdminPassword,
model::{
MutationContext, Revision,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
AccountGroupListQuery, AccountGroupMember, AccountGroupMembers, AccountGroupMutation,
AccountGroupPage, AccountGroupRecord, AccountGroupUsage, DeleteAccountGroup,
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountPageItem, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountListQuery, AccountPage, AccountPageItem, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -23,9 +30,9 @@ use gateway_admin::{
NewClientKey, SetClientKeyEnabled, UpdateClientKey,
},
observability::{
DashboardObservation, DiagnosticDimension, DiagnosticObservation, OpsError,
OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail, UsageFilter,
UsageOverview, UsagePage, UsageQuery, UsageRecord,
DashboardObservation, DecimalAmount, DiagnosticDimension, DiagnosticObservation,
OpsError, OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail,
UsageFilter, UsageOverview, UsagePage, UsageQuery, UsageRecord,
},
provider_credentials::{
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
@@ -45,8 +52,8 @@ use gateway_admin::{
ports::{
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
store::{
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult,
AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
AccountGroupStore, AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts,
AdminStoreResult, AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
},
system::{
SystemOperationError, SystemOperationErrorKind, SystemOperations,
@@ -66,6 +73,7 @@ use gateway_core::{
},
};
mod account_groups;
mod accounts;
mod auth;
mod client_keys;
@@ -96,6 +104,7 @@ impl AdminTestFixture {
let auth = Arc::new(MemoryAuthStore::new(api_key.clone()));
let settings = Arc::new(MemorySettingsStore::new(api_key));
let client_keys = Arc::new(MemoryClientKeyStore);
let account_groups = Arc::new(MemoryAccountGroupStore::new());
let usage_records = Arc::new(Mutex::new(Vec::new()));
let usage_detail = Arc::new(Mutex::new(None));
let diagnostics = Arc::new(Mutex::new(Vec::new()));
@@ -112,6 +121,7 @@ impl AdminTestFixture {
});
let stores = AdminStorePorts::new(
unused.clone(),
account_groups.clone(),
auth.clone(),
client_keys.clone(),
unused,
@@ -345,6 +355,258 @@ impl SettingsStore for MemorySettingsStore {
pub(super) struct MemoryClientKeyStore;
pub(super) const PRIMARY_GROUP_ID: &str = "grp_11111111111111111111111111111111";
pub(super) const SECONDARY_GROUP_ID: &str = "grp_22222222222222222222222222222222";
struct MemoryAccountGroupState {
revision: Revision,
groups: BTreeMap<gateway_core::routing::AccountGroupId, AccountGroupRecord>,
members: BTreeMap<gateway_core::routing::AccountGroupId, Vec<AccountGroupMember>>,
}
pub(super) struct MemoryAccountGroupStore {
state: Mutex<MemoryAccountGroupState>,
}
impl MemoryAccountGroupStore {
fn new() -> Self {
let now = Utc::now();
let primary_id = group_id(PRIMARY_GROUP_ID);
let secondary_id = group_id(SECONDARY_GROUP_ID);
let openai = AccountGroupMember {
id: "acct_openai".to_owned(),
name: "OpenAI primary".to_owned(),
provider_kind: ProviderKind::new("openai").expect("provider kind"),
email: Some("openai@example.invalid".to_owned()),
enabled: true,
};
let xai = AccountGroupMember {
id: "acct_xai".to_owned(),
name: "xAI primary".to_owned(),
provider_kind: ProviderKind::new("xai").expect("provider kind"),
email: None,
enabled: false,
};
let groups = BTreeMap::from([
(
primary_id.clone(),
AccountGroupRecord {
id: primary_id.clone(),
name: "Alpha routing".to_owned(),
description: Some("Primary traffic".to_owned()),
color: group_color("#2563ebff"),
enabled: true,
member_count: 2,
provider_counts: BTreeMap::from([
("openai".to_owned(), 1),
("xai".to_owned(), 1),
]),
client_key_count: 2,
account_summary: account_summary(1, 1, 2),
capacity: capacity(Some(0), 1),
usage: usage("1.25", "5.5"),
created_at: now,
updated_at: now,
},
),
(
secondary_id.clone(),
AccountGroupRecord {
id: secondary_id,
name: "Beta routing".to_owned(),
description: None,
color: group_color("#64748B80"),
enabled: false,
member_count: 0,
provider_counts: BTreeMap::new(),
client_key_count: 0,
account_summary: account_summary(0, 0, 0),
capacity: capacity(Some(0), 0),
usage: usage("0", "0"),
created_at: now,
updated_at: now,
},
),
]);
Self {
state: Mutex::new(MemoryAccountGroupState {
revision: Revision::new(7).expect("revision"),
groups,
members: BTreeMap::from([(primary_id, vec![openai, xai])]),
}),
}
}
}
#[async_trait]
impl AccountGroupStore for MemoryAccountGroupStore {
async fn list_account_groups(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage> {
let state = self.state.lock().expect("account groups");
let search = query.search.as_ref().map(|value| value.to_lowercase());
let mut matching = state
.groups
.values()
.filter(|record| {
query
.enabled
.is_none_or(|enabled| record.enabled == enabled)
})
.filter(|record| {
search.as_ref().is_none_or(|search| {
record.name.to_lowercase().contains(search)
|| record
.description
.as_ref()
.is_some_and(|description| description.to_lowercase().contains(search))
})
})
.cloned()
.collect::<Vec<_>>();
matching.sort_by(|left, right| left.name.cmp(&right.name));
let total = matching.len() as u64;
let page_size = usize::from(query.page_size.get());
let offset = usize::try_from(query.page.saturating_sub(1))
.unwrap_or(usize::MAX)
.saturating_mul(page_size);
let items = matching.into_iter().skip(offset).take(page_size).collect();
Ok(AccountGroupPage {
config_revision: state.revision,
items,
total,
page: query.page,
page_size: query.page_size.get(),
})
}
async fn account_group_members(
&self,
id: &gateway_core::routing::AccountGroupId,
) -> AdminStoreResult<AccountGroupMembers> {
let state = self.state.lock().expect("account groups");
if !state.groups.contains_key(id) {
return Err(not_found("account group"));
}
let items = state.members.get(id).cloned().unwrap_or_default();
Ok(AccountGroupMembers {
config_revision: state.revision,
id: id.clone(),
total: items.len() as u64,
items,
})
}
async fn create_account_group(
&self,
command: NewAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let mut state = self.state.lock().expect("account groups");
let now = Utc::now();
let record = AccountGroupRecord {
id: command.id.clone(),
name: command.name,
description: command.description,
color: command.color,
enabled: true,
member_count: 0,
provider_counts: BTreeMap::new(),
client_key_count: 0,
account_summary: account_summary(0, 0, 0),
capacity: capacity(Some(0), 0),
usage: usage("0", "0"),
created_at: now,
updated_at: now,
};
state.groups.insert(command.id.clone(), record);
state.members.insert(command.id.clone(), Vec::new());
refresh_group_counts(&mut state, &command.id);
mutation(&mut state, command.id, true)
}
async fn update_account_group(
&self,
command: UpdateAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let mut state = self.state.lock().expect("account groups");
let record = state
.groups
.get_mut(&command.id)
.ok_or_else(|| not_found("account group"))?;
record.name = command.name;
record.description = command.description;
record.color = command.color;
record.updated_at = Utc::now();
mutation(&mut state, command.id, true)
}
async fn set_account_group_enabled(
&self,
command: SetAccountGroupEnabled,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let mut state = self.state.lock().expect("account groups");
let record = state
.groups
.get_mut(&command.id)
.ok_or_else(|| not_found("account group"))?;
record.enabled = command.enabled;
record.updated_at = Utc::now();
mutation(&mut state, command.id, true)
}
async fn delete_account_group(
&self,
command: DeleteAccountGroup,
_: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let mut state = self.state.lock().expect("account groups");
if state.groups.remove(&command.id).is_none() {
return Err(not_found("account group"));
}
state.members.remove(&command.id);
mutation(&mut state, command.id, false)
}
}
fn group_id(value: &str) -> gateway_core::routing::AccountGroupId {
gateway_core::routing::AccountGroupId::new(value).expect("account group ID")
}
fn refresh_group_counts(
state: &mut MemoryAccountGroupState,
id: &gateway_core::routing::AccountGroupId,
) {
let members = state.members.get(id).map(Vec::as_slice).unwrap_or_default();
let provider_counts = members.iter().fold(BTreeMap::new(), |mut counts, member| {
*counts
.entry(member.provider_kind.as_str().to_owned())
.or_insert(0) += 1;
counts
});
if let Some(record) = state.groups.get_mut(id) {
record.member_count = members.len() as u64;
record.provider_counts = provider_counts;
record.updated_at = Utc::now();
}
}
fn mutation(
state: &mut MemoryAccountGroupState,
id: gateway_core::routing::AccountGroupId,
include_record: bool,
) -> AdminStoreResult<AccountGroupMutation> {
state.revision = next_revision(state.revision);
Ok(AccountGroupMutation {
config_revision: state.revision,
record: include_record.then(|| state.groups.get(&id).expect("group record").clone()),
id,
})
}
#[async_trait]
impl ClientKeyStore for MemoryClientKeyStore {
async fn list_client_keys(&self, _: ClientKeyListQuery) -> AdminStoreResult<ClientKeyPage> {
@@ -366,7 +628,8 @@ impl ClientKeyStore for MemoryClientKeyStore {
id: id.clone(),
name: "revealed".to_owned(),
label: None,
provider_kind: ProviderKind::new("openai").expect("provider kind"),
groups: Vec::new(),
provider_kinds: vec![ProviderKind::new("openai").expect("provider kind")],
prefix: "sk_aaaaaaaaa".to_owned(),
enabled: true,
limits: RateLimits::unlimited(),
@@ -501,14 +764,22 @@ impl AccountStore for UnusedStore {
Err(unavailable("credential refresh"))
}
async fn set_account_enabled(
async fn update_account(
&self,
_: SetAccountEnabled,
_: UpdateAccount,
_: &MutationContext,
) -> AdminStoreResult<Revision> {
) -> AdminStoreResult<AccountUpdateResult> {
Err(unavailable("account enabled"))
}
async fn batch_update_accounts(
&self,
_: BatchUpdateAccounts,
_: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
Err(unavailable("account batch update"))
}
async fn delete_accounts(
&self,
_: DeleteAccounts,
@@ -792,6 +1063,32 @@ fn next_revision(revision: Revision) -> Revision {
Revision::new(revision.get().saturating_add(1)).expect("next revision")
}
fn account_summary(available: u64, limited: u64, total: u64) -> AccountGroupAccountSummary {
AccountGroupAccountSummary {
available,
limited,
total,
}
}
fn group_color(value: &str) -> AccountGroupColor {
AccountGroupColor::parse(value).expect("group color")
}
fn capacity(used_slots: Option<u64>, total_slots: u64) -> AccountGroupCapacity {
AccountGroupCapacity {
used_slots,
total_slots,
}
}
fn usage(today_usd: &str, total_usd: &str) -> AccountGroupUsage {
AccountGroupUsage {
today_usd: today_usd.parse::<DecimalAmount>().expect("today cost"),
total_usd: total_usd.parse::<DecimalAmount>().expect("total cost"),
}
}
fn unavailable(resource: &'static str) -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
@@ -800,6 +1097,10 @@ fn unavailable(resource: &'static str) -> AdminStoreError {
)
}
fn not_found(resource: &'static str) -> AdminStoreError {
AdminStoreError::new(AdminStoreErrorKind::NotFound, resource, "not found")
}
fn unsupported_provider() -> ProviderAdminError {
ProviderAdminError::new(ProviderAdminErrorKind::Unsupported)
}
@@ -674,6 +674,9 @@ fn usage_record_with_account(
id: id.to_owned(),
client_api_key_ref: "key_detail".to_owned(),
config_revision: 1,
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: Vec::new(),
protocol: "openai".to_owned(),
operation: "responses".to_owned(),
endpoint: "/v1/responses".to_owned(),
@@ -800,6 +803,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
id: "request_endpoint".to_owned(),
client_api_key_ref: "key_endpoint".to_owned(),
config_revision: 1,
routing_scope: "groups".to_owned(),
routing_group_refs: vec!["grp_usage".to_owned()],
routing_group_names_snapshot: vec!["Usage group".to_owned()],
protocol: "openai".to_owned(),
operation: "generate".to_owned(),
endpoint: "/v1/responses".to_owned(),
@@ -919,6 +925,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
assert_eq!(
serde_json::json!({
"route": value["data"]["items"][0]["route"],
"routingScope": value["data"]["items"][0]["routingScope"],
"routingGroupRefs": value["data"]["items"][0]["routingGroupRefs"],
"routingGroupNamesSnapshot": value["data"]["items"][0]["routingGroupNamesSnapshot"],
"serviceTier": value["data"]["items"][0]["serviceTier"],
"accountId": value["data"]["items"][0]["accountId"],
"accountName": value["data"]["items"][0]["accountName"],
@@ -949,6 +958,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
}),
serde_json::json!({
"route": "/v1/responses",
"routingScope": "groups",
"routingGroupRefs": ["grp_usage"],
"routingGroupNamesSnapshot": ["Usage group"],
"serviceTier": "priority",
"accountId": "acct_snapshot",
"accountName": "Snapshot Alpha",
@@ -33,6 +33,7 @@ fn source_tree_should_match_frozen_machine_manifest() {
let mut actual = rust_files(&root.join("src"));
actual.sort();
let mut expected = vec![
"src/admin/account_groups.rs",
"src/admin/accounts.rs",
"src/admin/auth.rs",
"src/admin/backups.rs",
@@ -73,6 +74,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
let mut actual = rust_files(&root.join("tests"));
actual.sort();
let mut expected = vec![
"tests/admin/account_groups.rs",
"tests/admin/accounts.rs",
"tests/admin/auth.rs",
"tests/admin/client_keys.rs",
+16 -5
View File
@@ -4,7 +4,7 @@ mod models;
mod responses;
mod router;
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
@@ -18,7 +18,9 @@ use gateway_core::engine::admission::{
use gateway_core::engine::continuation::{
NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId,
};
use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy};
use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
};
use gateway_core::engine::execution::{
AuthenticatedClient, ClientApiKeyUsageSink, DefaultExecutionService, ExecutionService,
ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
@@ -34,8 +36,8 @@ use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLif
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, RuntimeSnapshot,
UpstreamModelId,
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
ProviderModel, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
};
pub(super) async fn api_router(execution: Arc<dyn ExecutionService>) -> axum::Router {
@@ -94,6 +96,10 @@ impl ClientApiKeyUsageSink for IgnoredClientApiKeyUsage {
fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
let provider = ProviderKind::new(provider_name).expect("provider");
let account_directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new("acct_api_test").expect("account ID"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)])));
let capabilities = ModelCapabilities::new(
BTreeSet::from([gateway_core::operation::OperationKind::Generate]),
Some(16_000),
@@ -119,12 +125,16 @@ fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
vec![ClientPolicy::new(
ClientApiKeyId::new("key_api_test").expect("key ID"),
PlaintextClientApiKey::new(plaintext).expect("plaintext key"),
provider,
Arc::new(FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::all_accounts(),
)),
true,
RateLimits::unlimited(),
)],
)
.expect("runtime snapshot")
.with_account_directory(account_directory)
}
#[derive(Default)]
@@ -260,6 +270,7 @@ struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async { unreachable!("authentication fixture does not execute") })
@@ -451,10 +451,12 @@ fn decoder_should_not_reject_large_bodies_using_catalog_context_limits() {
.to_string();
let decoded = decode_request(body.as_bytes()).expect("large request should decode");
let plan = super::snapshot("sk_context_test", "openai")
let snapshot = super::snapshot("sk_context_test", "openai");
let plan = snapshot
.plan(
&PublicModelId::new(decoded.metadata().public_model()).expect("public model"),
decoded.operation(),
snapshot.all_account_scope(),
&RoutingContext::default(),
)
.expect("large body must not be locally context-gated");
@@ -9,6 +9,7 @@ use futures::future::BoxFuture;
use crate::engine::credential::ProviderAccountId;
use crate::operation::ProviderSessionState;
use crate::policy::ClientApiKeyId;
use crate::routing::ProviderKind;
/// 客户端或 Provider 传递的 opaque response handle。
@@ -48,14 +49,15 @@ pub enum NativeContinuationScope {
/// 从可丢失会话亲和存储恢复的 native previous-response pin。
///
/// 该值不依赖调用方 API Key;它只阻止 native handle 与不透明 Provider state
/// 被发送到错误的 Provider 或账号。
/// 该值同时绑定调用方 API Key、Provider 与账号,阻止 native handle 和不透明
/// Provider state 被跨租户或跨上游目标复用。
#[derive(Clone, PartialEq)]
pub struct NativeContinuationPin {
/// 客户端提交、用于查找可丢失会话亲和记录的 response ID。
previous_response_id: PreviousResponseId,
/// Provider 原生 response handle。
upstream_response_id: PreviousResponseId,
client_api_key_id: ClientApiKeyId,
provider: ProviderKind,
account: ProviderAccountId,
scope: NativeContinuationScope,
@@ -67,12 +69,14 @@ impl NativeContinuationPin {
pub const fn new(
previous_response_id: PreviousResponseId,
upstream_response_id: PreviousResponseId,
client_api_key_id: ClientApiKeyId,
provider: ProviderKind,
account: ProviderAccountId,
) -> Self {
Self {
previous_response_id,
upstream_response_id,
client_api_key_id,
provider,
account,
scope: NativeContinuationScope::ConnectionLocal,
@@ -105,6 +109,16 @@ impl NativeContinuationPin {
&self.upstream_response_id
}
#[must_use]
pub const fn client_api_key_id(&self) -> &ClientApiKeyId {
&self.client_api_key_id
}
#[must_use]
pub fn matches_client(&self, client_api_key_id: &ClientApiKeyId) -> bool {
self.client_api_key_id == *client_api_key_id
}
#[must_use]
pub const fn provider(&self) -> &ProviderKind {
&self.provider
@@ -139,6 +153,7 @@ impl fmt::Debug for NativeContinuationPin {
.debug_struct("NativeContinuationPin")
.field("previous_response_id", &"<redacted>")
.field("upstream_response_id", &"<redacted>")
.field("client_api_key_id", &self.client_api_key_id)
.field("provider", &self.provider)
.field("account", &self.account)
.field("scope", &self.scope)
@@ -190,21 +205,57 @@ impl fmt::Debug for ContinuationBinding {
}
}
/// Native continuation lookup 的稳定失败分类。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum NativeContinuationStoreErrorKind {
/// Redis 等可丢失协调存储暂不可用;调用方可以按外部 continuation fail-open。
Unavailable,
/// 存储中的记录无法还原为可信 pin;不得将原始 handle 继续透传上游。
InvalidData,
/// 已存在的记录属于另一个客户端 API Key;必须 fail closed。
OwnershipMismatch,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("native continuation store is unavailable: {detail}")]
#[error("native continuation store failed: {kind:?}: {detail}")]
pub struct NativeContinuationStoreError {
kind: NativeContinuationStoreErrorKind,
detail: String,
}
impl NativeContinuationStoreError {
/// 构造带可排查细节的亲和存储错误;`detail` 不应包含 response handle。
/// 构造可 fail-open 的协调存储不可用错误;`detail` 不应包含 response handle。
#[must_use]
pub fn new(detail: impl Into<String>) -> Self {
pub fn unavailable(detail: impl Into<String>) -> Self {
Self {
kind: NativeContinuationStoreErrorKind::Unavailable,
detail: detail.into(),
}
}
/// 构造必须 fail closed 的无效记录错误;`detail` 不应包含 response handle。
#[must_use]
pub fn invalid_data(detail: impl Into<String>) -> Self {
Self {
kind: NativeContinuationStoreErrorKind::InvalidData,
detail: detail.into(),
}
}
/// 构造不泄露记录归属细节的跨调用方错误。
#[must_use]
pub fn ownership_mismatch() -> Self {
Self {
kind: NativeContinuationStoreErrorKind::OwnershipMismatch,
detail: "continuation belongs to a different client API key".to_owned(),
}
}
#[must_use]
pub const fn kind(&self) -> NativeContinuationStoreErrorKind {
self.kind
}
#[must_use]
pub fn detail(&self) -> &str {
&self.detail
@@ -213,11 +264,12 @@ impl NativeContinuationStoreError {
/// 可丢失的 previous-response 亲和存储端口。
///
/// 亲和记录不按调用方 API Key 隔离:与旧版一样,response ID 是会话级 opaque
/// handle;Redis 不可用或超时必须由调用方 fail-open。
/// 亲和记录按调用方 API Key 隔离;Redis 不可用或超时必须由调用方 fail-open,
/// 但命中其他 Key 的记录必须 fail closed。
pub trait NativeContinuationPort: Send + Sync {
fn resolve<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
previous_response_id: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>>;
@@ -46,6 +46,10 @@ impl AccountSelection {
Self::Diagnostic(account) => Some(account),
}
}
const fn is_diagnostic(&self) -> bool {
matches!(self, Self::Diagnostic(_))
}
}
impl<S: ?Sized> AttemptCoordinator<S>
@@ -126,6 +130,15 @@ where
.map(ProviderAccountStateOwner::from_continuation);
let continuation_attempt =
initial_continuation_attempt(&operation, &plan, continuation.as_ref());
let candidate_index = continuation
.as_ref()
.and_then(ContinuationBinding::pinned)
.map_or(Ok(0), |pin| {
plan.candidates()
.iter()
.position(|candidate| candidate.provider() == pin.provider())
.ok_or(EngineError::ContinuationPinMismatch)
})?;
let image_generation_requested = operation.image_generation_requested();
let mut session = ResponseExecutionSession {
engine: Arc::clone(&self.engine),
@@ -149,6 +162,7 @@ where
account_state_owner,
cancellation,
attempts: 0,
candidate_index,
excluded_accounts: BTreeSet::new(),
credential_recovery_attempted_accounts: BTreeSet::new(),
recovery_account: None,
@@ -225,6 +239,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
account_state_owner: Option<ProviderAccountStateOwner>,
cancellation: CancellationToken,
attempts: u32,
candidate_index: usize,
excluded_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
credential_recovery_attempted_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
/// 凭据恢复后的一次性同账号钉选;只约束紧随其后的 replay attempt,
@@ -461,6 +476,7 @@ where
NativeContinuationPin::new(
previous_response_id,
upstream_response_id,
self.client_api_key_ref.clone(),
provider,
account,
)
@@ -565,7 +581,7 @@ where
if self.attempts >= self.plan.max_attempts().get() {
return Err(EngineError::EmptyRoutingPlan);
}
let Some(candidate) = self.plan.candidates().first().cloned() else {
let Some(candidate) = self.plan.candidates().get(self.candidate_index).cloned() else {
let error = GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream Provider is available",
@@ -606,7 +622,8 @@ where
self.excluded_accounts.clone(),
pinned_account.clone(),
self.account_state_owner.clone(),
),
)
.with_account_scope(Arc::clone(self.plan.account_scope())),
}
.with_credential_recovery_attempted(pinned_account.as_ref().is_some_and(|account| {
self.credential_recovery_attempted_accounts
@@ -663,6 +680,20 @@ where
if self.prepare_unavailable_native_continuation_replay(&error) {
return Ok(Some(PullOutcome::AttemptDiscarded));
}
if matches!(
error.kind(),
ProviderErrorKind::AccountCapacityUnavailable
| ProviderErrorKind::NoEligibleAccount
| ProviderErrorKind::ProviderInfrastructureUnavailable
) && error.send_state() == UpstreamSendState::NotSent
&& matches!(
self.continuation_attempt,
ContinuationAttempt::None | ContinuationAttempt::ReplayAny
)
&& self.advance_provider_candidate()
{
return Ok(Some(PullOutcome::AttemptDiscarded));
}
if matches!(
error.kind(),
ProviderErrorKind::AccountCapacityUnavailable
@@ -716,6 +747,27 @@ where
}
let metadata = stream.metadata().clone();
if !self.account_selection.is_diagnostic()
&& metadata
.provider_account_id()
.is_none_or(|account| !candidate.account_scope().allows(account))
{
let error = GatewayError::new(
GatewayErrorKind::Internal,
"provider selected an account outside the frozen client scope",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::AccountOutsideClientScope);
}
if pinned_account
.as_ref()
.is_some_and(|required| metadata.provider_account_id() != Some(required))
@@ -819,6 +871,17 @@ where
Ok(None)
}
fn advance_provider_candidate(&mut self) -> bool {
let Some(next) = self.candidate_index.checked_add(1) else {
return false;
};
if next >= self.plan.candidates().len() {
return false;
}
self.candidate_index = next;
true
}
async fn observe_event(&mut self, event: &GatewayEvent) {
self.observe_event_facts(event);
self.mark_send_observed().await;
@@ -3,7 +3,7 @@
use std::cmp::Reverse;
use std::collections::{BTreeSet, HashMap};
use std::fmt;
use std::num::{NonZeroU32, NonZeroU64};
use std::num::{NonZeroU16, NonZeroU32, NonZeroU64};
use std::sync::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime};
@@ -38,6 +38,58 @@ impl ProviderAccountId {
}
}
/// Per-account scheduling concurrency override.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountConcurrencyLimit(NonZeroU32);
impl AccountConcurrencyLimit {
#[must_use]
pub const fn new(value: u32) -> Option<Self> {
match NonZeroU32::new(value) {
Some(value) => Some(Self(value)),
None => None,
}
}
#[must_use]
pub const fn get(self) -> u32 {
self.0.get()
}
#[must_use]
pub const fn into_non_zero(self) -> NonZeroU32 {
self.0
}
}
/// Relative scheduling priority for an account.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountWeight(NonZeroU16);
impl AccountWeight {
pub const DEFAULT: Self = Self(NonZeroU16::MIN);
pub const MAX: u16 = 100;
#[must_use]
pub const fn new(value: u16) -> Option<Self> {
match NonZeroU16::new(value) {
Some(value) if value.get() <= Self::MAX => Some(Self(value)),
_ => None,
}
}
#[must_use]
pub const fn get(self) -> u16 {
self.0.get()
}
}
impl Default for AccountWeight {
fn default() -> Self {
Self::DEFAULT
}
}
impl fmt::Display for ProviderAccountId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
@@ -619,6 +671,8 @@ pub struct ProviderAccount {
authentication_kind: String,
revision: CredentialRevision,
enabled: bool,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
credential_state: CredentialState,
quota: QuotaState,
last_error_reason: Option<AccountErrorReason>,
@@ -651,6 +705,8 @@ impl ProviderAccount {
authentication_kind,
revision,
enabled: true,
concurrency_limit: None,
weight: AccountWeight::DEFAULT,
credential_state: CredentialState::Unknown,
quota: QuotaState::unknown(),
last_error_reason: None,
@@ -695,6 +751,17 @@ impl ProviderAccount {
self
}
#[must_use]
pub const fn with_scheduling(
mut self,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
) -> Self {
self.concurrency_limit = concurrency_limit;
self.weight = weight;
self
}
/// 设置 RT 存在性与失败后的最早重试时刻。
///
/// 正常 OAuth 预刷新由 worker 使用 AT 原始过期时间与当前运行时策略动态判断,
@@ -780,6 +847,24 @@ impl ProviderAccount {
self.enabled
}
#[must_use]
pub const fn concurrency_limit(&self) -> Option<AccountConcurrencyLimit> {
self.concurrency_limit
}
#[must_use]
pub const fn weight(&self) -> AccountWeight {
self.weight
}
#[must_use]
pub const fn effective_concurrency(&self, default: NonZeroU32) -> NonZeroU32 {
match self.concurrency_limit {
Some(limit) => limit.into_non_zero(),
None => default,
}
}
#[must_use]
pub const fn access_token_expires_at(&self) -> Option<SystemTime> {
self.access_token_expires_at
@@ -1393,6 +1478,7 @@ pub struct AccountSelectionContext {
pub preferred_account: Option<ProviderAccountId>,
pub round_robin_cursor: u64,
pub eligibility: AccountEligibilityPolicy,
pub account_scope: Option<std::sync::Arc<crate::routing::FrozenAccountScope>>,
}
/// 选择账号时是否执行本地调度资格投影。
@@ -1417,10 +1503,12 @@ impl AccountEligibilityPolicy {
/// 候选账号未进入调度池的约束。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountSchedulingBlocker {
OutsideClientScope,
LocalAvailability,
Excluded,
ConcurrencyLimit,
RequestInterval,
LowerWeight,
}
/// 优先账号在本次选择中的处理结果。
@@ -1463,26 +1551,6 @@ impl AccountSelector {
candidates: &'a [AccountCandidate],
context: &AccountSelectionContext,
) -> Option<AccountSelection<'a>> {
let preferred = if let Some(preferred) = context.preferred_account.as_ref() {
match candidates
.iter()
.find(|candidate| candidate.account.id() == preferred)
{
Some(candidate) => match self.scheduling_blocker(candidate, context) {
None => {
return Some(AccountSelection {
candidate,
preferred: PreferredAccountSelection::Hit,
});
}
Some(blocker) => PreferredAccountSelection::Blocked(blocker),
},
None => PreferredAccountSelection::Missing,
}
} else {
PreferredAccountSelection::NotRequested
};
let mut eligible = candidates
.iter()
.filter(|candidate| self.scheduling_blocker(candidate, context).is_none())
@@ -1490,6 +1558,33 @@ impl AccountSelector {
if eligible.is_empty() {
return None;
}
let highest_weight = eligible
.iter()
.map(|candidate| candidate.account.weight())
.max()?;
let preferred = if let Some(preferred) = context.preferred_account.as_ref() {
match candidates
.iter()
.find(|candidate| candidate.account.id() == preferred)
{
Some(candidate) => match self.scheduling_blocker(candidate, context) {
Some(blocker) => PreferredAccountSelection::Blocked(blocker),
None if candidate.account.weight() < highest_weight => {
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight)
}
None => {
return Some(AccountSelection {
candidate,
preferred: PreferredAccountSelection::Hit,
});
}
},
None => PreferredAccountSelection::Missing,
}
} else {
PreferredAccountSelection::NotRequested
};
eligible.retain(|candidate| candidate.account.weight() == highest_weight);
let candidate = match context.policy.strategy() {
RotationStrategy::QuotaResetPriority => {
@@ -1533,6 +1628,13 @@ impl AccountSelector {
candidate: &AccountCandidate,
context: &AccountSelectionContext,
) -> Option<AccountSchedulingBlocker> {
if context
.account_scope
.as_ref()
.is_some_and(|scope| !scope.allows(candidate.account.id()))
{
return Some(AccountSchedulingBlocker::OutsideClientScope);
}
if !context.eligibility.bypasses_local_eligibility() {
let status = candidate
.account
@@ -1545,7 +1647,12 @@ impl AccountSelector {
if context.excluded_accounts.contains(candidate.account.id()) {
return Some(AccountSchedulingBlocker::Excluded);
}
if candidate.signals.in_flight >= context.policy.max_concurrent_per_account().get() {
if candidate.signals.in_flight
>= candidate
.account
.effective_concurrency(context.policy.max_concurrent_per_account())
.get()
{
return Some(AccountSchedulingBlocker::ConcurrencyLimit);
}
if candidate
@@ -15,7 +15,8 @@ use crate::engine::admission::{
ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest,
};
use crate::engine::continuation::{
ContinuationBinding, NativeContinuationPin, NativeContinuationPort, PreviousResponseId,
ContinuationBinding, NativeContinuationPin, NativeContinuationPort,
NativeContinuationStoreErrorKind, PreviousResponseId,
};
use crate::engine::coordinator::ResponseExecutionSession;
use crate::engine::probe::{
@@ -261,22 +262,32 @@ impl DefaultExecutionService {
})?;
let request_id = new_request_id()?;
let routing_context = self
.route_context(request.client.policy.provider_kind())
.route_context(request.client.policy.account_scope().provider_kinds())
.await?;
let plan = request
.client
.snapshot
.plan(&request.public_model, &request.operation, &routing_context)
.plan(
&request.public_model,
&request.operation,
Arc::clone(request.client.policy.account_scope()),
&routing_context,
)
.map_err(map_routing_error)?;
let continuation = match request.metadata.previous_response_id.as_ref() {
Some(previous) => {
let resolve = self.continuation.resolve(previous).fuse();
let resolve = self
.continuation
.resolve(request.client.policy.key_id(), previous)
.fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(resolve, timeout);
let pin = select_biased! {
result = resolve => match result {
Ok(pin) => pin,
Err(error) => {
Err(error)
if error.kind() == NativeContinuationStoreErrorKind::Unavailable =>
{
tracing::debug!(
request_id = request_id.as_str(),
%error,
@@ -284,10 +295,53 @@ impl DefaultExecutionService {
);
None
}
Err(error)
if error.kind()
== NativeContinuationStoreErrorKind::OwnershipMismatch =>
{
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"continuation does not belong to this client API key",
));
}
Err(error) => {
tracing::warn!(
request_id = request_id.as_str(),
%error,
"Continuation affinity 记录无效,已拒绝续接"
);
return Err(GatewayError::new(
GatewayErrorKind::Internal,
"continuation state is invalid",
));
}
},
_ = timeout => None,
};
match pin {
Some(pin) if !pin.matches_client(request.client.policy.key_id()) => {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"continuation does not belong to this client API key",
));
}
Some(pin) if !request.client.policy.account_scope().allows(pin.account()) => {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"continuation account is outside the client account scope",
));
}
Some(pin)
if !plan
.candidates()
.iter()
.any(|candidate| candidate.provider() == pin.provider()) =>
{
return Err(GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"continuation provider is not available",
));
}
Some(pin) => {
attach_continuation_session_state(&mut request.operation, &pin);
ContinuationBinding::Pinned(pin)
@@ -360,14 +414,19 @@ impl DefaultExecutionService {
client_api_key_id: request.client.policy.key_id().clone(),
model_request_id: request_id.clone(),
};
let observation = self
.providers
.request_observation(request.client.policy.provider_kind(), &request.operation);
let observation = plan
.candidates()
.first()
.map_or_else(Default::default, |candidate| {
self.providers
.request_observation(candidate.provider(), &request.operation)
});
let new_request = NewModelRequest {
id: request_id.clone(),
client_api_key_id: Some(request.client.policy.key_id().clone()),
client_api_key_ref: request.client.policy.key_id().clone(),
config_revision: plan.config_revision(),
routing: request.client.policy.account_scope().routing_snapshot(),
protocol: request.metadata.protocol,
operation: request.operation.kind(),
endpoint: request.metadata.endpoint,
@@ -417,38 +476,42 @@ impl DefaultExecutionService {
async fn route_context(
&self,
provider_kind: &ProviderKind,
provider_kinds: &BTreeSet<ProviderKind>,
) -> Result<RoutingContext, GatewayError> {
let decision = self.circuits.decision(provider_kind).fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(decision, timeout);
let decision = select_biased! {
result = decision => Some(result),
_ = timeout => None,
};
let blocked_providers = match decision {
Some(Ok(ProviderCircuitDecision::Allow)) => BTreeSet::new(),
Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => {
BTreeSet::from([provider_kind.clone()])
let decisions = futures::future::join_all(provider_kinds.iter().map(|provider_kind| {
let circuits = Arc::clone(&self.circuits);
async move {
let decision = circuits.decision(provider_kind).fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(decision, timeout);
let decision = select_biased! {
result = decision => Some(result),
_ = timeout => None,
};
(provider_kind, decision)
}
Some(Err(error)) => {
tracing::warn!(
}))
.await;
let mut blocked_providers = BTreeSet::new();
for (provider_kind, decision) in decisions {
match decision {
Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => {
blocked_providers.insert(provider_kind.clone());
}
Some(Err(error)) => tracing::warn!(
provider = provider_kind.as_str(),
%error,
"Provider circuit 读取失败,按可重建协调状态 fail-open"
);
BTreeSet::new()
}
None => {
tracing::warn!(
),
None => tracing::warn!(
provider = provider_kind.as_str(),
"Provider circuit 读取超时,按可重建协调状态 fail-open"
);
BTreeSet::new()
),
Some(Ok(ProviderCircuitDecision::Allow)) => {}
}
};
}
Ok(RoutingContext {
provider_kind: Some(provider_kind.clone()),
required_provider: None,
blocked_providers,
})
}
@@ -479,11 +542,16 @@ impl DefaultExecutionService {
GatewayError::new(GatewayErrorKind::Unsupported, "requested model is invalid")
})?;
let routing_context = RoutingContext {
provider_kind: Some(provider_kind),
required_provider: Some(provider_kind),
..RoutingContext::default()
};
let plan = snapshot
.plan(&public_model, &operation, &routing_context)
.plan(
&public_model,
&operation,
snapshot.all_account_scope(),
&routing_context,
)
.map_err(map_routing_error)?;
let started_at = SystemTime::now();
let deadline_at = started_at
@@ -499,6 +567,7 @@ impl DefaultExecutionService {
client_api_key_id: None,
client_api_key_ref: actor,
config_revision: plan.config_revision(),
routing: crate::routing::AccountRoutingSnapshot::all(),
protocol: "admin_connection_test".to_owned(),
operation: operation.kind(),
endpoint: "/api/admin/accounts/connection-test".to_owned(),
@@ -693,19 +762,19 @@ impl ExecutionService for DefaultExecutionService {
fn public_models(&self, client: &AuthenticatedClient) -> Vec<PublicModelId> {
client
.snapshot
.public_models_for_provider(client.policy.provider_kind())
.public_models_for_scope(client.policy.account_scope())
}
fn public_model_profiles(&self, client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
client
.snapshot
.public_model_profiles_for_provider(client.policy.provider_kind())
.public_model_profiles_for_scope(client.policy.account_scope())
}
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool {
client
.snapshot
.contains_public_model_for_provider(model, client.policy.provider_kind())
.contains_public_model_for_scope(model, client.policy.account_scope())
}
fn start(
@@ -1006,7 +1075,8 @@ fn new_request_id() -> Result<ModelRequestId, GatewayError> {
fn map_routing_error(error: crate::error::RoutingError) -> GatewayError {
match error {
crate::error::RoutingError::NoCapableProvider { .. } => GatewayError::new(
crate::error::RoutingError::NoCapableProvider { .. }
| crate::error::RoutingError::EmptyAccountScope => GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no provider can execute this request",
),
@@ -233,6 +233,7 @@ pub struct AccountAttemptContext {
state_owner: Option<ProviderAccountStateOwner>,
credential_recovery_attempted: bool,
diagnostic_required_account: bool,
account_scope: Option<Arc<crate::routing::FrozenAccountScope>>,
}
impl AccountAttemptContext {
@@ -248,6 +249,7 @@ impl AccountAttemptContext {
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: false,
account_scope: None,
}
}
@@ -266,9 +268,17 @@ impl AccountAttemptContext {
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: true,
account_scope: None,
}
}
/// 附着普通请求认证时冻结的账号范围。
#[must_use]
pub fn with_account_scope(mut self, scope: Arc<crate::routing::FrozenAccountScope>) -> Self {
self.account_scope = Some(scope);
self
}
/// 标记本请求已对即将选择的固定账号执行过一次凭据恢复。
#[must_use]
pub const fn with_credential_recovery_attempted(mut self, attempted: bool) -> Self {
@@ -301,6 +311,11 @@ impl AccountAttemptContext {
pub const fn is_diagnostic_required_account(&self) -> bool {
self.diagnostic_required_account
}
#[must_use]
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
self.account_scope.as_ref()
}
}
/// 请求中 Provider 账号绑定状态的唯一归属。
@@ -506,6 +521,12 @@ impl AttemptContext {
self.account.is_diagnostic_required_account()
}
/// 普通请求认证时冻结的账号范围;管理端诊断为 `None`。
#[must_use]
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
self.account.account_scope()
}
#[must_use]
pub const fn continuation(&self) -> Option<&ContinuationBinding> {
self.continuation.as_ref()
@@ -529,6 +550,7 @@ pub struct NewModelRequest {
pub client_api_key_id: Option<ClientApiKeyId>,
pub client_api_key_ref: ClientApiKeyId,
pub config_revision: ConfigRevision,
pub routing: crate::routing::AccountRoutingSnapshot,
pub protocol: String,
pub operation: OperationKind,
pub endpoint: String,
@@ -749,6 +771,8 @@ pub enum EngineError {
ContinuationPinMismatch,
#[error("provider did not use the account required by this execution")]
RequiredAccountMismatch,
#[error("provider selected an account outside the frozen client scope")]
AccountOutsideClientScope,
#[error("provider execution failed")]
Provider(ProviderError),
#[error("request was cancelled")]
+9
View File
@@ -33,6 +33,9 @@ pub enum IdentifierError {
/// 规定前缀。
expected: &'static str,
},
/// 标识不满足该领域值对象的固定格式。
#[error("identifier has an invalid format")]
InvalidFormat,
}
/// Operation 构造或校验失败。
@@ -67,6 +70,12 @@ pub enum RoutingError {
/// 配置 revision 必须为正数。
#[error("config revision must be greater than zero")]
InvalidRevision,
/// Restricted scope 必须保留至少一个持久化分组 binding。
#[error("account routing scope is invalid")]
InvalidAccountScope,
/// Key 的账号范围中当前没有任何可参与路由的账号。
#[error("account routing scope is empty")]
EmptyAccountScope,
/// 快照中存在重复实体。
#[error("duplicate {entity} `{id}`")]
DuplicateEntity {
@@ -1,11 +1,12 @@
//! 下游 Client API Key 的准入策略。
//!
//! Client API Key 绑定一个 Provider 平台;模型名称不参与改写或 allowlist 判断。
//! Client API Key 冻结账号分组权限;模型名称不参与权限判断。
use std::fmt;
use std::sync::Arc;
use crate::error::{IdentifierError, PolicyError, validate_text};
use crate::routing::ProviderKind;
use crate::routing::FrozenAccountScope;
/// `client_api_keys.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -88,7 +89,7 @@ impl RateLimits {
pub struct ClientPolicy {
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: ProviderKind,
account_scope: Arc<FrozenAccountScope>,
enabled: bool,
limits: RateLimits,
}
@@ -98,14 +99,14 @@ impl ClientPolicy {
pub const fn new(
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: ProviderKind,
account_scope: Arc<FrozenAccountScope>,
enabled: bool,
limits: RateLimits,
) -> Self {
Self {
key_id,
plaintext_key,
provider_kind,
account_scope,
enabled,
limits,
}
@@ -122,8 +123,8 @@ impl ClientPolicy {
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
#[must_use]
@@ -13,6 +13,7 @@ use crate::engine::credential::{
OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
};
use crate::error::{IdentifierError, validate_text};
use crate::policy::ClientApiKeyId;
use crate::routing::{ProviderKind, UpstreamModelId};
const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60);
@@ -170,6 +171,7 @@ pub enum ProviderLeaseRequest {
pub trait ProviderLeasePort: Send + Sync {
fn load_state<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
provider_kind: &'a ProviderKind,
accounts: &'a [ProviderAccountId],
) -> BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>>;
+296 -3
View File
@@ -9,12 +9,293 @@ use std::fmt;
use std::num::{NonZeroU32, NonZeroU64};
use std::sync::Arc;
use crate::engine::credential::AccountSelectionPolicy;
use crate::engine::credential::{AccountSelectionPolicy, ProviderAccountId};
use crate::error::{IdentifierError, RoutingError, validate_text};
use crate::operation::{CapabilityRequirements, Feature, OperationKind};
const MAX_REQUEST_ATTEMPTS: u32 = 32;
/// `account_groups.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountGroupId(String);
impl AccountGroupId {
/// 校验并创建账号分组 ID。
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
let value = value.into();
let Some(suffix) = value.strip_prefix("grp_") else {
return Err(IdentifierError::MissingPrefix { expected: "grp_" });
};
if suffix.len() != 32
|| !suffix
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
{
return Err(IdentifierError::InvalidFormat);
}
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for AccountGroupId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
/// 快照中一个账号的 Provider 与分组归属。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeAccount {
provider_kind: ProviderKind,
group_ids: Arc<BTreeSet<AccountGroupId>>,
}
impl RuntimeAccount {
#[must_use]
pub fn new(provider_kind: ProviderKind, group_ids: BTreeSet<AccountGroupId>) -> Self {
Self {
provider_kind,
group_ids: Arc::new(group_ids),
}
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
}
#[must_use]
pub fn group_ids(&self) -> &BTreeSet<AccountGroupId> {
&self.group_ids
}
}
/// 全快照共享的账号、Provider 与分组反向索引。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RuntimeAccountDirectory {
accounts: BTreeMap<ProviderAccountId, RuntimeAccount>,
providers_with_accounts: BTreeSet<ProviderKind>,
providers_by_group: BTreeMap<AccountGroupId, BTreeSet<ProviderKind>>,
}
impl RuntimeAccountDirectory {
#[must_use]
pub fn new(accounts: BTreeMap<ProviderAccountId, RuntimeAccount>) -> Self {
let mut providers_with_accounts = BTreeSet::new();
let mut providers_by_group = BTreeMap::<AccountGroupId, BTreeSet<ProviderKind>>::new();
for account in accounts.values() {
providers_with_accounts.insert(account.provider_kind.clone());
for group_id in account.group_ids.iter() {
providers_by_group
.entry(group_id.clone())
.or_default()
.insert(account.provider_kind.clone());
}
}
Self {
accounts,
providers_with_accounts,
providers_by_group,
}
}
#[must_use]
pub fn account(&self, account_id: &ProviderAccountId) -> Option<&RuntimeAccount> {
self.accounts.get(account_id)
}
#[must_use]
pub fn providers_with_accounts(&self) -> &BTreeSet<ProviderKind> {
&self.providers_with_accounts
}
#[must_use]
pub fn providers_for_groups<'a>(
&self,
group_ids: impl IntoIterator<Item = &'a AccountGroupId>,
) -> BTreeSet<ProviderKind> {
group_ids
.into_iter()
.filter_map(|group_id| self.providers_by_group.get(group_id))
.flatten()
.cloned()
.collect()
}
}
/// 历史请求保存的账号范围种类。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountRoutingScopeKind {
All,
Groups,
}
impl AccountRoutingScopeKind {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::All => "all",
Self::Groups => "groups",
}
}
}
/// 请求开始时冻结的分组 ID 与名称。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RoutingGroupSnapshot {
id: AccountGroupId,
name: String,
}
impl RoutingGroupSnapshot {
#[must_use]
pub fn new(id: AccountGroupId, name: String) -> Self {
Self { id, name }
}
#[must_use]
pub const fn id(&self) -> &AccountGroupId {
&self.id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
}
/// 请求历史所需的完整、稳定账号范围快照。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountRoutingSnapshot {
kind: AccountRoutingScopeKind,
groups: Arc<[RoutingGroupSnapshot]>,
}
impl AccountRoutingSnapshot {
#[must_use]
pub fn all() -> Self {
Self {
kind: AccountRoutingScopeKind::All,
groups: Arc::from([]),
}
}
#[must_use]
pub fn groups(groups: Vec<RoutingGroupSnapshot>) -> Self {
Self {
kind: AccountRoutingScopeKind::Groups,
groups: Arc::from(groups),
}
}
#[must_use]
pub const fn kind(&self) -> AccountRoutingScopeKind {
self.kind
}
#[must_use]
pub fn groups_snapshot(&self) -> &[RoutingGroupSnapshot] {
&self.groups
}
}
/// Key 持久 binding 编译出的账号权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClientRoutingScope {
AllAccounts,
Restricted {
bound_groups: Arc<[RoutingGroupSnapshot]>,
enabled_group_ids: Arc<BTreeSet<AccountGroupId>>,
provider_kinds: Arc<BTreeSet<ProviderKind>>,
},
}
impl ClientRoutingScope {
#[must_use]
pub fn all_accounts() -> Self {
Self::AllAccounts
}
pub fn restricted(
bound_groups: Vec<RoutingGroupSnapshot>,
enabled_group_ids: BTreeSet<AccountGroupId>,
provider_kinds: BTreeSet<ProviderKind>,
) -> Result<Self, RoutingError> {
if bound_groups.is_empty() {
return Err(RoutingError::InvalidAccountScope);
}
Ok(Self::Restricted {
bound_groups: Arc::from(bound_groups),
enabled_group_ids: Arc::new(enabled_group_ids),
provider_kinds: Arc::new(provider_kinds),
})
}
}
/// 一次认证随 RuntimeSnapshot 冻结的账号目录与 Key 权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FrozenAccountScope {
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
}
impl FrozenAccountScope {
#[must_use]
pub const fn new(
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
) -> Self {
Self {
directory,
client_scope,
}
}
#[must_use]
pub fn allows(&self, account_id: &ProviderAccountId) -> bool {
let Some(account) = self.directory.account(account_id) else {
return false;
};
match &self.client_scope {
ClientRoutingScope::AllAccounts => true,
ClientRoutingScope::Restricted {
enabled_group_ids, ..
} => account
.group_ids()
.iter()
.any(|group_id| enabled_group_ids.contains(group_id)),
}
}
#[must_use]
pub fn provider_kinds(&self) -> &BTreeSet<ProviderKind> {
match &self.client_scope {
ClientRoutingScope::AllAccounts => self.directory.providers_with_accounts(),
ClientRoutingScope::Restricted { provider_kinds, .. } => provider_kinds,
}
}
#[must_use]
pub fn routing_snapshot(&self) -> AccountRoutingSnapshot {
match &self.client_scope {
ClientRoutingScope::AllAccounts => AccountRoutingSnapshot::all(),
ClientRoutingScope::Restricted { bound_groups, .. } => {
AccountRoutingSnapshot::groups(bound_groups.to_vec())
}
}
}
#[must_use]
pub const fn directory(&self) -> &Arc<RuntimeAccountDirectory> {
&self.directory
}
}
/// 编译进二进制的 Provider adapter slug。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProviderKind(String);
@@ -494,8 +775,8 @@ impl ProviderModel {
/// 本次请求选择 Provider 时使用的动态过滤事实。
#[derive(Debug, Clone, Default)]
pub struct RoutingContext {
/// 已认证 Client API Key 绑定的平台;模型名称不参与平台猜测。
pub provider_kind: Option<ProviderKind>,
/// 管理端 connection test 显式限制的 Provider;普通请求留空。
pub required_provider: Option<ProviderKind>,
pub blocked_providers: BTreeSet<ProviderKind>,
}
@@ -505,6 +786,7 @@ pub struct ProviderCandidate {
provider: ProviderKind,
upstream_model: UpstreamModelId,
emulated_features: BTreeSet<Feature>,
account_scope: Arc<FrozenAccountScope>,
}
impl ProviderCandidate {
@@ -522,6 +804,11 @@ impl ProviderCandidate {
pub const fn emulated_features(&self) -> &BTreeSet<Feature> {
&self.emulated_features
}
#[must_use]
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
}
/// 一次请求冻结的 Provider 尝试顺序。
@@ -532,6 +819,7 @@ pub struct RoutingPlan {
public_model: PublicModelId,
operation: OperationKind,
max_attempts: NonZeroU32,
account_scope: Arc<FrozenAccountScope>,
candidates: Arc<[ProviderCandidate]>,
}
@@ -561,6 +849,11 @@ impl RoutingPlan {
self.max_attempts
}
#[must_use]
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
#[must_use]
pub fn candidates(&self) -> &[ProviderCandidate] {
&self.candidates
@@ -11,7 +11,7 @@ use futures::{FutureExt as _, Stream, StreamExt as _, pin_mut, select_biased};
use futures_timer::Delay;
use crate::engine::CancellationToken;
use crate::engine::credential::{AccountSelectionPolicy, RotationStrategy};
use crate::engine::credential::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
use crate::engine::provider::{ProviderCatalogGeneration, ProviderRegistry};
use crate::error::RoutingError;
use crate::health::{HealthProbe, HealthState};
@@ -24,8 +24,9 @@ use crate::task::{
};
use super::{
ConfigRevision, ModelCapabilities, ProviderCandidate, ProviderKind, ProviderModel,
PublicModelId, RoutingContext, RoutingPlan, UpstreamModelId,
AccountGroupId, ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
ProviderCandidate, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RoutingGroupSnapshot, RoutingPlan, RuntimeAccount, RuntimeAccountDirectory, UpstreamModelId,
};
const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5);
@@ -66,7 +67,7 @@ impl SnapshotSettingsFacts {
pub struct SnapshotClientPolicyFacts {
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: String,
group_ids: Vec<AccountGroupId>,
limits: RateLimits,
}
@@ -75,18 +76,67 @@ impl SnapshotClientPolicyFacts {
pub fn new(
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: impl Into<String>,
group_ids: Vec<AccountGroupId>,
limits: RateLimits,
) -> Self {
Self {
key_id,
plaintext_key,
provider_kind: provider_kind.into(),
group_ids,
limits,
}
}
}
/// Store 读取到的账号分组事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotAccountGroupFacts {
id: AccountGroupId,
name: String,
enabled: bool,
}
impl SnapshotAccountGroupFacts {
#[must_use]
pub fn new(id: AccountGroupId, name: String, enabled: bool) -> Self {
Self { id, name, enabled }
}
}
/// Store 读取到的账号及其固有 Provider 事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotProviderAccountFacts {
account_id: ProviderAccountId,
provider_kind: String,
}
impl SnapshotProviderAccountFacts {
#[must_use]
pub fn new(account_id: ProviderAccountId, provider_kind: impl Into<String>) -> Self {
Self {
account_id,
provider_kind: provider_kind.into(),
}
}
}
/// Store 读取到的一条分组成员关系。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotAccountGroupMemberFacts {
group_id: AccountGroupId,
account_id: ProviderAccountId,
}
impl SnapshotAccountGroupMemberFacts {
#[must_use]
pub const fn new(group_id: AccountGroupId, account_id: ProviderAccountId) -> Self {
Self {
group_id,
account_id,
}
}
}
/// 一次一致性读取产生的全部 RuntimeSnapshot 持久事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotFacts {
@@ -94,6 +144,9 @@ pub struct SnapshotFacts {
observed_current_revision: ConfigRevision,
settings: SnapshotSettingsFacts,
client_policies: Vec<SnapshotClientPolicyFacts>,
account_groups: Vec<SnapshotAccountGroupFacts>,
provider_accounts: Vec<SnapshotProviderAccountFacts>,
group_memberships: Vec<SnapshotAccountGroupMemberFacts>,
}
impl SnapshotFacts {
@@ -103,12 +156,18 @@ impl SnapshotFacts {
observed_current_revision: ConfigRevision,
settings: SnapshotSettingsFacts,
client_policies: Vec<SnapshotClientPolicyFacts>,
account_groups: Vec<SnapshotAccountGroupFacts>,
provider_accounts: Vec<SnapshotProviderAccountFacts>,
group_memberships: Vec<SnapshotAccountGroupMemberFacts>,
) -> Self {
Self {
config_revision,
observed_current_revision,
settings,
client_policies,
account_groups,
provider_accounts,
group_memberships,
}
}
@@ -223,13 +282,16 @@ async fn compile_runtime_snapshot(
providers: &ProviderRegistry,
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
let provider_kinds = providers.provider_kinds().cloned().collect::<Vec<_>>();
let registered_providers = provider_kinds.iter().cloned().collect::<BTreeSet<_>>();
// 实时目录只提供公开模型与能力提示;查询失败时保留 Provider 透传语义。
// 目录查询失败表示未知;查询成功后,即使为空,也必须与“已知缺少模型”区分。
let mut provider_models = Vec::new();
let mut known_provider_catalogs = BTreeSet::new();
for provider in &provider_kinds {
let Ok(models) = providers.query_model_capabilities(provider).await else {
continue;
};
known_provider_catalogs.insert(provider.clone());
provider_models.extend(models.into_iter().map(|model| {
let compiled = ProviderModel::new(
provider.clone(),
@@ -243,6 +305,56 @@ async fn compile_runtime_snapshot(
}));
}
let mut groups = BTreeMap::new();
for group in facts.account_groups {
if group.name.trim() != group.name
|| group.name.is_empty()
|| group.name.chars().count() > 100
|| group.name.chars().any(char::is_control)
|| groups.insert(group.id.clone(), group).is_some()
{
return Err(RuntimeSnapshotCompileError::InvalidData);
}
}
let mut account_groups = BTreeMap::<ProviderAccountId, BTreeSet<AccountGroupId>>::new();
let mut accounts = BTreeMap::new();
for account in facts.provider_accounts {
let provider_kind = ProviderKind::new(account.provider_kind)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
if !registered_providers.contains(&provider_kind)
|| accounts
.insert(
account.account_id.clone(),
RuntimeAccount::new(provider_kind, BTreeSet::new()),
)
.is_some()
{
return Err(RuntimeSnapshotCompileError::InvalidData);
}
account_groups.insert(account.account_id, BTreeSet::new());
}
let mut memberships = BTreeSet::new();
for membership in facts.group_memberships {
if !groups.contains_key(&membership.group_id)
|| !accounts.contains_key(&membership.account_id)
|| !memberships.insert((membership.group_id.clone(), membership.account_id.clone()))
{
return Err(RuntimeSnapshotCompileError::InvalidData);
}
account_groups
.get_mut(&membership.account_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?
.insert(membership.group_id);
}
for (account_id, group_ids) in account_groups {
let account = accounts
.get_mut(&account_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
*account = RuntimeAccount::new(account.provider_kind().clone(), group_ids);
}
let account_directory = Arc::new(RuntimeAccountDirectory::new(accounts));
let model_mappings = facts.settings.model_mappings;
let rotation_strategy = RotationStrategy::parse(facts.settings.rotation_strategy.as_str())
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
@@ -252,21 +364,48 @@ async fn compile_runtime_snapshot(
.ok_or(RuntimeSnapshotCompileError::InvalidData)?,
Duration::from_millis(facts.settings.request_interval_ms),
);
let client_policies = facts
.client_policies
.into_iter()
.map(|policy| {
let provider = ProviderKind::new(policy.provider_kind)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
Ok(ClientPolicy::new(
policy.key_id,
policy.plaintext_key,
provider,
true,
policy.limits,
))
})
.collect::<Result<Vec<_>, RuntimeSnapshotCompileError>>()?;
let mut client_policies = Vec::with_capacity(facts.client_policies.len());
for policy in facts.client_policies {
let account_scope = if policy.group_ids.is_empty() {
FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::all_accounts(),
)
} else {
let mut seen = BTreeSet::new();
let mut bound_groups = Vec::with_capacity(policy.group_ids.len());
let mut enabled_group_ids = BTreeSet::new();
for group_id in policy.group_ids {
if !seen.insert(group_id.clone()) {
return Err(RuntimeSnapshotCompileError::InvalidData);
}
let group = groups
.get(&group_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
bound_groups.push(RoutingGroupSnapshot::new(
group.id.clone(),
group.name.clone(),
));
if group.enabled {
enabled_group_ids.insert(group_id);
}
}
bound_groups.sort_by(|left, right| left.id().cmp(right.id()));
let provider_kinds = account_directory.providers_for_groups(&enabled_group_ids);
FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::restricted(bound_groups, enabled_group_ids, provider_kinds)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?,
)
};
client_policies.push(ClientPolicy::new(
policy.key_id,
policy.plaintext_key,
Arc::new(account_scope),
true,
policy.limits,
));
}
RuntimeSnapshot::new(
facts.config_revision,
@@ -276,7 +415,12 @@ async fn compile_runtime_snapshot(
client_policies,
)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)
.map(|snapshot| snapshot.with_model_mappings(model_mappings))
.map(|snapshot| {
snapshot
.with_model_mappings(model_mappings)
.with_account_directory(account_directory)
.with_known_provider_catalogs(known_provider_catalogs)
})
}
/// 数据面使用的不可变配置快照。
@@ -290,6 +434,8 @@ pub struct RuntimeSnapshot {
Arc<BTreeMap<ProviderKind, BTreeMap<UpstreamModelId, super::ModelPresentation>>>,
model_mappings: Arc<BTreeMap<String, String>>,
provider_catalog_generations: Arc<BTreeMap<ProviderKind, ProviderCatalogGeneration>>,
known_provider_catalogs: Arc<BTreeSet<ProviderKind>>,
account_directory: Arc<RuntimeAccountDirectory>,
client_policies: Arc<BTreeMap<ClientApiKeyId, ClientPolicy>>,
}
@@ -312,6 +458,7 @@ impl RuntimeSnapshot {
}
}
let mut known_provider_catalogs = BTreeSet::new();
let mut model_map =
BTreeMap::<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>::new();
let mut presentation_map =
@@ -323,6 +470,7 @@ impl RuntimeSnapshot {
capabilities,
presentation,
} = model;
known_provider_catalogs.insert(provider.clone());
if !provider_set.contains(&provider) {
return Err(RoutingError::NotFound {
entity: "provider",
@@ -367,6 +515,8 @@ impl RuntimeSnapshot {
provider_model_presentations: Arc::new(presentation_map),
model_mappings: Arc::new(BTreeMap::new()),
provider_catalog_generations: Arc::new(BTreeMap::new()),
known_provider_catalogs: Arc::new(known_provider_catalogs),
account_directory: Arc::new(RuntimeAccountDirectory::default()),
client_policies: Arc::new(client_policy_map),
})
}
@@ -377,6 +527,26 @@ impl RuntimeSnapshot {
self
}
#[must_use]
pub fn with_account_directory(mut self, directory: Arc<RuntimeAccountDirectory>) -> Self {
self.account_directory = directory;
self
}
#[must_use]
fn with_known_provider_catalogs(mut self, providers: BTreeSet<ProviderKind>) -> Self {
self.known_provider_catalogs = Arc::new(providers);
self
}
#[must_use]
pub fn all_account_scope(&self) -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
Arc::clone(&self.account_directory),
ClientRoutingScope::all_accounts(),
))
}
#[must_use]
fn with_provider_catalog_generations(
mut self,
@@ -459,8 +629,39 @@ impl RuntimeSnapshot {
.collect()
}
/// 已取得目录时按目录与全局别名判断模型可用性;目录暂不可用时沿用旧版
/// fail-open 语义,把具体 wire 可接受性留给上游。
/// 合并冻结账号范围内实际存在 Provider 的公开模型。
#[must_use]
pub fn public_models_for_scope(&self, scope: &FrozenAccountScope) -> Vec<PublicModelId> {
scope
.provider_kinds()
.iter()
.flat_map(|provider| self.public_models_for_provider(provider))
.collect::<BTreeSet<_>>()
.into_iter()
.collect()
}
/// 合并冻结账号范围内实际存在 Provider 的公开模型画像。
#[must_use]
pub fn public_model_profiles_for_scope(
&self,
scope: &FrozenAccountScope,
) -> Vec<super::PublicModelProfile> {
let mut profiles = BTreeMap::new();
for provider in scope.provider_kinds() {
for profile in self.public_model_profiles_for_provider(provider) {
profiles
.entry(profile.model().clone())
.or_insert_with(|| profile.presentation().clone());
}
}
profiles
.into_iter()
.map(|(model, presentation)| super::PublicModelProfile::new(model, presentation))
.collect()
}
/// 已取得目录时以映射后的上游模型为准;目录不可用时由 Provider 在发送前判定。
#[must_use]
pub fn contains_public_model_for_provider(
&self,
@@ -470,17 +671,25 @@ impl RuntimeSnapshot {
if !self.providers.contains(provider) {
return false;
}
if self.model_mappings.contains_key(public_model.as_str()) {
return true;
}
let upstream_model = self.mapped_model(public_model.as_str());
match self.provider_models.get(provider) {
Some(models) if !models.is_empty() => models
.keys()
.any(|model| model.as_str() == public_model.as_str()),
Some(_) | None => true,
Some(models) => models.keys().any(|model| model.as_str() == upstream_model),
None => !self.known_provider_catalogs.contains(provider),
}
}
#[must_use]
pub fn contains_public_model_for_scope(
&self,
public_model: &PublicModelId,
scope: &FrozenAccountScope,
) -> bool {
scope
.provider_kinds()
.iter()
.any(|provider| self.contains_public_model_for_provider(public_model, provider))
}
#[must_use]
pub fn mapped_model(&self, requested: &str) -> String {
let original = requested;
@@ -506,14 +715,26 @@ impl RuntimeSnapshot {
&self,
public_model: &PublicModelId,
operation: &Operation,
account_scope: Arc<FrozenAccountScope>,
context: &RoutingContext,
) -> Result<RoutingPlan, RoutingError> {
let requirements = operation.capability_requirements();
let mut candidates = Vec::new();
for provider in self.providers.iter() {
if context.required_provider.is_none() && account_scope.provider_kinds().is_empty() {
return Err(RoutingError::EmptyAccountScope);
}
let providers = context.required_provider.as_ref().map_or_else(
|| account_scope.provider_kinds().clone(),
|provider| BTreeSet::from([provider.clone()]),
);
for provider in &providers {
if !self.providers.contains(provider) {
continue;
}
if context
.provider_kind
.required_provider
.as_ref()
.is_some_and(|expected| expected != provider)
|| context.blocked_providers.contains(provider)
@@ -539,12 +760,14 @@ impl RuntimeSnapshot {
};
emulated
}
None if self.known_provider_catalogs.contains(provider) => continue,
None => BTreeSet::new(),
};
candidates.push(ProviderCandidate {
provider: provider.clone(),
upstream_model,
emulated_features,
account_scope: Arc::clone(&account_scope),
});
}
@@ -561,6 +784,7 @@ impl RuntimeSnapshot {
operation: operation.kind(),
max_attempts: NonZeroU32::new(super::MAX_REQUEST_ATTEMPTS)
.expect("constant request attempt limit is non-zero"),
account_scope,
candidates: Arc::from(candidates),
})
}
@@ -2,12 +2,14 @@ use gateway_core::engine::continuation::{
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::ProviderKind;
fn pin() -> NativeContinuationPin {
NativeContinuationPin::new(
PreviousResponseId::new("response-private"),
PreviousResponseId::new("upstream-private"),
ClientApiKeyId::new("key_client").expect("valid client key"),
ProviderKind::new("openai").expect("valid provider"),
ProviderAccountId::new("acct_codex").expect("valid account"),
)
@@ -43,3 +45,8 @@ fn native_pin_should_reject_different_account() {
&ProviderAccountId::new("acct_other").expect("valid account"),
));
}
#[test]
fn native_pin_should_reject_different_client_api_key() {
assert!(!pin().matches_client(&ClientApiKeyId::new("key_other").expect("valid client key")));
}
@@ -43,8 +43,9 @@ use gateway_core::operation::{
};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RoutingPlan, RuntimeSnapshot, UpstreamModelId,
AccountRoutingSnapshot, ClientRoutingScope, ConfigRevision, FrozenAccountScope,
ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan,
RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
};
use serde_json::{Map, Value, json};
@@ -531,7 +532,32 @@ fn plan_with_policy(
let public_model = PublicModelId::new("gpt-5").expect("public model");
let capabilities = ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000))
.with_upstream_feature_validation();
RuntimeSnapshot::new(
let directory = Arc::new(RuntimeAccountDirectory::new(
[
"acct_failed",
"acct_first",
"acct_image",
"acct_observation_mismatch",
"acct_observed",
"acct_one",
"acct_only",
"acct_other",
"acct_required",
"acct_second",
"acct_tool",
"acct_two",
"acct_wrong",
]
.into_iter()
.map(|id| {
(
ProviderAccountId::new(id).expect("account"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)
})
.collect(),
));
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
account_selection_policy,
vec![provider.clone()],
@@ -543,15 +569,21 @@ fn plan_with_policy(
Vec::new(),
)
.expect("snapshot")
.plan(
&public_model,
operation,
&RoutingContext {
provider_kind: Some(provider),
..RoutingContext::default()
},
)
.expect("routing plan")
.with_account_directory(Arc::clone(&directory));
snapshot
.plan(
&public_model,
operation,
Arc::new(FrozenAccountScope::new(
directory,
ClientRoutingScope::all_accounts(),
)),
&RoutingContext {
required_provider: Some(provider),
..RoutingContext::default()
},
)
.expect("routing plan")
}
fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest {
@@ -561,6 +593,7 @@ fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest
client_api_key_id: Some(client_key.clone()),
client_api_key_ref: client_key,
config_revision: ConfigRevision::new(1).expect("config revision"),
routing: AccountRoutingSnapshot::all(),
protocol: "openai".to_owned(),
operation: operation.kind(),
endpoint: "responses".to_owned(),
@@ -1354,6 +1387,7 @@ fn authenticated_native_continuation_reaches_every_attempt_context() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-secret-id"),
PreviousResponseId::new("provider-native-id"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -1411,6 +1445,7 @@ fn streaming_native_continuation_delivers_started_event_before_later_output() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -1487,6 +1522,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-secret-id"),
PreviousResponseId::new("provider-native-id"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -1552,6 +1588,7 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
original.clone(),
);
@@ -2498,6 +2535,7 @@ fn native_continuation_explicit_429_is_not_retried() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-secret-id"),
PreviousResponseId::new("provider-native-id"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_first").expect("account"),
);
@@ -1,18 +1,21 @@
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use serde_json::{Map, Value};
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountCandidate, AccountEligibilityPolicy, AccountFeedbackStats,
AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker, AccountSelectionContext,
AccountSelectionPolicy, AccountSelector, AccountStatus, CredentialCasUpdate,
CredentialRevision, CredentialState, OpaqueProviderData, PlaintextCredential,
PreferredAccountSelection, ProviderAccount, ProviderAccountId, ProviderAccountIdentity,
ProviderAccountUpdate, QuotaEvidence, QuotaState, RotationStrategy,
AccountAttemptFeedback, AccountCandidate, AccountConcurrencyLimit, AccountEligibilityPolicy,
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker,
AccountSelectionContext, AccountSelectionPolicy, AccountSelector, AccountStatus, AccountWeight,
CredentialCasUpdate, CredentialRevision, CredentialState, OpaqueProviderData,
PlaintextCredential, PreferredAccountSelection, ProviderAccount, ProviderAccountId,
ProviderAccountIdentity, ProviderAccountUpdate, QuotaEvidence, QuotaState, RotationStrategy,
};
use gateway_core::routing::{
ClientRoutingScope, FrozenAccountScope, ProviderKind, RuntimeAccount, RuntimeAccountDirectory,
};
use gateway_core::routing::ProviderKind;
fn account(id: &str) -> ProviderAccount {
ProviderAccount::new(
@@ -48,6 +51,34 @@ fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandida
}
}
fn weighted_candidate(id: &str, weight: u16, in_flight: u32) -> AccountCandidate {
let mut candidate = candidate(id, in_flight, None);
candidate.account = candidate.account.with_scheduling(
None,
AccountWeight::new(weight).expect("valid account weight"),
);
candidate
}
#[test]
fn account_scheduling_should_default_and_resolve_concurrency_override() {
let default = NonZeroU32::new(12).expect("default concurrency");
let account = account("acct_scheduling_default");
assert_eq!(account.concurrency_limit(), None);
assert_eq!(account.weight(), AccountWeight::DEFAULT);
assert_eq!(account.effective_concurrency(default), default);
let limit = AccountConcurrencyLimit::new(7).expect("concurrency override");
let weight = AccountWeight::new(80).expect("weight");
let account = account.with_scheduling(Some(limit), weight);
assert_eq!(account.concurrency_limit(), Some(limit));
assert_eq!(account.weight(), weight);
assert_eq!(account.effective_concurrency(default).get(), 7);
assert!(AccountConcurrencyLimit::new(0).is_none());
assert!(AccountWeight::new(0).is_none());
assert!(AccountWeight::new(101).is_none());
}
fn context(strategy: RotationStrategy) -> AccountSelectionContext {
AccountSelectionContext {
policy: AccountSelectionPolicy::new(
@@ -60,6 +91,7 @@ fn context(strategy: RotationStrategy) -> AccountSelectionContext {
preferred_account: None,
round_robin_cursor: 0,
eligibility: AccountEligibilityPolicy::Enforce,
account_scope: None,
}
}
@@ -215,6 +247,35 @@ fn diagnostic_selection_bypasses_all_local_account_eligibility() {
);
}
#[test]
fn selector_should_reject_a_preferred_account_outside_the_frozen_client_scope() {
let allowed = candidate("acct_allowed", 0, None);
let outside = candidate("acct_outside", 0, None);
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
allowed.account.id().clone(),
RuntimeAccount::new(
ProviderKind::new("openai").expect("provider"),
BTreeSet::new(),
),
)])));
let mut context = context(RotationStrategy::Sticky);
context.preferred_account = Some(outside.account.id().clone());
context.account_scope = Some(Arc::new(FrozenAccountScope::new(
directory,
ClientRoutingScope::all_accounts(),
)));
let candidates = [outside, allowed];
let selection = AccountSelector
.select(&candidates, &context)
.expect("allowed account should remain eligible");
assert_eq!(selection.candidate().account.id().as_str(), "acct_allowed");
assert_eq!(
selection.preferred(),
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::OutsideClientScope)
);
}
#[test]
fn disabled_account_should_not_be_schedulable() {
let account = account("acct_disabled").with_account_facts(
@@ -287,6 +348,70 @@ fn smart_selector_should_prefer_lower_inflight_count() {
assert_eq!(selected.candidate().account.id().as_str(), "acct_idle");
}
#[test]
fn selector_should_prioritize_the_highest_weight_for_every_rotation_strategy() {
let candidates = vec![
weighted_candidate("acct_low_weight", 1, 0),
weighted_candidate("acct_high_weight", 100, 2),
];
for strategy in [
RotationStrategy::Smart,
RotationStrategy::QuotaResetPriority,
RotationStrategy::RoundRobin,
RotationStrategy::Sticky,
] {
let selected = AccountSelector
.select(&candidates, &context(strategy))
.expect("highest weight account should be eligible");
assert_eq!(
selected.candidate().account.id().as_str(),
"acct_high_weight",
"strategy {strategy:?} must select inside the highest weight tier"
);
}
}
#[test]
fn selector_should_not_let_soft_affinity_bypass_a_higher_weight_account() {
let preferred = weighted_candidate("acct_preferred", 1, 0);
let preferred_id = preferred.account.id().clone();
let higher = weighted_candidate("acct_higher", 10, 0);
let mut selection = context(RotationStrategy::Smart);
selection.preferred_account = Some(preferred_id);
let candidates = [preferred, higher];
let selected = AccountSelector
.select(&candidates, &selection)
.expect("higher weight account should be eligible");
assert_eq!(selected.candidate().account.id().as_str(), "acct_higher");
assert_eq!(
selected.preferred(),
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight)
);
}
#[test]
fn selector_should_fall_back_when_a_higher_weight_account_reaches_its_override() {
let mut saturated = weighted_candidate("acct_high_weight", 100, 1);
saturated.account = saturated.account.with_scheduling(
Some(AccountConcurrencyLimit::new(1).expect("concurrency override")),
AccountWeight::new(100).expect("weight"),
);
let fallback = weighted_candidate("acct_low_weight", 1, 0);
let candidates = [saturated, fallback];
let selected = AccountSelector
.select(&candidates, &context(RotationStrategy::Smart))
.expect("lower weight account should be used after saturation");
assert_eq!(
selected.candidate().account.id().as_str(),
"acct_low_weight"
);
}
#[test]
fn smart_selector_should_balance_signals_instead_of_using_lexicographic_load() {
let mut healthy = candidate("acct_healthy", 1, Some(100));
@@ -1,5 +1,5 @@
use std::{
collections::BTreeSet,
collections::{BTreeMap, BTreeSet},
sync::{
Arc, Mutex,
atomic::{AtomicBool, Ordering},
@@ -42,7 +42,8 @@ use gateway_core::operation::{GenerateRequest, Operation, OperationKind, Protoco
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RuntimeSnapshot,
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
ProviderModel, PublicModelId, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot,
UpstreamModelId,
};
use serde_json::json;
@@ -283,7 +284,7 @@ fn slow_circuit_store_should_time_out_and_fail_open_during_request_start() {
}
#[test]
fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
fn known_catalog_should_reject_a_model_that_the_provider_did_not_publish() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
@@ -298,7 +299,7 @@ fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
.expect("authenticated client");
let model = "gpt-future-not-in-catalog";
let started = block_on(service.start(StartExecution {
let result = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::from_client_wire(model).expect("client model"),
operation: start_operation_for_model(model),
@@ -311,10 +312,80 @@ fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
user_agent: None,
previous_response_id: None,
},
}))
.expect("stale model catalog must not reject the request");
}));
let Err(error) = result else {
panic!("a known provider catalog must be authoritative for model availability");
};
assert!(!started.session.is_finalized());
assert_eq!(error.kind(), GatewayErrorKind::NoAvailableProvider);
}
#[test]
fn continuation_owned_by_another_client_api_key_should_fail_closed() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(RejectedContinuation::OwnershipMismatch),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_start_test")
.expect("authenticated client");
let result = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::new("gpt-start").expect("public model"),
operation: start_operation(),
metadata: execution_metadata_with_continuation(),
}));
let Err(error) = result else {
panic!("cross-client continuation must be rejected");
};
assert_eq!(error.kind(), GatewayErrorKind::PolicyDenied);
}
#[test]
fn invalid_continuation_record_should_not_be_forwarded_as_an_external_handle() {
let service = DefaultExecutionService::new(
RuntimeSnapshotHandle::new(start_snapshot()),
Arc::new(TrackingExecutionStore::default()),
ProviderRegistry::default(),
Arc::new(UnusedAdmissions),
Arc::new(UnusedCircuits),
Arc::new(RejectedContinuation::InvalidData),
Arc::new(RecordingClientApiKeyUsage::default()),
);
let client = service
.authenticate("sk_start_test")
.expect("authenticated client");
let result = block_on(service.start(StartExecution {
client,
public_model: PublicModelId::new("gpt-start").expect("public model"),
operation: start_operation(),
metadata: execution_metadata_with_continuation(),
}));
let Err(error) = result else {
panic!("invalid continuation state must be rejected");
};
assert_eq!(error.kind(), GatewayErrorKind::Internal);
}
fn execution_metadata_with_continuation() -> ExecutionRequestMetadata {
ExecutionRequestMetadata {
protocol: "openai".to_owned(),
endpoint: "/v1/responses".to_owned(),
transport: ClientTransport::HttpJson,
stream: false,
client_ip: None,
user_agent: None,
previous_response_id: Some(PreviousResponseId::new("response-private")),
}
}
#[derive(Default)]
@@ -526,6 +597,7 @@ struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async { Ok(None) })
@@ -539,6 +611,43 @@ impl NativeContinuationPort for UnusedContinuation {
}
}
enum RejectedContinuation {
OwnershipMismatch,
InvalidData,
}
impl NativeContinuationPort for RejectedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async move {
Err(match self {
Self::OwnershipMismatch => NativeContinuationStoreError::ownership_mismatch(),
Self::InvalidData => NativeContinuationStoreError::invalid_data("invalid record"),
})
})
}
fn record<'a>(
&'a self,
_: NativeContinuationPin,
) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> {
Box::pin(async { Ok(()) })
}
}
fn account_scope(provider: &ProviderKind, account_id: &str) -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new(account_id).expect("account ID"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)]))),
ClientRoutingScope::all_accounts(),
))
}
fn probe_snapshot() -> RuntimeSnapshot {
let provider = ProviderKind::new("openai").expect("provider kind");
let capabilities =
@@ -575,7 +684,7 @@ fn client_snapshot() -> RuntimeSnapshot {
vec![ClientPolicy::new(
ClientApiKeyId::new("key_usage_test").expect("client API key ID"),
PlaintextClientApiKey::new("sk_usage_test").expect("plaintext client API key"),
provider,
account_scope(&provider, "acct_usage"),
true,
RateLimits::unlimited(),
)],
@@ -603,7 +712,7 @@ fn start_snapshot() -> RuntimeSnapshot {
vec![ClientPolicy::new(
ClientApiKeyId::new("key_start_test").expect("client API key ID"),
PlaintextClientApiKey::new("sk_start_test").expect("plaintext client API key"),
provider,
account_scope(&provider, "acct_start"),
true,
RateLimits::unlimited(),
)],
@@ -1,16 +1,25 @@
use std::sync::Arc;
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::ProviderKind;
use gateway_core::routing::{ClientRoutingScope, FrozenAccountScope, RuntimeAccountDirectory};
fn plaintext(value: &str) -> PlaintextClientApiKey {
PlaintextClientApiKey::new(value).expect("valid plaintext client key")
}
fn account_scope() -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::default()),
ClientRoutingScope::all_accounts(),
))
}
#[test]
fn disabled_client_key_should_be_denied() {
let policy = ClientPolicy::new(
ClientApiKeyId::new("key_disabled").expect("valid key ID"),
plaintext("sk_disabled_secret"),
ProviderKind::new("openai").expect("valid provider"),
account_scope(),
false,
RateLimits::unlimited(),
);
@@ -23,7 +32,7 @@ fn enabled_client_key_should_be_authorized() {
let policy = ClientPolicy::new(
ClientApiKeyId::new("key_enabled").expect("valid key ID"),
plaintext("sk_enabled_secret"),
ProviderKind::new("openai").expect("valid provider"),
account_scope(),
true,
RateLimits::unlimited(),
);
+168 -16
View File
@@ -1,15 +1,20 @@
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::Duration;
use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy};
use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
};
use gateway_core::operation::{
CapabilityRequirements, Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload,
};
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RuntimeSnapshot, UpstreamModelId,
AccountGroupId, AccountRoutingScopeKind, ClientRoutingScope, ConfigRevision,
FrozenAccountScope, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId,
RoutingContext, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot,
UpstreamModelId,
};
mod snapshot;
@@ -45,11 +50,34 @@ fn operation() -> Operation {
))
}
fn account_directory() -> Arc<RuntimeAccountDirectory> {
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
(
ProviderAccountId::new("acct_openai").expect("account"),
RuntimeAccount::new(
ProviderKind::new("openai").expect("provider"),
BTreeSet::new(),
),
),
(
ProviderAccountId::new("acct_xai").expect("account"),
RuntimeAccount::new(ProviderKind::new("xai").expect("provider"), BTreeSet::new()),
),
])))
}
fn account_scope() -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
account_directory(),
ClientRoutingScope::all_accounts(),
))
}
fn client_policy(id: &str, plaintext: &str, enabled: bool) -> ClientPolicy {
ClientPolicy::new(
ClientApiKeyId::new(id).expect("client key ID"),
PlaintextClientApiKey::new(plaintext).expect("plaintext client key"),
ProviderKind::new("openai").expect("valid provider"),
account_scope(),
enabled,
RateLimits::unlimited(),
)
@@ -70,6 +98,7 @@ fn snapshot() -> RuntimeSnapshot {
Vec::new(),
)
.expect("snapshot")
.with_account_directory(account_directory())
.with_model_mappings(BTreeMap::from([
("gpt-5.4".to_owned(), "gpt-5.5".to_owned()),
("grok-latest".to_owned(), "grok-4.5".to_owned()),
@@ -81,6 +110,92 @@ fn config_revision_should_reject_zero() {
assert!(ConfigRevision::new(0).is_err());
}
#[test]
fn account_group_id_should_match_the_database_contract_exactly() {
let valid = "grp_0123456789abcdef0123456789abcdef";
assert_eq!(
AccountGroupId::new(valid).expect("valid group ID").as_str(),
valid
);
for invalid in [
"0123456789abcdef0123456789abcdef",
"grp_0123456789abcdef0123456789abcde",
"grp_0123456789abcdef0123456789abcdef0",
"grp_0123456789abcdef0123456789abcdeF",
"grp_0123456789abcdef0123456789abcdeg",
] {
assert!(
AccountGroupId::new(invalid).is_err(),
"invalid group ID was accepted: {invalid}"
);
}
}
#[test]
fn restricted_scope_should_use_enabled_group_union_and_preserve_all_bindings_for_history() {
let enabled =
AccountGroupId::new("grp_11111111111111111111111111111111").expect("enabled group");
let disabled =
AccountGroupId::new("grp_22222222222222222222222222222222").expect("disabled group");
let provider = ProviderKind::new("openai").expect("provider");
let allowed_account = ProviderAccountId::new("acct_allowed").expect("account");
let disabled_only_account = ProviderAccountId::new("acct_disabled").expect("account");
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
(
allowed_account.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([enabled.clone()])),
),
(
disabled_only_account.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([disabled.clone()])),
),
])));
let scope = FrozenAccountScope::new(
directory,
ClientRoutingScope::restricted(
vec![
RoutingGroupSnapshot::new(enabled.clone(), "Enabled".to_owned()),
RoutingGroupSnapshot::new(disabled, "Disabled".to_owned()),
],
BTreeSet::from([enabled]),
BTreeSet::from([provider]),
)
.expect("restricted scope"),
);
assert!(scope.allows(&allowed_account));
assert!(!scope.allows(&disabled_only_account));
let history = scope.routing_snapshot();
assert_eq!(history.kind(), AccountRoutingScopeKind::Groups);
assert_eq!(history.groups_snapshot().len(), 2);
assert_eq!(history.groups_snapshot()[1].name(), "Disabled");
}
#[test]
fn restricted_scope_with_no_enabled_group_should_fail_closed() {
let group = AccountGroupId::new("grp_33333333333333333333333333333333").expect("group");
let account = ProviderAccountId::new("acct_group_member").expect("account");
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
account.clone(),
RuntimeAccount::new(
ProviderKind::new("openai").expect("provider"),
BTreeSet::from([group.clone()]),
),
)])));
let scope = FrozenAccountScope::new(
directory,
ClientRoutingScope::restricted(
vec![RoutingGroupSnapshot::new(group, "Disabled".to_owned())],
BTreeSet::new(),
BTreeSet::new(),
)
.expect("restricted scope"),
);
assert!(scope.provider_kinds().is_empty());
assert!(!scope.allows(&account));
}
#[test]
fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
let snapshot = RuntimeSnapshot::new(
@@ -98,7 +213,12 @@ fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
let policies = snapshot.client_policies().collect::<Vec<_>>();
assert_eq!(policies.len(), 1);
assert_eq!(policies[0].key_id().as_str(), "key_enabled");
assert_eq!(policies[0].provider_kind().as_str(), "openai");
assert!(
policies[0]
.account_scope()
.provider_kinds()
.contains(&ProviderKind::new("openai").expect("provider"))
);
}
#[test]
@@ -137,8 +257,9 @@ fn selected_provider_should_use_global_model_mapping() {
.plan(
&PublicModelId::new("gpt-5.4").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
@@ -170,26 +291,37 @@ fn cyclic_model_mapping_should_fall_back_to_the_original_name() {
}
#[test]
fn global_mapping_should_apply_to_every_provider() {
fn known_provider_catalog_missing_mapped_model_should_be_filtered() {
let snapshot = snapshot();
let plan = snapshot
let error = snapshot
.plan(
&PublicModelId::new("gpt-5.4").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("xai").expect("provider")),
required_provider: Some(ProviderKind::new("xai").expect("provider")),
..RoutingContext::default()
},
)
.expect("transparent plan");
.expect_err("known xAI catalog does not contain the mapped OpenAI model");
assert_eq!(plan.candidates()[0].provider().as_str(), "xai");
assert_eq!(plan.candidates()[0].upstream_model().as_str(), "gpt-5.5");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProvider { .. }
));
}
#[test]
fn unmapped_model_should_pass_through_unchanged() {
let snapshot = snapshot();
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
scheduling(),
vec![ProviderKind::new("openai").expect("provider")],
Vec::new(),
Vec::new(),
)
.expect("snapshot")
.with_account_directory(account_directory());
for requested in [
"future-openai-model".to_owned(),
format!("future-{}", "x".repeat(512)),
@@ -200,8 +332,9 @@ fn unmapped_model_should_pass_through_unchanged() {
.plan(
&PublicModelId::from_client_wire(requested.clone()).expect("client model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
@@ -229,8 +362,9 @@ fn blocked_provider_should_be_filtered() {
.plan(
&PublicModelId::new("gpt-5.5").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
},
)
@@ -262,8 +396,9 @@ fn known_unsupported_operation_should_not_be_bypassed() {
.plan(
&PublicModelId::new("gpt-known-unsupported").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
@@ -316,3 +451,20 @@ fn unknown_model_should_be_hidden_when_provider_catalog_is_available() {
&ProviderKind::new("openai").expect("provider"),
));
}
#[test]
fn alias_should_only_be_available_from_a_provider_with_its_mapped_model() {
let snapshot = snapshot();
let alias = PublicModelId::new("grok-latest").expect("model");
assert!(!snapshot.contains_public_model_for_provider(
&alias,
&ProviderKind::new("openai").expect("provider"),
));
assert!(
snapshot.contains_public_model_for_provider(
&alias,
&ProviderKind::new("xai").expect("provider"),
)
);
}
@@ -19,7 +19,8 @@ use gateway_core::operation::OperationKind;
use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::{
RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, RuntimeSnapshotHandle,
RuntimeSnapshotPublisher, SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts,
RuntimeSnapshotPublisher, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts,
SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts, SnapshotProviderAccountFacts,
SnapshotRevisionStream, SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort,
SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh,
};
@@ -348,9 +349,12 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
vec![SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_one").expect("key ID"),
PlaintextClientApiKey::new("sk_test").expect("plaintext key"),
"alpha",
Vec::new(),
RateLimits::unlimited(),
)],
Vec::<SnapshotAccountGroupFacts>::new(),
Vec::<SnapshotProviderAccountFacts>::new(),
Vec::<SnapshotAccountGroupMemberFacts>::new(),
)
}
@@ -90,6 +90,11 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
pool.clone(),
Some(Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>),
)),
Arc::new(postgres::PgAccountGroupRepository::with_runtime_state(
pool.clone(),
credential_leases.clone(),
Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>,
)),
Arc::new(AdminAuthStoreAdapter {
security: postgres::PgAdminSecurityAuditRepository::new(pool.clone()),
settings: postgres::PgRuntimeSettingsRepository::new(pool.clone()),
@@ -0,0 +1,783 @@
//! PostgreSQL owner for provider-neutral account groups and memberships.
use std::{
collections::{BTreeMap, BTreeSet},
str::FromStr as _,
sync::Arc,
};
use async_trait::async_trait;
use futures::future::BoxFuture;
use gateway_admin::{
model::{
MutationContext,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
AccountGroupListQuery, AccountGroupMember, AccountGroupMembers, AccountGroupMutation,
AccountGroupPage, AccountGroupRecord, AccountGroupUsage, DeleteAccountGroup,
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
observability::DecimalAmount,
},
ports::store::{AccountGroupStore, AdminStoreError, AdminStoreResult},
};
use gateway_core::{
engine::credential::AccountStatus,
provider_ports::ProviderCooldownPort,
routing::{AccountGroupId, ProviderKind},
};
use sqlx::{PgPool, Postgres, QueryBuilder, Row as _, Transaction};
use crate::redis::CredentialLeaseRepository as _;
use crate::{
ConflictKind, StoreError, StoreResult, admin_revision, admin_store_error, mutation_audit,
postgres_unavailable,
};
use super::{
PgProviderAccountRepository, PgRuntimeSettingsRepository, ProviderAccountRepository,
RuntimeSettingsRepository, account_status_projection, append_admin_audit_event_in_transaction,
bump_config_revision_in_transaction, load_rate_limited_until,
};
const ENTITY: &str = "account group";
/// Account group store with transactional revision and audit ownership.
#[derive(Clone)]
pub struct PgAccountGroupRepository {
pool: PgPool,
runtime_signals: Option<crate::redis::RedisCredentialLeaseRepository>,
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
}
impl PgAccountGroupRepository {
#[must_use]
pub const fn new(pool: PgPool) -> Self {
Self {
pool,
runtime_signals: None,
cooldowns: None,
}
}
#[must_use]
pub fn with_runtime_state(
pool: PgPool,
runtime_signals: crate::redis::RedisCredentialLeaseRepository,
cooldowns: Arc<dyn ProviderCooldownPort>,
) -> Self {
Self {
pool,
runtime_signals: Some(runtime_signals),
cooldowns: Some(cooldowns),
}
}
async fn current_revision(&self) -> AdminStoreResult<gateway_admin::model::Revision> {
RuntimeSettingsRepository::load_runtime_settings(&PgRuntimeSettingsRepository::new(
self.pool.clone(),
))
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(|settings| admin_revision(settings.config_revision))
}
async fn required_record(&self, id: &AccountGroupId) -> AdminStoreResult<AccountGroupRecord> {
let record = load_record(&self.pool, id.as_str())
.await
.map_err(|error| admin_store_error(ENTITY, error))?
.ok_or_else(|| not_found(id.as_str()))?;
let mut records = vec![record];
self.enrich_records(&mut records).await?;
records
.pop()
.ok_or_else(|| invalid_admin("account group record disappeared"))
}
async fn enrich_records(&self, records: &mut [AccountGroupRecord]) -> AdminStoreResult<()> {
enrich_group_records(
&self.pool,
self.runtime_signals.as_ref(),
self.cooldowns.as_deref(),
records,
)
.await
.map_err(|error| admin_store_error(ENTITY, error))
}
async fn mutate<F>(
&self,
audit: super::AdminAuditEvent,
mutation: F,
) -> AdminStoreResult<gateway_admin::model::Revision>
where
F: for<'a> FnOnce(&'a mut Transaction<'_, Postgres>) -> BoxFuture<'a, StoreResult<()>>,
{
let mut transaction =
self.pool.begin().await.map_err(|_| {
admin_store_error(ENTITY, unavailable("begin account group mutation"))
})?;
let result = async {
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
mutation(&mut transaction).await?;
append_admin_audit_event_in_transaction(&mut transaction, audit, revision).await?;
Ok(revision)
}
.await;
match result {
Ok(revision) => {
transaction.commit().await.map_err(|_| {
admin_store_error(ENTITY, unavailable("commit account group mutation"))
})?;
admin_revision(revision)
}
Err(error) => {
transaction.rollback().await.map_err(|_| {
admin_store_error(ENTITY, unavailable("rollback account group mutation"))
})?;
Err(admin_store_error(ENTITY, error))
}
}
}
}
#[async_trait]
impl AccountGroupStore for PgAccountGroupRepository {
async fn list_account_groups(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage> {
validate_page_query(&query)?;
let total = count_groups(&self.pool, &query)
.await
.map_err(|error| admin_store_error(ENTITY, error))?;
let offset = u64::from(query.page.saturating_sub(1)) * u64::from(query.page_size.get());
let mut statement = group_select();
push_group_filter(&mut statement, &query);
statement.push(" order by g.created_at desc, g.id desc limit ");
statement.push_bind(i64::from(query.page_size.get()));
statement.push(" offset ");
statement.push_bind(i64::try_from(offset).map_err(|_| invalid_admin("page is too large"))?);
let rows = statement
.build()
.fetch_all(&self.pool)
.await
.map_err(|_| admin_store_error(ENTITY, unavailable("list account groups")))?;
let mut items = rows
.iter()
.map(group_record)
.collect::<StoreResult<Vec<_>>>()
.map_err(|error| admin_store_error(ENTITY, error))?;
self.enrich_records(&mut items).await?;
Ok(AccountGroupPage {
config_revision: self.current_revision().await?,
items,
total,
page: query.page,
page_size: query.page_size.get(),
})
}
async fn account_group_members(
&self,
id: &AccountGroupId,
) -> AdminStoreResult<AccountGroupMembers> {
self.required_record(id).await?;
let rows = sqlx::query(
"select a.id, a.name, a.provider_kind, a.email, a.enabled
from account_group_accounts m
join provider_accounts a on a.id = m.provider_account_id
where m.account_group_id = $1
order by a.provider_kind, lower(a.name), a.id",
)
.bind(id.as_str())
.fetch_all(&self.pool)
.await
.map_err(|_| admin_store_error(ENTITY, unavailable("load account group members")))?;
let items = rows
.iter()
.map(group_member)
.collect::<StoreResult<Vec<_>>>()
.map_err(|error| admin_store_error(ENTITY, error))?;
Ok(AccountGroupMembers {
config_revision: self.current_revision().await?,
id: id.clone(),
total: u64::try_from(items.len())
.map_err(|_| invalid_admin("member count overflow"))?,
items,
})
}
async fn create_account_group(
&self,
command: NewAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
validate_group_fields(&command.name, command.description.as_deref())?;
let id = command.id.clone();
let audit = mutation_audit(
context,
"create",
"account_group",
id.as_str(),
vec!["name".to_owned(), "description".to_owned()],
);
let revision = self
.mutate(audit, |transaction| {
Box::pin(async move {
sqlx::query(
"insert into account_groups
(id, name, description, color, enabled, created_at, updated_at)
values ($1, $2, $3, $4, true, now(), now())",
)
.bind(command.id.as_str())
.bind(command.name)
.bind(command.description)
.bind(command.color.as_str())
.execute(&mut **transaction)
.await
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
Ok(())
})
})
.await?;
Ok(AccountGroupMutation {
config_revision: revision,
id: id.clone(),
record: Some(self.required_record(&id).await?),
})
}
async fn update_account_group(
&self,
command: UpdateAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
validate_group_fields(&command.name, command.description.as_deref())?;
let id = command.id.clone();
let audit = mutation_audit(
context,
"update",
"account_group",
id.as_str(),
vec!["name".to_owned(), "description".to_owned()],
);
let revision = self
.mutate(audit, |transaction| {
Box::pin(async move {
let result = sqlx::query(
"update account_groups
set name = $2, description = $3, color = $4, updated_at = now()
where id = $1",
)
.bind(command.id.as_str())
.bind(command.name)
.bind(command.description)
.bind(command.color.as_str())
.execute(&mut **transaction)
.await
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
require_one(result.rows_affected(), command.id.as_str())
})
})
.await?;
Ok(AccountGroupMutation {
config_revision: revision,
id: id.clone(),
record: Some(self.required_record(&id).await?),
})
}
async fn set_account_group_enabled(
&self,
command: SetAccountGroupEnabled,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let id = command.id.clone();
let audit = mutation_audit(
context,
if command.enabled { "enable" } else { "disable" },
"account_group",
id.as_str(),
vec!["enabled".to_owned()],
);
let revision = self
.mutate(audit, |transaction| {
Box::pin(async move {
let result = sqlx::query(
"update account_groups set enabled = $2, updated_at = now() where id = $1",
)
.bind(command.id.as_str())
.bind(command.enabled)
.execute(&mut **transaction)
.await
.map_err(|_| unavailable("set account group state"))?;
require_one(result.rows_affected(), command.id.as_str())
})
})
.await?;
Ok(AccountGroupMutation {
config_revision: revision,
id: id.clone(),
record: Some(self.required_record(&id).await?),
})
}
async fn delete_account_group(
&self,
command: DeleteAccountGroup,
context: &MutationContext,
) -> AdminStoreResult<AccountGroupMutation> {
let id = command.id.clone();
let audit = mutation_audit(context, "delete", "account_group", id.as_str(), Vec::new());
let revision = self
.mutate(audit, |transaction| {
Box::pin(async move {
let result = sqlx::query(
"delete from account_groups g
where g.id = $1
and not exists (
select 1 from client_api_key_groups k where k.account_group_id = g.id
)",
)
.bind(command.id.as_str())
.execute(&mut **transaction)
.await
.map_err(|_| unavailable("delete account group"))?;
if result.rows_affected() == 1 {
return Ok(());
}
let exists = sqlx::query_scalar::<_, bool>(
"select exists(select 1 from account_groups where id = $1)",
)
.bind(command.id.as_str())
.fetch_one(&mut **transaction)
.await
.map_err(|_| unavailable("check account group delete conflict"))?;
if exists {
Err(conflict(command.id.as_str()))
} else {
Err(not_found_store(command.id.as_str()))
}
})
})
.await?;
Ok(AccountGroupMutation {
config_revision: revision,
id,
record: None,
})
}
}
fn group_select() -> QueryBuilder<Postgres> {
QueryBuilder::new(
"select g.id, g.name, g.description, g.color, g.enabled, g.created_at, g.updated_at,
coalesce(members.member_count, 0)::bigint as member_count,
coalesce(keys.client_key_count, 0)::bigint as client_key_count,
coalesce(members.provider_counts, '{}'::jsonb) as provider_counts
from account_groups g
left join lateral (
select coalesce(sum(counts.provider_count), 0)::bigint as member_count,
jsonb_object_agg(counts.provider_kind, counts.provider_count)
as provider_counts
from (
select a.provider_kind, count(*)::bigint as provider_count
from account_group_accounts gm
join provider_accounts a on a.id = gm.provider_account_id
where gm.account_group_id = g.id
group by a.provider_kind
) counts
) members on true
left join lateral (
select count(*)::bigint as client_key_count
from client_api_key_groups kg
where kg.account_group_id = g.id
) keys on true
where true",
)
}
fn push_group_filter(statement: &mut QueryBuilder<Postgres>, query: &AccountGroupListQuery) {
if let Some(search) = &query.search {
statement.push(" and lower(g.name) like ");
statement.push_bind(format!("%{}%", search.to_lowercase()));
}
if let Some(enabled) = query.enabled {
statement.push(" and g.enabled = ");
statement.push_bind(enabled);
}
}
async fn count_groups(pool: &PgPool, query: &AccountGroupListQuery) -> StoreResult<u64> {
let mut statement =
QueryBuilder::<Postgres>::new("select count(*)::bigint from account_groups g where true");
if let Some(search) = &query.search {
statement.push(" and lower(g.name) like ");
statement.push_bind(format!("%{}%", search.to_lowercase()));
}
if let Some(enabled) = query.enabled {
statement.push(" and g.enabled = ");
statement.push_bind(enabled);
}
let count = statement
.build_query_scalar::<i64>()
.fetch_one(pool)
.await
.map_err(|_| unavailable("count account groups"))?;
u64::try_from(count).map_err(|_| invalid("negative account group count"))
}
async fn load_record(pool: &PgPool, id: &str) -> StoreResult<Option<AccountGroupRecord>> {
let mut statement = group_select();
statement.push(" and g.id = ");
statement.push_bind(id.to_owned());
statement
.build()
.fetch_optional(pool)
.await
.map_err(|_| unavailable("load account group"))?
.as_ref()
.map(group_record)
.transpose()
}
fn group_record(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupRecord> {
let provider_counts: serde_json::Value = row
.try_get("provider_counts")
.map_err(|_| invalid("invalid provider counts"))?;
let provider_counts = serde_json::from_value::<BTreeMap<String, u64>>(provider_counts)
.map_err(|_| invalid("invalid provider counts"))?;
Ok(AccountGroupRecord {
id: AccountGroupId::new(
row.try_get::<String, _>("id")
.map_err(|_| invalid("invalid id"))?,
)
.map_err(|_| invalid("invalid id"))?,
name: row.try_get("name").map_err(|_| invalid("invalid name"))?,
description: row
.try_get("description")
.map_err(|_| invalid("invalid description"))?,
color: AccountGroupColor::parse(
row.try_get::<String, _>("color")
.map_err(|_| invalid("invalid color"))?
.as_str(),
)
.ok_or_else(|| invalid("invalid color"))?,
enabled: row
.try_get("enabled")
.map_err(|_| invalid("invalid enabled"))?,
member_count: count_value(row, "member_count")?,
provider_counts,
client_key_count: count_value(row, "client_key_count")?,
account_summary: AccountGroupAccountSummary {
available: 0,
limited: 0,
total: 0,
},
capacity: AccountGroupCapacity {
used_slots: None,
total_slots: 0,
},
usage: AccountGroupUsage {
today_usd: DecimalAmount::from_str("0").map_err(|_| invalid("invalid zero cost"))?,
total_usd: DecimalAmount::from_str("0").map_err(|_| invalid("invalid zero cost"))?,
},
created_at: row
.try_get("created_at")
.map_err(|_| invalid("invalid created_at"))?,
updated_at: row
.try_get("updated_at")
.map_err(|_| invalid("invalid updated_at"))?,
})
}
async fn enrich_group_records(
pool: &PgPool,
runtime_signals: Option<&crate::redis::RedisCredentialLeaseRepository>,
cooldowns: Option<&dyn ProviderCooldownPort>,
records: &mut [AccountGroupRecord],
) -> StoreResult<()> {
if records.is_empty() {
return Ok(());
}
let settings = RuntimeSettingsRepository::load_runtime_settings(
&PgRuntimeSettingsRepository::new(pool.clone()),
)
.await?;
let accounts = PgProviderAccountRepository::new(pool.clone())
.list_provider_accounts(None, true)
.await?;
let now = chrono::Utc::now();
let cooldowns = load_rate_limited_until(cooldowns, &accounts, now.into()).await;
let mut status_by_account = BTreeMap::new();
let mut capacity_by_account = BTreeMap::new();
for account in &accounts {
let status =
account_status_projection(account, now.into(), cooldowns.get(&account.id).copied())
.status;
status_by_account.insert(account.id.as_str(), status);
capacity_by_account.insert(
account.id.as_str(),
u64::from(
account
.concurrency_limit
.map_or(settings.max_concurrent_per_account, |limit| limit.get()),
),
);
}
let group_ids = records
.iter()
.map(|record| record.id.as_str().to_owned())
.collect::<Vec<_>>();
let memberships = sqlx::query(
"select account_group_id, provider_account_id
from account_group_accounts
where account_group_id = any($1::text[])",
)
.bind(&group_ids)
.fetch_all(pool)
.await
.map_err(|_| unavailable("load account group availability"))?;
let mut accounts_by_group = BTreeMap::<String, Vec<String>>::new();
for row in memberships {
accounts_by_group
.entry(
row.try_get("account_group_id")
.map_err(|_| invalid("invalid group ID"))?,
)
.or_default()
.push(
row.try_get("provider_account_id")
.map_err(|_| invalid("invalid account ID"))?,
);
}
let normal_ids = accounts_by_group
.values()
.flatten()
.filter(|account_id| {
status_by_account.get(account_id.as_str()) == Some(&AccountStatus::Normal)
})
.cloned()
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let used_by_account = if normal_ids.is_empty() {
Some(BTreeMap::new())
} else if let Some(runtime_signals) = runtime_signals {
runtime_signals
.credential_runtime_signals(&normal_ids)
.await
.ok()
.map(|signals| {
signals
.into_iter()
.map(|signal| (signal.resource_id, u64::from(signal.in_flight)))
.collect::<BTreeMap<_, _>>()
})
} else {
None
};
let costs = group_costs(pool, &group_ids).await?;
for record in records {
let account_ids = accounts_by_group
.get(record.id.as_str())
.map_or(&[][..], Vec::as_slice);
let available = account_ids
.iter()
.filter(|id| status_by_account.get(id.as_str()) == Some(&AccountStatus::Normal))
.count() as u64;
let total = u64::try_from(account_ids.len())
.map_err(|_| invalid("group account count overflow"))?;
record.account_summary = AccountGroupAccountSummary {
available,
limited: total.saturating_sub(available),
total,
};
record.capacity = AccountGroupCapacity {
used_slots: if available == 0 {
Some(0)
} else {
used_by_account.as_ref().map(|used| {
account_ids.iter().fold(0_u64, |sum, id| {
sum.saturating_add(used.get(id).copied().unwrap_or(0))
})
})
},
total_slots: account_ids.iter().fold(0_u64, |total, account_id| {
if status_by_account.get(account_id.as_str()) == Some(&AccountStatus::Normal) {
total.saturating_add(
capacity_by_account
.get(account_id.as_str())
.copied()
.unwrap_or(0),
)
} else {
total
}
}),
};
if let Some(usage) = costs.get(record.id.as_str()) {
record.usage = usage.clone();
}
}
Ok(())
}
async fn group_costs(
pool: &PgPool,
group_ids: &[String],
) -> StoreResult<BTreeMap<String, AccountGroupUsage>> {
let rows = sqlx::query(
"select gm.account_group_id,
coalesce(sum(mr.cost_amount) filter (
where mr.started_at >= date_trunc('day', now() at time zone 'Asia/Shanghai')
at time zone 'Asia/Shanghai'
), 0)::text as today_usd,
coalesce(sum(mr.cost_amount), 0)::text as total_usd
from account_group_accounts gm
left join model_requests mr
on mr.provider_account_ref = gm.provider_account_id
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
and mr.cost_currency = 'USD'
and mr.cost_amount is not null
where gm.account_group_id = any($1::text[])
group by gm.account_group_id",
)
.bind(group_ids)
.fetch_all(pool)
.await
.map_err(|_| unavailable("load account group costs"))?;
rows.into_iter()
.map(|row| {
let group_id = row
.try_get("account_group_id")
.map_err(|_| invalid("invalid group ID"))?;
let today = row
.try_get::<String, _>("today_usd")
.map_err(|_| invalid("invalid today cost"))?;
let total = row
.try_get::<String, _>("total_usd")
.map_err(|_| invalid("invalid total cost"))?;
Ok((
group_id,
AccountGroupUsage {
today_usd: DecimalAmount::from_str(&today)
.map_err(|_| invalid("invalid today cost"))?,
total_usd: DecimalAmount::from_str(&total)
.map_err(|_| invalid("invalid total cost"))?,
},
))
})
.collect()
}
fn group_member(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupMember> {
Ok(AccountGroupMember {
id: row
.try_get("id")
.map_err(|_| invalid("invalid member id"))?,
name: row
.try_get("name")
.map_err(|_| invalid("invalid member name"))?,
provider_kind: ProviderKind::new(
row.try_get::<String, _>("provider_kind")
.map_err(|_| invalid("invalid member provider"))?,
)
.map_err(|_| invalid("invalid member provider"))?,
email: row
.try_get("email")
.map_err(|_| invalid("invalid member email"))?,
enabled: row
.try_get("enabled")
.map_err(|_| invalid("invalid member enabled"))?,
})
}
fn count_value(row: &sqlx::postgres::PgRow, field: &str) -> StoreResult<u64> {
u64::try_from(
row.try_get::<i64, _>(field)
.map_err(|_| invalid("invalid count"))?,
)
.map_err(|_| invalid("negative count"))
}
fn validate_page_query(query: &AccountGroupListQuery) -> AdminStoreResult<()> {
if query.page == 0 {
return Err(invalid_admin("page must be positive"));
}
if query.search.as_deref().is_some_and(|search| {
search.trim().is_empty() || search.len() > 256 || search.chars().any(char::is_control)
}) {
return Err(invalid_admin("invalid search"));
}
Ok(())
}
fn validate_group_fields(name: &str, description: Option<&str>) -> AdminStoreResult<()> {
if name.trim() != name
|| name.is_empty()
|| name.chars().count() > 100
|| name.chars().any(char::is_control)
|| description
.is_some_and(|value| value.len() > 4096 || value.chars().any(char::is_control))
{
return Err(invalid_admin("invalid account group fields"));
}
Ok(())
}
fn require_one(rows: u64, id: &str) -> StoreResult<()> {
if rows == 1 {
Ok(())
} else {
Err(not_found_store(id))
}
}
fn map_group_write_error(error: sqlx::Error, id: &str) -> StoreError {
if error
.as_database_error()
.is_some_and(|database| database.is_unique_violation())
{
conflict(id)
} else {
unavailable("write account group")
}
}
fn not_found(id: &str) -> AdminStoreError {
admin_store_error(ENTITY, not_found_store(id))
}
fn not_found_store(id: &str) -> StoreError {
StoreError::NotFound {
entity: ENTITY,
id: id.to_owned(),
}
}
fn conflict(id: &str) -> StoreError {
StoreError::Conflict {
entity: ENTITY,
id: id.to_owned(),
kind: ConflictKind::InvalidTransition,
}
}
fn invalid_admin(message: &str) -> AdminStoreError {
admin_store_error(ENTITY, invalid(message))
}
fn invalid(message: &str) -> StoreError {
StoreError::InvalidData {
entity: ENTITY,
message: message.to_owned(),
}
}
fn unavailable(message: &'static str) -> StoreError {
postgres_unavailable(message)
}
@@ -1,7 +1,7 @@
//! 明文 `client_api_keys` 的 PostgreSQL owner。
use std::{
collections::BTreeMap,
collections::{BTreeMap, BTreeSet},
fmt,
sync::{
Arc, Mutex,
@@ -15,6 +15,7 @@ use chrono::{DateTime, Utc};
use gateway_admin::{
model::{
MutationContext,
account_groups::AccountGroupRef as AdminAccountGroupRef,
client_keys::{
ClientKeyCursor as AdminClientKeyCursor,
ClientKeyCursorValue as AdminClientKeyCursorValue,
@@ -27,11 +28,12 @@ use gateway_admin::{
},
ports::store::{AdminStoreResult, ClientKeyStore},
};
use gateway_core::routing::ProviderKind;
use gateway_core::{
engine::execution::ClientApiKeyUsageSink,
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
routing::{AccountGroupId, ProviderKind},
};
use serde::Deserialize;
use sqlx::{PgPool, Postgres, QueryBuilder, Transaction};
use crate::{
@@ -49,7 +51,7 @@ const CLIENT_API_KEY_LAST_USED_FLUSH_DELAY: Duration = Duration::from_secs(1);
pub struct ClientApiKeySnapshot {
pub id: ClientApiKeyId,
pub plaintext_key: PlaintextClientApiKey,
pub provider_kind: String,
pub group_ids: Vec<AccountGroupId>,
pub limits: RateLimits,
}
@@ -57,7 +59,7 @@ impl ClientApiKeySnapshot {
pub(crate) fn from_persisted(
id: String,
key: String,
provider_kind: String,
group_ids: Vec<String>,
max_concurrency: i64,
requests_per_minute: i64,
) -> StoreResult<Self> {
@@ -65,7 +67,12 @@ impl ClientApiKeySnapshot {
id: ClientApiKeyId::new(id).map_err(|_| invalid("persisted key ID is invalid"))?,
plaintext_key: PlaintextClientApiKey::new(key)
.map_err(|_| invalid("persisted plaintext key is invalid"))?,
provider_kind,
group_ids: group_ids
.into_iter()
.map(|id| {
AccountGroupId::new(id).map_err(|_| invalid("persisted group ID is invalid"))
})
.collect::<StoreResult<Vec<_>>>()?,
limits: RateLimits {
max_concurrency: to_u64(max_concurrency)?,
requests_per_minute: to_u64(requests_per_minute)?,
@@ -78,7 +85,6 @@ impl ClientApiKeySnapshot {
pub struct ClientApiKeySecret {
pub id: String,
pub key: String,
pub provider_kind: String,
pub enabled: bool,
pub max_concurrency: u64,
pub requests_per_minute: u64,
@@ -100,7 +106,8 @@ pub struct ClientApiKeyRecord {
pub id: String,
pub name: String,
pub label: Option<String>,
pub provider_kind: String,
pub groups: Vec<ClientApiKeyGroupRecord>,
pub provider_kinds: Vec<String>,
pub prefix: String,
pub enabled: bool,
pub max_concurrency: u64,
@@ -110,6 +117,14 @@ pub struct ClientApiKeyRecord {
pub updated_at: DateTime<Utc>,
}
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct ClientApiKeyGroupRecord {
pub id: String,
pub name: String,
pub color: String,
pub enabled: bool,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum ClientApiKeySortField {
Name,
@@ -250,7 +265,7 @@ pub struct NewClientApiKey {
pub id: String,
pub name: String,
pub label: Option<String>,
pub provider_kind: String,
pub group_ids: Vec<String>,
pub key: String,
pub max_concurrency: u64,
pub requests_per_minute: u64,
@@ -271,28 +286,17 @@ impl NewClientApiKey {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "id", &self.id)?;
require_nonempty(ENTITY, "name", &self.name)?;
require_nonempty(ENTITY, "provider_kind", &self.provider_kind)?;
validate_group_ids(&self.group_ids)?;
validate_key(&self.key)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateClientApiKey {
pub id: String,
pub name: String,
pub label: Option<String>,
pub provider_kind: String,
pub enabled: bool,
pub max_concurrency: u64,
pub requests_per_minute: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UpdateClientApiKeyDetails {
pub id: String,
pub name: String,
pub label: Option<String>,
pub provider_kind: String,
pub group_ids: Vec<String>,
pub max_concurrency: u64,
pub requests_per_minute: u64,
}
@@ -301,7 +305,7 @@ impl UpdateClientApiKeyDetails {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "id", &self.id)?;
require_nonempty(ENTITY, "name", &self.name)?;
require_nonempty(ENTITY, "provider_kind", &self.provider_kind)?;
validate_group_ids(&self.group_ids)?;
to_i64(self.max_concurrency)?;
to_i64(self.requests_per_minute)?;
Ok(())
@@ -316,7 +320,6 @@ pub trait ClientApiKeyRepository: Send + Sync {
) -> StoreResult<ClientApiKeyPage>;
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>>;
async fn reveal_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeySecret>>;
async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult<bool>;
async fn touch_client_api_keys(
&self,
touched_at: &BTreeMap<String, DateTime<Utc>>,
@@ -344,9 +347,38 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
query.validate()?;
let total = count_client_api_keys(&self.pool, query.search.as_deref()).await?;
let mut statement = QueryBuilder::<Postgres>::new(
"select id, name, label, provider_kind, left(key, 10) as prefix, enabled, max_concurrency,
requests_per_minute, last_used_at, created_at, updated_at
from client_api_keys where true",
"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,
case
when groups.binding_count = 0 then coalesce(
(select array_agg(distinct a.provider_kind order by a.provider_kind)
from provider_accounts a),
'{}'
)
else coalesce(groups.provider_kinds, '{}')
end as provider_kinds
from client_api_keys k
left join lateral (
select
(select count(*)::bigint
from client_api_key_groups kg
where kg.client_api_key_id = k.id) as binding_count,
(select jsonb_agg(jsonb_build_object(
'id', g.id, 'name', g.name, 'color', g.color, 'enabled', g.enabled
) order by g.id)
from client_api_key_groups kg
join account_groups g on g.id = kg.account_group_id
where kg.client_api_key_id = k.id) as groups,
(select array_agg(distinct a.provider_kind order by a.provider_kind)
from client_api_key_groups kg
join account_groups g on g.id = kg.account_group_id and g.enabled
join account_group_accounts gm on gm.account_group_id = g.id
join provider_accounts a on a.id = gm.provider_account_id
where kg.client_api_key_id = k.id) as provider_kinds
) groups on true
where true",
);
push_client_key_search(&mut statement, query.search.as_deref());
if let Some(cursor) = &query.cursor {
@@ -356,12 +388,12 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
statement.push(" limit ");
statement.push_bind(i64::from(query.page_size) + 1);
let rows = statement
.build_query_as::<ClientRecordRow>()
.build()
.fetch_all(&self.pool)
.await
.map_err(|_| postgres_unavailable("list client API keys"))?;
let mut items = rows
.into_iter()
.iter()
.map(client_record_from_row)
.collect::<StoreResult<Vec<_>>>()?;
let has_more = items.len() > usize::from(query.page_size);
@@ -384,8 +416,8 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
async fn reveal_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeySecret>> {
require_nonempty(ENTITY, "id", id)?;
sqlx::query_as::<_, (String, String, String, bool, i64, i64)>(
"select id, key, provider_kind, enabled, max_concurrency, requests_per_minute
sqlx::query_as::<_, (String, String, bool, i64, i64)>(
"select id, key, enabled, max_concurrency, requests_per_minute
from client_api_keys where id = $1",
)
.bind(id)
@@ -398,41 +430,48 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>> {
require_nonempty(ENTITY, "id", id)?;
sqlx::query_as::<_, ClientRecordRow>(
"select id, name, label, provider_kind, left(key, 10) as prefix, enabled, max_concurrency,
requests_per_minute, last_used_at, created_at, updated_at
from client_api_keys where id = $1",
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,
case
when groups.binding_count = 0 then coalesce(
(select array_agg(distinct a.provider_kind order by a.provider_kind)
from provider_accounts a),
'{}'
)
else coalesce(groups.provider_kinds, '{}')
end as provider_kinds
from client_api_keys k
left join lateral (
select
(select count(*)::bigint
from client_api_key_groups kg
where kg.client_api_key_id = k.id) as binding_count,
(select jsonb_agg(jsonb_build_object(
'id', g.id, 'name', g.name, 'color', g.color, 'enabled', g.enabled
) order by g.id)
from client_api_key_groups kg
join account_groups g on g.id = kg.account_group_id
where kg.client_api_key_id = k.id) as groups,
(select array_agg(distinct a.provider_kind order by a.provider_kind)
from client_api_key_groups kg
join account_groups g on g.id = kg.account_group_id and g.enabled
join account_group_accounts gm on gm.account_group_id = g.id
join provider_accounts a on a.id = gm.provider_account_id
where kg.client_api_key_id = k.id) as provider_kinds
) groups on true
where k.id = $1",
)
.bind(id)
.fetch_optional(&self.pool)
.await
.map_err(|_| postgres_unavailable("get client API key"))?
.as_ref()
.map(client_record_from_row)
.transpose()
}
async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult<bool> {
require_nonempty(ENTITY, "id", &key.id)?;
require_nonempty(ENTITY, "name", &key.name)?;
let result = sqlx::query(
"update client_api_keys
set name = $2, label = $3, provider_kind = $4, enabled = $5, max_concurrency = $6,
requests_per_minute = $7, updated_at = now()
where id = $1",
)
.bind(key.id)
.bind(key.name)
.bind(key.label)
.bind(key.provider_kind)
.bind(key.enabled)
.bind(to_i64(key.max_concurrency)?)
.bind(to_i64(key.requests_per_minute)?)
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("update client API key"))?;
Ok(result.rows_affected() == 1)
}
async fn touch_client_api_keys(
&self,
touched_at: &BTreeMap<String, DateTime<Utc>>,
@@ -637,7 +676,11 @@ impl ClientKeyStore for PgAdminClientKeyStore {
id: id.as_str().to_owned(),
name: command.name,
label: command.label,
provider_kind: command.provider_kind.as_str().to_owned(),
group_ids: command
.group_ids
.iter()
.map(|id| id.as_str().to_owned())
.collect(),
key: command.plaintext,
max_concurrency: command.limits.max_concurrency,
requests_per_minute: command.limits.requests_per_minute,
@@ -650,7 +693,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
[
"name",
"label",
"provider_kind",
"group_ids",
"key",
"enabled",
"max_concurrency",
@@ -679,7 +722,11 @@ impl ClientKeyStore for PgAdminClientKeyStore {
id: id.as_str().to_owned(),
name: command.name,
label: command.label,
provider_kind: command.provider_kind.as_str().to_owned(),
group_ids: command
.group_ids
.iter()
.map(|id| id.as_str().to_owned())
.collect(),
max_concurrency: command.limits.max_concurrency,
requests_per_minute: command.limits.requests_per_minute,
},
@@ -691,7 +738,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
[
"name",
"label",
"provider_kind",
"group_ids",
"max_concurrency",
"requests_per_minute",
]
@@ -828,8 +875,30 @@ fn admin_client_key_record(record: ClientApiKeyRecord) -> AdminStoreResult<Admin
.map_err(|_| admin_store_error(ENTITY, invalid("invalid client key id")))?,
name: record.name,
label: record.label,
provider_kind: ProviderKind::new(record.provider_kind)
.map_err(|_| admin_store_error(ENTITY, invalid("invalid provider kind")))?,
groups: record
.groups
.into_iter()
.map(|group| {
Ok(AdminAccountGroupRef {
id: AccountGroupId::new(group.id)
.map_err(|_| admin_store_error(ENTITY, invalid("invalid group id")))?,
name: group.name,
color: gateway_admin::model::account_groups::AccountGroupColor::parse(
&group.color,
)
.ok_or_else(|| admin_store_error(ENTITY, invalid("invalid group color")))?,
enabled: group.enabled,
})
})
.collect::<AdminStoreResult<Vec<_>>>()?,
provider_kinds: record
.provider_kinds
.into_iter()
.map(|provider| {
ProviderKind::new(provider)
.map_err(|_| admin_store_error(ENTITY, invalid("invalid provider kind")))
})
.collect::<AdminStoreResult<Vec<_>>>()?,
prefix: record.prefix,
enabled: record.enabled,
limits: RateLimits {
@@ -849,20 +918,20 @@ pub(crate) async fn insert_client_api_key_in_transaction(
key.validate()?;
sqlx::query(
"insert into client_api_keys (
id, name, label, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, label, key, enabled, max_concurrency, requests_per_minute,
last_used_at, created_at, updated_at
) values ($1, $2, $3, $4, $5, true, $6, $7, null, now(), now())",
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now())",
)
.bind(&key.id)
.bind(&key.name)
.bind(&key.label)
.bind(&key.provider_kind)
.bind(&key.key)
.bind(to_i64(key.max_concurrency)?)
.bind(to_i64(key.requests_per_minute)?)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("insert client API key in transaction"))?;
replace_client_api_key_groups_in_transaction(transaction, &key.id, &key.group_ids).await?;
Ok(())
}
@@ -873,20 +942,20 @@ pub(crate) async fn update_client_api_key_in_transaction(
key.validate()?;
let result = sqlx::query(
"update client_api_keys
set name = $2, label = $3, provider_kind = $4, max_concurrency = $5,
requests_per_minute = $6, updated_at = now()
set name = $2, label = $3, max_concurrency = $4,
requests_per_minute = $5, updated_at = now()
where id = $1",
)
.bind(&key.id)
.bind(&key.name)
.bind(&key.label)
.bind(&key.provider_kind)
.bind(to_i64(key.max_concurrency)?)
.bind(to_i64(key.requests_per_minute)?)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("update client API key in transaction"))?;
require_changed(result.rows_affected(), &key.id)
require_changed(result.rows_affected(), &key.id)?;
replace_client_api_key_groups_in_transaction(transaction, &key.id, &key.group_ids).await
}
pub(crate) async fn set_client_api_key_enabled_in_transaction(
@@ -929,6 +998,59 @@ fn require_changed(rows_affected: u64, id: &str) -> StoreResult<()> {
}
}
async fn replace_client_api_key_groups_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
key_id: &str,
group_ids: &[String],
) -> StoreResult<()> {
validate_group_ids(group_ids)?;
if !group_ids.is_empty() {
let count = sqlx::query_scalar::<_, i64>(
"select count(*)::bigint from account_groups where id = any($1::text[])",
)
.bind(group_ids)
.fetch_one(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("validate client API key groups"))?;
if usize::try_from(count).ok() != Some(group_ids.len()) {
return Err(StoreError::NotFound {
entity: "account group",
id: "client API key group selection".to_owned(),
});
}
}
sqlx::query("delete from client_api_key_groups where client_api_key_id = $1")
.bind(key_id)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("delete client API key groups"))?;
if !group_ids.is_empty() {
sqlx::query(
"insert into client_api_key_groups
(client_api_key_id, account_group_id, created_at)
select $1, group_id, now() from unnest($2::text[]) group_id",
)
.bind(key_id)
.bind(group_ids)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("insert client API key groups"))?;
}
Ok(())
}
fn validate_group_ids(group_ids: &[String]) -> StoreResult<()> {
if group_ids.len() > 1000
|| group_ids.iter().collect::<BTreeSet<_>>().len() != group_ids.len()
|| group_ids
.iter()
.any(|id| AccountGroupId::new(id.clone()).is_err())
{
return Err(invalid("group IDs are invalid or duplicated"));
}
Ok(())
}
fn validate_key(key: &str) -> StoreResult<()> {
let valid = key.len() == KEY_LENGTH
&& key.starts_with("sk_")
@@ -945,45 +1067,53 @@ fn validate_key(key: &str) -> StoreResult<()> {
}
fn client_secret_from_row(
row: (String, String, String, bool, i64, i64),
row: (String, String, bool, i64, i64),
) -> StoreResult<ClientApiKeySecret> {
Ok(ClientApiKeySecret {
id: row.0,
key: row.1,
provider_kind: row.2,
enabled: row.3,
max_concurrency: to_u64(row.4)?,
requests_per_minute: to_u64(row.5)?,
enabled: row.2,
max_concurrency: to_u64(row.3)?,
requests_per_minute: to_u64(row.4)?,
})
}
type ClientRecordRow = (
String,
String,
Option<String>,
String,
String,
bool,
i64,
i64,
Option<DateTime<Utc>>,
DateTime<Utc>,
DateTime<Utc>,
);
fn client_record_from_row(row: ClientRecordRow) -> StoreResult<ClientApiKeyRecord> {
fn client_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<ClientApiKeyRecord> {
use sqlx::Row as _;
let groups: serde_json::Value = row
.try_get("groups")
.map_err(|_| invalid("invalid groups"))?;
Ok(ClientApiKeyRecord {
id: row.0,
name: row.1,
label: row.2,
provider_kind: row.3,
prefix: row.4,
enabled: row.5,
max_concurrency: to_u64(row.6)?,
requests_per_minute: to_u64(row.7)?,
last_used_at: row.8,
created_at: row.9,
updated_at: row.10,
id: row.try_get("id").map_err(|_| invalid("invalid id"))?,
name: row.try_get("name").map_err(|_| invalid("invalid name"))?,
label: row.try_get("label").map_err(|_| invalid("invalid label"))?,
groups: serde_json::from_value(groups).map_err(|_| invalid("invalid groups"))?,
provider_kinds: row
.try_get("provider_kinds")
.map_err(|_| invalid("invalid provider kinds"))?,
prefix: row
.try_get("prefix")
.map_err(|_| invalid("invalid prefix"))?,
enabled: row
.try_get("enabled")
.map_err(|_| invalid("invalid enabled"))?,
max_concurrency: to_u64(
row.try_get("max_concurrency")
.map_err(|_| invalid("invalid max concurrency"))?,
)?,
requests_per_minute: to_u64(
row.try_get("requests_per_minute")
.map_err(|_| invalid("invalid requests per minute"))?,
)?,
last_used_at: row
.try_get("last_used_at")
.map_err(|_| invalid("invalid last used at"))?,
created_at: row
.try_get("created_at")
.map_err(|_| invalid("invalid created at"))?,
updated_at: row
.try_get("updated_at")
.map_err(|_| invalid("invalid updated at"))?,
})
}
@@ -14,6 +14,7 @@ use gateway_core::engine::{
UpstreamSendState as CoreUpstreamSendState,
};
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
use gateway_core::routing::{AccountRoutingScopeKind, AccountRoutingSnapshot};
use crate::{
ConflictKind, DecimalAmount, StoreError, StoreResult, postgres_unavailable, require_nonempty,
@@ -83,6 +84,9 @@ pub struct NewModelRequest {
pub client_api_key_id: Option<String>,
pub client_api_key_ref: String,
pub config_revision: u64,
pub routing_scope: String,
pub routing_group_refs: Vec<String>,
pub routing_group_names_snapshot: Value,
pub protocol: String,
pub operation: String,
pub endpoint: String,
@@ -109,6 +113,11 @@ impl NewModelRequest {
require_nonempty(ENTITY, "endpoint", &self.endpoint)?;
require_nonempty(ENTITY, "client_transport", &self.client_transport)?;
require_nonempty(ENTITY, "requested_model_id", &self.requested_model_id)?;
validate_routing_snapshot(
&self.routing_scope,
&self.routing_group_refs,
&self.routing_group_names_snapshot,
)?;
if self.config_revision == 0 || self.started_at > self.deadline_at {
return Err(invalid(
"revision and deadline violate the frozen constraints",
@@ -362,13 +371,15 @@ impl ModelRequestRepository for PgExecutionStore {
sqlx::query(
"insert into model_requests (
id, client_api_key_id, client_api_key_ref, config_revision, protocol,
routing_scope, routing_group_refs, routing_group_names_snapshot,
operation, endpoint, client_transport, requested_model_id,
client_ip, user_agent, reasoning_effort,
reasoning_preset, request_kind, subagent_kind, compact,
image_generation_requested, started_at, deadline_at
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10::inet, $11, $12,
$13, $14, $15, $16, $17, $18, $19
$1, $2, $3, $4, $5, $6, $7, $8,
$9, $10, $11, $12, $13::inet, $14, $15,
$16, $17, $18, $19, $20, $21, $22
)",
)
.bind(request.id)
@@ -376,6 +387,9 @@ impl ModelRequestRepository for PgExecutionStore {
.bind(request.client_api_key_ref)
.bind(to_i64(request.config_revision, "config_revision")?)
.bind(request.protocol)
.bind(request.routing_scope)
.bind(request.routing_group_refs)
.bind(sqlx::types::Json(request.routing_group_names_snapshot))
.bind(request.operation)
.bind(request.endpoint)
.bind(request.client_transport)
@@ -409,6 +423,7 @@ impl ModelRequestRepository for PgExecutionStore {
sqlx::query(
"insert into model_requests (
id, client_api_key_id, client_api_key_ref, config_revision, protocol,
routing_scope, routing_group_refs, routing_group_names_snapshot,
operation, endpoint, client_transport, requested_model_id,
client_ip, user_agent, reasoning_effort,
reasoning_preset, request_kind, subagent_kind, compact,
@@ -419,13 +434,14 @@ impl ModelRequestRepository for PgExecutionStore {
upstream_model_id, upstream_transport, http_version,
attempt_count, upstream_send_state
) values (
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10::inet, $11, $12,
$13, $14, $15, $16, $17, $18, $19,
$20, $21, $22,
(select name from provider_accounts where id = $21),
(select email from provider_accounts where id = $21),
(select authentication_kind from provider_accounts where id = $21),
$23, $24, $25, 1, 'not_sent'
$1, $2, $3, $4, $5, $6, $7, $8,
$9, $10, $11, $12, $13::inet, $14, $15,
$16, $17, $18, $19, $20, $21, $22,
$23, $24, $25,
(select name from provider_accounts where id = $24),
(select email from provider_accounts where id = $24),
(select authentication_kind from provider_accounts where id = $24),
$26, $27, $28, 1, 'not_sent'
)",
)
.bind(request.id)
@@ -433,6 +449,9 @@ impl ModelRequestRepository for PgExecutionStore {
.bind(request.client_api_key_ref)
.bind(to_i64(request.config_revision, "config_revision")?)
.bind(request.protocol)
.bind(request.routing_scope)
.bind(request.routing_group_refs)
.bind(sqlx::types::Json(request.routing_group_names_snapshot))
.bind(request.operation)
.bind(request.endpoint)
.bind(request.client_transport)
@@ -1045,6 +1064,8 @@ fn core_store_error(error: StoreError) -> CoreStoreError {
}
fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
let (routing_scope, routing_group_refs, routing_group_names_snapshot) =
routing_snapshot_row(&request.routing);
NewModelRequest {
id: request.id.as_str().to_owned(),
client_api_key_id: request
@@ -1053,6 +1074,9 @@ fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
.map(|id| id.as_str().to_owned()),
client_api_key_ref: request.client_api_key_ref.as_str().to_owned(),
config_revision: request.config_revision.get(),
routing_scope,
routing_group_refs,
routing_group_names_snapshot,
protocol: request.protocol,
operation: request.operation.as_str().to_owned(),
endpoint: request.endpoint,
@@ -1071,6 +1095,48 @@ fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
}
}
fn routing_snapshot_row(snapshot: &AccountRoutingSnapshot) -> (String, Vec<String>, Value) {
let groups = snapshot.groups_snapshot();
(
snapshot.kind().as_str().to_owned(),
groups
.iter()
.map(|group| group.id().as_str().to_owned())
.collect(),
Value::Array(
groups
.iter()
.map(|group| Value::String(group.name().to_owned()))
.collect(),
),
)
}
fn validate_routing_snapshot(scope: &str, refs: &[String], names: &Value) -> StoreResult<()> {
let Some(names) = names.as_array() else {
return Err(invalid("routing group names snapshot must be an array"));
};
let valid = match scope {
value if value == AccountRoutingScopeKind::All.as_str() => {
refs.is_empty() && names.is_empty()
}
value if value == AccountRoutingScopeKind::Groups.as_str() => {
!refs.is_empty()
&& refs.len() == names.len()
&& refs.iter().all(|id| !id.is_empty())
&& names
.iter()
.all(|name| name.as_str().is_some_and(|name| !name.is_empty()))
}
_ => false,
};
if valid {
Ok(())
} else {
Err(invalid("invalid routing scope snapshot"))
}
}
fn attempt_start_row(attempt: CoreAttemptRecord) -> ModelRequestAttemptStart {
ModelRequestAttemptStart {
model_request_id: attempt.request_id.as_str().to_owned(),
@@ -7,6 +7,7 @@ use crate::{
Revision, StoreBackend, StoreError, StorePoolConfig, StoreResult, postgres_unavailable,
};
mod account_groups;
mod admin_security_audit;
mod admission_recovery;
mod backup;
@@ -20,6 +21,7 @@ mod retention;
mod runtime_settings;
mod snapshot;
pub use account_groups::*;
pub use admin_security_audit::*;
pub use admission_recovery::*;
pub use backup::*;
@@ -251,7 +253,7 @@ impl PgControlPlaneRepository {
async fn apply_targeted_mutation(
&self,
mutation: ControlPlaneMutation,
audit: AdminAuditEvent,
mut audit: AdminAuditEvent,
) -> StoreResult<Revision> {
let mut transaction = self
.pool
@@ -262,9 +264,34 @@ impl PgControlPlaneRepository {
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
match mutation {
ControlPlaneMutation::CreateClientApiKey(key) => {
audit.changed_fields.push(if key.group_ids.is_empty() {
"routing_scope:all".to_owned()
} else {
"routing_scope:groups".to_owned()
});
insert_client_api_key_in_transaction(&mut transaction, &key).await?;
}
ControlPlaneMutation::UpdateClientApiKey(key) => {
let previously_restricted = sqlx::query_scalar::<_, bool>(
"select exists(
select 1 from client_api_key_groups where client_api_key_id = $1
)",
)
.bind(&key.id)
.fetch_one(&mut *transaction)
.await
.map_err(|_| {
postgres_unavailable("load client API key routing scope for audit")
})?;
if previously_restricted && key.group_ids.is_empty() {
audit
.changed_fields
.push("routing_scope:groups->all".to_owned());
} else if !previously_restricted && !key.group_ids.is_empty() {
audit
.changed_fields
.push("routing_scope:all->groups".to_owned());
}
update_client_api_key_in_transaction(&mut transaction, &key).await?;
}
ControlPlaneMutation::SetClientApiKeyEnabled { id, enabled } => {
@@ -6,8 +6,8 @@ use std::sync::Arc;
use super::*;
use crate::postgres::{
PgProviderAccountRepository, ProviderAccountRepository, account_status_projection,
load_rate_limited_until,
PgProviderAccountRepository, ProviderAccountRepository, ProviderAccountSummary,
account_status_projection, load_rate_limited_until,
};
use crate::redis::{CredentialLeaseRepository as _, RedisCredentialLeaseRepository};
@@ -28,7 +28,7 @@ impl PgObservabilityRepository {
async fn account_status_snapshot(
&self,
observed_at: DateTime<Utc>,
) -> StoreResult<(ProviderAccountMetrics, Vec<String>)> {
) -> StoreResult<(ProviderAccountMetrics, Vec<ProviderAccountSummary>)> {
let accounts = PgProviderAccountRepository::new(self.pool.clone())
.list_provider_accounts(None, true)
.await?;
@@ -39,7 +39,7 @@ impl PgObservabilityRepository {
total: u64::try_from(accounts.len()).unwrap_or(u64::MAX),
..ProviderAccountMetrics::default()
};
let mut normal_account_ids = Vec::new();
let mut normal_accounts = Vec::new();
for account in &accounts {
let projection = account_status_projection(
account,
@@ -49,7 +49,7 @@ impl PgObservabilityRepository {
match projection.status {
gateway_core::engine::credential::AccountStatus::Normal => {
metrics.normal = metrics.normal.saturating_add(1);
normal_account_ids.push(account.id.clone());
normal_accounts.push(account.clone());
}
gateway_core::engine::credential::AccountStatus::QuotaExhausted => {
metrics.quota_exhausted = metrics.quota_exhausted.saturating_add(1);
@@ -65,7 +65,7 @@ impl PgObservabilityRepository {
}
}
}
Ok((metrics, normal_account_ids))
Ok((metrics, normal_accounts))
}
}
@@ -230,21 +230,42 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
&self,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<Option<admin_observability::DashboardRuntimeSlots>> {
let (_, normal_account_ids) = self
let (_, normal_accounts) = self
.repository
.account_status_snapshot(observed_at)
.await
.map_err(observability_error)?;
let active_accounts = u64::try_from(normal_account_ids.len())
.map_err(|_| observability_error(invalid("active account count overflows u64")))?;
if normal_account_ids.is_empty() {
let inherited_accounts = u64::try_from(
normal_accounts
.iter()
.filter(|account| account.concurrency_limit.is_none())
.count(),
)
.map_err(|_| observability_error(invalid("inherited account count overflows u64")))?;
let overridden_slots = normal_accounts.iter().fold(0_u64, |total, account| {
total.saturating_add(
account
.concurrency_limit
.map_or(0, |limit| u64::from(limit.get())),
)
});
if normal_accounts.is_empty() {
return Ok(Some(admin_observability::DashboardRuntimeSlots {
active_accounts,
inherited_accounts,
overridden_slots,
used_slots: Some(0),
}));
}
let normal_account_ids = normal_accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let Some(runtime_signals) = &self.runtime_signals else {
return Ok(None);
return Ok(Some(admin_observability::DashboardRuntimeSlots {
inherited_accounts,
overridden_slots,
used_slots: None,
}));
};
let signals = match runtime_signals
.credential_runtime_signals(&normal_account_ids)
@@ -253,7 +274,8 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
Ok(signals) => signals,
Err(_) => {
return Ok(Some(admin_observability::DashboardRuntimeSlots {
active_accounts,
inherited_accounts,
overridden_slots,
used_slots: None,
}));
}
@@ -262,7 +284,8 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
total.saturating_add(u64::from(signal.in_flight))
});
Ok(Some(admin_observability::DashboardRuntimeSlots {
active_accounts,
inherited_accounts,
overridden_slots,
used_slots: Some(used_slots),
}))
}
@@ -436,6 +436,9 @@ pub(crate) fn admin_usage_record(
id: record.id,
client_api_key_ref: record.client_api_key_ref,
config_revision: record.config_revision,
routing_scope: record.routing_scope,
routing_group_refs: record.routing_group_refs,
routing_group_names_snapshot: record.routing_group_names_snapshot,
protocol: record.protocol,
operation: record.operation,
endpoint: record.endpoint,
@@ -655,6 +658,13 @@ pub(crate) fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<
id: get(row, "id")?,
client_api_key_ref: get(row, "client_api_key_ref")?,
config_revision: unsigned(row, "config_revision")?,
routing_scope: get(row, "routing_scope")?,
routing_group_refs: get(row, "routing_group_refs")?,
routing_group_names_snapshot: get::<sqlx::types::Json<Vec<String>>>(
row,
"routing_group_names_snapshot",
)?
.0,
protocol: get(row, "protocol")?,
operation: get(row, "operation")?,
endpoint: get(row, "endpoint")?,
@@ -497,6 +497,9 @@ pub struct UsageRecord {
pub id: String,
pub client_api_key_ref: String,
pub config_revision: u64,
pub routing_scope: String,
pub routing_group_refs: Vec<String>,
pub routing_group_names_snapshot: Vec<String>,
pub protocol: String,
pub operation: String,
pub endpoint: String,
@@ -103,7 +103,9 @@ pub(crate) fn push_completed_usage_fact_filter(query: &mut QueryBuilder<Postgres
}
pub(crate) const USAGE_RECORD_SELECT: &str =
"select mr.id, mr.client_api_key_ref, mr.config_revision, mr.protocol, mr.operation,
"select mr.id, mr.client_api_key_ref, mr.config_revision,
mr.routing_scope, mr.routing_group_refs, mr.routing_group_names_snapshot,
mr.protocol, mr.operation,
mr.endpoint, mr.client_transport, mr.requested_model_id,
mr.provider_kind, mr.provider_account_ref,
mr.provider_account_name_snapshot as provider_account_name,
@@ -231,6 +231,111 @@ impl PgAdminAccountStore {
credential_revision: Some(admin_revision(rotation.credential_revision)?),
})
}
async fn account_groups_by_account(
&self,
account_ids: &[String],
) -> AdminStoreResult<BTreeMap<String, Vec<AccountGroupRef>>> {
if account_ids.is_empty() {
return Ok(BTreeMap::new());
}
let rows = sqlx::query_as::<_, (String, String, String, String, bool)>(
"select m.provider_account_id, g.id, g.name, g.color, g.enabled
from account_group_accounts m
join account_groups g on g.id = m.account_group_id
where m.provider_account_id = any($1::text[])
order by m.provider_account_id, g.id",
)
.bind(account_ids)
.fetch_all(&self.pool)
.await
.map_err(|_| {
admin_store_error(
ENTITY,
postgres_unavailable("load account group references"),
)
})?;
let mut groups = BTreeMap::<String, Vec<AccountGroupRef>>::new();
for (account_id, group_id, name, color, enabled) in rows {
let id = AccountGroupId::new(group_id).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"persisted account group ID is invalid",
)
})?;
groups.entry(account_id).or_default().push(AccountGroupRef {
id,
name,
color: gateway_admin::model::account_groups::AccountGroupColor::parse(&color)
.ok_or_else(|| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"persisted account group color is invalid",
)
})?,
enabled,
});
}
Ok(groups)
}
async fn account_ids_matching_group(
&self,
filter: Option<&AccountGroupFilter>,
) -> AdminStoreResult<Option<BTreeSet<String>>> {
let Some(filter) = filter else {
return Ok(None);
};
let rows = match filter {
AccountGroupFilter::Group(group_id) => {
let exists = sqlx::query_scalar::<_, bool>(
"select exists(select 1 from account_groups where id = $1)",
)
.bind(group_id.as_str())
.fetch_one(&self.pool)
.await
.map_err(|_| {
admin_store_error(
ENTITY,
postgres_unavailable("validate account group filter"),
)
})?;
if !exists {
return Err(admin_store_error(
ENTITY,
StoreError::NotFound {
entity: "account group",
id: group_id.as_str().to_owned(),
},
));
}
sqlx::query_scalar::<_, String>(
"select provider_account_id from account_group_accounts
where account_group_id = $1",
)
.bind(group_id.as_str())
.fetch_all(&self.pool)
.await
}
AccountGroupFilter::Ungrouped => {
sqlx::query_scalar::<_, String>(
"select a.id from provider_accounts a
where not exists (
select 1 from account_group_accounts m
where m.provider_account_id = a.id
)",
)
.fetch_all(&self.pool)
.await
}
}
.map_err(|_| {
admin_store_error(ENTITY, postgres_unavailable("load account group filter"))
})?;
Ok(Some(rows.into_iter().collect()))
}
}
#[async_trait]
@@ -252,18 +357,23 @@ impl AccountStore for PgAdminAccountStore {
let rate_limited_until =
load_rate_limited_until(self.cooldowns.as_deref(), &accounts, now.into()).await;
let summary = admin_account_summary(&accounts, now, &rate_limited_until);
let group_matches = self
.account_ids_matching_group(query.group_filter.as_ref())
.await?;
let mut items = accounts
.into_iter()
.filter_map(|account| {
let until = rate_limited_until.get(&account.id).copied();
let projection = account_status_projection(&account, now.into(), until);
account_matches_admin_query(&account, projection.status, &query).then_some(
AdminAccountListItem {
account,
projection,
usage: None,
},
)
(account_matches_admin_query(&account, projection.status, &query)
&& group_matches
.as_ref()
.is_none_or(|matches| matches.contains(&account.id)))
.then_some(AdminAccountListItem {
account,
projection,
usage: None,
})
})
.collect::<Vec<_>>();
@@ -298,9 +408,20 @@ impl AccountStore for PgAdminAccountStore {
.into_iter()
.skip(offset)
.take(page_size)
.collect::<Vec<_>>();
let item_ids = items
.iter()
.map(|item| item.account.id.clone())
.collect::<Vec<_>>();
let mut groups_by_account = self.account_groups_by_account(&item_ids).await?;
let items = items
.into_iter()
.map(|item| {
let account_id = item.account.id.clone();
let mut account = admin_account_record(item.account)?;
account.groups = groups_by_account.remove(&account_id).unwrap_or_default();
Ok(AccountPageItem {
account: admin_account_record(item.account)?,
account,
projection: item.projection,
})
})
@@ -332,8 +453,14 @@ impl AccountStore for PgAdminAccountStore {
.get(account_id)
.copied();
let projection = account_status_projection(&record.summary, now.into(), rate_limited_until);
let account_id = record.summary.id.clone();
let mut groups = self
.account_groups_by_account(std::slice::from_ref(&account_id))
.await?;
let mut account = admin_account_record(record.summary)?;
account.groups = groups.remove(&account_id).unwrap_or_default();
Ok(Some(AccountPageItem {
account: admin_account_record(record.summary)?,
account,
projection,
}))
}
@@ -572,29 +699,99 @@ impl AccountStore for PgAdminAccountStore {
.await
}
async fn set_account_enabled(
async fn update_account(
&self,
command: SetAccountEnabled,
command: UpdateAccount,
context: &MutationContext,
) -> AdminStoreResult<AdminRevision> {
let scope = self.required_scope(&command.account_id).await?;
let action = if command.enabled { "enable" } else { "disable" };
self.accounts
.set_provider_account_enabled_admin(SetProviderAccountEnabled {
scope,
account_id: command.account_id.clone(),
) -> AdminStoreResult<AccountUpdateResult> {
let account_id = CoreProviderAccountId::new(command.account_id.clone()).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"invalid provider account ID",
)
})?;
let config_revision = self
.accounts
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
account_ids: vec![command.account_id.clone()],
enabled: command.enabled,
concurrency_limit: command.concurrency_limit,
weight: command.weight,
group_ids: command.group_ids,
audit: mutation_audit(
context,
action,
"update",
"provider_account",
&command.account_id,
vec!["enabled".to_owned()],
vec![
"enabled".to_owned(),
"concurrency_limit".to_owned(),
"weight".to_owned(),
"groups".to_owned(),
],
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(admin_revision)
.and_then(admin_revision)?;
Ok(AccountUpdateResult {
config_revision,
account_id,
})
}
async fn batch_update_accounts(
&self,
command: BatchUpdateAccounts,
context: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
let account_ids = command
.account_ids
.iter()
.map(|id| {
CoreProviderAccountId::new(id.clone()).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"invalid provider account ID",
)
})
})
.collect::<AdminStoreResult<Vec<_>>>()?;
let audit_target = if command.account_ids.len() == 1 {
command.account_ids[0].clone()
} else {
"provider_accounts".to_owned()
};
let config_revision = self
.accounts
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
account_ids: command.account_ids,
enabled: command.enabled,
concurrency_limit: command.concurrency_limit,
weight: command.weight,
group_ids: command.group_ids,
audit: mutation_audit(
context,
"batch_update",
"provider_account",
&audit_target,
vec![
"enabled".to_owned(),
"concurrency_limit".to_owned(),
"weight".to_owned(),
"groups".to_owned(),
],
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(admin_revision)?;
Ok(AccountsUpdateResult {
config_revision,
account_ids,
})
}
async fn delete_accounts(
@@ -31,6 +31,8 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.map(DateTime::<Utc>::from),
next_refresh_at: account.account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.account.enabled(),
concurrency_limit: account.account.concurrency_limit(),
weight: account.account.weight(),
credential_state: account.account.credential_state(),
credential_observed_at: Utc::now(),
})
@@ -146,6 +146,7 @@ pub(crate) fn admin_account_record(
"persisted Provider kind is invalid",
)
})?,
groups: Vec::new(),
name: summary.name,
email: summary.email,
upstream_user_id: summary.upstream_user_id,
@@ -157,6 +158,8 @@ pub(crate) fn admin_account_record(
access_token_expires_at: summary.access_token_expires_at,
next_refresh_at: summary.next_refresh_at,
enabled: summary.enabled,
concurrency_limit: summary.concurrency_limit,
weight: summary.weight,
credential_state: summary.credential_state,
credential_observed_at: summary.credential_observed_at,
quota: summary.quota,
@@ -184,6 +187,8 @@ pub(crate) fn prepared_account(
access_token_expires_at: credential.access_token_expires_at,
next_refresh_at: credential.next_refresh_at,
enabled: credential.enabled,
concurrency_limit: None,
weight: AccountWeight::DEFAULT,
credential_state: credential.credential_state,
credential_observed_at: credential.credential_observed_at,
})
@@ -11,12 +11,15 @@ use chrono::{DateTime, TimeDelta, Utc};
use gateway_admin::{
model::{
MutationContext, Revision as AdminRevision,
account_groups::AccountGroupRef,
accounts::{
AccountCost, AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
AccountPageItem, AccountRecord, AccountRequestBucket, AccountSort as AdminAccountSort,
AccountSortField as AdminAccountSortField, AccountStatus as AdminAccountStatus,
AccountSummary, AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult,
DeleteAccounts, SetAccountEnabled, SortDirection as AdminSortDirection,
AccountCost, AccountGroupFilter, AccountListQuery as AdminAccountListQuery,
AccountModelUsage, AccountPage, AccountPageItem, AccountRecord, AccountRequestBucket,
AccountSort as AdminAccountSort, AccountSortField as AdminAccountSortField,
AccountStatus as AdminAccountStatus, AccountSummary, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, SortDirection as AdminSortDirection,
UpdateAccount,
},
observability::{
CostCoverage as AdminCostCoverage, DecimalAmount as AdminDecimalAmount, TimeRange,
@@ -34,16 +37,16 @@ use gateway_admin::{
use sqlx::{PgPool, Postgres, Row, Transaction};
use gateway_core::engine::credential::{
AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision as CoreCredentialRevision, CredentialState, LoadedCredential,
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
ProviderAccountIdentity, ProviderAccountStore,
AccountConcurrencyLimit, AccountErrorReason, AccountStateChange, AccountWeight,
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision as CoreCredentialRevision,
CredentialState, LoadedCredential, NewProviderAccount as CoreNewProviderAccount,
OpaqueProviderData, PlaintextCredential, ProviderAccount as CoreProviderAccount,
ProviderAccountId as CoreProviderAccountId, ProviderAccountIdentity, ProviderAccountStore,
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaAccessChange, QuotaAccessState,
QuotaEvidence, QuotaObservation, QuotaState, QuotaWriteOutcome,
};
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
use gateway_core::routing::ProviderKind;
use gateway_core::routing::{AccountGroupId, ProviderKind};
use crate::{
ConflictKind, JsonObject, Revision, StoreError, StoreResult, admin_revision, admin_store_error,
@@ -56,9 +56,9 @@ pub trait ProviderAccountAdminRepository: Send + Sync {
command: RotateProviderAccount,
) -> StoreResult<ProviderAccountAdminRotation>;
async fn set_provider_account_enabled_admin(
async fn batch_update_provider_accounts_admin(
&self,
command: SetProviderAccountEnabled,
command: BatchUpdateProviderAccountsAdmin,
) -> StoreResult<Revision>;
async fn delete_provider_accounts_admin(
@@ -99,7 +99,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
let rows = sqlx::query(
"select 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, credential_state,
access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
credential_observed_at, quota_access_state, quota_evidence,
quota_access_observed_at, quota_reset_at,
quota_observed_at, last_error_reason, last_error_message, created_at, updated_at
@@ -127,11 +127,11 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
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,
credential_state, provider_quota_json,
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,
$14, null, $15, null, null, now(), greatest(now(), $15)
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
)",
)
.bind(account.id)
@@ -147,6 +147,8 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
.bind(account.access_token_expires_at)
.bind(account.next_refresh_at)
.bind(account.enabled)
.bind(account.concurrency_limit.map(|limit| i64::from(limit.get())))
.bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?)
.bind(credential_state.as_str())
.bind(account.credential_observed_at)
.execute(&self.pool)
@@ -471,12 +473,12 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
finish_admin_transaction(transaction, result, "provider account admin rotation").await
}
async fn set_provider_account_enabled_admin(
async fn batch_update_provider_accounts_admin(
&self,
command: SetProviderAccountEnabled,
command: BatchUpdateProviderAccountsAdmin,
) -> StoreResult<Revision> {
command.scope.validate()?;
require_nonempty(ENTITY, "account_id", &command.account_id)?;
validate_batch_update_account_ids(&command.account_ids)?;
validate_batch_update_group_ids(&command.group_ids)?;
let mut transaction = self
.pool
.begin()
@@ -484,11 +486,18 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
.map_err(|_| postgres_unavailable("begin provider account admin state change"))?;
let result = async {
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
set_provider_account_enabled_in_transaction(
update_provider_accounts_scheduling_in_transaction(
&mut transaction,
&command.scope,
&command.account_id,
&command.account_ids,
command.enabled,
command.concurrency_limit,
command.weight,
)
.await?;
replace_account_group_assignments_in_transaction(
&mut transaction,
&command.account_ids,
&command.group_ids,
)
.await?;
append_admin_audit_event_in_transaction(&mut transaction, command.audit, revision)
@@ -527,6 +536,53 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
}
}
async fn replace_account_group_assignments_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
account_ids: &[String],
group_ids: &[AccountGroupId],
) -> StoreResult<()> {
let group_ids = group_ids
.iter()
.map(|group_id| group_id.as_str().to_owned())
.collect::<Vec<_>>();
if !group_ids.is_empty() {
let known_group_count = sqlx::query_scalar::<_, i64>(
"select count(*)::bigint from account_groups where id = any($1::text[])",
)
.bind(&group_ids)
.fetch_one(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("validate account group assignment"))?;
if usize::try_from(known_group_count).ok() != Some(group_ids.len()) {
return Err(StoreError::NotFound {
entity: "account group",
id: "one or more group IDs".to_owned(),
});
}
}
sqlx::query("delete from account_group_accounts where provider_account_id = any($1::text[])")
.bind(account_ids)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("clear account group assignments"))?;
if group_ids.is_empty() {
return Ok(());
}
sqlx::query(
"insert into account_group_accounts
(account_group_id, provider_account_id, created_at)
select group_id, account_id, now()
from unnest($1::text[]) group_id
cross join unnest($2::text[]) account_id",
)
.bind(group_ids)
.bind(account_ids)
.execute(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("assign accounts to groups"))?;
Ok(())
}
pub(crate) async fn upsert_provider_account_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
account: &NewProviderAccount,
@@ -542,11 +598,11 @@ pub(crate) async fn upsert_provider_account_in_transaction(
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,
credential_state, provider_quota_json,
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,
$14, null, $15, null, null, now(), greatest(now(), $15)
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
)
on conflict (
provider_kind,
@@ -589,6 +645,8 @@ pub(crate) async fn upsert_provider_account_in_transaction(
.bind(account.access_token_expires_at)
.bind(account.next_refresh_at)
.bind(account.enabled)
.bind(account.concurrency_limit.map(|limit| i64::from(limit.get())))
.bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?)
.bind(credential_state.as_str())
.bind(account.credential_observed_at)
.fetch_optional(&mut **transaction)
@@ -691,23 +749,37 @@ pub(crate) async fn rotate_provider_account_in_transaction(
Revision::new(to_u64(next)?)
}
pub(crate) async fn set_provider_account_enabled_in_transaction(
pub(crate) async fn update_provider_accounts_scheduling_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
scope: &ProviderAccountAdminScope,
account_id: &str,
account_ids: &[String],
enabled: bool,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
) -> StoreResult<()> {
let result = sqlx::query(
"update provider_accounts set enabled = $3, updated_at = now()
where id = $1 and provider_kind = $2",
let updated = sqlx::query_scalar::<_, String>(
"update provider_accounts
set enabled = $2, concurrency_limit = $3, weight = $4, updated_at = now()
where id = any($1::text[])
returning id",
)
.bind(account_id)
.bind(&scope.provider_kind)
.bind(account_ids)
.bind(enabled)
.execute(&mut **transaction)
.bind(concurrency_limit.map(|limit| i64::from(limit.get())))
.bind(i16::try_from(weight.get()).map_err(|_| invalid("invalid weight"))?)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("set provider account state in admin transaction"))?;
require_admin_account_changed(result.rows_affected(), account_id)
.map_err(|_| postgres_unavailable("set provider accounts state in admin transaction"))?
.into_iter()
.collect::<BTreeSet<_>>();
let expected = account_ids.iter().cloned().collect::<BTreeSet<_>>();
if updated == expected {
Ok(())
} else {
Err(StoreError::NotFound {
entity: ENTITY,
id: "one or more provider account IDs".to_owned(),
})
}
}
pub(crate) async fn delete_provider_accounts_in_transaction(
@@ -749,20 +821,6 @@ pub(crate) fn validate_credential_update(update: &ProviderCredentialUpdate) -> S
Ok(())
}
pub(crate) fn require_admin_account_changed(
rows_affected: u64,
account_id: &str,
) -> StoreResult<()> {
if rows_affected == 1 {
Ok(())
} else {
Err(StoreError::NotFound {
entity: ENTITY,
id: account_id.to_owned(),
})
}
}
pub(crate) fn validate_admin_account_ids(account_ids: &[String]) -> StoreResult<()> {
if account_ids.is_empty() || account_ids.len() > MAX_ADMIN_IMPORT_BATCH {
return Err(invalid(
@@ -779,6 +837,38 @@ pub(crate) fn validate_admin_account_ids(account_ids: &[String]) -> StoreResult<
Ok(())
}
fn validate_batch_update_account_ids(account_ids: &[String]) -> StoreResult<()> {
const MAX_BATCH_UPDATE_ACCOUNTS: usize = 1000;
if account_ids.is_empty() || account_ids.len() > MAX_BATCH_UPDATE_ACCOUNTS {
return Err(invalid(
"account batch update must contain between 1 and 1000 IDs",
));
}
let mut unique = BTreeSet::new();
for account_id in account_ids {
require_nonempty(ENTITY, "account_id", account_id)?;
if !unique.insert(account_id.as_str()) {
return Err(invalid("account batch update contains duplicate IDs"));
}
}
Ok(())
}
fn validate_batch_update_group_ids(group_ids: &[AccountGroupId]) -> StoreResult<()> {
const MAX_BATCH_UPDATE_GROUPS: usize = 1000;
if group_ids.len() > MAX_BATCH_UPDATE_GROUPS {
return Err(invalid("account batch update contains too many group IDs"));
}
let mut unique = BTreeSet::new();
if group_ids
.iter()
.any(|group_id| !unique.insert(group_id.as_str()))
{
return Err(invalid("account batch update contains duplicate group IDs"));
}
Ok(())
}
pub(crate) async fn finish_admin_transaction<T>(
transaction: Transaction<'_, Postgres>,
result: StoreResult<T>,
@@ -1,6 +1,7 @@
//! `provider_accounts` row 模型、解码与校验。
use super::*;
use gateway_core::routing::AccountGroupId;
pub(crate) const ENTITY: &str = "provider account";
pub(crate) const CREDENTIALS_MAX_BYTES: usize = 256 * 1024;
@@ -105,6 +106,8 @@ pub struct ProviderAccountSummary {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
pub quota: QuotaState,
@@ -150,6 +153,8 @@ pub struct NewProviderAccount {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
}
@@ -274,10 +279,12 @@ impl fmt::Debug for RotateProviderAccount {
}
#[derive(Debug, Clone)]
pub struct SetProviderAccountEnabled {
pub scope: ProviderAccountAdminScope,
pub account_id: String,
pub struct BatchUpdateProviderAccountsAdmin {
pub account_ids: Vec<String>,
pub enabled: bool,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub group_ids: Vec<AccountGroupId>,
pub audit: AdminAuditEvent,
}
@@ -334,7 +341,7 @@ impl ProviderAccountStateUpdate {
pub(crate) const ACCOUNT_SELECT: &str = "select 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, credential_state,
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,
last_error_reason, last_error_message,
credential_observed_at, quota_observed_at, created_at, updated_at
@@ -342,7 +349,7 @@ pub(crate) const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email,
pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select 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, credential_state,
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,
last_error_reason, last_error_message,
credential_observed_at, quota_observed_at, created_at, updated_at
@@ -401,6 +408,7 @@ pub(crate) fn core_account_from_summary(
summary.last_error_reason,
summary.last_error_message,
)
.with_scheduling(summary.concurrency_limit, summary.weight)
.with_refresh_schedule(
summary.has_refresh_token,
summary.next_refresh_at.map(Into::into),
@@ -446,6 +454,18 @@ pub(crate) fn account_summary_from_row(
quota_reset_at.map(Into::into),
)
.ok_or_else(|| invalid("invalid persisted quota fact"))?;
let concurrency_limit = get::<Option<i64>>(&row, "concurrency_limit")?
.map(|value| {
u32::try_from(value)
.ok()
.and_then(AccountConcurrencyLimit::new)
.ok_or_else(|| invalid("invalid concurrency_limit"))
})
.transpose()?;
let weight = u16::try_from(get::<i16>(&row, "weight")?)
.ok()
.and_then(AccountWeight::new)
.ok_or_else(|| invalid("invalid weight"))?;
Ok(ProviderAccountSummary {
id: get(&row, "id")?,
provider_kind: get(&row, "provider_kind")?,
@@ -460,6 +480,8 @@ pub(crate) fn account_summary_from_row(
access_token_expires_at: get(&row, "access_token_expires_at")?,
next_refresh_at: get(&row, "next_refresh_at")?,
enabled: get(&row, "enabled")?,
concurrency_limit,
weight,
credential_state: parse_credential_state(&credential_state)?,
credential_observed_at: get(&row, "credential_observed_at")?,
quota,
@@ -4,10 +4,12 @@ use std::collections::BTreeMap;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::routing::{
ConfigRevision,
AccountGroupId, ConfigRevision,
snapshot::{
SnapshotClientPolicyFacts, SnapshotFacts, SnapshotSettingsFacts, SnapshotStoreError,
SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts, SnapshotClientPolicyFacts,
SnapshotFacts, SnapshotProviderAccountFacts, SnapshotSettingsFacts, SnapshotStoreError,
SnapshotStorePort,
},
};
@@ -41,6 +43,28 @@ pub struct RuntimeSnapshotData {
pub observed_current_revision: Revision,
pub settings: SnapshotRuntimeSettings,
pub client_api_keys: Vec<ClientApiKeySnapshot>,
pub account_groups: Vec<SnapshotAccountGroupData>,
pub provider_accounts: Vec<SnapshotProviderAccountData>,
pub group_memberships: Vec<SnapshotGroupMembershipData>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotAccountGroupData {
pub id: AccountGroupId,
pub name: String,
pub enabled: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotProviderAccountData {
pub id: String,
pub provider_kind: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotGroupMembershipData {
pub group_id: AccountGroupId,
pub account_id: String,
}
#[async_trait]
@@ -77,6 +101,9 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
let (config_revision, settings) = load_settings(&mut transaction).await?;
let client_api_keys = load_client_keys(&mut transaction).await?;
let account_groups = load_account_groups(&mut transaction).await?;
let provider_accounts = load_provider_accounts(&mut transaction).await?;
let group_memberships = load_group_memberships(&mut transaction).await?;
transaction
.commit()
.await
@@ -89,6 +116,9 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
observed_current_revision,
settings,
client_api_keys,
account_groups,
provider_accounts,
group_memberships,
})
}
@@ -151,16 +181,44 @@ impl SnapshotStorePort for PgRuntimeSnapshotRepository {
SnapshotClientPolicyFacts::new(
key.id,
key.plaintext_key,
key.provider_kind,
key.group_ids,
key.limits,
)
})
.collect();
let account_groups = data
.account_groups
.into_iter()
.map(|group| SnapshotAccountGroupFacts::new(group.id, group.name, group.enabled))
.collect();
let provider_accounts = data
.provider_accounts
.into_iter()
.map(|account| {
ProviderAccountId::new(account.id)
.map(|id| SnapshotProviderAccountFacts::new(id, account.provider_kind))
.map_err(|_| SnapshotStoreError::unavailable())
})
.collect::<Result<Vec<_>, _>>()?;
let group_memberships = data
.group_memberships
.into_iter()
.map(|membership| {
ProviderAccountId::new(membership.account_id)
.map(|account_id| {
SnapshotAccountGroupMemberFacts::new(membership.group_id, account_id)
})
.map_err(|_| SnapshotStoreError::unavailable())
})
.collect::<Result<Vec<_>, _>>()?;
Ok(SnapshotFacts::new(
config_revision,
observed_current_revision,
settings,
client_policies,
account_groups,
provider_accounts,
group_memberships,
))
})
}
@@ -224,9 +282,16 @@ async fn load_settings(
async fn load_client_keys(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<ClientApiKeySnapshot>> {
let rows = sqlx::query_as::<_, (String, String, String, i64, i64)>(
"select id, key, provider_kind, max_concurrency, requests_per_minute
from client_api_keys where enabled order by id",
let rows = sqlx::query_as::<_, (String, String, Vec<String>, i64, i64)>(
"select k.id, k.key,
coalesce(array_agg(kg.account_group_id order by kg.account_group_id)
filter (where kg.account_group_id is not null), '{}') as group_ids,
k.max_concurrency, k.requests_per_minute
from client_api_keys k
left join client_api_key_groups kg on kg.client_api_key_id = k.id
where k.enabled
group by k.id
order by k.id",
)
.fetch_all(&mut **transaction)
.await
@@ -236,6 +301,63 @@ async fn load_client_keys(
.collect()
}
async fn load_account_groups(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<SnapshotAccountGroupData>> {
let rows = sqlx::query_as::<_, (String, String, bool)>(
"select id, name, enabled from account_groups order by id",
)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("load snapshot account groups"))?;
rows.into_iter()
.map(|(id, name, enabled)| {
Ok(SnapshotAccountGroupData {
id: AccountGroupId::new(id).map_err(|_| invalid("invalid account group id"))?,
name,
enabled,
})
})
.collect()
}
async fn load_provider_accounts(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<SnapshotProviderAccountData>> {
sqlx::query_as::<_, (String, String)>(
"select id, provider_kind from provider_accounts order by id",
)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("load snapshot provider accounts"))
.map(|rows| {
rows.into_iter()
.map(|(id, provider_kind)| SnapshotProviderAccountData { id, provider_kind })
.collect()
})
}
async fn load_group_memberships(
transaction: &mut Transaction<'_, Postgres>,
) -> StoreResult<Vec<SnapshotGroupMembershipData>> {
let rows = sqlx::query_as::<_, (String, String)>(
"select account_group_id, provider_account_id
from account_group_accounts order by account_group_id, provider_account_id",
)
.fetch_all(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("load snapshot group memberships"))?;
rows.into_iter()
.map(|(group_id, account_id)| {
Ok(SnapshotGroupMembershipData {
group_id: AccountGroupId::new(group_id)
.map_err(|_| invalid("invalid membership group id"))?,
account_id,
})
})
.collect()
}
fn revision_from_i64(value: i64) -> StoreResult<Revision> {
Revision::new(to_u64(value)?)
}
@@ -9,6 +9,7 @@ use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures::future::join_all;
use gateway_core::engine::credential::{AccountRuntimeSignals, ProviderAccountId};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
ProviderRefreshCapacityRequest, ProviderSchedulingLeaseRequest, ProviderSchedulingState,
@@ -329,10 +330,13 @@ impl RedisCredentialLeaseRepository {
}))
}
/// 原子推进跨进程共享的 Provider 调度游标。
pub async fn advance_scheduling_cursor(&self, provider_kind: &str) -> StoreResult<u64> {
require_nonempty("provider scheduling cursor", "provider_kind", provider_kind)?;
let key = self.scheduling_cursor_key(provider_kind)?;
/// 原子推进跨进程共享、按 Client Key 与 Provider 隔离的调度游标。
pub async fn advance_scheduling_cursor(
&self,
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> StoreResult<u64> {
let key = self.scheduling_cursor_key(client_api_key_id, provider_kind)?;
let mut connection = self.connection.clone();
let cursor = Script::new(ADVANCE_SCHEDULING_CURSOR_SCRIPT)
.key(key)
@@ -369,10 +373,17 @@ impl RedisCredentialLeaseRepository {
))
}
fn scheduling_cursor_key(&self, provider_kind: &str) -> StoreResult<String> {
let fingerprint = resource_fingerprint("provider scheduling cursor", provider_kind)?;
fn scheduling_cursor_key(
&self,
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> StoreResult<String> {
let fingerprint = resource_fingerprint(
"provider scheduling cursor",
&format!("{}\0{}", client_api_key_id.as_str(), provider_kind.as_str()),
)?;
Ok(format!(
"{}:scheduler:provider:{{{fingerprint}}}:cursor",
"{}:scheduler:client-provider:{{{fingerprint}}}:cursor",
self.namespace
))
}
@@ -475,10 +486,11 @@ impl RedisProviderLeaseCoordinator {
async fn next_scheduling_cursor(
&self,
client_api_key_id: &ClientApiKeyId,
provider_kind: &ProviderKind,
) -> Result<u64, ProviderStoreError> {
self.repository
.advance_scheduling_cursor(provider_kind.as_str())
.advance_scheduling_cursor(client_api_key_id, provider_kind)
.await
.map_err(|_| provider_unavailable("advance scheduling cursor"))
}
@@ -589,12 +601,15 @@ impl RedisProviderLeaseCoordinator {
impl ProviderLeasePort for RedisProviderLeaseCoordinator {
fn load_state<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
provider_kind: &'a ProviderKind,
accounts: &'a [ProviderAccountId],
) -> futures::future::BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
Box::pin(async move {
let signals = self.load_signals(accounts).await?;
let round_robin_cursor = self.next_scheduling_cursor(provider_kind).await?;
let round_robin_cursor = self
.next_scheduling_cursor(client_api_key_id, provider_kind)
.await?;
Ok(ProviderSchedulingState::new(signals, round_robin_cursor))
})
}
@@ -1,8 +1,9 @@
//! 旧版 Responses continuation 的 Redis 会话亲和记录。
//! Responses continuation 的 Redis 会话亲和记录。
//!
//! 该存储只保存可丢失的 response → Provider/account pin;PostgreSQL 审计记录
//! 不是 continuation 的可用性前提。键只含 response ID 的不可逆指纹,Provider
//! 私有状态作为不透明 JSON 保存,Store 不解释其内容。
//! 不是 continuation 的可用性前提。键只含 response ID 的不可逆指纹;Client Key
//! 所有权写入 payload 并在读取时严格校验,使跨 Key 复用能够 fail closed。
//! Provider 私有状态作为不透明 JSON 保存,Store 不解释其内容。
use std::time::{Duration, SystemTime, UNIX_EPOCH};
@@ -13,6 +14,7 @@ use gateway_core::engine::continuation::{
};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::operation::ProviderSessionState;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::ProviderKind;
use redis::aio::ConnectionManager;
use serde::{Deserialize, Serialize};
@@ -72,8 +74,9 @@ pub struct RedisNativeContinuationRepository {
}
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ContinuationWire {
client_api_key_id: String,
upstream_response_id: String,
provider: String,
account: String,
@@ -86,13 +89,13 @@ impl RedisNativeContinuationRepository {
pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult<Self> {
Ok(Self {
connection,
namespace: format!("{}:response-affinity:v1", namespace(key_namespace)?),
namespace: format!("{}:response-affinity:v2", namespace(key_namespace)?),
})
}
fn entry_fingerprint(&self, response_id: &str) -> Result<String, NativeContinuationStoreError> {
resource_fingerprint("native continuation", response_id)
.map_err(|error| NativeContinuationStoreError::new(error.to_string()))
.map_err(|error| NativeContinuationStoreError::invalid_data(error.to_string()))
}
fn entry_key(&self, fingerprint: &str) -> String {
@@ -111,6 +114,7 @@ impl RedisNativeContinuationRepository {
impl NativeContinuationPort for RedisNativeContinuationRepository {
fn resolve<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
previous_response_id: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async move {
@@ -124,10 +128,10 @@ impl NativeContinuationPort for RedisNativeContinuationRepository {
.query_async::<Option<String>>(&mut connection)
.await
.map_err(|error| {
NativeContinuationStoreError::new(format!("redis GET failed: {error}"))
NativeContinuationStoreError::unavailable(format!("redis GET failed: {error}"))
})?;
payload
.map(|payload| decode_pin(previous_response_id, &payload))
.map(|payload| decode_pin(client_api_key_id, previous_response_id, &payload))
.transpose()
})
}
@@ -143,20 +147,24 @@ impl NativeContinuationPort for RedisNativeContinuationRepository {
let fingerprint = self.entry_fingerprint(pin.previous_response_id().as_str())?;
let payload = encode_pin(&pin)?;
if payload.len() > MAX_RECORD_BYTES {
return Err(NativeContinuationStoreError::new(format!(
return Err(NativeContinuationStoreError::invalid_data(format!(
"pin payload is too large ({} bytes, limit {MAX_RECORD_BYTES})",
payload.len()
)));
}
let now_millis = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|_| NativeContinuationStoreError::new("system clock predates Unix epoch"))?
.map_err(|_| {
NativeContinuationStoreError::unavailable("system clock predates Unix epoch")
})?
.as_millis();
let now_millis = u64::try_from(now_millis).map_err(|_| {
NativeContinuationStoreError::new("system time is outside supported range")
NativeContinuationStoreError::unavailable("system time is outside supported range")
})?;
let ttl_millis = u64::try_from(CONTINUATION_TTL.as_millis()).map_err(|_| {
NativeContinuationStoreError::new("continuation TTL is outside supported range")
NativeContinuationStoreError::invalid_data(
"continuation TTL is outside supported range",
)
})?;
let cutoff = now_millis.saturating_sub(ttl_millis);
let ttl_seconds = CONTINUATION_TTL.as_secs();
@@ -175,7 +183,7 @@ impl NativeContinuationPort for RedisNativeContinuationRepository {
.invoke_async::<i64>(&mut connection)
.await
.map_err(|error| {
NativeContinuationStoreError::new(format!(
NativeContinuationStoreError::unavailable(format!(
"redis affinity script failed: {error}"
))
})?;
@@ -189,11 +197,12 @@ fn encode_pin(pin: &NativeContinuationPin) -> Result<String, NativeContinuationS
.session_state()
.is_some_and(|state| state.provider() != pin.provider().as_str())
{
return Err(NativeContinuationStoreError::new(
return Err(NativeContinuationStoreError::invalid_data(
"session state provider does not match pin provider",
));
}
serde_json::to_string(&ContinuationWire {
client_api_key_id: pin.client_api_key_id().as_str().to_owned(),
upstream_response_id: pin.upstream_response_id().as_str().to_owned(),
provider: pin.provider().as_str().to_owned(),
account: pin.account().as_str().to_owned(),
@@ -201,34 +210,44 @@ fn encode_pin(pin: &NativeContinuationPin) -> Result<String, NativeContinuationS
session_state: pin.session_state().cloned(),
})
.map_err(|error| {
NativeContinuationStoreError::new(format!("pin serialization failed: {error}"))
NativeContinuationStoreError::invalid_data(format!("pin serialization failed: {error}"))
})
}
fn decode_pin(
expected_client_api_key_id: &ClientApiKeyId,
previous_response_id: &PreviousResponseId,
payload: &str,
) -> Result<NativeContinuationPin, NativeContinuationStoreError> {
let wire: ContinuationWire = serde_json::from_str(payload)
.map_err(|error| NativeContinuationStoreError::new(format!("{error}")))?;
.map_err(|error| NativeContinuationStoreError::invalid_data(format!("{error}")))?;
let provider = ProviderKind::new(wire.provider).map_err(|_| {
NativeContinuationStoreError::new("pin payload contains an invalid provider")
NativeContinuationStoreError::invalid_data("pin payload contains an invalid provider")
})?;
let account = ProviderAccountId::new(wire.account).map_err(|_| {
NativeContinuationStoreError::new("pin payload contains an invalid account id")
NativeContinuationStoreError::invalid_data("pin payload contains an invalid account id")
})?;
let client_api_key_id = ClientApiKeyId::new(wire.client_api_key_id).map_err(|_| {
NativeContinuationStoreError::invalid_data(
"pin payload contains an invalid client API key id",
)
})?;
if client_api_key_id != *expected_client_api_key_id {
return Err(NativeContinuationStoreError::ownership_mismatch());
}
let upstream_response_id = PreviousResponseId::new(wire.upstream_response_id);
let scope = parse_scope(&wire.scope)?;
let mut pin = NativeContinuationPin::new(
previous_response_id.clone(),
upstream_response_id,
client_api_key_id,
provider.clone(),
account,
)
.with_scope(scope);
if let Some(state) = wire.session_state {
if state.provider() != provider.as_str() {
return Err(NativeContinuationStoreError::new(
return Err(NativeContinuationStoreError::invalid_data(
"pin payload session state provider does not match pin provider",
));
}
@@ -248,7 +267,7 @@ fn parse_scope(scope: &str) -> Result<NativeContinuationScope, NativeContinuatio
match scope {
"persisted" => Ok(NativeContinuationScope::Persisted),
"connection_local" => Ok(NativeContinuationScope::ConnectionLocal),
_ => Err(NativeContinuationStoreError::new(
_ => Err(NativeContinuationStoreError::invalid_data(
"pin payload contains an unsupported continuation scope",
)),
}
@@ -0,0 +1,313 @@
use std::collections::BTreeMap;
use gateway_admin::{
model::{
MutationActor, MutationContext, PageSize,
account_groups::{
AccountGroupColor, AccountGroupListQuery, DeleteAccountGroup, NewAccountGroup,
},
client_keys::NewClientKey,
client_keys::UpdateClientKey,
},
ports::store::{AccountGroupStore, AdminStoreErrorKind, ClientKeyStore},
};
use gateway_core::{
policy::{ClientApiKeyId, RateLimits},
routing::AccountGroupId,
};
use gateway_store::postgres::{PgAccountGroupRepository, PgAdminClientKeyStore};
use super::TestDatabase;
const MIXED_GROUP: &str = "grp_00000000000000000000000000000001";
const EMPTY_GROUP: &str = "grp_00000000000000000000000000000002";
#[tokio::test]
async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multiplication() {
let Some(database) = TestDatabase::create("account_group_aggregate").await else {
return;
};
seed_account(
&database.pool,
"acct_group_openai",
"openai",
"OpenAI Account",
)
.await;
seed_account(&database.pool, "acct_group_xai", "xai", "xAI Account").await;
sqlx::query(
"update provider_accounts set concurrency_limit = 4 where id = 'acct_group_openai'",
)
.execute(&database.pool)
.await
.expect("set account concurrency override");
let groups = PgAccountGroupRepository::new(database.pool.clone());
let keys = PgAdminClientKeyStore::new(database.pool.clone());
let mixed_group = group_id(MIXED_GROUP);
let empty_group = group_id(EMPTY_GROUP);
groups
.create_account_group(
NewAccountGroup {
id: mixed_group.clone(),
name: "Mixed Production".to_owned(),
description: Some("cross-provider".to_owned()),
color: group_color("#2563EBFF"),
},
&context("create-mixed"),
)
.await
.expect("create mixed account group");
groups
.create_account_group(
NewAccountGroup {
id: empty_group.clone(),
name: "Empty Pool".to_owned(),
description: None,
color: group_color("#06B6D4CC"),
},
&context("create-empty"),
)
.await
.expect("create empty account group");
assign_accounts(
&database.pool,
MIXED_GROUP,
&["acct_group_openai", "acct_group_xai"],
)
.await;
for (id, group_ids) in [
("key_group_one", vec![mixed_group.clone()]),
("key_group_two", vec![mixed_group.clone()]),
("key_empty_pool", vec![empty_group.clone()]),
("key_all_accounts", Vec::new()),
] {
keys.create_client_key(new_key(id, group_ids), &context(id))
.await
.expect("create scoped client key");
}
let page = groups
.list_account_groups(AccountGroupListQuery {
page: 1,
page_size: PageSize::new(20).expect("page size"),
search: None,
enabled: None,
})
.await
.expect("list account groups");
assert_eq!(page.total, 2);
let by_id = page
.items
.into_iter()
.map(|group| (group.id.to_string(), group))
.collect::<BTreeMap<_, _>>();
let mixed = by_id.get(MIXED_GROUP).expect("mixed group");
assert_eq!(mixed.member_count, 2);
assert_eq!(
mixed.provider_counts,
BTreeMap::from([("openai".to_owned(), 1), ("xai".to_owned(), 1)])
);
assert_eq!(mixed.client_key_count, 2);
assert_eq!(mixed.account_summary.available, 2);
assert_eq!(mixed.account_summary.limited, 0);
assert_eq!(mixed.account_summary.total, 2);
assert_eq!(mixed.capacity.used_slots, None);
assert_eq!(mixed.capacity.total_slots, 7);
assert_eq!(mixed.usage.today_usd.as_str(), "0");
assert_eq!(mixed.usage.total_usd.as_str(), "0");
let empty = by_id.get(EMPTY_GROUP).expect("empty group");
assert_eq!(empty.member_count, 0);
assert!(empty.provider_counts.is_empty());
assert_eq!(empty.client_key_count, 1);
let members = groups
.account_group_members(&mixed_group)
.await
.expect("load mixed group members");
assert_eq!(members.total, 2);
assert_eq!(
members
.items
.iter()
.map(|member| member.provider_kind.as_str())
.collect::<Vec<_>>(),
["openai", "xai"]
);
let all_key = keys
.reveal_client_key(&client_key_id("key_all_accounts"))
.await
.expect("reveal all-accounts key")
.expect("all-accounts key exists");
assert!(all_key.record.groups.is_empty());
assert_eq!(
all_key
.record
.provider_kinds
.iter()
.map(|kind| kind.as_str())
.collect::<Vec<_>>(),
["openai", "xai"]
);
let empty_pool_key = keys
.reveal_client_key(&client_key_id("key_empty_pool"))
.await
.expect("reveal empty-pool key")
.expect("empty-pool key exists");
assert_eq!(empty_pool_key.record.groups.len(), 1);
assert!(empty_pool_key.record.provider_kinds.is_empty());
let (scope_revision, widened) = keys
.update_client_key(
UpdateClientKey {
id: client_key_id("key_group_one"),
name: "key_group_one".to_owned(),
label: None,
group_ids: Vec::new(),
limits: RateLimits::unlimited(),
},
&context("widen-group-key"),
)
.await
.expect("widen restricted key to all accounts");
assert!(widened.groups.is_empty());
let scope_audit: Vec<String> = sqlx::query_scalar(
"select changed_fields from admin_audit_events
where admin_request_id = 'widen-group-key'",
)
.fetch_one(&database.pool)
.await
.expect("load scope widening audit");
assert!(scope_audit.contains(&"routing_scope:groups->all".to_owned()));
assert_eq!(current_revision(&database.pool).await, scope_revision.get());
let (restricted_revision, restricted) = keys
.update_client_key(
UpdateClientKey {
id: client_key_id("key_group_one"),
name: "key_group_one".to_owned(),
label: None,
group_ids: vec![empty_group],
limits: RateLimits::unlimited(),
},
&context("restrict-all-key"),
)
.await
.expect("restrict all-accounts key to groups");
assert_eq!(restricted.groups.len(), 1);
let restricted_audit: Vec<String> = sqlx::query_scalar(
"select changed_fields from admin_audit_events
where admin_request_id = 'restrict-all-key'",
)
.fetch_one(&database.pool)
.await
.expect("load scope restriction audit");
assert!(restricted_audit.contains(&"routing_scope:all->groups".to_owned()));
assert_eq!(
current_revision(&database.pool).await,
restricted_revision.get()
);
let revision_before_delete = current_revision(&database.pool).await;
let audit_before_delete = audit_count(&database.pool).await;
let error = groups
.delete_account_group(
DeleteAccountGroup { id: mixed_group },
&context("delete-referenced"),
)
.await
.expect_err("referenced group must not be deleted");
assert_eq!(error.kind(), AdminStoreErrorKind::Conflict);
assert_eq!(
current_revision(&database.pool).await,
revision_before_delete
);
assert_eq!(audit_count(&database.pool).await, audit_before_delete);
database.close().await;
}
fn new_key(id: &str, group_ids: Vec<AccountGroupId>) -> NewClientKey {
let marker = char::from(id.as_bytes().last().copied().unwrap_or(b'k'));
NewClientKey {
id: client_key_id(id),
name: id.to_owned(),
label: None,
group_ids,
limits: RateLimits::unlimited(),
plaintext: format!("sk_{}", marker.to_string().repeat(43)),
}
}
fn group_id(value: &str) -> AccountGroupId {
AccountGroupId::new(value).expect("valid account group ID")
}
fn group_color(value: &str) -> AccountGroupColor {
AccountGroupColor::parse(value).expect("valid account group color")
}
fn client_key_id(value: &str) -> ClientApiKeyId {
ClientApiKeyId::new(value).expect("valid client key ID")
}
fn context(request_id: &str) -> MutationContext {
MutationContext {
actor: MutationActor::System,
request_id: request_id.to_owned(),
}
}
async fn seed_account(pool: &sqlx::PgPool, id: &str, provider: &str, name: &str) {
sqlx::query(
"insert into provider_accounts (
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,
credential_state, credential_observed_at, created_at, updated_at
) values (
$1, $2, $3, null, $1 || '-user', null, null, 'oauth', '{}'::jsonb, 1,
false, null, null, true, 'ready', now(), now(), now()
)",
)
.bind(id)
.bind(provider)
.bind(name)
.execute(pool)
.await
.expect("seed provider account");
}
async fn assign_accounts(pool: &sqlx::PgPool, group_id: &str, account_ids: &[&str]) {
for account_id in account_ids {
sqlx::query(
"insert into account_group_accounts (
account_group_id, provider_account_id, created_at
)
values ($1, $2, now())",
)
.bind(group_id)
.bind(account_id)
.execute(pool)
.await
.expect("seed account group membership");
}
}
async fn current_revision(pool: &sqlx::PgPool) -> u64 {
let value =
sqlx::query_scalar::<_, i64>("select config_revision from runtime_settings where id = 1")
.fetch_one(pool)
.await
.expect("load config revision");
u64::try_from(value).expect("positive config revision")
}
async fn audit_count(pool: &sqlx::PgPool) -> u64 {
let value = sqlx::query_scalar::<_, i64>("select count(*) from admin_audit_events")
.fetch_one(pool)
.await
.expect("count audit rows");
u64::try_from(value).expect("non-negative audit count")
}
@@ -93,10 +93,12 @@ async fn seed_request(
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, outcome,
started_at, deadline_at, completed_at
started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
$1, 'key-recovery', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'coding', $2, $3, $4, $5
'http_sse', 'coding', $2, $3, $4, $5,
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(id)
@@ -23,7 +23,7 @@ fn client_key_requires_the_frozen_plaintext_format() {
id: "key-1".to_owned(),
name: "default".to_owned(),
label: None,
provider_kind: "openai".to_owned(),
group_ids: Vec::new(),
key: format!("sk_{}", "a".repeat(43)),
max_concurrency: 0,
requests_per_minute: 0,
@@ -44,9 +44,9 @@ async fn client_key_list_uses_safe_keyset_search_and_filtered_total() {
);
sqlx::query(
"insert into client_api_keys (
id, name, label, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, label, key, enabled, max_concurrency, requests_per_minute,
created_at, updated_at
) values ($1, $2, $2, 'openai', $3, true, 0, 0,
) values ($1, $2, $2, $3, true, 0, 0,
now() - ($4::bigint * interval '1 minute'), now())",
)
.bind(id)
@@ -186,9 +186,9 @@ async fn client_key_database_sort_is_stable_and_keeps_null_last_used_at_last() {
{
sqlx::query(
"insert into client_api_keys (
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, key, enabled, max_concurrency, requests_per_minute,
last_used_at, created_at, updated_at
) values ($1, $2, 'openai', $3, $4, 0, 0, $5::timestamptz, $6::timestamptz,
) values ($1, $2, $3, $4, 0, 0, $5::timestamptz, $6::timestamptz,
$6::timestamptz)",
)
.bind(id)
@@ -272,9 +272,9 @@ async fn client_key_usage_touch_should_preserve_the_latest_observed_timestamp()
};
sqlx::query(
"insert into client_api_keys (
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, key, enabled, max_concurrency, requests_per_minute,
created_at, updated_at
) values ('key_usage_touch', 'usage', 'openai', $1, true, 0, 0, now(), now())",
) values ('key_usage_touch', 'usage', $1, true, 0, 0, now(), now())",
)
.bind(format!("sk_{}", "t".repeat(43)))
.execute(&database.pool)
@@ -315,9 +315,9 @@ async fn dedicated_reveal_returns_plaintext_without_debug_exposure() {
let plaintext = format!("sk_{}", "r".repeat(43));
sqlx::query(
"insert into client_api_keys (
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
id, name, key, enabled, max_concurrency, requests_per_minute,
created_at, updated_at
) values ('key_reveal', 'reveal', 'openai', $1, true, 1, 2, now(), now())",
) values ('key_reveal', 'reveal', $1, true, 1, 2, now(), now())",
)
.bind(&plaintext)
.execute(&database.pool)
@@ -340,7 +340,7 @@ fn client_key_debug_redacts_plaintext() {
id: "key-1".to_owned(),
name: "default".to_owned(),
label: None,
provider_kind: "openai".to_owned(),
group_ids: Vec::new(),
key: secret.clone(),
max_concurrency: 0,
requests_per_minute: 0,
@@ -29,6 +29,9 @@ fn model_request_rejects_mismatched_client_key_live_id() {
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: serde_json::json!([]),
client_ip: None,
user_agent: None,
reasoning_effort: None,
@@ -76,6 +79,9 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: serde_json::json!([]),
client_ip: None,
user_agent: None,
reasoning_effort: None,
@@ -153,6 +159,66 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
database.close().await;
}
#[tokio::test]
async fn model_request_persists_group_routing_snapshot_without_live_group_foreign_keys() {
let Some(database) = TestDatabase::create("execution_group_routing_history").await else {
return;
};
let repository = PgExecutionStore::new(database.pool.clone());
let started_at = Utc::now();
repository
.insert_model_request(NewModelRequest {
id: "req_group_history".to_owned(),
client_api_key_id: None,
client_api_key_ref: "key_group_history".to_owned(),
config_revision: 7,
routing_scope: "groups".to_owned(),
routing_group_refs: vec![
"grp_00000000000000000000000000000001".to_owned(),
"grp_00000000000000000000000000000002".to_owned(),
],
routing_group_names_snapshot: serde_json::json!(["Production Pool", "Overflow Pool"]),
protocol: "openai".to_owned(),
operation: "responses".to_owned(),
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
client_ip: None,
user_agent: None,
reasoning_effort: None,
reasoning_preset: None,
request_kind: None,
subagent_kind: None,
compact: false,
image_generation_requested: false,
started_at,
deadline_at: started_at + Duration::seconds(30),
})
.await
.expect("insert grouped request history");
let stored: (String, Vec<String>, serde_json::Value) = sqlx::query_as(
"select routing_scope, routing_group_refs, routing_group_names_snapshot
from model_requests where id = 'req_group_history'",
)
.fetch_one(&database.pool)
.await
.expect("load grouped request history");
assert_eq!(stored.0, "groups");
assert_eq!(
stored.1,
[
"grp_00000000000000000000000000000001",
"grp_00000000000000000000000000000002",
]
);
assert_eq!(
stored.2,
serde_json::json!(["Production Pool", "Overflow Pool"])
);
database.close().await;
}
#[tokio::test]
async fn downstream_commit_should_atomically_record_http_status_once() {
let Some(database) = TestDatabase::create("execution_commit_status").await else {
@@ -483,9 +549,11 @@ async fn seed_running_request(pool: &sqlx::PgPool, id: &str) -> Result<(), sqlx:
"insert into model_requests (
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
client_transport, requested_model_id, provider_kind, provider_account_ref, cost_source,
started_at, deadline_at
started_at, deadline_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values ($1, 'key_status', 1, 'openai_responses', 'generate', '/v1/responses',
'http_json', 'status-model', 'openai', 'acct_status', 'unavailable', now(), now() + interval '1 minute')",
'http_json', 'status-model', 'openai', 'acct_status', 'unavailable', now(), now() + interval '1 minute',
'all', '{}'::text[], '[]'::jsonb)",
)
.bind(id)
.execute(pool)
@@ -7,6 +7,7 @@ use sqlx::{
};
use uuid::Uuid;
mod account_groups;
mod admin_security_audit;
mod admission_recovery;
mod backup;
@@ -146,6 +147,29 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
.fetch_one(&second)
.await
.expect("check removed raw response ID index");
let legacy_key_provider_column_exists = sqlx::query_scalar::<_, bool>(
"select exists (
select 1 from information_schema.columns
where table_schema = 'public'
and table_name = 'client_api_keys'
and column_name = 'provider_kind'
)",
)
.fetch_one(&second)
.await
.expect("check removed client key provider column");
let routing_history_columns = sqlx::query_scalar::<_, String>(
"select column_name from information_schema.columns
where table_schema = 'public'
and table_name = 'model_requests'
and column_name in (
'routing_scope', 'routing_group_refs', 'routing_group_names_snapshot'
)
order by column_name",
)
.fetch_all(&second)
.await
.expect("load routing history columns");
second.close().await;
sqlx::raw_sql(sqlx::AssertSqlSafe(format!(
@@ -156,7 +180,7 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
.expect("drop migration test database");
admin.close().await;
assert_eq!(first_table_count, 10);
assert_eq!(first_table_count, 13);
assert_eq!(
migration_count,
i64::try_from(TEST_MIGRATOR.iter().count())
@@ -164,6 +188,15 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
);
assert_eq!(response_id_types, ["bytea", "bytea"]);
assert!(!raw_response_id_index_exists);
assert!(!legacy_key_provider_column_exists);
assert_eq!(
routing_history_columns,
[
"routing_group_names_snapshot",
"routing_group_refs",
"routing_scope",
]
);
}
#[test]
@@ -550,6 +550,12 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
.await
.expect("fully filtered usage page");
assert_eq!(filtered.total, 1);
assert_eq!(filtered.items[0].routing_scope, "groups");
assert_eq!(filtered.items[0].routing_group_refs, ["grp_history"]);
assert_eq!(
filtered.items[0].routing_group_names_snapshot,
["Historical group"]
);
assert_eq!(filtered.items[0].service_tier.as_deref(), Some("priority"));
assert_eq!(
filtered.items[0].upstream_request_id.as_deref(),
@@ -584,6 +590,12 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
admin_observability::RequestOutcome::Succeeded
);
assert_eq!(detail.request.service_tier.as_deref(), Some("priority"));
assert_eq!(detail.request.routing_scope, "groups");
assert_eq!(detail.request.routing_group_refs, ["grp_history"]);
assert_eq!(
detail.request.routing_group_names_snapshot,
["Historical group"]
);
assert_eq!(detail.attempts.len(), 1);
assert_eq!(
detail.attempts[0].outcome,
@@ -1021,7 +1033,8 @@ async fn seed_observability_facts(
provider_account_authentication_kind_snapshot,
attempt_count, upstream_send_state, outcome, client_status_code,
input_tokens, output_tokens, total_tokens, cost_source, latency_ms,
started_at, deadline_at, completed_at
started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_uncommitted', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe',
@@ -1029,7 +1042,8 @@ async fn seed_observability_facts(
'primary', 'account@example.invalid', 'oauth',
1, 'sent', 'succeeded', 200,
900, 900, 1800, 'unavailable', 650,
$1 - interval '15 minutes', $1 + interval '10 minutes', $1 - interval '14 minutes'
$1 - interval '15 minutes', $1 + interval '10 minutes', $1 - interval '14 minutes',
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(now)
@@ -1051,7 +1065,8 @@ async fn seed_observability_facts(
image_input_tokens, image_output_tokens, total_tokens,
image_generation_requested, image_generation_succeeded,
cost_source, cost_amount, cost_currency, first_token_ms, latency_ms,
started_at, deadline_at, completed_at
started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_success', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe',
@@ -1062,7 +1077,8 @@ async fn seed_observability_facts(
'resp_observe_success', 'upstream_req_success', 'upstream_resp_success',
100, 20, 40, 3, 5, 31, 9, 120, true, true,
'provider_reported', 1.25, 'USD', 120, 900,
$1 - interval '20 minutes', $1 + interval '10 minutes', $1 - interval '19 minutes'
$1 - interval '20 minutes', $1 + interval '10 minutes', $1 - interval '19 minutes',
'groups', array['grp_history'], jsonb_build_array('Historical group')
)",
)
.bind(now)
@@ -1080,7 +1096,8 @@ async fn seed_observability_facts(
error_kind, provider_error_code, error_message, retry_after_ms,
input_tokens, cached_tokens, image_generation_requested,
image_generation_succeeded, cost_source, latency_ms,
started_at, deadline_at, completed_at
started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_failed', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe',
@@ -1089,7 +1106,8 @@ async fn seed_observability_facts(
'primary', 'account@example.invalid', 'oauth',
'sent', 'failed', 502, 429, 'rate_limited', 'rate_limit',
'upstream limited', 1000, 0, 0, true, false, 'unavailable', 700,
$1 - interval '10 minutes', $1 + interval '20 minutes', $1 - interval '9 minutes'
$1 - interval '10 minutes', $1 + interval '20 minutes', $1 - interval '9 minutes',
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(now)
@@ -1131,7 +1149,8 @@ async fn seed_calculated_billing_facts(
provider_account_authentication_kind_snapshot,
upstream_send_state, downstream_committed_at, outcome, client_status_code,
input_tokens, output_tokens, cached_tokens, cache_write_tokens, total_tokens, service_tier,
cost_source, cost_amount, cost_currency, started_at, deadline_at, completed_at
cost_source, cost_amount, cost_currency, started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_calculated', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
'http_sse', 'public-model', 'openai', 'acct_observe', 'acct_observe', 'gpt-5.5',
@@ -1139,7 +1158,8 @@ async fn seed_calculated_billing_facts(
'primary', 'account@example.invalid', 'oauth',
'sent', $1 - interval '29 minutes', 'succeeded', 200,
800, 200, 0, 0, 1000, 'priority', 'calculated', 1.25, 'USD',
$1 - interval '30 minutes', $1 + interval '10 minutes', $1 - interval '29 minutes'
$1 - interval '30 minutes', $1 + interval '10 minutes', $1 - interval '29 minutes',
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(now)
@@ -1154,7 +1174,7 @@ async fn seed_calculated_billing_facts(
provider_account_authentication_kind_snapshot,
upstream_send_state, outcome, client_status_code, input_tokens, output_tokens,
total_tokens, cost_source, cost_amount, cost_currency, started_at, deadline_at,
completed_at
completed_at, routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
'req_observe_calculated_uncommitted', 'key_observe', 1, 'openai', 'responses',
'/v1/responses', 'http_sse', 'public-model', 'openai', 'acct_observe',
@@ -1162,7 +1182,8 @@ async fn seed_calculated_billing_facts(
'primary', 'account@example.invalid', 'oauth',
'sent', 'succeeded', 200, 800, 200,
1000, 'calculated', 1.25, 'USD',
$1 - interval '40 minutes', $1 + interval '10 minutes', $1 - interval '39 minutes'
$1 - interval '40 minutes', $1 + interval '10 minutes', $1 - interval '39 minutes',
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(now)
@@ -6,8 +6,8 @@ use gateway_admin::{
MutationActor, MutationContext, PageSize,
accounts::{
AccountListQuery, AccountSort, AccountSortField, AccountStatus,
AccountUsageWindowQuery, DeleteAccounts, SetAccountEnabled as AdminSetAccountEnabled,
SortDirection,
AccountUsageWindowQuery, BatchUpdateAccounts, DeleteAccounts, SortDirection,
UpdateAccount,
},
observability::TimeRange,
provider_credentials::{CredentialListQuery, CredentialListWindow, CredentialStateFilter},
@@ -21,14 +21,14 @@ use gateway_core::engine::credential::{
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState,
QuotaWriteOutcome,
};
use gateway_core::routing::ProviderKind;
use gateway_core::routing::{AccountGroupId, ProviderKind};
use gateway_store::{
ConflictKind, JsonObject, Revision, StoreError,
postgres::{
AdminAuditActorKind, AdminAuditEvent, DeleteProviderAccounts, ImportProviderAccounts,
NewProviderAccount, PgAdminAccountStore, PgProviderAccountRepository,
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountRepository,
ProviderCredentialUpdate, RotateProviderAccount, SetProviderAccountEnabled,
AdminAuditActorKind, AdminAuditEvent, BatchUpdateProviderAccountsAdmin,
DeleteProviderAccounts, ImportProviderAccounts, NewProviderAccount, PgAdminAccountStore,
PgProviderAccountRepository, ProviderAccountAdminRepository, ProviderAccountAdminScope,
ProviderAccountRepository, ProviderCredentialUpdate, RotateProviderAccount,
UpdateProviderAccount,
},
};
@@ -392,6 +392,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
page: 1,
page_size: PageSize::new(2).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: Some(AccountSort {
@@ -431,6 +432,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
page: 1,
page_size: PageSize::new(2).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: None,
sort: Some(AccountSort {
@@ -454,6 +456,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
page: 1,
page_size: PageSize::new(10).expect("page size"),
provider_kind: Some(ProviderKind::new("openai").expect("Provider kind")),
group_filter: None,
search: Some("ALPHA@EXAMPLE".to_owned()),
status: Some(AccountStatus::Normal),
sort: None,
@@ -470,6 +473,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
page: 1,
page_size: PageSize::new(10).expect("page size"),
provider_kind: None,
group_filter: None,
search: None,
status: Some(AccountStatus::Error),
sort: None,
@@ -725,17 +729,20 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
request_id: "request_terminal_mutation".to_owned(),
};
let revision = store
.set_account_enabled(
AdminSetAccountEnabled {
let result = store
.update_account(
UpdateAccount {
account_id: "acct_terminal_mutation".to_owned(),
enabled: false,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
group_ids: Vec::new(),
},
&context,
)
.await
.expect("disable account atomically");
assert_eq!(revision.get(), 2);
assert_eq!(result.config_revision.get(), 2);
let enabled: bool = sqlx::query_scalar(
"select enabled from provider_accounts where id = 'acct_terminal_mutation'",
)
@@ -768,7 +775,16 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
assert_eq!(
audit_rows,
vec![
("disable".to_owned(), 2, vec!["enabled".to_owned()]),
(
"update".to_owned(),
2,
vec![
"enabled".to_owned(),
"concurrency_limit".to_owned(),
"weight".to_owned(),
"groups".to_owned(),
],
),
("delete".to_owned(), 3, Vec::new()),
]
);
@@ -776,6 +792,118 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
database.close().await;
}
#[tokio::test]
async fn terminal_batch_update_replaces_state_and_groups_once_or_rolls_back_everything() {
const GROUP_ID: &str = "grp_00000000000000000000000000000071";
let Some(database) = TestDatabase::create("provider_account_batch_update").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
for (account_id, upstream_user_id) in [
("acct_batch_a", "user-batch-a"),
("acct_batch_b", "user-batch-b"),
] {
repository
.insert_provider_account(account(account_id, upstream_user_id))
.await
.expect("insert batch account");
}
sqlx::query(
"insert into account_groups (id, name, description, color, enabled, created_at, updated_at)
values ($1, 'Batch group', null, '#2563EBFF', true, now(), now())",
)
.bind(GROUP_ID)
.execute(&database.pool)
.await
.expect("insert batch account group");
let store = PgAdminAccountStore::new(database.pool.clone(), None);
let account_ids = vec!["acct_batch_a".to_owned(), "acct_batch_b".to_owned()];
let context = MutationContext {
actor: MutationActor::System,
request_id: "request_batch_update".to_owned(),
};
let result = store
.batch_update_accounts(
BatchUpdateAccounts {
account_ids: account_ids.clone(),
enabled: false,
concurrency_limit: gateway_core::engine::credential::AccountConcurrencyLimit::new(
7,
),
weight: gateway_core::engine::credential::AccountWeight::new(25).expect("weight"),
group_ids: vec![AccountGroupId::new(GROUP_ID).expect("group ID")],
},
&context,
)
.await
.expect("batch update accounts");
assert_eq!(result.config_revision.get(), 2);
assert_eq!(result.account_ids.len(), 2);
let scheduling: Vec<(bool, Option<i64>, i16)> = sqlx::query_as(
"select enabled, concurrency_limit, weight
from provider_accounts where id = any($1::text[]) order by id",
)
.bind(&account_ids)
.fetch_all(&database.pool)
.await
.expect("load batch account state");
assert_eq!(scheduling, [(false, Some(7), 25), (false, Some(7), 25)]);
for account_id in &account_ids {
assert_eq!(
account_group_ids(&database.pool, account_id).await,
[GROUP_ID]
);
}
assert_eq!(audit_count(&database.pool, &context.request_id).await, 1);
let revision_before_failure = current_revision(&database.pool).await;
let audit_before_failure = audit_count(&database.pool, &context.request_id).await;
store
.batch_update_accounts(
BatchUpdateAccounts {
account_ids: account_ids.clone(),
enabled: true,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
group_ids: vec![
AccountGroupId::new("grp_00000000000000000000000000000072")
.expect("missing group ID"),
],
},
&context,
)
.await
.expect_err("missing group must roll back the batch");
assert_eq!(
current_revision(&database.pool).await,
revision_before_failure
);
assert_eq!(
audit_count(&database.pool, &context.request_id).await,
audit_before_failure
);
let scheduling: Vec<(bool, Option<i64>, i16)> = sqlx::query_as(
"select enabled, concurrency_limit, weight
from provider_accounts where id = any($1::text[]) order by id",
)
.bind(&account_ids)
.fetch_all(&database.pool)
.await
.expect("load rolled back account state");
assert_eq!(scheduling, [(false, Some(7), 25), (false, Some(7), 25)]);
for account_id in &account_ids {
assert_eq!(
account_group_ids(&database.pool, account_id).await,
[GROUP_ID]
);
}
database.close().await;
}
#[tokio::test]
async fn terminal_admin_delete_removes_enabled_accounts_in_one_transaction() {
let Some(database) = TestDatabase::create("provider_account_enabled_delete").await else {
@@ -916,6 +1044,8 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
access_token_expires_at: Some(Utc::now() + TimeDelta::minutes(5)),
next_refresh_at: None,
enabled: true,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
credential_state: CredentialState::Ready,
credential_observed_at: Utc::now(),
})
@@ -1190,10 +1320,12 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
assert_eq!(unchanged.4.as_deref(), Some("workspace-admin-rebound"));
let revision = repository
.set_provider_account_enabled_admin(SetProviderAccountEnabled {
scope: scope.clone(),
account_id: "acct_admin_a".to_owned(),
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
account_ids: vec!["acct_admin_a".to_owned()],
enabled: false,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
group_ids: Vec::new(),
audit: audit("audit_account_disable", "disable", "acct_admin_a"),
})
.await
@@ -1230,6 +1362,116 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
database.close().await;
}
#[tokio::test]
async fn provider_account_import_and_reauthorization_preserve_existing_memberships() {
const GROUP_ID: &str = "grp_00000000000000000000000000000091";
let Some(database) = TestDatabase::create("provider_account_group_assignment").await else {
return;
};
sqlx::query(
"insert into account_groups
(id, name, description, color, enabled, created_at, updated_at)
values ($1, 'Credential group', null, '#2563EBFF', true, now(), now())",
)
.bind(GROUP_ID)
.execute(&database.pool)
.await
.expect("seed account group");
let repository = PgProviderAccountRepository::new(database.pool.clone());
let scope = ProviderAccountAdminScope {
provider_kind: "openai".to_owned(),
};
let imported = repository
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![account("acct_grouped_import", "user-grouped-import")],
audit: audit("audit_grouped_import", "import", "acct_grouped_import"),
})
.await
.expect("import account without a group");
assert_eq!(imported.config_revision.get(), 2);
assert!(
account_group_ids(&database.pool, "acct_grouped_import")
.await
.is_empty()
);
assert_eq!(
audit_count(&database.pool, "request:audit_grouped_import").await,
1
);
sqlx::query(
"insert into account_group_accounts (account_group_id, provider_account_id, created_at)
values ($1, $2, now())",
)
.bind(GROUP_ID)
.bind("acct_grouped_import")
.execute(&database.pool)
.await
.expect("assign existing account group");
sqlx::query("update provider_accounts set concurrency_limit = 9, weight = 40 where id = $1")
.bind("acct_grouped_import")
.execute(&database.pool)
.await
.expect("set account scheduling before reimport");
repository
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![account("acct_reimport_candidate", "user-grouped-import")],
audit: audit("audit_grouped_reimport", "import", "acct_grouped_import"),
})
.await
.expect("reimport existing identity");
let scheduling: (Option<i64>, i16) =
sqlx::query_as("select concurrency_limit, weight from provider_accounts where id = $1")
.bind("acct_grouped_import")
.fetch_one(&database.pool)
.await
.expect("load preserved scheduling");
assert_eq!(scheduling, (Some(9), 40));
assert_eq!(
account_group_ids(&database.pool, "acct_grouped_import").await,
[GROUP_ID]
);
let revision_before_reauthorization = current_revision(&database.pool).await;
let reauthorized = repository
.rotate_provider_account(RotateProviderAccount {
scope,
profile: profile("acct_grouped_import", "reauthorized"),
replacement_identity: None,
credential: credential_update("acct_grouped_import", 2, "reauthorized-secret"),
audit: audit(
"audit_clear_groups_reauthorize",
"reauthorize",
"acct_grouped_import",
),
})
.await
.expect("reauthorization preserves groups atomically");
assert_eq!(
reauthorized.config_revision.get(),
u64::try_from(revision_before_reauthorization + 1).expect("revision")
);
assert_eq!(
account_group_ids(&database.pool, "acct_grouped_import").await,
[GROUP_ID]
);
assert_eq!(
audit_count(&database.pool, "request:audit_clear_groups_reauthorize").await,
1
);
let scheduling: (Option<i64>, i16) =
sqlx::query_as("select concurrency_limit, weight from provider_accounts where id = $1")
.bind("acct_grouped_import")
.fetch_one(&database.pool)
.await
.expect("load scheduling after reauthorization");
assert_eq!(scheduling, (Some(9), 40));
database.close().await;
}
#[tokio::test]
async fn verified_credential_rotation_preserves_quota_exhaustion() {
let Some(database) = TestDatabase::create("provider_account_rotation_quota").await else {
@@ -1450,6 +1692,8 @@ fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
access_token_expires_at: Some(Utc::now() + TimeDelta::hours(1)),
next_refresh_at: None,
enabled: true,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
credential_state: CredentialState::Ready,
credential_observed_at: Utc::now(),
}
@@ -1532,13 +1776,15 @@ async fn seed_model_request(
upstream_status_code,
input_tokens, output_tokens, cached_tokens, cache_write_tokens, reasoning_tokens,
total_tokens, cost_source, cost_amount, cost_currency,
started_at, deadline_at, completed_at
started_at, deadline_at, completed_at,
routing_scope, routing_group_refs, routing_group_names_snapshot
) values (
$1, 'key-provider-account-test', 1, 'openai', 'responses', '/v1/responses',
'http_sse', $4, $3, $2, $2, $4, 'http_sse', 1,
'sent', $7 + interval '1 second', 'succeeded', 200, 200, $5, 0, 0, 0, 0,
$5, 'provider_reported', $6::numeric, 'USD', $7,
$7 + interval '5 minutes', $7 + interval '1 second'
$7 + interval '5 minutes', $7 + interval '1 second',
'all', '{}'::text[], '[]'::jsonb
)",
)
.bind(seed.request_id)
@@ -1568,6 +1814,25 @@ async fn account_count(pool: &sqlx::PgPool, account_id: &str) -> i64 {
.expect("count provider account")
}
async fn account_group_ids(pool: &sqlx::PgPool, account_id: &str) -> Vec<String> {
sqlx::query_scalar(
"select account_group_id from account_group_accounts
where provider_account_id = $1 order by account_group_id",
)
.bind(account_id)
.fetch_all(pool)
.await
.expect("load account groups")
}
async fn audit_count(pool: &sqlx::PgPool, request_id: &str) -> i64 {
sqlx::query_scalar("select count(*) from admin_audit_events where admin_request_id = $1")
.bind(request_id)
.fetch_one(pool)
.await
.expect("count audit events")
}
#[test]
fn provider_credentials_are_redacted_from_debug() {
let secret = "secret-access-token";
@@ -10,14 +10,14 @@ fn snapshot_client_policy_contains_only_common_limits() {
let policy = ClientApiKeySnapshot {
id: ClientApiKeyId::new("key-1").expect("client key ID"),
plaintext_key: PlaintextClientApiKey::new("sk_snapshot_secret").expect("plaintext key"),
provider_kind: "openai".to_owned(),
group_ids: Vec::new(),
limits: RateLimits {
max_concurrency: 3,
requests_per_minute: 60,
},
};
assert_eq!(policy.limits.max_concurrency, 3);
assert_eq!(policy.provider_kind, "openai");
assert!(policy.group_ids.is_empty());
assert!(!format!("{policy:?}").contains("sk_snapshot_secret"));
}
@@ -29,9 +29,8 @@ async fn runtime_snapshot_loads_enabled_plaintext_key_without_debug_exposure() {
let plaintext = format!("sk_{}", "s".repeat(43));
sqlx::query(
"insert into client_api_keys (
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
created_at, updated_at
) values ('key_snapshot', 'snapshot', 'xai', $1, true, 2, 60, now(), now())",
id, name, key, enabled, max_concurrency, requests_per_minute, created_at, updated_at
) values ('key_snapshot', 'snapshot', $1, true, 2, 60, now(), now())",
)
.bind(&plaintext)
.execute(&database.pool)
@@ -42,7 +41,7 @@ async fn runtime_snapshot_loads_enabled_plaintext_key_without_debug_exposure() {
.await
.expect("load runtime snapshot");
assert_eq!(snapshot.client_api_keys.len(), 1);
assert_eq!(snapshot.client_api_keys[0].provider_kind, "xai");
assert!(snapshot.client_api_keys[0].group_ids.is_empty());
assert_eq!(
snapshot.client_api_keys[0].plaintext_key.expose_for_auth(),
plaintext
@@ -665,6 +665,9 @@ fn new_request(id: &str, started_at: DateTime<Utc>) -> NewModelRequest {
endpoint: "/v1/responses".to_owned(),
client_transport: "http_sse".to_owned(),
requested_model_id: "coding".to_owned(),
routing_scope: "all".to_owned(),
routing_group_refs: Vec::new(),
routing_group_names_snapshot: serde_json::json!([]),
client_ip: None,
user_agent: None,
reasoning_effort: None,
@@ -756,8 +759,8 @@ async fn finalize_request(pool: &PgPool, id: &str, started_at: DateTime<Utc>) {
async fn seed_api_key(pool: &PgPool, id: &str, name: &str, now: DateTime<Utc>) {
sqlx::query(
"insert into client_api_keys (id, name, provider_kind, key, enabled, created_at, updated_at)
values ($1, $2, 'openai', 'sk_' || $3, true, $4, $4)",
"insert into client_api_keys (id, name, key, enabled, created_at, updated_at)
values ($1, $2, 'sk_' || $3, true, $4, $4)",
)
.bind(id)
.bind(name)
@@ -211,30 +211,41 @@ async fn scheduling_lease_ttl_and_cache_loss_rebuild_empty_runtime_signal() {
}
#[tokio::test]
async fn scheduling_cursor_is_shared_across_gateway_processes() {
async fn scheduling_cursor_is_shared_across_processes_and_isolated_by_client_and_provider() {
let Some((repository, mut connection, namespace)) = repository().await else {
return;
};
let peer = RedisCredentialLeaseRepository::new(connection.clone(), &namespace)
.expect("peer repository");
let primary = gateway_core::policy::ClientApiKeyId::new("key_primary").expect("client key");
let secondary = gateway_core::policy::ClientApiKeyId::new("key_secondary").expect("client key");
let openai = gateway_core::routing::ProviderKind::new("openai").expect("provider");
let xai = gateway_core::routing::ProviderKind::new("xai").expect("provider");
assert_eq!(
repository
.advance_scheduling_cursor("route_openai_real")
.advance_scheduling_cursor(&primary, &openai)
.await
.expect("first cursor"),
0
);
assert_eq!(
peer.advance_scheduling_cursor("route_openai_real")
peer.advance_scheduling_cursor(&primary, &openai)
.await
.expect("shared cursor"),
1
);
assert_eq!(
peer.advance_scheduling_cursor("route_openai_backup")
peer.advance_scheduling_cursor(&secondary, &openai)
.await
.expect("independent cursor"),
.expect("client-isolated cursor"),
0
);
assert_eq!(
peer.advance_scheduling_cursor(&primary, &xai)
.await
.expect("provider-isolated cursor"),
0
);
@@ -243,12 +254,13 @@ async fn scheduling_cursor_is_shared_across_gateway_processes() {
.query_async::<Vec<String>>(&mut connection)
.await
.expect("list cursor keys");
assert_eq!(keys.len(), 2);
assert!(
keys.iter().all(|key| {
!key.contains("route_openai_real") && !key.contains("route_openai_backup")
})
);
assert_eq!(keys.len(), 3);
assert!(keys.iter().all(|key| {
!key.contains("key_primary")
&& !key.contains("key_secondary")
&& !key.contains("openai")
&& !key.contains("xai")
}));
redis::cmd("DEL")
.arg(keys)
.query_async::<i64>(&mut connection)
@@ -1,8 +1,10 @@
use gateway_core::engine::continuation::{
NativeContinuationPin, NativeContinuationPort, NativeContinuationScope, PreviousResponseId,
NativeContinuationPin, NativeContinuationPort, NativeContinuationScope,
NativeContinuationStoreErrorKind, PreviousResponseId,
};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::operation::ProviderSessionState;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::ProviderKind;
use gateway_store::redis::RedisNativeContinuationRepository;
use redis::aio::ConnectionManager;
@@ -30,6 +32,7 @@ async fn native_continuation_round_trips_opaque_state_without_exposing_response_
NativeContinuationPin::new(
previous_response_id.clone(),
previous_response_id.clone(),
client_key("key_primary"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_primary").expect("account"),
)
@@ -40,7 +43,7 @@ async fn native_continuation_round_trips_opaque_state_without_exposing_response_
.expect("record continuation");
let resolved = repository
.resolve(&previous_response_id)
.resolve(&client_key("key_primary"), &previous_response_id)
.await
.expect("resolve continuation")
.expect("stored continuation");
@@ -79,7 +82,7 @@ async fn empty_opaque_continuation_is_a_store_miss_without_a_redis_record() {
assert!(
repository
.resolve(&response_id)
.resolve(&client_key("key_primary"), &response_id)
.await
.expect("resolve")
.is_none()
@@ -88,6 +91,7 @@ async fn empty_opaque_continuation_is_a_store_miss_without_a_redis_record() {
.record(NativeContinuationPin::new(
response_id.clone(),
response_id,
client_key("key_primary"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_primary").expect("account"),
))
@@ -104,6 +108,7 @@ async fn native_continuation_rejects_provider_state_for_a_different_provider() {
let pin = NativeContinuationPin::new(
PreviousResponseId::new("resp-mismatched-state"),
PreviousResponseId::new("resp-upstream-handle"),
client_key("key_primary"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_primary").expect("account"),
)
@@ -115,6 +120,35 @@ async fn native_continuation_rejects_provider_state_for_a_different_provider() {
assert!(namespace_keys(&mut connection, &namespace).await.is_empty());
}
#[tokio::test]
async fn native_continuation_rejects_cross_client_reuse() {
let Some((repository, mut connection, namespace)) = repository().await else {
return;
};
let response_id = PreviousResponseId::new("resp-client-owned");
repository
.record(NativeContinuationPin::new(
response_id.clone(),
response_id.clone(),
client_key("key_primary"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_primary").expect("account"),
))
.await
.expect("record client-owned continuation");
let error = repository
.resolve(&client_key("key_secondary"), &response_id)
.await
.expect_err("another client must not reuse the continuation");
assert_eq!(
error.kind(),
NativeContinuationStoreErrorKind::OwnershipMismatch
);
delete_namespace_keys(&mut connection, &namespace).await;
}
#[tokio::test]
async fn native_continuation_record_should_bound_index_cleanup_work() {
const STALE_ENTRY_COUNT: usize = 100;
@@ -177,11 +211,16 @@ fn pin(response_id: &str) -> NativeContinuationPin {
NativeContinuationPin::new(
PreviousResponseId::new(response_id),
PreviousResponseId::new(response_id),
client_key("key_primary"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_primary").expect("account"),
)
}
fn client_key(id: &str) -> ClientApiKeyId {
ClientApiKeyId::new(id).expect("client API key id")
}
async fn namespace_keys(connection: &mut ConnectionManager, namespace: &str) -> Vec<String> {
redis::cmd("KEYS")
.arg(format!("{namespace}:*"))
@@ -9,6 +9,7 @@ use crate::transport::request::derive_conversation_anchor;
pub(crate) fn derive_codex_session_affinity_key(
request: &CodexResponsesRequest,
client_api_key_id: &ClientApiKeyId,
) -> Option<ProviderSessionAffinityKey> {
let session_key = if let Some(conversation_id) = request
.local_conversation_id
@@ -27,16 +28,25 @@ pub(crate) fn derive_codex_session_affinity_key(
};
let session_discriminator = request.subagent_kind();
session_key.and_then(|session_key| match session_discriminator.as_deref() {
None => Some(session_key),
Some(discriminator) => {
// 仅扩展原会话锚点的键空间;后续查找、绑定和调度仍使用同一套亲和逻辑。
opaque_affinity_key(
"subagent-session",
&format!("{discriminator}\0{}", session_key.expose_to_store()),
)
}
})
let session_key =
session_key.and_then(|session_key| match session_discriminator.as_deref() {
None => Some(session_key),
Some(discriminator) => {
// 仅扩展原会话锚点的键空间;后续查找、绑定和调度仍使用同一套亲和逻辑。
opaque_affinity_key(
"subagent-session",
&format!("{discriminator}\0{}", session_key.expose_to_store()),
)
}
})?;
opaque_affinity_key(
"client-session",
&format!(
"{}\0{}",
client_api_key_id.as_str(),
session_key.expose_to_store()
),
)
}
fn opaque_affinity_key(domain: &str, value: &str) -> Option<ProviderSessionAffinityKey> {
@@ -127,6 +127,7 @@ enum AffinityEscapeReason {
QuotaExhausted,
Cooldown,
LeaseSaturated,
HigherPriority,
PinnedAccount,
SelectionInvariant,
}
@@ -138,6 +139,7 @@ impl AffinityEscapeReason {
Self::QuotaExhausted => "quota_exhausted",
Self::Cooldown => "cooldown",
Self::LeaseSaturated => "lease_saturated",
Self::HigherPriority => "higher_priority",
Self::PinnedAccount => "pinned_account",
Self::SelectionInvariant => "selection_invariant",
}
@@ -195,8 +197,13 @@ impl AffinitySelection {
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::RequestInterval) => {
self.escape(AffinityEscapeReason::Cooldown);
}
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight) => {
self.escape(AffinityEscapeReason::HigherPriority);
}
PreferredAccountSelection::Blocked(
AccountSchedulingBlocker::LocalAvailability | AccountSchedulingBlocker::Excluded,
AccountSchedulingBlocker::LocalAvailability
| AccountSchedulingBlocker::Excluded
| AccountSchedulingBlocker::OutsideClientScope,
)
| PreferredAccountSelection::Missing => {
self.escape(AffinityEscapeReason::HardUnavailable);
@@ -282,6 +289,11 @@ impl CodexCredentialSelector {
.into_iter()
.filter(|account| {
account.provider() == &self.provider_kind
&& (diagnostic
|| request
.attempt
.account_scope()
.is_some_and(|scope| scope.allows(account.id())))
&& (diagnostic || {
let observed_support = self
.catalog
@@ -310,7 +322,11 @@ impl CodexCredentialSelector {
.collect::<Vec<_>>();
let scheduling = self
.leases
.load_state(&self.provider_kind, &account_ids)
.load_state(
request.attempt.client_api_key_ref(),
&self.provider_kind,
&account_ids,
)
.await?;
let round_robin_cursor = scheduling.round_robin_cursor();
let candidates = accounts
@@ -411,6 +427,7 @@ impl CodexCredentialSelector {
} else {
AccountEligibilityPolicy::Enforce
},
account_scope: request.attempt.account_scope().cloned(),
};
let Some(selection) = AccountSelector.select(&candidates, &context) else {
return match shortest_retry {
@@ -436,7 +453,7 @@ impl CodexCredentialSelector {
self.provider_kind.clone(),
account.id().clone(),
account.revision(),
policy.max_concurrent_per_account(),
account.effective_concurrency(policy.max_concurrent_per_account()),
policy.request_interval(),
request.attempt.deadline(),
),
@@ -253,7 +253,8 @@ impl Provider for CodexProvider {
.as_ref()
.and_then(|state| state.turn_state.clone());
}
let session_affinity_key = derive_codex_session_affinity_key(&upstream_request);
let session_affinity_key =
derive_codex_session_affinity_key(&upstream_request, context.client_api_key_ref());
let cyber_policy_session_key =
derive_codex_cyber_policy_session_key(&upstream_request, context.client_api_key_ref());
let transport = selected_transport(&upstream_request);
@@ -797,6 +797,7 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
AccountRecord {
id: account.id().to_string(),
provider_kind: account.provider().clone(),
groups: Vec::new(),
name: account.name().to_owned(),
email: account.email().map(str::to_owned),
upstream_user_id: account.upstream_user_id().map(str::to_owned),
@@ -808,6 +809,8 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
access_token_expires_at: account.access_token_expires_at().map(DateTime::<Utc>::from),
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.enabled(),
concurrency_limit: account.concurrency_limit(),
weight: account.weight(),
credential_state: account.credential_state(),
credential_observed_at: now,
quota: account.quota(),
@@ -1,4 +1,4 @@
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
@@ -8,10 +8,11 @@ use gateway_core::engine::continuation::{
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
};
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountErrorReason, AccountFeedbackStats, AccountSelectionPolicy,
AccountStateChange, CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId,
ProviderAccountStore as _, QuotaAccessChange, QuotaAccessState, QuotaEvidence,
QuotaObservation, QuotaState, QuotaWriteOutcome, RotationStrategy,
AccountAttemptFeedback, AccountConcurrencyLimit, AccountErrorReason, AccountFeedbackStats,
AccountSelectionPolicy, AccountStateChange, AccountWeight, CredentialState, OpaqueProviderData,
ProviderAccount, ProviderAccountId, ProviderAccountStore as _, QuotaAccessChange,
QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState, QuotaWriteOutcome,
RotationStrategy,
};
use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, RequestAttemptContext,
@@ -20,7 +21,9 @@ use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderCooldownPort, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
};
use gateway_core::routing::ProviderKind;
use gateway_core::routing::{
ClientRoutingScope, FrozenAccountScope, ProviderKind, RuntimeAccount, RuntimeAccountDirectory,
};
use provider_openai::OFFICIAL_CODEX_BASE_URL;
use provider_openai::credential::{
CodexAccountFailure, CodexCookiePolicy, CodexCredentialCatalogService, CodexCredentialCodec,
@@ -48,6 +51,34 @@ fn create_account(store: &Arc<MemoryAccountStore>, id: &str, token: &str) {
}));
}
fn contract_account_scope() -> Arc<FrozenAccountScope> {
let provider = ProviderKind::new("openai").expect("provider");
let accounts = [
"acct_available",
"acct_fallback",
"acct_fallback_signal",
"acct_first",
"acct_missing",
"acct_original",
"acct_original_signal",
"acct_other",
"acct_primary",
"acct_second",
]
.into_iter()
.map(|id| {
(
ProviderAccountId::new(id).expect("account"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)
})
.collect::<BTreeMap<_, _>>();
Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(accounts)),
ClientRoutingScope::all_accounts(),
))
}
fn attempt(excluded_accounts: BTreeSet<ProviderAccountId>) -> AttemptContext {
attempt_with_required(excluded_accounts, None)
}
@@ -64,7 +95,8 @@ fn attempt_with_required(
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::new(excluded_accounts, required_account, None),
AccountAttemptContext::new(excluded_accounts, required_account, None)
.with_account_scope(contract_account_scope()),
None,
CancellationToken::new(),
)
@@ -83,7 +115,8 @@ fn round_robin_attempt() -> AttemptContext {
NonZeroU32::new(2).expect("concurrency"),
Duration::ZERO,
),
AccountAttemptContext::new(BTreeSet::new(), None, None),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
None,
CancellationToken::new(),
)
@@ -358,6 +391,32 @@ fn selector_uses_frozen_global_account_policy_for_lease() {
assert_eq!(requests[0].deadline(), attempt.deadline());
}
#[test]
fn selector_uses_the_account_concurrency_override_for_the_redis_lease() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
store.set_scheduling(
"acct_primary",
Some(AccountConcurrencyLimit::new(7).expect("concurrency override")),
AccountWeight::DEFAULT,
);
let leases = Arc::new(TestLeaseCoordinator::default());
let selector = selector(&store, Arc::clone(&leases));
let attempt = attempt(BTreeSet::new());
block_on(selector.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url:
&Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL"),
attempt: &attempt,
session_affinity_key: None,
}))
.expect("select account");
let requests = leases.requests.lock().expect("lease requests lock");
assert_eq!(requests[0].max_concurrent().get(), 7);
}
#[test]
fn selector_round_robin_cursor_advances_across_requests() {
let store = Arc::new(MemoryAccountStore::default());
@@ -1043,6 +1102,7 @@ fn native_continuation_surfaces_the_original_accounts_quota_status_to_the_coordi
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_codex_contract").expect("client key id"),
ProviderKind::new("openai").expect("provider"),
original.id().clone(),
);
@@ -1054,7 +1114,8 @@ fn native_continuation_surfaces_the_original_accounts_quota_status_to_the_coordi
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, None),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
Some(ContinuationBinding::Pinned(continuation)),
CancellationToken::new(),
);
@@ -1112,6 +1173,7 @@ fn native_continuation_surfaces_the_original_accounts_quota_signal_to_the_coordi
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_codex_contract").expect("client key id"),
ProviderKind::new("openai").expect("provider"),
original.id().clone(),
);
@@ -1123,7 +1185,8 @@ fn native_continuation_surfaces_the_original_accounts_quota_signal_to_the_coordi
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, None),
AccountAttemptContext::new(BTreeSet::new(), None, None)
.with_account_scope(contract_account_scope()),
Some(ContinuationBinding::Pinned(continuation)),
CancellationToken::new(),
);
@@ -7,6 +7,7 @@ use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures::future::BoxFuture;
use gateway_core::engine::credential::{ProviderAccountId, ProviderAccountStore};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderCredentialState, ProviderCredentialStatePort, ProviderLeaseAcquisition,
ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshPolicy, ProviderRuntimePolicyPort,
@@ -86,6 +87,7 @@ struct RefreshLeases;
impl ProviderLeasePort for RefreshLeases {
fn load_state<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a ProviderKind,
_: &'a [ProviderAccountId],
) -> BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
@@ -1,4 +1,4 @@
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
@@ -21,7 +21,8 @@ use gateway_core::event::GatewayEvent;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
ProviderModel, PublicModelId, RoutingContext, RuntimeAccount, RuntimeAccountDirectory,
RuntimeSnapshot, UpstreamModelId,
};
use provider_openai::credential::{
@@ -252,6 +253,13 @@ fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest
let provider = ProviderKind::new(provider_name).expect("provider");
let upstream_model = UpstreamModelId::new("gpt-5.4").expect("upstream model");
let public_model = PublicModelId::new(upstream_model.as_str()).expect("public model");
let account_scope = Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new("acct_provider_contract").expect("account"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)]))),
ClientRoutingScope::all_accounts(),
));
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
account_policy(),
@@ -265,12 +273,67 @@ fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest
)
.expect("snapshot");
let plan = snapshot
.plan(&public_model, &operation, &RoutingContext::default())
.plan(
&public_model,
&operation,
account_scope,
&RoutingContext::default(),
)
.expect("routing plan");
ProviderRequest::new(operation, plan.candidates()[0].clone())
}
fn contract_account_scope() -> Arc<FrozenAccountScope> {
let provider = ProviderKind::new("openai").expect("provider");
let accounts = [
"acct_affinity",
"acct_atomic_failure",
"acct_bare_atomic_failure",
"acct_bounded_replay_grace",
"acct_bounded_session_state",
"acct_capacity_busy",
"acct_client_history",
"acct_completed_affinity",
"acct_continuation_prefetch",
"acct_disabled_scheduling",
"acct_first_event_latency",
"acct_header_new",
"acct_header_old",
"acct_header_same",
"acct_http_sse_exhausted",
"acct_local_affinity",
"acct_metadata_new",
"acct_metadata_old",
"acct_prefetch_limit",
"acct_presentation",
"acct_provider_contract",
"acct_scope_new",
"acct_scope_old",
"acct_scope_same",
"acct_semantic_failure",
"acct_session_affinity",
"acct_subagent",
"acct_success_exhausted",
"acct_thread_spawn_affinity",
"acct_truncated_stream",
"acct_usage_limit_request_path",
"acct_websocket_close",
]
.into_iter()
.map(|id| {
(
ProviderAccountId::new(id).expect("account"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)
})
.collect::<BTreeMap<_, _>>();
Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(accounts)),
ClientRoutingScope::all_accounts(),
))
}
fn context(request_id: &str, cancellation: CancellationToken) -> AttemptContext {
AttemptContext::new(
RequestAttemptContext::new(
@@ -280,7 +343,8 @@ fn context(request_id: &str, cancellation: CancellationToken) -> AttemptContext
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::new(BTreeSet::<ProviderAccountId>::new(), None, None),
AccountAttemptContext::new(BTreeSet::<ProviderAccountId>::new(), None, None)
.with_account_scope(contract_account_scope()),
None,
cancellation,
)
@@ -318,7 +382,8 @@ fn context_with_state_owner(request_id: &str, owner_account_id: &str) -> Attempt
NonZeroU32::new(1).expect("attempt"),
SystemTime::now() + Duration::from_secs(30),
account_policy(),
AccountAttemptContext::new(BTreeSet::new(), None, Some(owner)),
AccountAttemptContext::new(BTreeSet::new(), None, Some(owner))
.with_account_scope(contract_account_scope()),
None,
CancellationToken::new(),
)
@@ -9,12 +9,14 @@ use std::time::{Duration, SystemTime};
use async_trait::async_trait;
use futures::future::BoxFuture;
use gateway_core::engine::credential::{
AccountErrorReason, AccountRuntimeSignals, AccountStateChange, CredentialCasOutcome,
CredentialCasUpdate, CredentialRevision, CredentialState, LoadedCredential, NewProviderAccount,
OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore,
ProviderAccountUpdate, QuotaAccessChange, QuotaObservation, QuotaState, QuotaWriteOutcome,
AccountConcurrencyLimit, AccountErrorReason, AccountRuntimeSignals, AccountStateChange,
AccountWeight, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState,
LoadedCredential, NewProviderAccount, OpaqueProviderData, ProviderAccount, ProviderAccountId,
ProviderAccountStore, ProviderAccountUpdate, QuotaAccessChange, QuotaObservation, QuotaState,
QuotaWriteOutcome,
};
use gateway_core::error::{StoreError, StoreErrorKind};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
@@ -70,6 +72,21 @@ impl MemoryAccountStore {
.map(|stored| stored.account.clone())
}
pub(crate) fn set_scheduling(
&self,
id: &str,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
) {
let id = ProviderAccountId::new(id).expect("valid account ID");
let mut accounts = self.accounts.lock().expect("account store lock");
let stored = accounts.get_mut(&id).expect("seeded account");
stored.account = stored
.account
.clone()
.with_scheduling(concurrency_limit, weight);
}
pub(crate) fn quota_reads(&self) -> usize {
self.quota_reads.load(Ordering::SeqCst)
}
@@ -450,6 +467,7 @@ fn rebuild_account(current: &ProviderAccount, rebuild: AccountRebuild) -> Provid
rebuild.last_error_reason,
rebuild.last_error_message,
)
.with_scheduling(current.concurrency_limit(), current.weight())
.with_refresh_schedule(rebuild.has_refresh_token, rebuild.next_refresh_at)
}
@@ -487,6 +505,7 @@ pub(crate) struct TestLeaseCoordinator {
impl ProviderLeasePort for TestLeaseCoordinator {
fn load_state<'a>(
&'a self,
_client_api_key_id: &'a ClientApiKeyId,
_provider_kind: &'a ProviderKind,
accounts: &'a [ProviderAccountId],
) -> BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
@@ -89,9 +89,14 @@ impl GrokAccountSessionSelector {
.list_accounts_for_provider()
.await
.map_err(|_| GrokSessionSelectorError::Unavailable)?;
if !diagnostic {
self.quota.prepare_scheduling(&accounts).await;
}
let accounts = if diagnostic {
accounts
} else {
accounts
.into_iter()
.filter(|account| request.account_scope().allows(account.id()))
.collect()
};
let catalog_eligible = if diagnostic {
accounts
} else {
@@ -128,6 +133,9 @@ impl GrokAccountSessionSelector {
if catalog_eligible.is_empty() {
return Err(GrokSessionSelectorError::NoEligibleSession);
}
if !diagnostic {
self.quota.prepare_scheduling(&catalog_eligible).await;
}
let account_ids = catalog_eligible
.iter()
@@ -135,7 +143,11 @@ impl GrokAccountSessionSelector {
.collect::<Vec<_>>();
let scheduling = self
.scheduling
.load_state(&self.provider_kind, &account_ids)
.load_state(
request.client_api_key_id(),
&self.provider_kind,
&account_ids,
)
.await
.map_err(|_| GrokSessionSelectorError::Unavailable)?;
let runtime_cooldowns = if diagnostic {
@@ -179,6 +191,7 @@ impl GrokAccountSessionSelector {
preferred_account: request.required_account().cloned().or(affinity_account),
round_robin_cursor: scheduling.round_robin_cursor(),
eligibility: request.eligibility(),
account_scope: (!diagnostic).then(|| Arc::clone(request.account_scope())),
};
let mut capacity_denied = false;
let mut retry_after = None;
@@ -194,9 +207,11 @@ impl GrokAccountSessionSelector {
self.provider_kind.clone(),
selected_id.clone(),
selected_revision,
request
.account_selection_policy()
.max_concurrent_per_account(),
selected.account.effective_concurrency(
request
.account_selection_policy()
.max_concurrent_per_account(),
),
request.account_selection_policy().request_interval(),
request.deadline(),
),
@@ -391,6 +391,8 @@ async fn select_grok_session(
required_account.clone(),
context.account_selection_policy(),
context.deadline(),
Arc::clone(candidate.account_scope()),
context.client_api_key_ref().clone(),
)
.with_eligibility_policy(if context.is_diagnostic_required_account() {
AccountEligibilityPolicy::BypassForDiagnostic
@@ -2,13 +2,15 @@ use std::collections::BTreeSet;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use gateway_core::engine::credential::{
AccountEligibilityPolicy, AccountSelectionPolicy, CredentialRevision, ProviderAccountId,
};
use gateway_core::engine::provider::ProviderResource;
use gateway_core::routing::UpstreamModelId;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{FrozenAccountScope, UpstreamModelId};
use sha2::{Digest as _, Sha256};
use crate::SecretValue;
@@ -213,6 +215,8 @@ pub struct GrokSessionSelection {
eligibility: AccountEligibilityPolicy,
affinity: Option<GrokSessionAffinityKey>,
deadline: SystemTime,
account_scope: Arc<FrozenAccountScope>,
client_api_key_id: ClientApiKeyId,
}
impl GrokSessionSelection {
@@ -224,6 +228,8 @@ impl GrokSessionSelection {
required_account: Option<ProviderAccountId>,
account_selection_policy: AccountSelectionPolicy,
deadline: SystemTime,
account_scope: Arc<FrozenAccountScope>,
client_api_key_id: ClientApiKeyId,
) -> Self {
Self {
upstream_model,
@@ -233,6 +239,8 @@ impl GrokSessionSelection {
eligibility: AccountEligibilityPolicy::Enforce,
affinity: None,
deadline,
account_scope,
client_api_key_id,
}
}
@@ -293,6 +301,16 @@ impl GrokSessionSelection {
pub const fn deadline(&self) -> SystemTime {
self.deadline
}
#[must_use]
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
#[must_use]
pub const fn client_api_key_id(&self) -> &ClientApiKeyId {
&self.client_api_key_id
}
}
/// Grok session selector 返回的 future。
@@ -460,6 +460,7 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
AccountRecord {
id: account.id().to_string(),
provider_kind: account.provider().clone(),
groups: Vec::new(),
name: account.name().to_owned(),
email: account.email().map(str::to_owned),
upstream_user_id: account.upstream_user_id().map(str::to_owned),
@@ -471,6 +472,8 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
access_token_expires_at: account.access_token_expires_at().map(DateTime::<Utc>::from),
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.enabled(),
concurrency_limit: account.concurrency_limit(),
weight: account.weight(),
credential_state: account.credential_state(),
credential_observed_at: now,
quota: account.quota(),
@@ -486,6 +489,7 @@ struct TestLeases;
impl ProviderLeasePort for TestLeases {
fn load_state<'a>(
&'a self,
_: &'a gateway_core::policy::ClientApiKeyId,
_: &'a gateway_core::routing::ProviderKind,
accounts: &'a [ProviderAccountId],
) -> BoxFuture<
@@ -110,6 +110,7 @@ impl ProviderRuntimePolicyPort for MutableRuntimePolicy {
impl ProviderLeasePort for TestRefreshLeases {
fn load_state<'a>(
&'a self,
_: &'a gateway_core::policy::ClientApiKeyId,
_: &'a gateway_core::routing::ProviderKind,
_: &'a [gateway_core::engine::credential::ProviderAccountId],
) -> futures::future::BoxFuture<
@@ -4,16 +4,22 @@ use std::time::{Duration, SystemTime};
use chrono::Utc;
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountFeedbackStats, AccountRuntimeSignals, AccountSelectionPolicy,
CredentialCasOutcome, CredentialRevision, CredentialState, OpaqueProviderData,
ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, QuotaAccessState,
QuotaObservation, QuotaState, QuotaWriteOutcome, RotationStrategy,
AccountAttemptFeedback, AccountConcurrencyLimit, AccountFeedbackStats, AccountRuntimeSignals,
AccountSelectionPolicy, AccountWeight, CredentialCasOutcome, CredentialRevision,
CredentialState, OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
ProviderAccountUpdate, QuotaAccessState, QuotaObservation, QuotaState, QuotaWriteOutcome,
RotationStrategy,
};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::provider_ports::{
ProviderCooldownPort, ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort,
ProviderLeaseRequest, ProviderSchedulingState, ProviderStoreError,
ProviderLeaseRequest, ProviderSchedulingLeaseRequest, ProviderSchedulingState,
ProviderStoreError,
};
use gateway_core::routing::{
ClientRoutingScope, FrozenAccountScope, RuntimeAccount, RuntimeAccountDirectory,
UpstreamModelId,
};
use gateway_core::routing::UpstreamModelId;
use provider_xai::{
GrokAccountSessionSelector, GrokBillingRequest, GrokBillingTransport,
GrokBillingTransportError, GrokBillingTransportErrorKind, GrokBillingTransportFuture,
@@ -31,6 +37,7 @@ use crate::support::{
struct SchedulingCoordinator {
signals: Mutex<BTreeMap<ProviderAccountId, AccountRuntimeSignals>>,
denied: Mutex<BTreeSet<ProviderAccountId>>,
requests: Mutex<Vec<ProviderSchedulingLeaseRequest>>,
}
struct UnavailableBillingTransport;
@@ -48,6 +55,7 @@ impl GrokBillingTransport for UnavailableBillingTransport {
impl ProviderLeasePort for SchedulingCoordinator {
fn load_state<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a gateway_core::routing::ProviderKind,
_: &'a [ProviderAccountId],
) -> futures::future::BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
@@ -67,20 +75,19 @@ impl ProviderLeasePort for SchedulingCoordinator {
let ProviderLeaseRequest::Scheduling(request) = request else {
panic!("expected scheduling lease request");
};
Ok(
if self
.denied
.lock()
.expect("denied")
.contains(request.account_id())
{
ProviderLeaseAcquisition::Busy {
retry_after: Some(Duration::from_millis(25)),
}
} else {
ProviderLeaseAcquisition::Acquired(Box::new(()))
},
)
let denied = self
.denied
.lock()
.expect("denied")
.contains(request.account_id());
self.requests.lock().expect("requests").push(request);
Ok(if denied {
ProviderLeaseAcquisition::Busy {
retry_after: Some(Duration::from_millis(25)),
}
} else {
ProviderLeaseAcquisition::Acquired(Box::new(()))
})
})
}
}
@@ -141,6 +148,7 @@ impl SelectorFixture {
let coordinator = Arc::new(SchedulingCoordinator {
signals: Mutex::new(signals),
denied: Mutex::new(BTreeSet::new()),
requests: Mutex::new(Vec::new()),
});
let cooldowns = Arc::new(MemoryCooldownPort::default());
let catalog_cache: Arc<dyn GrokCredentialCatalogCache> = cache.clone();
@@ -210,6 +218,21 @@ impl SelectorFixture {
required_account: Option<ProviderAccountId>,
strategy: RotationStrategy,
) -> GrokSessionSelection {
let provider = gateway_core::routing::ProviderKind::new("xai").expect("provider");
let accounts = self
.coordinator
.signals
.lock()
.expect("signals")
.keys()
.cloned()
.map(|account_id| {
(
account_id,
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)
})
.collect::<BTreeMap<_, _>>();
GrokSessionSelection::new(
UpstreamModelId::new(upstream_model).expect("model"),
excluded,
@@ -220,6 +243,11 @@ impl SelectorFixture {
Duration::ZERO,
),
SystemTime::now() + Duration::from_secs(30),
Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(accounts)),
ClientRoutingScope::all_accounts(),
)),
ClientApiKeyId::new("key_xai_selector").expect("client key"),
)
}
@@ -286,6 +314,26 @@ async fn required_account_overrides_smart_selection_without_fallback() {
));
}
#[tokio::test]
async fn selector_uses_the_account_concurrency_override_for_the_redis_lease() {
let fixture = SelectorFixture::new(&["concurrency-override"]).await;
let account = account_id("concurrency-override");
fixture.store.set_scheduling(
&account,
Some(AccountConcurrencyLimit::new(7).expect("concurrency override")),
AccountWeight::DEFAULT,
);
fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("select account");
let requests = fixture.coordinator.requests.lock().expect("requests");
assert_eq!(requests[0].max_concurrent().get(), 7);
}
#[tokio::test]
async fn unauthorized_feedback_records_runtime_cooldown_without_persisting_account_state() {
let fixture = SelectorFixture::new(&["feedback-a", "feedback-b"]).await;
@@ -1,4 +1,4 @@
use std::collections::{BTreeSet, VecDeque};
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::num::NonZeroU32;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
@@ -26,7 +26,8 @@ use gateway_core::operation::{
Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload, ProviderSessionState,
};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
ProviderModel, PublicModelId, RoutingContext, RuntimeAccount, RuntimeAccountDirectory,
RuntimeSnapshot, SupportLevel, UpstreamModelId,
};
use provider_xai::{
@@ -810,6 +811,13 @@ fn provider_request_with_operation(provider_kind: &str, operation: Operation) ->
.with_feature(Feature::Tools, SupportLevel::Native)
.with_feature(Feature::NativeContinuation, SupportLevel::Native),
);
let account_scope = Arc::new(FrozenAccountScope::new(
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
account_id("provider"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)]))),
ClientRoutingScope::all_accounts(),
));
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
selection_policy(),
@@ -822,6 +830,7 @@ fn provider_request_with_operation(provider_kind: &str, operation: Operation) ->
.plan(
&PublicModelId::new(MODEL).expect("model"),
&operation,
account_scope,
&RoutingContext::default(),
)
.expect("routing plan");
@@ -1858,6 +1867,7 @@ async fn native_previous_response_is_rejected_before_selection() {
let pin = NativeContinuationPin::new(
PreviousResponseId::new("resp_previous"),
PreviousResponseId::new("resp_upstream_previous"),
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
ProviderKind::new("openai").expect("provider"),
account_id("provider"),
);
@@ -1891,6 +1901,7 @@ async fn native_previous_response_pins_account_and_sends_upstream_handle() {
let pin = NativeContinuationPin::new(
PreviousResponseId::new("resp_previous"),
PreviousResponseId::new("resp_upstream_previous"),
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
ProviderKind::new("xai").expect("provider"),
account_id("provider"),
);
@@ -1936,6 +1947,7 @@ async fn native_previous_response_does_not_allow_quota_or_rate_limit_account_rot
let pin = NativeContinuationPin::new(
PreviousResponseId::new("resp_previous"),
PreviousResponseId::new("resp_upstream_previous"),
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
ProviderKind::new("xai").expect("provider"),
account_id("provider"),
);
@@ -2257,6 +2269,7 @@ async fn connection_state_inherits_session_and_recovers_reasoning_on_pinned_acco
let pin = NativeContinuationPin::new(
PreviousResponseId::new("resp_gateway_first"),
PreviousResponseId::new("resp_state"),
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
ProviderKind::new("xai").expect("provider"),
account_id("provider"),
);
@@ -2381,6 +2394,7 @@ async fn replay_owner_should_reencode_custom_apply_patch_call_for_grok() {
let pin = NativeContinuationPin::new(
PreviousResponseId::new("resp_gateway_patch"),
PreviousResponseId::new("resp_custom_patch"),
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
ProviderKind::new("xai").expect("provider"),
account_id("provider"),
);
@@ -2456,6 +2470,7 @@ async fn missing_native_response_should_be_replay_safe_for_the_same_account() {
let pin = NativeContinuationPin::new(
PreviousResponseId::new("resp_gateway_first"),
PreviousResponseId::new("resp_upstream_first"),
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
ProviderKind::new("xai").expect("provider"),
account_id("provider"),
);
+20 -4
View File
@@ -8,10 +8,11 @@ use std::time::{Duration, SystemTime};
use async_trait::async_trait;
use chrono::Utc;
use gateway_core::engine::credential::{
AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision, CredentialState, LoadedCredential, NewProviderAccount, PlaintextCredential,
ProviderAccount, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
QuotaAccessChange, QuotaObservation, QuotaState, QuotaWriteOutcome,
AccountConcurrencyLimit, AccountErrorReason, AccountStateChange, AccountWeight,
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState,
LoadedCredential, NewProviderAccount, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountStore, ProviderAccountUpdate, QuotaAccessChange, QuotaObservation, QuotaState,
QuotaWriteOutcome,
};
use gateway_core::error::{StoreError, StoreErrorKind};
use gateway_core::provider_ports::{
@@ -213,6 +214,20 @@ impl MemoryProviderAccountStore {
.map(|stored| stored.account.clone())
}
pub fn set_scheduling(
&self,
id: &ProviderAccountId,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
) {
let mut accounts = lock(&self.accounts);
let stored = accounts.get_mut(id).expect("seeded account");
stored.account = stored
.account
.clone()
.with_scheduling(concurrency_limit, weight);
}
pub fn last_error_message(&self, id: &ProviderAccountId) -> Option<String> {
lock(&self.accounts)
.get(id)
@@ -543,6 +558,7 @@ fn rebuild_account(previous: &ProviderAccount, replacement: AccountReplacement)
replacement.last_error_reason,
replacement.last_error_message,
)
.with_scheduling(previous.concurrency_limit(), previous.weight())
.with_refresh_schedule(replacement.has_refresh_token, replacement.next_refresh_at)
}
+4
View File
@@ -2,3 +2,7 @@
f4168f7e23d42cf9d37096339e931b0e0c823f7a8dbfc6580a04f096c4c3d8e3 0002_provider_accounts_nullable_upstream_user_id.sql
4163f61183d9e90f1368c08f614040d6a1a0048de3886a77dd708ce21430b2d0 0003_backfill_openai_access_token_expiry.sql
b41bb3be5a84b95852158d82267456f8415811a1137025581eefd1c96c0882d1 0004_normalize_account_state_facts.sql
00ef40c9cd404a09afe916262e894eec361039a60a93d6c46e6c96104b1346b9 0005_account_groups.sql
73846775d2c533fc17dc12164efd9ce6beaa99e2a67a064589ebadd50cd9b349 0006_account_group_color.sql
4e3a367b1bf75de82e7706127080627e0bb8f190b4d3bbc69be954de38eb1ead 0007_account_group_color_alpha.sql
3f07a905e7413f0be519825a10c73effb392df4e6069a5df7901bcd0ec2ff1af 0008_provider_account_scheduling.sql
+135
View File
@@ -0,0 +1,135 @@
-- Account groups are provider-neutral account sets. Existing provider-scoped
-- keys are placed in ordinary migration groups before their provider column is
-- removed, preserving their pre-upgrade account scope.
create table account_groups (
id text primary key,
name text not null,
description text,
enabled boolean not null default true,
created_at timestamptz not null,
updated_at timestamptz not null,
constraint account_groups_id_ck check (
id ~ '^grp_[0-9a-f]{32}$'
),
constraint account_groups_name_ck check (
char_length(btrim(name)) between 1 and 100
and name = btrim(name)
and name !~ '[[:cntrl:]]'
),
constraint account_groups_description_ck check (
description is null or (
octet_length(description) <= 4096
and description !~ '[[:cntrl:]]'
)
),
constraint account_groups_time_ck check (
created_at <= updated_at
)
);
create unique index account_groups_name_uq
on account_groups (lower(name));
create index account_groups_list_idx
on account_groups (enabled, created_at desc, id desc);
create table account_group_accounts (
account_group_id text not null,
provider_account_id text not null,
created_at timestamptz not null,
primary key (account_group_id, provider_account_id),
constraint account_group_accounts_group_fk foreign key (account_group_id)
references account_groups (id)
on update restrict
on delete cascade,
constraint account_group_accounts_account_fk foreign key (provider_account_id)
references provider_accounts (id)
on update restrict
on delete cascade
);
create index account_group_accounts_account_idx
on account_group_accounts (provider_account_id, account_group_id);
create table client_api_key_groups (
client_api_key_id text not null,
account_group_id text not null,
created_at timestamptz not null,
primary key (client_api_key_id, account_group_id),
constraint client_api_key_groups_key_fk foreign key (client_api_key_id)
references client_api_keys (id)
on update restrict
on delete cascade,
constraint client_api_key_groups_group_fk foreign key (account_group_id)
references account_groups (id)
on update restrict
on delete restrict
);
create index client_api_key_groups_group_idx
on client_api_key_groups (account_group_id, client_api_key_id);
with provider_kinds as (
select provider_kind from provider_accounts
union
select provider_kind from client_api_keys
)
insert into account_groups (id, name, description, enabled, created_at, updated_at)
select
'grp_' || md5('codex-proxy-rs:migration-group:' || provider_kind),
'迁移池-' || provider_kind,
'升级时按原 Provider 权限自动生成',
true,
now(),
now()
from provider_kinds;
insert into account_group_accounts (account_group_id, provider_account_id, created_at)
select
'grp_' || md5('codex-proxy-rs:migration-group:' || provider_kind),
id,
clock_timestamp()
from provider_accounts;
insert into client_api_key_groups (client_api_key_id, account_group_id, created_at)
select
id,
'grp_' || md5('codex-proxy-rs:migration-group:' || provider_kind),
clock_timestamp()
from client_api_keys;
alter table model_requests
add column routing_scope text not null default 'legacy_provider',
add column routing_group_refs text[] not null default '{}',
add column routing_group_names_snapshot jsonb not null default '[]'::jsonb,
add constraint model_requests_routing_scope_ck check (
routing_scope in ('legacy_provider', 'all', 'groups')
),
add constraint model_requests_routing_group_names_ck check (
jsonb_typeof(routing_group_names_snapshot) = 'array'
and (
(
routing_scope in ('legacy_provider', 'all')
and cardinality(routing_group_refs) = 0
and jsonb_array_length(routing_group_names_snapshot) = 0
)
or (
routing_scope = 'groups'
and cardinality(routing_group_refs) > 0
and array_position(routing_group_refs, null) is null
and jsonb_array_length(routing_group_names_snapshot)
= cardinality(routing_group_refs)
)
)
);
create index model_requests_routing_group_refs_idx
on model_requests using gin (routing_group_refs);
alter table model_requests
alter column routing_scope drop default;
alter table client_api_keys
drop constraint client_api_keys_provider_kind_ck,
drop column provider_kind;
@@ -0,0 +1,9 @@
alter table account_groups
add column color text;
update account_groups
set color = '#2563EB';
alter table account_groups
alter column color set not null,
add constraint account_groups_color_ck check (color ~ '^#[0-9A-F]{6}$');
@@ -0,0 +1,8 @@
alter table account_groups
drop constraint account_groups_color_ck;
update account_groups
set color = color || 'FF';
alter table account_groups
add constraint account_groups_color_ck check (color ~ '^#[0-9A-F]{8}$');

Some files were not shown because too many files have changed in this diff Show More