diff --git a/backend/crates/gateway-admin/src/lib.rs b/backend/crates/gateway-admin/src/lib.rs index 9c3e8981a..d8e893303 100644 --- a/backend/crates/gateway-admin/src/lib.rs +++ b/backend/crates/gateway-admin/src/lib.rs @@ -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, accounts: Arc, + account_groups: Arc, client_keys: Arc, observability: Arc, settings: Arc, @@ -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(), diff --git a/backend/crates/gateway-admin/src/model/account_groups.rs b/backend/crates/gateway-admin/src/model/account_groups.rs new file mode 100644 index 000000000..36814a513 --- /dev/null +++ b/backend/crates/gateway-admin/src/model/account_groups.rs @@ -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 { + (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, + 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, + pub enabled: Option, +} + +/// Complete account group summary. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountGroupRecord { + pub id: AccountGroupId, + pub name: String, + pub description: Option, + pub color: AccountGroupColor, + pub enabled: bool, + pub member_count: u64, + pub provider_counts: BTreeMap, + pub client_key_count: u64, + pub account_summary: AccountGroupAccountSummary, + pub capacity: AccountGroupCapacity, + pub usage: AccountGroupUsage, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// Paginated account groups. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountGroupPage { + pub config_revision: Revision, + pub items: Vec, + 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, + 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, + pub total: u64, +} + +/// Create an account group. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CreateAccountGroup { + pub name: String, + pub description: Option, + 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, + 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, + 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, +} diff --git a/backend/crates/gateway-admin/src/model/accounts.rs b/backend/crates/gateway-admin/src/model/accounts.rs index 84a44696e..0bd95eb58 100644 --- a/backend/crates/gateway-admin/src/model/accounts.rs +++ b/backend/crates/gateway-admin/src/model/accounts.rs @@ -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, + pub group_filter: Option, pub search: Option, pub status: Option, pub sort: Option, } +/// 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, pub name: String, pub email: Option, pub upstream_user_id: Option, @@ -66,6 +76,8 @@ pub struct AccountRecord { pub access_token_expires_at: Option>, pub next_refresh_at: Option>, pub enabled: bool, + pub concurrency_limit: Option, + pub weight: AccountWeight, pub credential_state: CredentialState, pub credential_observed_at: DateTime, 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, + pub weight: AccountWeight, + pub group_ids: Vec, +} + +/// 账号更新结果。 +#[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, + pub enabled: bool, + pub concurrency_limit: Option, + pub weight: AccountWeight, + pub group_ids: Vec, +} + +/// 批量账号更新结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountsUpdateResult { + pub config_revision: Revision, + pub account_ids: Vec, } /// 账号批量删除命令。 diff --git a/backend/crates/gateway-admin/src/model/client_keys.rs b/backend/crates/gateway-admin/src/model/client_keys.rs index d29919c48..3176833f9 100644 --- a/backend/crates/gateway-admin/src/model/client_keys.rs +++ b/backend/crates/gateway-admin/src/model/client_keys.rs @@ -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, - pub provider_kind: ProviderKind, + pub groups: Vec, + pub provider_kinds: Vec, 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, - pub provider_kind: ProviderKind, + pub group_ids: Vec, pub limits: RateLimits, } @@ -153,7 +154,7 @@ pub struct NewClientKey { pub id: ClientApiKeyId, pub name: String, pub label: Option, - pub provider_kind: ProviderKind, + pub group_ids: Vec, 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, - pub provider_kind: ProviderKind, + pub group_ids: Vec, pub limits: RateLimits, } diff --git a/backend/crates/gateway-admin/src/model/mod.rs b/backend/crates/gateway-admin/src/model/mod.rs index 9f4fe1cb3..d6f50d3da 100644 --- a/backend/crates/gateway-admin/src/model/mod.rs +++ b/backend/crates/gateway-admin/src/model/mod.rs @@ -2,6 +2,7 @@ use std::num::{NonZeroU16, NonZeroU64}; +pub mod account_groups; pub mod accounts; pub mod auth; pub mod backup; diff --git a/backend/crates/gateway-admin/src/model/observability.rs b/backend/crates/gateway-admin/src/model/observability.rs index 00740dc09..a49ec483b 100644 --- a/backend/crates/gateway-admin/src/model/observability.rs +++ b/backend/crates/gateway-admin/src/model/observability.rs @@ -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, } @@ -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, + pub routing_group_names_snapshot: Vec, pub protocol: String, pub operation: String, pub endpoint: String, diff --git a/backend/crates/gateway-admin/src/ports/store.rs b/backend/crates/gateway-admin/src/ports/store.rs index 7a5f6eda9..abbe0c2d2 100644 --- a/backend/crates/gateway-admin/src/ports/store.rs +++ b/backend/crates/gateway-admin/src/ports/store.rs @@ -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; - async fn set_account_enabled( + async fn update_account( &self, - command: SetAccountEnabled, + command: UpdateAccount, context: &MutationContext, - ) -> AdminStoreResult; + ) -> AdminStoreResult; + + async fn batch_update_accounts( + &self, + command: BatchUpdateAccounts, + context: &MutationContext, + ) -> AdminStoreResult; async fn delete_accounts( &self, @@ -215,6 +226,44 @@ pub trait ClientKeyStore: Send + Sync { ) -> AdminStoreResult; } +/// Provider-neutral account group management and membership transactions. +#[async_trait] +pub trait AccountGroupStore: Send + Sync { + async fn list_account_groups( + &self, + query: AccountGroupListQuery, + ) -> AdminStoreResult; + + async fn account_group_members( + &self, + id: &gateway_core::routing::AccountGroupId, + ) -> AdminStoreResult; + + async fn create_account_group( + &self, + command: NewAccountGroup, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn update_account_group( + &self, + command: UpdateAccountGroup, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn set_account_group_enabled( + &self, + command: SetAccountGroupEnabled, + context: &MutationContext, + ) -> AdminStoreResult; + + async fn delete_account_group( + &self, + command: DeleteAccountGroup, + context: &MutationContext, + ) -> AdminStoreResult; +} + /// 用量、趋势、诊断与运维错误的只读能力。 #[async_trait] pub trait ObservabilityStore: Send + Sync { @@ -302,6 +351,7 @@ pub trait SettingsStore: Send + Sync { #[derive(Clone)] pub struct AdminStorePorts { accounts: Arc, + account_groups: Arc, auth: Arc, client_keys: Arc, observability: Arc, @@ -313,6 +363,7 @@ impl AdminStorePorts { #[must_use] pub fn new( accounts: Arc, + account_groups: Arc, auth: Arc, client_keys: Arc, observability: Arc, @@ -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 { + self.account_groups.clone() + } + #[must_use] pub fn auth(&self) -> Arc { self.auth.clone() diff --git a/backend/crates/gateway-admin/src/use_case/account_groups.rs b/backend/crates/gateway-admin/src/use_case/account_groups.rs new file mode 100644 index 000000000..c3d15b7fb --- /dev/null +++ b/backend/crates/gateway-admin/src/use_case/account_groups.rs @@ -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; + async fn members(&self, id: &AccountGroupId) -> Result; + async fn create( + &self, + context: &MutationContext, + command: CreateAccountGroup, + ) -> Result; + async fn update( + &self, + context: &MutationContext, + command: UpdateAccountGroup, + ) -> Result; + async fn set_enabled( + &self, + context: &MutationContext, + command: SetAccountGroupEnabled, + ) -> Result; + async fn delete( + &self, + context: &MutationContext, + command: DeleteAccountGroup, + ) -> Result; +} + +pub(crate) struct DefaultAccountGroupService { + store: Arc, + snapshot: Arc, +} + +impl DefaultAccountGroupService { + #[must_use] + pub(crate) fn new( + store: Arc, + snapshot: Arc, + ) -> Self { + Self { store, snapshot } + } + + async fn publish( + &self, + result: Result, + ) -> Result { + 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 { + self.store + .list_account_groups(query) + .await + .map_err(|error| map_store_error(error, "account group")) + } + + async fn members(&self, id: &AccountGroupId) -> Result { + 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 { + 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 { + self.publish(self.store.update_account_group(command, context).await) + .await + } + + async fn set_enabled( + &self, + context: &MutationContext, + command: SetAccountGroupEnabled, + ) -> Result { + self.publish(self.store.set_account_group_enabled(command, context).await) + .await + } + + async fn delete( + &self, + context: &MutationContext, + command: DeleteAccountGroup, + ) -> Result { + self.publish(self.store.delete_account_group(command, context).await) + .await + } +} diff --git a/backend/crates/gateway-admin/src/use_case/accounts.rs b/backend/crates/gateway-admin/src/use_case/accounts.rs index d57df6043..c929f69f7 100644 --- a/backend/crates/gateway-admin/src/use_case/accounts.rs +++ b/backend/crates/gateway-admin/src/use_case/accounts.rs @@ -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; + async fn update( + &self, + context: &MutationContext, + command: UpdateAccount, + ) -> Result; + + async fn batch_update( + &self, + context: &MutationContext, + command: BatchUpdateAccounts, + ) -> Result; + async fn quota( &self, account_id: &ProviderAccountId, @@ -383,6 +396,76 @@ impl AccountsService for DefaultAccountsService { }) } + async fn update( + &self, + context: &MutationContext, + command: UpdateAccount, + ) -> Result { + 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 { + let account_ids = command + .account_ids + .iter() + .map(|id| { + ProviderAccountId::new(id.clone()) + .map_err(|_| AdminError::invalid("Invalid provider account ID")) + }) + .collect::, _>>()?; + let mut providers = BTreeMap::< + ProviderKind, + ( + Arc, + Vec, + ), + >::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, diff --git a/backend/crates/gateway-admin/src/use_case/client_keys.rs b/backend/crates/gateway-admin/src/use_case/client_keys.rs index 2bcc2e2c2..271b9c141 100644 --- a/backend/crates/gateway-admin/src/use_case/client_keys.rs +++ b/backend/crates/gateway-admin/src/use_case/client_keys.rs @@ -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(), }, diff --git a/backend/crates/gateway-admin/src/use_case/mod.rs b/backend/crates/gateway-admin/src/use_case/mod.rs index 9e048b561..ac47f4b7d 100644 --- a/backend/crates/gateway-admin/src/use_case/mod.rs +++ b/backend/crates/gateway-admin/src/use_case/mod.rs @@ -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 { - 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, diff --git a/backend/crates/gateway-admin/src/use_case/observability.rs b/backend/crates/gateway-admin/src/use_case/observability.rs index 8939e6426..ea9296023 100644 --- a/backend/crates/gateway-admin/src/use_case/observability.rs +++ b/backend/crates/gateway-admin/src/use_case/observability.rs @@ -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 { diff --git a/backend/crates/gateway-admin/src/use_case/openai.rs b/backend/crates/gateway-admin/src/use_case/openai.rs index af9e20834..18fa7dc26 100644 --- a/backend/crates/gateway-admin/src/use_case/openai.rs +++ b/backend/crates/gateway-admin/src/use_case/openai.rs @@ -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; - async fn enable( - &self, - command: CredentialMutation, - ) -> Result; - async fn disable( - &self, - command: CredentialMutation, - ) -> Result; async fn delete( &self, command: CredentialDeletion, @@ -252,38 +243,6 @@ impl OpenAiService for DefaultOpenAiService { Ok(result) } - async fn enable( - &self, - command: CredentialMutation, - ) -> Result { - 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 { - 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, diff --git a/backend/crates/gateway-admin/src/use_case/xai.rs b/backend/crates/gateway-admin/src/use_case/xai.rs index a6eb7ea3b..b32e044aa 100644 --- a/backend/crates/gateway-admin/src/use_case/xai.rs +++ b/backend/crates/gateway-admin/src/use_case/xai.rs @@ -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; - async fn enable( - &self, - command: CredentialMutation, - ) -> Result; - async fn disable( - &self, - command: CredentialMutation, - ) -> Result; async fn delete( &self, command: CredentialDeletion, @@ -166,38 +158,6 @@ impl XaiService for DefaultXaiService { Ok(result) } - async fn enable( - &self, - command: CredentialMutation, - ) -> Result { - 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 { - 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, diff --git a/backend/crates/gateway-admin/tests/use_case/accounts.rs b/backend/crates/gateway-admin/tests/use_case/accounts.rs index 43b4d1fec..a3bd31c4e 100644 --- a/backend/crates/gateway-admin/tests/use_case/accounts.rs +++ b/backend/crates/gateway-admin/tests/use_case/accounts.rs @@ -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 { - self.record("store.set_enabled"); + ) -> AdminStoreResult { + 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 { + 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()), diff --git a/backend/crates/gateway-admin/tests/use_case/mod.rs b/backend/crates/gateway-admin/tests/use_case/mod.rs index cda8de9db..eda5bfd58 100644 --- a/backend/crates/gateway-admin/tests/use_case/mod.rs +++ b/backend/crates/gateway-admin/tests/use_case/mod.rs @@ -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 { + Err(unavailable("account groups")) + } + + async fn account_group_members( + &self, + _: &gateway_core::routing::AccountGroupId, + ) -> AdminStoreResult { + Err(unavailable("account group members")) + } + + async fn create_account_group( + &self, + _: NewAccountGroup, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("account group create")) + } + + async fn update_account_group( + &self, + _: UpdateAccountGroup, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("account group update")) + } + + async fn set_account_group_enabled( + &self, + _: SetAccountGroupEnabled, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("account group state")) + } + + async fn delete_account_group( + &self, + _: DeleteAccountGroup, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("account group delete")) + } +} + #[async_trait] impl AccountStore for UnavailableStore { async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult { @@ -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 { + ) -> AdminStoreResult { Err(unavailable("account enabled")) } + async fn batch_update_accounts( + &self, + _: BatchUpdateAccounts, + _: &MutationContext, + ) -> AdminStoreResult { + Err(unavailable("account batch update")) + } + async fn delete_accounts( &self, _: DeleteAccounts, diff --git a/backend/crates/gateway-admin/tests/use_case/observability.rs b/backend/crates/gateway-admin/tests/use_case/observability.rs index 15a91e5d6..dd0271e28 100644 --- a/backend/crates/gateway-admin/tests/use_case/observability.rs +++ b/backend/crates/gateway-admin/tests/use_case/observability.rs @@ -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(), diff --git a/backend/crates/gateway-admin/tests/use_case/openai.rs b/backend/crates/gateway-admin/tests/use_case/openai.rs index 2f352ad84..89c2a6243 100644 --- a/backend/crates/gateway-admin/tests/use_case/openai.rs +++ b/backend/crates/gateway-admin/tests/use_case/openai.rs @@ -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"), diff --git a/backend/crates/gateway-admin/tests/use_case/xai.rs b/backend/crates/gateway-admin/tests/use_case/xai.rs index 2b12d635e..adf73a857 100644 --- a/backend/crates/gateway-admin/tests/use_case/xai.rs +++ b/backend/crates/gateway-admin/tests/use_case/xai.rs @@ -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, store: Arc) .build() .await } - -fn mutation(account_id: &str) -> CredentialMutation { - CredentialMutation { - context: context("request-xai"), - account_id: ProviderAccountId::new(account_id).expect("account ID"), - } -} diff --git a/backend/crates/gateway-api/src/admin/account_groups.rs b/backend/crates/gateway-api/src/admin/account_groups.rs new file mode 100644 index 000000000..db365f79b --- /dev/null +++ b/backend/crates/gateway-api/src/admin/account_groups.rs @@ -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, + page_size: Option, + search: Option, + enabled: Option, +} + +impl ListAccountGroupsQuery { + fn into_command(self) -> Result { + 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, + color: String, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct UpdateAccountGroupRequest { + id: String, + name: String, + description: Option, + 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, + color: String, + enabled: bool, + member_count: u64, + provider_counts: BTreeMap, + client_key_count: u64, + account_summary: AccountGroupAccountSummaryView, + capacity: AccountGroupCapacityView, + usage: AccountGroupUsageView, + created_at: DateTime, + updated_at: DateTime, +} + +#[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, + total_slots: u64, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountGroupUsageView { + today_usd: String, + total_usd: String, +} + +impl From 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, + page: PageMeta, + config_revision: u64, +} + +impl From 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, + enabled: bool, +} + +impl From 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, + total: u64, + config_revision: u64, +} + +impl From 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, + config_revision: u64, +} + +impl From 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() -> Router +where + S: AdminSessionState + Clone + Send + Sync + 'static, +{ + Router::new() + .route("/api/admin/account-groups", get(list::)) + .route("/api/admin/account-groups/create", post(create::)) + .route("/api/admin/account-groups/update", post(update::)) + .route("/api/admin/account-groups/enable", post(enable::)) + .route("/api/admin/account-groups/disable", post(disable::)) + .route("/api/admin/account-groups/delete", post(delete::)) + .route("/api/admin/account-groups/members", get(members::)) +} + +async fn list( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +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( + _auth: AdminAuth, + State(state): State, + Query(query): Query, +) -> Result +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( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +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( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +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( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + set_enabled(auth, state, request, true).await +} + +async fn disable( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +where + S: AdminSessionState + Send + Sync, +{ + set_enabled(auth, state, request, false).await +} + +async fn set_enabled( + auth: AdminAuth, + state: S, + request: AccountGroupIdRequest, + enabled: bool, +) -> Result>, 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( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +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, +) -> Result>, 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::new(value).map_err(|_| AdminError::bad_request("Invalid account group ID")) +} + +fn group_color(value: &str) -> Result { + 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") +} diff --git a/backend/crates/gateway-api/src/admin/accounts.rs b/backend/crates/gateway-api/src/admin/accounts.rs index 30c4a1bcc..a06c8f19c 100644 --- a/backend/crates/gateway-api/src/admin/accounts.rs +++ b/backend/crates/gateway-api/src/admin/accounts.rs @@ -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, pub page_size: Option, pub provider: Option, + pub group_id: Option, pub search: Option, pub status: Option, pub sort_by: Option, pub sort_direction: Option, } +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct BatchUpdateAccountsRequest { + pub account_ids: Vec, + pub enabled: bool, + #[serde(deserialize_with = "deserialize_required_nullable")] + pub concurrency_limit: Option, + pub weight: u64, + pub group_ids: Vec, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct BatchUpdatedAccountsData { + account_ids: Vec, + config_revision: u64, +} + +impl From 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::>().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 { + 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 { @@ -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, pub resource_ref: String, pub email: Option, pub account_id: Option, @@ -191,6 +265,8 @@ pub struct AccountView { /// 最近一次失败的上游错误描述；仅错误状态存在。 pub error_message: Option, pub enabled: bool, + pub concurrency_limit: Option, + pub weight: u16, pub access_token_expires_at: Option, pub access_token_expires_at_display: Option, pub refresh_token_expires_at: Option, @@ -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, + pub weight: u64, + pub group_ids: Vec, } -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 { 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 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, } -impl From 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, +) -> Result, 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, D::Error> +where + D: Deserializer<'de>, +{ + Option::::deserialize(deserializer) +} + +fn parse_account_weight(value: u64) -> Result { + u16::try_from(value) + .ok() + .and_then(AccountWeight::new) + .ok_or_else(|| WireValidationError::new("weight")) +} + +fn validate_wire_group_ids(values: &[String]) -> Result, WireValidationError> { + if values.len() > MAX_ACCOUNT_GROUP_BATCH + || values.iter().collect::>().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() -> Router where @@ -900,9 +1043,12 @@ where .route("/api/admin/accounts/import", post(import_accounts::)) .route("/api/admin/accounts/refresh", post(refresh_account::)) .route("/api/admin/accounts/rotate", post(rotate_account::)) - .route("/api/admin/accounts/enable", post(enable_account::)) - .route("/api/admin/accounts/disable", post(disable_account::)) + .route("/api/admin/accounts/update", post(update_account::)) .route("/api/admin/accounts/delete", post(delete_accounts::)) + .route( + "/api/admin/accounts/batch-update", + post(batch_update_accounts::), + ) .route("/api/admin/accounts/quota", get(account_quota::)) .route( "/api/admin/accounts/quota/refresh", @@ -927,6 +1073,27 @@ where ) } +async fn batch_update_accounts( + auth: AdminAuth, + State(state): State, + Json(request): Json, +) -> Result +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( _auth: AdminAuth, State(state): State, @@ -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( +async fn update_account( auth: AdminAuth, State(state): State, - Json(request): Json, + Json(request): Json, ) -> Result where S: AdminSessionState + Send + Sync, { - mutate_account(auth, state, request, AccountMutationAction::Enable).await -} - -async fn disable_account( - auth: AdminAuth, - State(state): State, - Json(request): Json, -) -> Result -where - S: AdminSessionState + Send + Sync, -{ - mutate_account(auth, state, request, AccountMutationAction::Disable).await -} - -#[derive(Clone, Copy)] -enum AccountMutationAction { - Enable, - Disable, -} - -async fn mutate_account( - auth: AdminAuth, - state: S, - request: AccountMutationRequest, - action: AccountMutationAction, -) -> Result -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) -> 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) -> 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 diff --git a/backend/crates/gateway-api/src/admin/client_keys.rs b/backend/crates/gateway-api/src/admin/client_keys.rs index 6244c18b4..57e6b32b8 100644 --- a/backend/crates/gateway-api/src/admin/client_keys.rs +++ b/backend/crates/gateway-api/src/admin/client_keys.rs @@ -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, - provider_kind: String, + group_ids: Vec, max_concurrency: u64, requests_per_minute: u64, } @@ -147,14 +147,13 @@ impl CreateClientKeyRequest { pub fn into_command(self) -> Result { 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, - provider_kind: String, + group_ids: Vec, 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, - provider_kind: String, + routing_scope: &'static str, + groups: Vec, + provider_kinds: Vec, prefix: String, enabled: bool, max_concurrency: u64, @@ -252,13 +252,42 @@ pub struct ClientKeyView { last_used_at: Option>, } +#[derive(Debug, Clone, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ClientKeyGroupView { + id: String, + name: String, + color: String, + enabled: bool, +} + impl From 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) -> Result, WireValidationError> { + if values.len() > 1000 + || values + .iter() + .collect::>() + .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( diff --git a/backend/crates/gateway-api/src/admin/mod.rs b/backend/crates/gateway-api/src/admin/mod.rs index 1734c557e..bc3c334ee 100644 --- a/backend/crates/gateway-api/src/admin/mod.rs +++ b/backend/crates/gateway-api/src/admin/mod.rs @@ -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::()) .merge(accounts::router::()) .merge(auth::router::()) .merge(backups::router::()) diff --git a/backend/crates/gateway-api/src/admin/observability/presenter.rs b/backend/crates/gateway-api/src/admin/observability/presenter.rs index 17dbaebb0..9dfafe402 100644 --- a/backend/crates/gateway-api/src/admin/observability/presenter.rs +++ b/backend/crates/gateway-api/src/admin/observability/presenter.rs @@ -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, diff --git a/backend/crates/gateway-api/src/admin/observability/wire.rs b/backend/crates/gateway-api/src/admin/observability/wire.rs index f91fbc86f..5f982b65d 100644 --- a/backend/crates/gateway-api/src/admin/observability/wire.rs +++ b/backend/crates/gateway-api/src/admin/observability/wire.rs @@ -84,6 +84,9 @@ pub struct UsageRecordView { pub id: String, pub request_id: String, pub client_api_key_id: Option, + pub routing_scope: String, + pub routing_group_refs: Vec, + pub routing_group_names_snapshot: Vec, pub kind: String, pub provider: Option, pub authentication_kind: Option, diff --git a/backend/crates/gateway-api/src/openai/error.rs b/backend/crates/gateway-api/src/openai/error.rs index e2151359b..16f8144b2 100644 --- a/backend/crates/gateway-api/src/openai/error.rs +++ b/backend/crates/gateway-api/src/openai/error.rs @@ -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 => { diff --git a/backend/crates/gateway-api/tests/admin/account_groups.rs b/backend/crates/gateway-api/tests/admin/account_groups.rs new file mode 100644 index 000000000..3b5a86f58 --- /dev/null +++ b/backend/crates/gateway-api/tests/admin/account_groups.rs @@ -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::().with_state(fixture.state()) +} + +async fn request( + router: Router, + method: Method, + uri: &str, + body: Option, + 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") +} diff --git a/backend/crates/gateway-api/tests/admin/accounts.rs b/backend/crates/gateway-api/tests/admin/accounts.rs index 3efc5ec1d..c41fc5f4b 100644 --- a/backend/crates/gateway-api/tests/admin/accounts.rs +++ b/backend/crates/gateway-api/tests/admin/accounts.rs @@ -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::(json!({ + "accountIds": ["acct_test"], + "concurrencyLimit": null, + "weight": 1, + "groupIds": [] + })) + .is_err() + ); + assert!( + serde_json::from_value::(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::(json!({ + "accountIds": ["acct_test"], + "enabled": true, + "weight": 1, + "groupIds": [] + })) + .is_err() + ); + assert!( + serde_json::from_value::(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::(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::(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::(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")], }); diff --git a/backend/crates/gateway-api/tests/admin/client_keys.rs b/backend/crates/gateway-api/tests/admin/client_keys.rs index 44d0baba3..97c9b5234 100644 --- a/backend/crates/gateway-api/tests/admin/client_keys.rs +++ b/backend/crates/gateway-api/tests/admin/client_keys.rs @@ -105,7 +105,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() { let valid = serde_json::from_value::(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::(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()); diff --git a/backend/crates/gateway-api/tests/admin/mod.rs b/backend/crates/gateway-api/tests/admin/mod.rs index eab9d45d5..b6dc48b53 100644 --- a/backend/crates/gateway-api/tests/admin/mod.rs +++ b/backend/crates/gateway-api/tests/admin/mod.rs @@ -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, + members: BTreeMap>, +} + +pub(super) struct MemoryAccountGroupStore { + state: Mutex, +} + +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 { + 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::>(); + 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 { + 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 { + 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 { + 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 { + 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 { + 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 { + 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 { @@ -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 { + ) -> AdminStoreResult { Err(unavailable("account enabled")) } + async fn batch_update_accounts( + &self, + _: BatchUpdateAccounts, + _: &MutationContext, + ) -> AdminStoreResult { + 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, total_slots: u64) -> AccountGroupCapacity { + AccountGroupCapacity { + used_slots, + total_slots, + } +} + +fn usage(today_usd: &str, total_usd: &str) -> AccountGroupUsage { + AccountGroupUsage { + today_usd: today_usd.parse::().expect("today cost"), + total_usd: total_usd.parse::().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) } diff --git a/backend/crates/gateway-api/tests/admin/observability/response.rs b/backend/crates/gateway-api/tests/admin/observability/response.rs index 2b2b26c4a..4191bd0f6 100644 --- a/backend/crates/gateway-api/tests/admin/observability/response.rs +++ b/backend/crates/gateway-api/tests/admin/observability/response.rs @@ -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", diff --git a/backend/crates/gateway-api/tests/architecture.rs b/backend/crates/gateway-api/tests/architecture.rs index 07c25307c..e25850982 100644 --- a/backend/crates/gateway-api/tests/architecture.rs +++ b/backend/crates/gateway-api/tests/architecture.rs @@ -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", diff --git a/backend/crates/gateway-api/tests/openai/mod.rs b/backend/crates/gateway-api/tests/openai/mod.rs index cdc866eaa..e43204e32 100644 --- a/backend/crates/gateway-api/tests/openai/mod.rs +++ b/backend/crates/gateway-api/tests/openai/mod.rs @@ -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) -> 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, NativeContinuationStoreError>> { Box::pin(async { unreachable!("authentication fixture does not execute") }) diff --git a/backend/crates/gateway-api/tests/openai/responses/mod.rs b/backend/crates/gateway-api/tests/openai/responses/mod.rs index f55eb3f74..97a0e27fc 100644 --- a/backend/crates/gateway-api/tests/openai/responses/mod.rs +++ b/backend/crates/gateway-api/tests/openai/responses/mod.rs @@ -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"); diff --git a/backend/crates/gateway-core/src/engine/continuation.rs b/backend/crates/gateway-core/src/engine/continuation.rs index 4a7368a75..e678822f8 100644 --- a/backend/crates/gateway-core/src/engine/continuation.rs +++ b/backend/crates/gateway-core/src/engine/continuation.rs @@ -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", &"") .field("upstream_response_id", &"") + .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) -> Self { + pub fn unavailable(detail: impl Into) -> Self { Self { + kind: NativeContinuationStoreErrorKind::Unavailable, detail: detail.into(), } } + /// 构造必须 fail closed 的无效记录错误；`detail` 不应包含 response handle。 + #[must_use] + pub fn invalid_data(detail: impl Into) -> 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, NativeContinuationStoreError>>; diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index 6bddd6f83..104bcb524 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -46,6 +46,10 @@ impl AccountSelection { Self::Diagnostic(account) => Some(account), } } + + const fn is_diagnostic(&self) -> bool { + matches!(self, Self::Diagnostic(_)) + } } impl AttemptCoordinator @@ -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 { account_state_owner: Option, cancellation: CancellationToken, attempts: u32, + candidate_index: usize, excluded_accounts: BTreeSet, credential_recovery_attempted_accounts: BTreeSet, /// 凭据恢复后的一次性同账号钉选；只约束紧随其后的 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; diff --git a/backend/crates/gateway-core/src/engine/credential.rs b/backend/crates/gateway-core/src/engine/credential.rs index 1c44eac15..cf53ff4e4 100644 --- a/backend/crates/gateway-core/src/engine/credential.rs +++ b/backend/crates/gateway-core/src/engine/credential.rs @@ -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 { + 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 { + 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, + weight: AccountWeight, credential_state: CredentialState, quota: QuotaState, last_error_reason: Option, @@ -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, + 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 { + 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 { self.access_token_expires_at @@ -1393,6 +1478,7 @@ pub struct AccountSelectionContext { pub preferred_account: Option, pub round_robin_cursor: u64, pub eligibility: AccountEligibilityPolicy, + pub account_scope: Option>, } /// 选择账号时是否执行本地调度资格投影。 @@ -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> { - 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 { + 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 diff --git a/backend/crates/gateway-core/src/engine/execution.rs b/backend/crates/gateway-core/src/engine/execution.rs index 275060cb1..4c4a6ca64 100644 --- a/backend/crates/gateway-core/src/engine/execution.rs +++ b/backend/crates/gateway-core/src/engine/execution.rs @@ -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, ) -> Result { - 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 { 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 { 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 { 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", ), diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index a071804ad..7cae7179a 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -233,6 +233,7 @@ pub struct AccountAttemptContext { state_owner: Option, credential_recovery_attempted: bool, diagnostic_required_account: bool, + account_scope: Option>, } 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) -> 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> { + 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> { + 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, 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")] diff --git a/backend/crates/gateway-core/src/error.rs b/backend/crates/gateway-core/src/error.rs index 9d948ae83..df6c5296a 100644 --- a/backend/crates/gateway-core/src/error.rs +++ b/backend/crates/gateway-core/src/error.rs @@ -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 { diff --git a/backend/crates/gateway-core/src/policy/mod.rs b/backend/crates/gateway-core/src/policy/mod.rs index ab328f523..3f56e54b2 100644 --- a/backend/crates/gateway-core/src/policy/mod.rs +++ b/backend/crates/gateway-core/src/policy/mod.rs @@ -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, 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, 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 { + &self.account_scope } #[must_use] diff --git a/backend/crates/gateway-core/src/provider_ports/mod.rs b/backend/crates/gateway-core/src/provider_ports/mod.rs index 2bd85a32f..45ff071db 100644 --- a/backend/crates/gateway-core/src/provider_ports/mod.rs +++ b/backend/crates/gateway-core/src/provider_ports/mod.rs @@ -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>; diff --git a/backend/crates/gateway-core/src/routing/mod.rs b/backend/crates/gateway-core/src/routing/mod.rs index 683ff1e3c..775587215 100644 --- a/backend/crates/gateway-core/src/routing/mod.rs +++ b/backend/crates/gateway-core/src/routing/mod.rs @@ -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) -> Result { + 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>, +} + +impl RuntimeAccount { + #[must_use] + pub fn new(provider_kind: ProviderKind, group_ids: BTreeSet) -> 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 { + &self.group_ids + } +} + +/// 全快照共享的账号、Provider 与分组反向索引。 +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct RuntimeAccountDirectory { + accounts: BTreeMap, + providers_with_accounts: BTreeSet, + providers_by_group: BTreeMap>, +} + +impl RuntimeAccountDirectory { + #[must_use] + pub fn new(accounts: BTreeMap) -> Self { + let mut providers_with_accounts = BTreeSet::new(); + let mut providers_by_group = BTreeMap::>::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 { + &self.providers_with_accounts + } + + #[must_use] + pub fn providers_for_groups<'a>( + &self, + group_ids: impl IntoIterator, + ) -> BTreeSet { + 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) -> 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>, + provider_kinds: Arc>, + }, +} + +impl ClientRoutingScope { + #[must_use] + pub fn all_accounts() -> Self { + Self::AllAccounts + } + + pub fn restricted( + bound_groups: Vec, + enabled_group_ids: BTreeSet, + provider_kinds: BTreeSet, + ) -> Result { + 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, + client_scope: ClientRoutingScope, +} + +impl FrozenAccountScope { + #[must_use] + pub const fn new( + directory: Arc, + 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 { + 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 { + &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, + /// 管理端 connection test 显式限制的 Provider；普通请求留空。 + pub required_provider: Option, pub blocked_providers: BTreeSet, } @@ -505,6 +786,7 @@ pub struct ProviderCandidate { provider: ProviderKind, upstream_model: UpstreamModelId, emulated_features: BTreeSet, + account_scope: Arc, } impl ProviderCandidate { @@ -522,6 +804,11 @@ impl ProviderCandidate { pub const fn emulated_features(&self) -> &BTreeSet { &self.emulated_features } + + #[must_use] + pub const fn account_scope(&self) -> &Arc { + &self.account_scope + } } /// 一次请求冻结的 Provider 尝试顺序。 @@ -532,6 +819,7 @@ pub struct RoutingPlan { public_model: PublicModelId, operation: OperationKind, max_attempts: NonZeroU32, + account_scope: Arc, candidates: Arc<[ProviderCandidate]>, } @@ -561,6 +849,11 @@ impl RoutingPlan { self.max_attempts } + #[must_use] + pub const fn account_scope(&self) -> &Arc { + &self.account_scope + } + #[must_use] pub fn candidates(&self) -> &[ProviderCandidate] { &self.candidates diff --git a/backend/crates/gateway-core/src/routing/snapshot.rs b/backend/crates/gateway-core/src/routing/snapshot.rs index e74c1626c..937807557 100644 --- a/backend/crates/gateway-core/src/routing/snapshot.rs +++ b/backend/crates/gateway-core/src/routing/snapshot.rs @@ -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, limits: RateLimits, } @@ -75,18 +76,67 @@ impl SnapshotClientPolicyFacts { pub fn new( key_id: ClientApiKeyId, plaintext_key: PlaintextClientApiKey, - provider_kind: impl Into, + group_ids: Vec, 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) -> 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, + account_groups: Vec, + provider_accounts: Vec, + group_memberships: Vec, } impl SnapshotFacts { @@ -103,12 +156,18 @@ impl SnapshotFacts { observed_current_revision: ConfigRevision, settings: SnapshotSettingsFacts, client_policies: Vec, + account_groups: Vec, + provider_accounts: Vec, + group_memberships: Vec, ) -> 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 { let provider_kinds = providers.provider_kinds().cloned().collect::>(); + let registered_providers = provider_kinds.iter().cloned().collect::>(); - // 实时目录只提供公开模型与能力提示；查询失败时保留 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::>::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::, 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>>, model_mappings: Arc>, provider_catalog_generations: Arc>, + known_provider_catalogs: Arc>, + account_directory: Arc, client_policies: Arc>, } @@ -312,6 +458,7 @@ impl RuntimeSnapshot { } } + let mut known_provider_catalogs = BTreeSet::new(); let mut model_map = BTreeMap::>::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) -> Self { + self.account_directory = directory; + self + } + + #[must_use] + fn with_known_provider_catalogs(mut self, providers: BTreeSet) -> Self { + self.known_provider_catalogs = Arc::new(providers); + self + } + + #[must_use] + pub fn all_account_scope(&self) -> Arc { + 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 { + scope + .provider_kinds() + .iter() + .flat_map(|provider| self.public_models_for_provider(provider)) + .collect::>() + .into_iter() + .collect() + } + + /// 合并冻结账号范围内实际存在 Provider 的公开模型画像。 + #[must_use] + pub fn public_model_profiles_for_scope( + &self, + scope: &FrozenAccountScope, + ) -> Vec { + 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, context: &RoutingContext, ) -> Result { 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), }) } diff --git a/backend/crates/gateway-core/tests/engine/continuation.rs b/backend/crates/gateway-core/tests/engine/continuation.rs index 2bc420f54..f6e9651b2 100644 --- a/backend/crates/gateway-core/tests/engine/continuation.rs +++ b/backend/crates/gateway-core/tests/engine/continuation.rs @@ -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"))); +} diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index 0e6ac914f..efef8ffa9 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -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"), ); diff --git a/backend/crates/gateway-core/tests/engine/credential.rs b/backend/crates/gateway-core/tests/engine/credential.rs index 3afbedac7..878e3da04 100644 --- a/backend/crates/gateway-core/tests/engine/credential.rs +++ b/backend/crates/gateway-core/tests/engine/credential.rs @@ -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) -> 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)); diff --git a/backend/crates/gateway-core/tests/engine/execution.rs b/backend/crates/gateway-core/tests/engine/execution.rs index 5dd90c82b..e73012af1 100644 --- a/backend/crates/gateway-core/tests/engine/execution.rs +++ b/backend/crates/gateway-core/tests/engine/execution.rs @@ -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, 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, 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 { + 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(), )], diff --git a/backend/crates/gateway-core/tests/policy/mod.rs b/backend/crates/gateway-core/tests/policy/mod.rs index ee37ee5a8..912608c10 100644 --- a/backend/crates/gateway-core/tests/policy/mod.rs +++ b/backend/crates/gateway-core/tests/policy/mod.rs @@ -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 { + 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(), ); diff --git a/backend/crates/gateway-core/tests/routing/mod.rs b/backend/crates/gateway-core/tests/routing/mod.rs index 428a3d6e9..db475aeec 100644 --- a/backend/crates/gateway-core/tests/routing/mod.rs +++ b/backend/crates/gateway-core/tests/routing/mod.rs @@ -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 { + 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 { + 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::>(); 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"), + ) + ); +} diff --git a/backend/crates/gateway-core/tests/routing/snapshot.rs b/backend/crates/gateway-core/tests/routing/snapshot.rs index 4edb1b974..00f4e02f9 100644 --- a/backend/crates/gateway-core/tests/routing/snapshot.rs +++ b/backend/crates/gateway-core/tests/routing/snapshot.rs @@ -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::::new(), + Vec::::new(), + Vec::::new(), ) } diff --git a/backend/crates/gateway-store/src/bundle.rs b/backend/crates/gateway-store/src/bundle.rs index 1ea76459c..e0cd316a4 100644 --- a/backend/crates/gateway-store/src/bundle.rs +++ b/backend/crates/gateway-store/src/bundle.rs @@ -90,6 +90,11 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult { pool.clone(), Some(Arc::clone(&cooldowns) as Arc), )), + Arc::new(postgres::PgAccountGroupRepository::with_runtime_state( + pool.clone(), + credential_leases.clone(), + Arc::clone(&cooldowns) as Arc, + )), Arc::new(AdminAuthStoreAdapter { security: postgres::PgAdminSecurityAuditRepository::new(pool.clone()), settings: postgres::PgRuntimeSettingsRepository::new(pool.clone()), diff --git a/backend/crates/gateway-store/src/postgres/account_groups.rs b/backend/crates/gateway-store/src/postgres/account_groups.rs new file mode 100644 index 000000000..5b6ddf777 --- /dev/null +++ b/backend/crates/gateway-store/src/postgres/account_groups.rs @@ -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, + cooldowns: Option>, +} + +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, + ) -> Self { + Self { + pool, + runtime_signals: Some(runtime_signals), + cooldowns: Some(cooldowns), + } + } + + async fn current_revision(&self) -> AdminStoreResult { + 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 { + 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( + &self, + audit: super::AdminAuditEvent, + mutation: F, + ) -> AdminStoreResult + 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 { + 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::>>() + .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 { + 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::>>() + .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 { + 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 { + 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 { + 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 { + 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 { + 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, 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 { + let mut statement = + QueryBuilder::::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::() + .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> { + 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 { + let provider_counts: serde_json::Value = row + .try_get("provider_counts") + .map_err(|_| invalid("invalid provider counts"))?; + let provider_counts = serde_json::from_value::>(provider_counts) + .map_err(|_| invalid("invalid provider counts"))?; + Ok(AccountGroupRecord { + id: AccountGroupId::new( + row.try_get::("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::("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::>(); + 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::>::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::>() + .into_iter() + .collect::>(); + 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::>() + }) + } 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> { + 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::("today_usd") + .map_err(|_| invalid("invalid today cost"))?; + let total = row + .try_get::("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 { + 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::("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::try_from( + row.try_get::(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) +} diff --git a/backend/crates/gateway-store/src/postgres/client_keys.rs b/backend/crates/gateway-store/src/postgres/client_keys.rs index 8581cb13d..16bc210d3 100644 --- a/backend/crates/gateway-store/src/postgres/client_keys.rs +++ b/backend/crates/gateway-store/src/postgres/client_keys.rs @@ -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, pub limits: RateLimits, } @@ -57,7 +59,7 @@ impl ClientApiKeySnapshot { pub(crate) fn from_persisted( id: String, key: String, - provider_kind: String, + group_ids: Vec, max_concurrency: i64, requests_per_minute: i64, ) -> StoreResult { @@ -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::>>()?, 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, - pub provider_kind: String, + pub groups: Vec, + pub provider_kinds: Vec, pub prefix: String, pub enabled: bool, pub max_concurrency: u64, @@ -110,6 +117,14 @@ pub struct ClientApiKeyRecord { pub updated_at: DateTime, } +#[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, - pub provider_kind: String, + pub group_ids: Vec, 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, - 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, - pub provider_kind: String, + pub group_ids: Vec, 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; async fn get_client_api_key(&self, id: &str) -> StoreResult>; async fn reveal_client_api_key(&self, id: &str) -> StoreResult>; - async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult; async fn touch_client_api_keys( &self, touched_at: &BTreeMap>, @@ -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::::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::() + .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::>>()?; 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> { 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> { 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 { - 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>, @@ -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>>()?, + provider_kinds: record + .provider_kinds + .into_iter() + .map(|provider| { + ProviderKind::new(provider) + .map_err(|_| admin_store_error(ENTITY, invalid("invalid provider kind"))) + }) + .collect::>>()?, 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::>().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 { 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, - bool, - i64, - i64, - Option>, - DateTime, - DateTime, -); - -fn client_record_from_row(row: ClientRecordRow) -> StoreResult { +fn client_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult { + 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"))?, }) } diff --git a/backend/crates/gateway-store/src/postgres/execution.rs b/backend/crates/gateway-store/src/postgres/execution.rs index 355dbed46..e726edb46 100644 --- a/backend/crates/gateway-store/src/postgres/execution.rs +++ b/backend/crates/gateway-store/src/postgres/execution.rs @@ -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, pub client_api_key_ref: String, pub config_revision: u64, + pub routing_scope: String, + pub routing_group_refs: Vec, + 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, 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(), diff --git a/backend/crates/gateway-store/src/postgres/mod.rs b/backend/crates/gateway-store/src/postgres/mod.rs index e8f731e54..ffd81f526 100644 --- a/backend/crates/gateway-store/src/postgres/mod.rs +++ b/backend/crates/gateway-store/src/postgres/mod.rs @@ -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 { 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 } => { diff --git a/backend/crates/gateway-store/src/postgres/observability/admin_adapter.rs b/backend/crates/gateway-store/src/postgres/observability/admin_adapter.rs index fa4171470..ba3930cfd 100644 --- a/backend/crates/gateway-store/src/postgres/observability/admin_adapter.rs +++ b/backend/crates/gateway-store/src/postgres/observability/admin_adapter.rs @@ -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, - ) -> StoreResult<(ProviderAccountMetrics, Vec)> { + ) -> StoreResult<(ProviderAccountMetrics, Vec)> { 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, ) -> AdminStoreResult> { - 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::>(); 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), })) } diff --git a/backend/crates/gateway-store/src/postgres/observability/mapping.rs b/backend/crates/gateway-store/src/postgres/observability/mapping.rs index 4e692d371..6d4b4ecda 100644 --- a/backend/crates/gateway-store/src/postgres/observability/mapping.rs +++ b/backend/crates/gateway-store/src/postgres/observability/mapping.rs @@ -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::>>( + row, + "routing_group_names_snapshot", + )? + .0, protocol: get(row, "protocol")?, operation: get(row, "operation")?, endpoint: get(row, "endpoint")?, diff --git a/backend/crates/gateway-store/src/postgres/observability/model.rs b/backend/crates/gateway-store/src/postgres/observability/model.rs index cb43aac33..26eceff75 100644 --- a/backend/crates/gateway-store/src/postgres/observability/model.rs +++ b/backend/crates/gateway-store/src/postgres/observability/model.rs @@ -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, + pub routing_group_names_snapshot: Vec, pub protocol: String, pub operation: String, pub endpoint: String, diff --git a/backend/crates/gateway-store/src/postgres/observability/queries/usage.rs b/backend/crates/gateway-store/src/postgres/observability/queries/usage.rs index 7157bcf73..607297062 100644 --- a/backend/crates/gateway-store/src/postgres/observability/queries/usage.rs +++ b/backend/crates/gateway-store/src/postgres/observability/queries/usage.rs @@ -103,7 +103,9 @@ pub(crate) fn push_completed_usage_fact_filter(query: &mut QueryBuilder AdminStoreResult>> { + 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::>::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>> { + 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::>(); @@ -298,9 +408,20 @@ impl AccountStore for PgAdminAccountStore { .into_iter() .skip(offset) .take(page_size) + .collect::>(); + let item_ids = items + .iter() + .map(|item| item.account.id.clone()) + .collect::>(); + 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 { - 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 { + 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 { + 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::>>()?; + 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( diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs index a4a09f1a2..4448f0fe9 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/core_adapter.rs @@ -31,6 +31,8 @@ impl ProviderAccountStore for PgProviderAccountRepository { .map(DateTime::::from), next_refresh_at: account.account.next_refresh_at().map(DateTime::::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(), }) diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs index 882c2adef..bd5e221fa 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/mapping.rs @@ -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, }) diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs index a2f89ebd3..18691d309 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/mod.rs @@ -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, diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs index f08137e1f..c4a828e1b 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/repository.rs @@ -56,9 +56,9 @@ pub trait ProviderAccountAdminRepository: Send + Sync { command: RotateProviderAccount, ) -> StoreResult; - async fn set_provider_account_enabled_admin( + async fn batch_update_provider_accounts_admin( &self, - command: SetProviderAccountEnabled, + command: BatchUpdateProviderAccountsAdmin, ) -> StoreResult; 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 { - 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::>(); + 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, + 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::>(); + let expected = account_ids.iter().cloned().collect::>(); + 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( transaction: Transaction<'_, Postgres>, result: StoreResult, diff --git a/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs b/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs index ce213471a..5799b3dac 100644 --- a/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs +++ b/backend/crates/gateway-store/src/postgres/provider_accounts/rows.rs @@ -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>, pub next_refresh_at: Option>, pub enabled: bool, + pub concurrency_limit: Option, + pub weight: AccountWeight, pub credential_state: CredentialState, pub credential_observed_at: DateTime, pub quota: QuotaState, @@ -150,6 +153,8 @@ pub struct NewProviderAccount { pub access_token_expires_at: Option>, pub next_refresh_at: Option>, pub enabled: bool, + pub concurrency_limit: Option, + pub weight: AccountWeight, pub credential_state: CredentialState, pub credential_observed_at: DateTime, } @@ -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, pub enabled: bool, + pub concurrency_limit: Option, + pub weight: AccountWeight, + pub group_ids: Vec, 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::>(&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::(&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, diff --git a/backend/crates/gateway-store/src/postgres/snapshot.rs b/backend/crates/gateway-store/src/postgres/snapshot.rs index 158ff150d..8d0dd4000 100644 --- a/backend/crates/gateway-store/src/postgres/snapshot.rs +++ b/backend/crates/gateway-store/src/postgres/snapshot.rs @@ -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, + pub account_groups: Vec, + pub provider_accounts: Vec, + pub group_memberships: Vec, +} + +#[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::, _>>()?; + 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::, _>>()?; 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> { - 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, 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> { + 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> { + 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> { + 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::new(to_u64(value)?) } diff --git a/backend/crates/gateway-store/src/redis/credential_leases.rs b/backend/crates/gateway-store/src/redis/credential_leases.rs index e4090ef66..119a14768 100644 --- a/backend/crates/gateway-store/src/redis/credential_leases.rs +++ b/backend/crates/gateway-store/src/redis/credential_leases.rs @@ -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 { - 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 { + 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 { - let fingerprint = resource_fingerprint("provider scheduling cursor", provider_kind)?; + fn scheduling_cursor_key( + &self, + client_api_key_id: &ClientApiKeyId, + provider_kind: &ProviderKind, + ) -> StoreResult { + 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 { 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> { 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)) }) } diff --git a/backend/crates/gateway-store/src/redis/native_continuation.rs b/backend/crates/gateway-store/src/redis/native_continuation.rs index c2cee05f0..cf4b9f8b6 100644 --- a/backend/crates/gateway-store/src/redis/native_continuation.rs +++ b/backend/crates/gateway-store/src/redis/native_continuation.rs @@ -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 { 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 { 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, NativeContinuationStoreError>> { Box::pin(async move { @@ -124,10 +128,10 @@ impl NativeContinuationPort for RedisNativeContinuationRepository { .query_async::>(&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::(&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 Result Result { 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 Ok(NativeContinuationScope::Persisted), "connection_local" => Ok(NativeContinuationScope::ConnectionLocal), - _ => Err(NativeContinuationStoreError::new( + _ => Err(NativeContinuationStoreError::invalid_data( "pin payload contains an unsupported continuation scope", )), } diff --git a/backend/crates/gateway-store/tests/postgres/account_groups.rs b/backend/crates/gateway-store/tests/postgres/account_groups.rs new file mode 100644 index 000000000..484b8332b --- /dev/null +++ b/backend/crates/gateway-store/tests/postgres/account_groups.rs @@ -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::>(); + 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::>(), + ["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::>(), + ["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 = 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 = 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) -> 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") +} diff --git a/backend/crates/gateway-store/tests/postgres/admission_recovery.rs b/backend/crates/gateway-store/tests/postgres/admission_recovery.rs index 9937ec7de..56d7e5d27 100644 --- a/backend/crates/gateway-store/tests/postgres/admission_recovery.rs +++ b/backend/crates/gateway-store/tests/postgres/admission_recovery.rs @@ -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) diff --git a/backend/crates/gateway-store/tests/postgres/client_keys.rs b/backend/crates/gateway-store/tests/postgres/client_keys.rs index 12abad905..b4a340f63 100644 --- a/backend/crates/gateway-store/tests/postgres/client_keys.rs +++ b/backend/crates/gateway-store/tests/postgres/client_keys.rs @@ -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, diff --git a/backend/crates/gateway-store/tests/postgres/execution.rs b/backend/crates/gateway-store/tests/postgres/execution.rs index 1a6e8d3fd..5c277278f 100644 --- a/backend/crates/gateway-store/tests/postgres/execution.rs +++ b/backend/crates/gateway-store/tests/postgres/execution.rs @@ -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, 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) diff --git a/backend/crates/gateway-store/tests/postgres/mod.rs b/backend/crates/gateway-store/tests/postgres/mod.rs index ca21dcc35..23a23a30c 100644 --- a/backend/crates/gateway-store/tests/postgres/mod.rs +++ b/backend/crates/gateway-store/tests/postgres/mod.rs @@ -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] diff --git a/backend/crates/gateway-store/tests/postgres/observability.rs b/backend/crates/gateway-store/tests/postgres/observability.rs index ab23dc1fa..fa2877c82 100644 --- a/backend/crates/gateway-store/tests/postgres/observability.rs +++ b/backend/crates/gateway-store/tests/postgres/observability.rs @@ -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) diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs index 9b4d7855b..9a81c37d1 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts.rs @@ -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, 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, 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, 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, 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 { + 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"; diff --git a/backend/crates/gateway-store/tests/postgres/snapshot.rs b/backend/crates/gateway-store/tests/postgres/snapshot.rs index 4a6dc2efd..74e506744 100644 --- a/backend/crates/gateway-store/tests/postgres/snapshot.rs +++ b/backend/crates/gateway-store/tests/postgres/snapshot.rs @@ -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 diff --git a/backend/crates/gateway-store/tests/postgres/snapshots.rs b/backend/crates/gateway-store/tests/postgres/snapshots.rs index 7b43a7a83..63208b10a 100644 --- a/backend/crates/gateway-store/tests/postgres/snapshots.rs +++ b/backend/crates/gateway-store/tests/postgres/snapshots.rs @@ -665,6 +665,9 @@ fn new_request(id: &str, started_at: DateTime) -> 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) { async fn seed_api_key(pool: &PgPool, id: &str, name: &str, now: DateTime) { 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) diff --git a/backend/crates/gateway-store/tests/redis/credential_leases.rs b/backend/crates/gateway-store/tests/redis/credential_leases.rs index c31bed6a5..675258acf 100644 --- a/backend/crates/gateway-store/tests/redis/credential_leases.rs +++ b/backend/crates/gateway-store/tests/redis/credential_leases.rs @@ -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::>(&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::(&mut connection) diff --git a/backend/crates/gateway-store/tests/redis/native_continuation.rs b/backend/crates/gateway-store/tests/redis/native_continuation.rs index 3a7f9d6c3..4ec23c607 100644 --- a/backend/crates/gateway-store/tests/redis/native_continuation.rs +++ b/backend/crates/gateway-store/tests/redis/native_continuation.rs @@ -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 { redis::cmd("KEYS") .arg(format!("{namespace}:*")) diff --git a/backend/crates/providers/openai/src/credential/affinity.rs b/backend/crates/providers/openai/src/credential/affinity.rs index d7a8def2b..5ec43dbb3 100644 --- a/backend/crates/providers/openai/src/credential/affinity.rs +++ b/backend/crates/providers/openai/src/credential/affinity.rs @@ -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 { 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 { diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index b6c2c18ff..45786dd1c 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -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::>(); 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(), ), diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index 974cf6332..4f6f4e91d 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -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); diff --git a/backend/crates/providers/openai/tests/admin.rs b/backend/crates/providers/openai/tests/admin.rs index 19f7f2d29..55415d038 100644 --- a/backend/crates/providers/openai/tests/admin.rs +++ b/backend/crates/providers/openai/tests/admin.rs @@ -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::::from), next_refresh_at: account.next_refresh_at().map(DateTime::::from), enabled: account.enabled(), + concurrency_limit: account.concurrency_limit(), + weight: account.weight(), credential_state: account.credential_state(), credential_observed_at: now, quota: account.quota(), diff --git a/backend/crates/providers/openai/tests/credential/contract.rs b/backend/crates/providers/openai/tests/credential/contract.rs index 823ddde5b..b8e58e695 100644 --- a/backend/crates/providers/openai/tests/credential/contract.rs +++ b/backend/crates/providers/openai/tests/credential/contract.rs @@ -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, id: &str, token: &str) { })); } +fn contract_account_scope() -> Arc { + 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::>(); + Arc::new(FrozenAccountScope::new( + Arc::new(RuntimeAccountDirectory::new(accounts)), + ClientRoutingScope::all_accounts(), + )) +} + fn attempt(excluded_accounts: BTreeSet) -> 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(), ); diff --git a/backend/crates/providers/openai/tests/credential/refresh.rs b/backend/crates/providers/openai/tests/credential/refresh.rs index 11b67a9ac..acd0a6512 100644 --- a/backend/crates/providers/openai/tests/credential/refresh.rs +++ b/backend/crates/providers/openai/tests/credential/refresh.rs @@ -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> { diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index c5591e06f..654373bf5 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -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 { + 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::>(); + 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::::new(), None, None), + AccountAttemptContext::new(BTreeSet::::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(), ) diff --git a/backend/crates/providers/openai/tests/support.rs b/backend/crates/providers/openai/tests/support.rs index 879120f46..09dc3204c 100644 --- a/backend/crates/providers/openai/tests/support.rs +++ b/backend/crates/providers/openai/tests/support.rs @@ -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, + 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> { diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index e9a4a0a88..27d2a76f8 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -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::>(); 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(), ), diff --git a/backend/crates/providers/xai/src/provider.rs b/backend/crates/providers/xai/src/provider.rs index edc11babf..988456ce0 100644 --- a/backend/crates/providers/xai/src/provider.rs +++ b/backend/crates/providers/xai/src/provider.rs @@ -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 diff --git a/backend/crates/providers/xai/src/transport/session.rs b/backend/crates/providers/xai/src/transport/session.rs index e2b6b644a..d4211506a 100644 --- a/backend/crates/providers/xai/src/transport/session.rs +++ b/backend/crates/providers/xai/src/transport/session.rs @@ -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, deadline: SystemTime, + account_scope: Arc, + client_api_key_id: ClientApiKeyId, } impl GrokSessionSelection { @@ -224,6 +228,8 @@ impl GrokSessionSelection { required_account: Option, account_selection_policy: AccountSelectionPolicy, deadline: SystemTime, + account_scope: Arc, + 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 { + &self.account_scope + } + + #[must_use] + pub const fn client_api_key_id(&self) -> &ClientApiKeyId { + &self.client_api_key_id + } } /// Grok session selector 返回的 future。 diff --git a/backend/crates/providers/xai/tests/admin.rs b/backend/crates/providers/xai/tests/admin.rs index c7eed7f06..234c02f4b 100644 --- a/backend/crates/providers/xai/tests/admin.rs +++ b/backend/crates/providers/xai/tests/admin.rs @@ -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::::from), next_refresh_at: account.next_refresh_at().map(DateTime::::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< diff --git a/backend/crates/providers/xai/tests/credential/refresh.rs b/backend/crates/providers/xai/tests/credential/refresh.rs index 69446604b..bdce7e493 100644 --- a/backend/crates/providers/xai/tests/credential/refresh.rs +++ b/backend/crates/providers/xai/tests/credential/refresh.rs @@ -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< diff --git a/backend/crates/providers/xai/tests/credential/selector.rs b/backend/crates/providers/xai/tests/credential/selector.rs index 3a6b749d2..fbb50a2d9 100644 --- a/backend/crates/providers/xai/tests/credential/selector.rs +++ b/backend/crates/providers/xai/tests/credential/selector.rs @@ -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>, denied: Mutex>, + requests: Mutex>, } 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> { @@ -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 = cache.clone(); @@ -210,6 +218,21 @@ impl SelectorFixture { required_account: Option, 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::>(); 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; diff --git a/backend/crates/providers/xai/tests/provider/contract.rs b/backend/crates/providers/xai/tests/provider/contract.rs index 4e3b3dc3e..c31d271c2 100644 --- a/backend/crates/providers/xai/tests/provider/contract.rs +++ b/backend/crates/providers/xai/tests/provider/contract.rs @@ -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"), ); diff --git a/backend/crates/providers/xai/tests/support.rs b/backend/crates/providers/xai/tests/support.rs index 9b2062cad..9dd20b259 100644 --- a/backend/crates/providers/xai/tests/support.rs +++ b/backend/crates/providers/xai/tests/support.rs @@ -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, + 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 { 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) } diff --git a/backend/migrations/.frozen-sha256 b/backend/migrations/.frozen-sha256 index bdcefc8ac..86c6c0b16 100644 --- a/backend/migrations/.frozen-sha256 +++ b/backend/migrations/.frozen-sha256 @@ -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 diff --git a/backend/migrations/0005_account_groups.sql b/backend/migrations/0005_account_groups.sql new file mode 100644 index 000000000..9f61726cb --- /dev/null +++ b/backend/migrations/0005_account_groups.sql @@ -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; diff --git a/backend/migrations/0006_account_group_color.sql b/backend/migrations/0006_account_group_color.sql new file mode 100644 index 000000000..b9c333fed --- /dev/null +++ b/backend/migrations/0006_account_group_color.sql @@ -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}$'); diff --git a/backend/migrations/0007_account_group_color_alpha.sql b/backend/migrations/0007_account_group_color_alpha.sql new file mode 100644 index 000000000..07a51a21e --- /dev/null +++ b/backend/migrations/0007_account_group_color_alpha.sql @@ -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}$'); diff --git a/backend/migrations/0008_provider_account_scheduling.sql b/backend/migrations/0008_provider_account_scheduling.sql new file mode 100644 index 000000000..bf296ac3c --- /dev/null +++ b/backend/migrations/0008_provider_account_scheduling.sql @@ -0,0 +1,7 @@ +alter table provider_accounts + add column concurrency_limit bigint, + add column weight smallint not null default 1, + add constraint provider_accounts_concurrency_limit_ck + check (concurrency_limit is null or concurrency_limit between 1 and 4294967295), + add constraint provider_accounts_weight_ck + check (weight between 1 and 100);