mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
feat(routing): add account groups and weighted scheduling
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -21,9 +21,9 @@ pub mod ports;
|
||||
mod use_case;
|
||||
|
||||
pub use use_case::{
|
||||
accounts::AccountsService, auth::AuthService, backup::BackupService,
|
||||
client_keys::ClientKeyService, observability::ObservabilityService, openai::OpenAiService,
|
||||
settings::SettingsService, system::SystemService, xai::XaiService,
|
||||
account_groups::AccountGroupService, accounts::AccountsService, auth::AuthService,
|
||||
backup::BackupService, client_keys::ClientKeyService, observability::ObservabilityService,
|
||||
openai::OpenAiService, settings::SettingsService, system::SystemService, xai::XaiService,
|
||||
};
|
||||
|
||||
use model::{AdminError, AdminErrorKind};
|
||||
@@ -33,10 +33,10 @@ use ports::{
|
||||
system::SystemOperations,
|
||||
};
|
||||
use use_case::{
|
||||
accounts::DefaultAccountsService, auth::DefaultAuthService, backup::DefaultBackupService,
|
||||
client_keys::DefaultClientKeyService, observability::DefaultObservabilityService,
|
||||
openai::DefaultOpenAiService, settings::DefaultSettingsService, system::DefaultSystemService,
|
||||
xai::DefaultXaiService,
|
||||
account_groups::DefaultAccountGroupService, accounts::DefaultAccountsService,
|
||||
auth::DefaultAuthService, backup::DefaultBackupService, client_keys::DefaultClientKeyService,
|
||||
observability::DefaultObservabilityService, openai::DefaultOpenAiService,
|
||||
settings::DefaultSettingsService, system::DefaultSystemService, xai::DefaultXaiService,
|
||||
};
|
||||
|
||||
const OPENAI_PROVIDER_KIND: &str = "openai";
|
||||
@@ -147,6 +147,7 @@ pub enum AdminConfigError {
|
||||
pub struct AdminServices {
|
||||
auth: Arc<dyn AuthService>,
|
||||
accounts: Arc<dyn AccountsService>,
|
||||
account_groups: Arc<dyn AccountGroupService>,
|
||||
client_keys: Arc<dyn ClientKeyService>,
|
||||
observability: Arc<dyn ObservabilityService>,
|
||||
settings: Arc<dyn SettingsService>,
|
||||
@@ -167,6 +168,11 @@ impl AdminServices {
|
||||
self.accounts.as_ref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn account_groups(&self) -> &dyn AccountGroupService {
|
||||
self.account_groups.as_ref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn client_keys(&self) -> &dyn ClientKeyService {
|
||||
self.client_keys.as_ref()
|
||||
@@ -274,6 +280,10 @@ pub async fn initialize(
|
||||
let services = AdminServices {
|
||||
auth,
|
||||
accounts,
|
||||
account_groups: Arc::new(DefaultAccountGroupService::new(
|
||||
store.account_groups(),
|
||||
snapshot.clone(),
|
||||
)),
|
||||
client_keys: Arc::new(DefaultClientKeyService::new(
|
||||
store.client_keys(),
|
||||
snapshot.clone(),
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
//! Provider-neutral account group commands and projections.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_core::routing::{AccountGroupId, ProviderKind};
|
||||
|
||||
use super::{PageSize, Revision, observability::DecimalAmount};
|
||||
|
||||
/// Canonical `#RRGGBBAA` color persisted with an account group.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupColor(String);
|
||||
|
||||
impl AccountGroupColor {
|
||||
#[must_use]
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn parse(value: &str) -> Option<Self> {
|
||||
(value.len() == 9
|
||||
&& value.starts_with('#')
|
||||
&& value[1..].bytes().all(|byte| byte.is_ascii_hexdigit()))
|
||||
.then(|| Self(value.to_ascii_uppercase()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Group member availability at the current observation time.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AccountGroupAccountSummary {
|
||||
pub available: u64,
|
||||
pub limited: u64,
|
||||
pub total: u64,
|
||||
}
|
||||
|
||||
/// Group scheduling slots derived from available accounts and runtime leases.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct AccountGroupCapacity {
|
||||
pub used_slots: Option<u64>,
|
||||
pub total_slots: u64,
|
||||
}
|
||||
|
||||
/// Successful, downstream-committed USD request costs for group accounts.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupUsage {
|
||||
pub today_usd: DecimalAmount,
|
||||
pub total_usd: DecimalAmount,
|
||||
}
|
||||
|
||||
/// Lightweight group reference embedded in account and client-key views.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupRef {
|
||||
pub id: AccountGroupId,
|
||||
pub name: String,
|
||||
pub color: AccountGroupColor,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
/// Account group list query.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupListQuery {
|
||||
pub page: u32,
|
||||
pub page_size: PageSize,
|
||||
pub search: Option<String>,
|
||||
pub enabled: Option<bool>,
|
||||
}
|
||||
|
||||
/// Complete account group summary.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupRecord {
|
||||
pub id: AccountGroupId,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub color: AccountGroupColor,
|
||||
pub enabled: bool,
|
||||
pub member_count: u64,
|
||||
pub provider_counts: BTreeMap<String, u64>,
|
||||
pub client_key_count: u64,
|
||||
pub account_summary: AccountGroupAccountSummary,
|
||||
pub capacity: AccountGroupCapacity,
|
||||
pub usage: AccountGroupUsage,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
/// Paginated account groups.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupPage {
|
||||
pub config_revision: Revision,
|
||||
pub items: Vec<AccountGroupRecord>,
|
||||
pub total: u64,
|
||||
pub page: u32,
|
||||
pub page_size: u16,
|
||||
}
|
||||
|
||||
/// Lightweight account shown by the membership editor.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupMember {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub provider_kind: ProviderKind,
|
||||
pub email: Option<String>,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
/// Complete membership projection for one group.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupMembers {
|
||||
pub config_revision: Revision,
|
||||
pub id: AccountGroupId,
|
||||
pub items: Vec<AccountGroupMember>,
|
||||
pub total: u64,
|
||||
}
|
||||
|
||||
/// Create an account group.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateAccountGroup {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub color: AccountGroupColor,
|
||||
}
|
||||
|
||||
/// Store-ready create command with a generated stable ID.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct NewAccountGroup {
|
||||
pub id: AccountGroupId,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub color: AccountGroupColor,
|
||||
}
|
||||
|
||||
/// Update an account group's descriptive fields.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpdateAccountGroup {
|
||||
pub id: AccountGroupId,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub color: AccountGroupColor,
|
||||
}
|
||||
|
||||
/// Enable or disable an account group.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SetAccountGroupEnabled {
|
||||
pub id: AccountGroupId,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
/// Delete an account group that is not referenced by a client key.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct DeleteAccountGroup {
|
||||
pub id: AccountGroupId,
|
||||
}
|
||||
|
||||
/// Result of any account-group mutation.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountGroupMutation {
|
||||
pub config_revision: Revision,
|
||||
pub id: AccountGroupId,
|
||||
pub record: Option<AccountGroupRecord>,
|
||||
}
|
||||
@@ -7,11 +7,12 @@ use futures::Stream;
|
||||
|
||||
use gateway_core::routing::ProviderKind;
|
||||
|
||||
use super::{PageSize, Revision, observability::TimeRange};
|
||||
use super::{PageSize, Revision, account_groups::AccountGroupRef, observability::TimeRange};
|
||||
|
||||
pub use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountStatus, AccountStatusFacts, AccountStatusProjection,
|
||||
CredentialState, QuotaAccessState, QuotaEvidence, QuotaState, resolve_account_status,
|
||||
AccountConcurrencyLimit, AccountErrorReason, AccountStatus, AccountStatusFacts,
|
||||
AccountStatusProjection, AccountWeight, CredentialState, QuotaAccessState, QuotaEvidence,
|
||||
QuotaState, resolve_account_status,
|
||||
};
|
||||
|
||||
/// 账号列表排序字段。
|
||||
@@ -45,16 +46,25 @@ pub struct AccountListQuery {
|
||||
pub page: u32,
|
||||
pub page_size: PageSize,
|
||||
pub provider_kind: Option<ProviderKind>,
|
||||
pub group_filter: Option<AccountGroupFilter>,
|
||||
pub search: Option<String>,
|
||||
pub status: Option<AccountStatus>,
|
||||
pub sort: Option<AccountSort>,
|
||||
}
|
||||
|
||||
/// Optional account membership filter.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum AccountGroupFilter {
|
||||
Group(gateway_core::routing::AccountGroupId),
|
||||
Ungrouped,
|
||||
}
|
||||
|
||||
/// 账号公共存储投影;Provider 专属字段不进入此结构。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountRecord {
|
||||
pub id: String,
|
||||
pub provider_kind: ProviderKind,
|
||||
pub groups: Vec<AccountGroupRef>,
|
||||
pub name: String,
|
||||
pub email: Option<String>,
|
||||
pub upstream_user_id: Option<String>,
|
||||
@@ -66,6 +76,8 @@ pub struct AccountRecord {
|
||||
pub access_token_expires_at: Option<DateTime<Utc>>,
|
||||
pub next_refresh_at: Option<DateTime<Utc>>,
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub credential_state: CredentialState,
|
||||
pub credential_observed_at: DateTime<Utc>,
|
||||
pub quota: QuotaState,
|
||||
@@ -180,11 +192,38 @@ pub struct AccountSummary {
|
||||
pub error: u64,
|
||||
}
|
||||
|
||||
/// 账号启停写入命令。
|
||||
/// 账号可编辑事实的一次性替换命令。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SetAccountEnabled {
|
||||
pub struct UpdateAccount {
|
||||
pub account_id: String,
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
|
||||
}
|
||||
|
||||
/// 账号更新结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountUpdateResult {
|
||||
pub config_revision: Revision,
|
||||
pub account_id: gateway_core::engine::credential::ProviderAccountId,
|
||||
}
|
||||
|
||||
/// 一批账号可编辑事实的一次性替换命令。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct BatchUpdateAccounts {
|
||||
pub account_ids: Vec<String>,
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub group_ids: Vec<gateway_core::routing::AccountGroupId>,
|
||||
}
|
||||
|
||||
/// 批量账号更新结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountsUpdateResult {
|
||||
pub config_revision: Revision,
|
||||
pub account_ids: Vec<gateway_core::engine::credential::ProviderAccountId>,
|
||||
}
|
||||
|
||||
/// 账号批量删除命令。
|
||||
|
||||
@@ -6,10 +6,10 @@ use chrono::{DateTime, Utc};
|
||||
|
||||
use gateway_core::{
|
||||
policy::{ClientApiKeyId, RateLimits},
|
||||
routing::ProviderKind,
|
||||
routing::{AccountGroupId, ProviderKind},
|
||||
};
|
||||
|
||||
use super::{AdminModelError, Revision};
|
||||
use super::{AdminModelError, Revision, account_groups::AccountGroupRef};
|
||||
|
||||
/// Client Key 列表保持旧 HTTP 合同允许的完整非零 `u16` 页大小。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
@@ -88,7 +88,8 @@ pub struct ClientKeyRecord {
|
||||
pub id: ClientApiKeyId,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub provider_kind: ProviderKind,
|
||||
pub groups: Vec<AccountGroupRef>,
|
||||
pub provider_kinds: Vec<ProviderKind>,
|
||||
pub prefix: String,
|
||||
pub enabled: bool,
|
||||
pub limits: RateLimits,
|
||||
@@ -143,7 +144,7 @@ impl fmt::Debug for ClientKeySecret {
|
||||
pub struct CreateClientKey {
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub provider_kind: ProviderKind,
|
||||
pub group_ids: Vec<AccountGroupId>,
|
||||
pub limits: RateLimits,
|
||||
}
|
||||
|
||||
@@ -153,7 +154,7 @@ pub struct NewClientKey {
|
||||
pub id: ClientApiKeyId,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub provider_kind: ProviderKind,
|
||||
pub group_ids: Vec<AccountGroupId>,
|
||||
pub limits: RateLimits,
|
||||
pub plaintext: String,
|
||||
}
|
||||
@@ -164,7 +165,7 @@ impl fmt::Debug for NewClientKey {
|
||||
.debug_struct("NewClientKey")
|
||||
.field("id", &self.id)
|
||||
.field("name", &self.name)
|
||||
.field("provider_kind", &self.provider_kind)
|
||||
.field("group_ids", &self.group_ids)
|
||||
.field("plaintext", &"[REDACTED]")
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
@@ -176,7 +177,7 @@ pub struct UpdateClientKey {
|
||||
pub id: ClientApiKeyId,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub provider_kind: ProviderKind,
|
||||
pub group_ids: Vec<AccountGroupId>,
|
||||
pub limits: RateLimits,
|
||||
}
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
use std::num::{NonZeroU16, NonZeroU64};
|
||||
|
||||
pub mod account_groups;
|
||||
pub mod accounts;
|
||||
pub mod auth;
|
||||
pub mod backup;
|
||||
|
||||
@@ -513,8 +513,10 @@ pub struct DashboardAccountUsage {
|
||||
/// Dashboard 当前账号池的可重建运行时槽位事实。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct DashboardRuntimeSlots {
|
||||
/// 与持久账号池指标相同谓词下可参与调度的账号数。
|
||||
pub active_accounts: u64,
|
||||
/// 继承运行参数默认并发上限的可调度账号数。
|
||||
pub inherited_accounts: u64,
|
||||
/// 已配置账号级并发上限的可调度账号槽位总数。
|
||||
pub overridden_slots: u64,
|
||||
/// Redis 可用时的实时 in-flight 槽位总数;不可用时为 `None`。
|
||||
pub used_slots: Option<u64>,
|
||||
}
|
||||
@@ -548,6 +550,9 @@ pub struct UsageRecord {
|
||||
pub id: String,
|
||||
pub client_api_key_ref: String,
|
||||
pub config_revision: u64,
|
||||
pub routing_scope: String,
|
||||
pub routing_group_refs: Vec<String>,
|
||||
pub routing_group_names_snapshot: Vec<String>,
|
||||
pub protocol: String,
|
||||
pub operation: String,
|
||||
pub endpoint: String,
|
||||
|
||||
@@ -10,9 +10,14 @@ use chrono::{DateTime, Utc};
|
||||
use super::backup::BackupStorePorts;
|
||||
use crate::model::{
|
||||
MutationContext, Revision,
|
||||
account_groups::{
|
||||
AccountGroupListQuery, AccountGroupMembers, AccountGroupMutation, AccountGroupPage,
|
||||
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
},
|
||||
accounts::{
|
||||
AccountListQuery, AccountPage, AccountPageItem, AccountUsage, AccountUsageWindowQuery,
|
||||
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
|
||||
AccountListQuery, AccountPage, AccountPageItem, AccountUpdateResult, AccountUsage,
|
||||
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
},
|
||||
auth::{AdminAuditEvent, AdminSession},
|
||||
client_keys::{
|
||||
@@ -138,11 +143,17 @@ pub trait AccountStore: Send + Sync {
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<CredentialMutationResult>;
|
||||
|
||||
async fn set_account_enabled(
|
||||
async fn update_account(
|
||||
&self,
|
||||
command: SetAccountEnabled,
|
||||
command: UpdateAccount,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<Revision>;
|
||||
) -> AdminStoreResult<AccountUpdateResult>;
|
||||
|
||||
async fn batch_update_accounts(
|
||||
&self,
|
||||
command: BatchUpdateAccounts,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountsUpdateResult>;
|
||||
|
||||
async fn delete_accounts(
|
||||
&self,
|
||||
@@ -215,6 +226,44 @@ pub trait ClientKeyStore: Send + Sync {
|
||||
) -> AdminStoreResult<Revision>;
|
||||
}
|
||||
|
||||
/// Provider-neutral account group management and membership transactions.
|
||||
#[async_trait]
|
||||
pub trait AccountGroupStore: Send + Sync {
|
||||
async fn list_account_groups(
|
||||
&self,
|
||||
query: AccountGroupListQuery,
|
||||
) -> AdminStoreResult<AccountGroupPage>;
|
||||
|
||||
async fn account_group_members(
|
||||
&self,
|
||||
id: &gateway_core::routing::AccountGroupId,
|
||||
) -> AdminStoreResult<AccountGroupMembers>;
|
||||
|
||||
async fn create_account_group(
|
||||
&self,
|
||||
command: NewAccountGroup,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation>;
|
||||
|
||||
async fn update_account_group(
|
||||
&self,
|
||||
command: UpdateAccountGroup,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation>;
|
||||
|
||||
async fn set_account_group_enabled(
|
||||
&self,
|
||||
command: SetAccountGroupEnabled,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation>;
|
||||
|
||||
async fn delete_account_group(
|
||||
&self,
|
||||
command: DeleteAccountGroup,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation>;
|
||||
}
|
||||
|
||||
/// 用量、趋势、诊断与运维错误的只读能力。
|
||||
#[async_trait]
|
||||
pub trait ObservabilityStore: Send + Sync {
|
||||
@@ -302,6 +351,7 @@ pub trait SettingsStore: Send + Sync {
|
||||
#[derive(Clone)]
|
||||
pub struct AdminStorePorts {
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
account_groups: Arc<dyn AccountGroupStore>,
|
||||
auth: Arc<dyn AuthStore>,
|
||||
client_keys: Arc<dyn ClientKeyStore>,
|
||||
observability: Arc<dyn ObservabilityStore>,
|
||||
@@ -313,6 +363,7 @@ impl AdminStorePorts {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
account_groups: Arc<dyn AccountGroupStore>,
|
||||
auth: Arc<dyn AuthStore>,
|
||||
client_keys: Arc<dyn ClientKeyStore>,
|
||||
observability: Arc<dyn ObservabilityStore>,
|
||||
@@ -321,6 +372,7 @@ impl AdminStorePorts {
|
||||
) -> Self {
|
||||
Self {
|
||||
accounts,
|
||||
account_groups,
|
||||
auth,
|
||||
client_keys,
|
||||
observability,
|
||||
@@ -334,6 +386,11 @@ impl AdminStorePorts {
|
||||
self.accounts.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn account_groups(&self) -> Arc<dyn AccountGroupStore> {
|
||||
self.account_groups.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn auth(&self) -> Arc<dyn AuthStore> {
|
||||
self.auth.clone()
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
//! Account group management use cases.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use gateway_core::routing::{AccountGroupId, snapshot::SnapshotControl};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError, MutationContext,
|
||||
account_groups::{
|
||||
AccountGroupListQuery, AccountGroupMembers, AccountGroupMutation, AccountGroupPage,
|
||||
CreateAccountGroup, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
|
||||
UpdateAccountGroup,
|
||||
},
|
||||
},
|
||||
ports::store::AccountGroupStore,
|
||||
};
|
||||
|
||||
use super::{map_store_error, publish_committed};
|
||||
|
||||
/// API-facing account group management service.
|
||||
#[async_trait]
|
||||
pub trait AccountGroupService: Send + Sync {
|
||||
async fn list(&self, query: AccountGroupListQuery) -> Result<AccountGroupPage, AdminError>;
|
||||
async fn members(&self, id: &AccountGroupId) -> Result<AccountGroupMembers, AdminError>;
|
||||
async fn create(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: CreateAccountGroup,
|
||||
) -> Result<AccountGroupMutation, AdminError>;
|
||||
async fn update(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: UpdateAccountGroup,
|
||||
) -> Result<AccountGroupMutation, AdminError>;
|
||||
async fn set_enabled(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: SetAccountGroupEnabled,
|
||||
) -> Result<AccountGroupMutation, AdminError>;
|
||||
async fn delete(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: DeleteAccountGroup,
|
||||
) -> Result<AccountGroupMutation, AdminError>;
|
||||
}
|
||||
|
||||
pub(crate) struct DefaultAccountGroupService {
|
||||
store: Arc<dyn AccountGroupStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
}
|
||||
|
||||
impl DefaultAccountGroupService {
|
||||
#[must_use]
|
||||
pub(crate) fn new(
|
||||
store: Arc<dyn AccountGroupStore>,
|
||||
snapshot: Arc<dyn SnapshotControl>,
|
||||
) -> Self {
|
||||
Self { store, snapshot }
|
||||
}
|
||||
|
||||
async fn publish(
|
||||
&self,
|
||||
result: Result<AccountGroupMutation, crate::ports::store::AdminStoreError>,
|
||||
) -> Result<AccountGroupMutation, AdminError> {
|
||||
let result = result.map_err(|error| map_store_error(error, "account group"))?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AccountGroupService for DefaultAccountGroupService {
|
||||
async fn list(&self, query: AccountGroupListQuery) -> Result<AccountGroupPage, AdminError> {
|
||||
self.store
|
||||
.list_account_groups(query)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "account group"))
|
||||
}
|
||||
|
||||
async fn members(&self, id: &AccountGroupId) -> Result<AccountGroupMembers, AdminError> {
|
||||
self.store
|
||||
.account_group_members(id)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "account group"))
|
||||
}
|
||||
|
||||
async fn create(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: CreateAccountGroup,
|
||||
) -> Result<AccountGroupMutation, AdminError> {
|
||||
let id = AccountGroupId::new(format!("grp_{}", Uuid::now_v7().simple()))
|
||||
.map_err(|_| AdminError::internal("Failed to create account group ID"))?;
|
||||
self.publish(
|
||||
self.store
|
||||
.create_account_group(
|
||||
NewAccountGroup {
|
||||
id,
|
||||
name: command.name,
|
||||
description: command.description,
|
||||
color: command.color,
|
||||
},
|
||||
context,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn update(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: UpdateAccountGroup,
|
||||
) -> Result<AccountGroupMutation, AdminError> {
|
||||
self.publish(self.store.update_account_group(command, context).await)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn set_enabled(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: SetAccountGroupEnabled,
|
||||
) -> Result<AccountGroupMutation, AdminError> {
|
||||
self.publish(self.store.set_account_group_enabled(command, context).await)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: DeleteAccountGroup,
|
||||
) -> Result<AccountGroupMutation, AdminError> {
|
||||
self.publish(self.store.delete_account_group(command, context).await)
|
||||
.await
|
||||
}
|
||||
}
|
||||
@@ -18,7 +18,8 @@ use crate::{
|
||||
AdminError, MutationContext,
|
||||
accounts::{
|
||||
AccountConnectionTestEvent, AccountConnectionTestEventStream, AccountListQuery,
|
||||
AccountPageItem, AccountUsageWindowQuery,
|
||||
AccountPageItem, AccountUpdateResult, AccountUsageWindowQuery, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, UpdateAccount,
|
||||
},
|
||||
observability::TimeRange,
|
||||
provider_credentials::{
|
||||
@@ -54,6 +55,18 @@ pub trait AccountsService: Send + Sync {
|
||||
account_id: ProviderAccountId,
|
||||
) -> Result<AccountRefreshResult, AdminError>;
|
||||
|
||||
async fn update(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: UpdateAccount,
|
||||
) -> Result<AccountUpdateResult, AdminError>;
|
||||
|
||||
async fn batch_update(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: BatchUpdateAccounts,
|
||||
) -> Result<AccountsUpdateResult, AdminError>;
|
||||
|
||||
async fn quota(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
@@ -383,6 +396,76 @@ impl AccountsService for DefaultAccountsService {
|
||||
})
|
||||
}
|
||||
|
||||
async fn update(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: UpdateAccount,
|
||||
) -> Result<AccountUpdateResult, AdminError> {
|
||||
let account_id = ProviderAccountId::new(command.account_id.clone())
|
||||
.map_err(|_| AdminError::invalid("Invalid provider account ID"))?;
|
||||
let (_, provider) = self.provider_for_account(&account_id).await?;
|
||||
let enabled = command.enabled;
|
||||
let result = self
|
||||
.accounts
|
||||
.update_account(command, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "provider account"))?;
|
||||
if !enabled {
|
||||
provider.account_unavailable(&account_id).await;
|
||||
}
|
||||
provider
|
||||
.account_facts_changed(std::slice::from_ref(&account_id))
|
||||
.await;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn batch_update(
|
||||
&self,
|
||||
context: &MutationContext,
|
||||
command: BatchUpdateAccounts,
|
||||
) -> Result<AccountsUpdateResult, AdminError> {
|
||||
let account_ids = command
|
||||
.account_ids
|
||||
.iter()
|
||||
.map(|id| {
|
||||
ProviderAccountId::new(id.clone())
|
||||
.map_err(|_| AdminError::invalid("Invalid provider account ID"))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let mut providers = BTreeMap::<
|
||||
ProviderKind,
|
||||
(
|
||||
Arc<dyn crate::ports::provider::ProviderAdmin>,
|
||||
Vec<ProviderAccountId>,
|
||||
),
|
||||
>::new();
|
||||
for account_id in &account_ids {
|
||||
let (item, provider) = self.provider_for_account(account_id).await?;
|
||||
providers
|
||||
.entry(item.account.provider_kind)
|
||||
.or_insert_with(|| (provider, Vec::new()))
|
||||
.1
|
||||
.push(account_id.clone());
|
||||
}
|
||||
let enabled = command.enabled;
|
||||
let result = self
|
||||
.accounts
|
||||
.batch_update_accounts(command, context)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, "provider accounts"))?;
|
||||
for (provider, provider_ids) in providers.values() {
|
||||
if !enabled {
|
||||
for account_id in provider_ids {
|
||||
provider.account_unavailable(account_id).await;
|
||||
}
|
||||
}
|
||||
provider.account_facts_changed(provider_ids).await;
|
||||
}
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn quota(
|
||||
&self,
|
||||
account_id: &ProviderAccountId,
|
||||
|
||||
@@ -97,7 +97,7 @@ impl ClientKeyService for DefaultClientKeyService {
|
||||
id,
|
||||
name: command.name,
|
||||
label: command.label,
|
||||
provider_kind: command.provider_kind,
|
||||
group_ids: command.group_ids,
|
||||
limits: command.limits,
|
||||
plaintext: plaintext.clone(),
|
||||
},
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
//! 管理控制面的用例实现。
|
||||
|
||||
pub mod account_groups;
|
||||
pub mod accounts;
|
||||
pub mod auth;
|
||||
pub mod backup;
|
||||
@@ -13,11 +14,11 @@ pub mod xai;
|
||||
use crate::{
|
||||
model::{
|
||||
AdminError, AdminErrorKind, MutationContext,
|
||||
accounts::{DeleteAccounts, SetAccountEnabled},
|
||||
accounts::DeleteAccounts,
|
||||
provider_credentials::{
|
||||
AuthorizationMutationTarget, CredentialDeletion, CredentialDeletionResult,
|
||||
CredentialDetails, CredentialMutation, CredentialMutationResult,
|
||||
CredentialRotationCommit, PendingAuthorizationMutation, PreparedAuthorizationCommit,
|
||||
CredentialDetails, CredentialMutationResult, CredentialRotationCommit,
|
||||
PendingAuthorizationMutation, PreparedAuthorizationCommit,
|
||||
PreparedAuthorizationCredential, PreparedCredentialImport, PreparedCredentialRotation,
|
||||
StartAuthorization,
|
||||
},
|
||||
@@ -280,44 +281,6 @@ async fn commit_credential_refresh(
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_credential_enabled(
|
||||
accounts: &dyn AccountStore,
|
||||
provider: &dyn ProviderAdmin,
|
||||
command: CredentialMutation,
|
||||
enabled: bool,
|
||||
resource: &'static str,
|
||||
) -> Result<CredentialMutationResult, AdminError> {
|
||||
required_credential(
|
||||
accounts,
|
||||
provider.provider_kind(),
|
||||
&command.account_id,
|
||||
resource,
|
||||
)
|
||||
.await?;
|
||||
let account_id = command.account_id;
|
||||
let revision = accounts
|
||||
.set_account_enabled(
|
||||
SetAccountEnabled {
|
||||
account_id: account_id.as_str().to_owned(),
|
||||
enabled,
|
||||
},
|
||||
&command.context,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| map_store_error(error, resource))?;
|
||||
if !enabled {
|
||||
provider.account_unavailable(&account_id).await;
|
||||
}
|
||||
provider
|
||||
.account_facts_changed(std::slice::from_ref(&account_id))
|
||||
.await;
|
||||
Ok(CredentialMutationResult {
|
||||
config_revision: revision,
|
||||
account_id,
|
||||
credential_revision: None,
|
||||
})
|
||||
}
|
||||
|
||||
async fn delete_credentials(
|
||||
accounts: &dyn AccountStore,
|
||||
provider: &dyn ProviderAdmin,
|
||||
|
||||
@@ -126,12 +126,20 @@ impl ObservabilityService for DefaultObservabilityService {
|
||||
let health_timeline = health_timeline_at(&observation.trend, Utc::now());
|
||||
let wire_profiles = self.providers.dashboard_wire_profiles();
|
||||
let max_concurrent_per_account = u64::from(settings.max_concurrent_per_account);
|
||||
let active_accounts = runtime_slots
|
||||
.as_ref()
|
||||
.map_or(observation.provider_accounts.normal, |slots| {
|
||||
slots.active_accounts
|
||||
});
|
||||
let total_slots = active_accounts.saturating_mul(max_concurrent_per_account);
|
||||
let total_slots = runtime_slots.as_ref().map_or_else(
|
||||
|| {
|
||||
observation
|
||||
.provider_accounts
|
||||
.normal
|
||||
.saturating_mul(max_concurrent_per_account)
|
||||
},
|
||||
|slots| {
|
||||
slots
|
||||
.inherited_accounts
|
||||
.saturating_mul(max_concurrent_per_account)
|
||||
.saturating_add(slots.overridden_slots)
|
||||
},
|
||||
);
|
||||
let used_slots = runtime_slots.and_then(|slots| slots.used_slots);
|
||||
Ok(DashboardResult {
|
||||
capacity: DashboardCapacity {
|
||||
|
||||
@@ -11,9 +11,9 @@ use crate::{
|
||||
provider_credentials::{
|
||||
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
|
||||
CredentialDeletionResult, CredentialDetails, CredentialImportCommit,
|
||||
CredentialImportResult, CredentialListQuery, CredentialMutation,
|
||||
CredentialMutationResult, CredentialPage, ImportCredentials, PrepareCredentialImport,
|
||||
PrepareCredentialRotation, ProviderQuotaRequest, RotateCredential, StartAuthorization,
|
||||
CredentialImportResult, CredentialListQuery, CredentialMutationResult, CredentialPage,
|
||||
ImportCredentials, PrepareCredentialImport, PrepareCredentialRotation,
|
||||
ProviderQuotaRequest, RotateCredential, StartAuthorization,
|
||||
},
|
||||
},
|
||||
ports::{provider::ProviderAdmin, store::AccountStore},
|
||||
@@ -22,8 +22,7 @@ use crate::{
|
||||
use super::{
|
||||
commit_authorization, commit_credential_rotation, delete_credentials, map_provider_error,
|
||||
map_store_error, pending_authorization, publish_committed, required_credential,
|
||||
set_credential_enabled, validate_authorization_commit, validate_prepared_import,
|
||||
validate_prepared_rotation,
|
||||
validate_authorization_commit, validate_prepared_import, validate_prepared_rotation,
|
||||
};
|
||||
|
||||
/// OpenAI 固定管理路由消费的服务。
|
||||
@@ -50,14 +49,6 @@ pub trait OpenAiService: Send + Sync {
|
||||
&self,
|
||||
command: RotateCredential,
|
||||
) -> Result<CredentialMutationResult, AdminError>;
|
||||
async fn enable(
|
||||
&self,
|
||||
command: CredentialMutation,
|
||||
) -> Result<CredentialMutationResult, AdminError>;
|
||||
async fn disable(
|
||||
&self,
|
||||
command: CredentialMutation,
|
||||
) -> Result<CredentialMutationResult, AdminError>;
|
||||
async fn delete(
|
||||
&self,
|
||||
command: CredentialDeletion,
|
||||
@@ -252,38 +243,6 @@ impl OpenAiService for DefaultOpenAiService {
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn enable(
|
||||
&self,
|
||||
command: CredentialMutation,
|
||||
) -> Result<CredentialMutationResult, AdminError> {
|
||||
let result = set_credential_enabled(
|
||||
self.accounts.as_ref(),
|
||||
self.provider.as_ref(),
|
||||
command,
|
||||
true,
|
||||
"OpenAI credential",
|
||||
)
|
||||
.await?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn disable(
|
||||
&self,
|
||||
command: CredentialMutation,
|
||||
) -> Result<CredentialMutationResult, AdminError> {
|
||||
let result = set_credential_enabled(
|
||||
self.accounts.as_ref(),
|
||||
self.provider.as_ref(),
|
||||
command,
|
||||
false,
|
||||
"OpenAI credential",
|
||||
)
|
||||
.await?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
command: CredentialDeletion,
|
||||
|
||||
@@ -11,8 +11,8 @@ use crate::{
|
||||
provider_credentials::{
|
||||
AuthorizationStarted, CompleteAuthorization, CredentialDeletion,
|
||||
CredentialDeletionResult, CredentialImportCommit, CredentialImportResult,
|
||||
CredentialListQuery, CredentialMutation, CredentialMutationResult, CredentialPage,
|
||||
ImportCredentials, PrepareCredentialImport, ProviderQuotaRequest, StartAuthorization,
|
||||
CredentialListQuery, CredentialMutationResult, CredentialPage, ImportCredentials,
|
||||
PrepareCredentialImport, ProviderQuotaRequest, StartAuthorization,
|
||||
},
|
||||
},
|
||||
ports::{provider::ProviderAdmin, store::AccountStore},
|
||||
@@ -20,8 +20,8 @@ use crate::{
|
||||
|
||||
use super::{
|
||||
commit_authorization, delete_credentials, map_provider_error, map_store_error,
|
||||
pending_authorization, publish_committed, set_credential_enabled,
|
||||
validate_authorization_commit, validate_prepared_import,
|
||||
pending_authorization, publish_committed, validate_authorization_commit,
|
||||
validate_prepared_import,
|
||||
};
|
||||
|
||||
/// xAI 固定管理路由消费的服务。
|
||||
@@ -40,14 +40,6 @@ pub trait XaiService: Send + Sync {
|
||||
&self,
|
||||
command: CompleteAuthorization,
|
||||
) -> Result<CredentialMutationResult, AdminError>;
|
||||
async fn enable(
|
||||
&self,
|
||||
command: CredentialMutation,
|
||||
) -> Result<CredentialMutationResult, AdminError>;
|
||||
async fn disable(
|
||||
&self,
|
||||
command: CredentialMutation,
|
||||
) -> Result<CredentialMutationResult, AdminError>;
|
||||
async fn delete(
|
||||
&self,
|
||||
command: CredentialDeletion,
|
||||
@@ -166,38 +158,6 @@ impl XaiService for DefaultXaiService {
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn enable(
|
||||
&self,
|
||||
command: CredentialMutation,
|
||||
) -> Result<CredentialMutationResult, AdminError> {
|
||||
let result = set_credential_enabled(
|
||||
self.accounts.as_ref(),
|
||||
self.provider.as_ref(),
|
||||
command,
|
||||
true,
|
||||
"xAI credential",
|
||||
)
|
||||
.await?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn disable(
|
||||
&self,
|
||||
command: CredentialMutation,
|
||||
) -> Result<CredentialMutationResult, AdminError> {
|
||||
let result = set_credential_enabled(
|
||||
self.accounts.as_ref(),
|
||||
self.provider.as_ref(),
|
||||
command,
|
||||
false,
|
||||
"xAI credential",
|
||||
)
|
||||
.await?;
|
||||
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn delete(
|
||||
&self,
|
||||
command: CredentialDeletion,
|
||||
|
||||
@@ -22,8 +22,9 @@ use gateway_admin::{
|
||||
AdminError, MutationContext, Revision,
|
||||
accounts::{
|
||||
AccountConnectionTestEvent, AccountListQuery, AccountPage, AccountPageItem,
|
||||
AccountRecord, AccountSummary, AccountUsage, AccountUsageWindowQuery,
|
||||
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
|
||||
AccountRecord, AccountSummary, AccountUpdateResult, AccountUsage,
|
||||
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
},
|
||||
observability::TimeRange,
|
||||
provider_credentials::{
|
||||
@@ -661,15 +662,36 @@ impl AccountStore for FakeAccountStore {
|
||||
Ok(rotation_result(command))
|
||||
}
|
||||
|
||||
async fn set_account_enabled(
|
||||
async fn update_account(
|
||||
&self,
|
||||
_: SetAccountEnabled,
|
||||
command: UpdateAccount,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
self.record("store.set_enabled");
|
||||
) -> AdminStoreResult<AccountUpdateResult> {
|
||||
self.record("store.update_account");
|
||||
self.record_context(context);
|
||||
self.require_commit()?;
|
||||
Ok(revision(2))
|
||||
Ok(AccountUpdateResult {
|
||||
config_revision: revision(2),
|
||||
account_id: ProviderAccountId::new(command.account_id).expect("account ID"),
|
||||
})
|
||||
}
|
||||
|
||||
async fn batch_update_accounts(
|
||||
&self,
|
||||
command: BatchUpdateAccounts,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountsUpdateResult> {
|
||||
self.record("store.batch_update_accounts");
|
||||
self.record_context(context);
|
||||
self.require_commit()?;
|
||||
Ok(AccountsUpdateResult {
|
||||
config_revision: revision(2),
|
||||
account_ids: command
|
||||
.account_ids
|
||||
.into_iter()
|
||||
.map(|id| ProviderAccountId::new(id).expect("account ID"))
|
||||
.collect(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn delete_accounts(
|
||||
@@ -1001,6 +1023,122 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
|
||||
assert_eq!(store.audit_requests(), ["refresh-request"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_update_should_commit_then_release_disabled_account_and_publish_facts() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("openai", events.clone());
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
let services = accounts_service(provider, store.clone()).await;
|
||||
|
||||
let result = services
|
||||
.accounts()
|
||||
.update(
|
||||
&context("update-request"),
|
||||
UpdateAccount {
|
||||
account_id: "acct_test".to_owned(),
|
||||
enabled: false,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
|
||||
group_ids: Vec::new(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("update account");
|
||||
|
||||
assert_eq!(result.config_revision, revision(2));
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
[
|
||||
"store.load_account",
|
||||
"store.update_account",
|
||||
"provider.account_unavailable",
|
||||
"provider.account_facts_changed",
|
||||
]
|
||||
);
|
||||
assert_eq!(store.audit_requests(), ["update-request"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_update_should_not_notify_provider_when_store_commit_fails() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("openai", events.clone());
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
store.fail_next_commit();
|
||||
let services = accounts_service(provider, store).await;
|
||||
|
||||
services
|
||||
.accounts()
|
||||
.update(
|
||||
&context("update-failure"),
|
||||
UpdateAccount {
|
||||
account_id: "acct_test".to_owned(),
|
||||
enabled: false,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
|
||||
group_ids: Vec::new(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("failed account update");
|
||||
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
["store.load_account", "store.update_account"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_batch_update_should_commit_once_and_notify_each_provider() {
|
||||
let events = events();
|
||||
let openai = FakeProviderAdmin::new("openai", events.clone());
|
||||
let xai = FakeProviderAdmin::new("xai", events.clone());
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
let mut openai_account = account_record("openai");
|
||||
openai_account.id = "acct_openai".to_owned();
|
||||
let mut xai_account = account_record("xai");
|
||||
xai_account.id = "acct_xai".to_owned();
|
||||
store.set_accounts(vec![openai_account, xai_account]);
|
||||
let services = super::AdminHarness::new()
|
||||
.accounts(store.clone())
|
||||
.settings(Arc::new(StaticSettingsStore))
|
||||
.provider(openai)
|
||||
.provider(xai)
|
||||
.probe(Arc::new(SuccessfulAccountProbe))
|
||||
.build()
|
||||
.await;
|
||||
|
||||
let result = services
|
||||
.accounts()
|
||||
.batch_update(
|
||||
&context("batch-update-request"),
|
||||
BatchUpdateAccounts {
|
||||
account_ids: vec!["acct_openai".to_owned(), "acct_xai".to_owned()],
|
||||
enabled: false,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
|
||||
group_ids: Vec::new(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("batch update accounts");
|
||||
|
||||
assert_eq!(result.config_revision, revision(2));
|
||||
assert_eq!(result.account_ids.len(), 2);
|
||||
assert_eq!(store.audit_requests(), ["batch-update-request"]);
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
[
|
||||
"store.load_account",
|
||||
"store.load_account",
|
||||
"store.batch_update_accounts",
|
||||
"provider.account_unavailable",
|
||||
"provider.account_facts_changed",
|
||||
"provider.account_unavailable",
|
||||
"provider.account_facts_changed",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn accounts_list_should_return_complete_directory_semantics() {
|
||||
let provider = FakeProviderAdmin::new("openai", events());
|
||||
@@ -1012,6 +1150,7 @@ async fn accounts_list_should_return_complete_directory_semantics() {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
@@ -1073,6 +1212,7 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
@@ -1115,6 +1255,7 @@ async fn accounts_list_should_prefer_credential_error_over_quota_exhaustion() {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
@@ -1143,6 +1284,7 @@ async fn accounts_list_should_map_unknown_credential_to_error_not_normal() {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
@@ -1187,6 +1329,7 @@ async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view()
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
@@ -1231,6 +1374,7 @@ async fn accounts_list_should_not_derive_exhaustion_from_provider_quota_view() {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
@@ -1281,6 +1425,7 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
@@ -1366,6 +1511,7 @@ async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_w
|
||||
page: 1,
|
||||
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: None,
|
||||
@@ -1513,6 +1659,7 @@ pub(super) fn account_record(kind: &str) -> AccountRecord {
|
||||
AccountRecord {
|
||||
id: "acct_test".to_owned(),
|
||||
provider_kind: ProviderKind::new(kind).expect("provider kind"),
|
||||
groups: Vec::new(),
|
||||
name: "test account".to_owned(),
|
||||
email: Some("test@example.invalid".to_owned()),
|
||||
upstream_user_id: Some("upstream-user".to_owned()),
|
||||
@@ -1524,6 +1671,8 @@ pub(super) fn account_record(kind: &str) -> AccountRecord {
|
||||
access_token_expires_at: Some(now + TimeDelta::hours(1)),
|
||||
next_refresh_at: Some(now + TimeDelta::minutes(30)),
|
||||
enabled: true,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
|
||||
credential_state: CredentialState::Ready,
|
||||
credential_observed_at: now,
|
||||
quota: QuotaState::allowed(now.into()),
|
||||
|
||||
@@ -20,9 +20,14 @@ use gateway_admin::{
|
||||
AdminConfig, AdminServices, InitialAdminPassword,
|
||||
model::{
|
||||
MutationContext, Revision,
|
||||
account_groups::{
|
||||
AccountGroupListQuery, AccountGroupMembers, AccountGroupMutation, AccountGroupPage,
|
||||
DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
},
|
||||
accounts::{
|
||||
AccountListQuery, AccountPage, AccountUsage, AccountUsageWindowQuery,
|
||||
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
|
||||
AccountListQuery, AccountPage, AccountUpdateResult, AccountUsage,
|
||||
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
},
|
||||
auth::{AdminAuditEvent, AdminSession},
|
||||
client_keys::{
|
||||
@@ -51,8 +56,8 @@ use gateway_admin::{
|
||||
backup::BackupStorePorts,
|
||||
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
|
||||
store::{
|
||||
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult,
|
||||
AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
|
||||
AccountGroupStore, AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts,
|
||||
AdminStoreResult, AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
|
||||
},
|
||||
system::{
|
||||
SystemOperationError, SystemOperationErrorKind, SystemOperations,
|
||||
@@ -171,6 +176,7 @@ impl AdminHarness {
|
||||
},
|
||||
AdminStorePorts::new(
|
||||
self.accounts,
|
||||
Arc::new(UnavailableAccountGroupStore),
|
||||
self.auth,
|
||||
self.client_keys,
|
||||
self.observability,
|
||||
@@ -235,6 +241,57 @@ impl AuthStore for BootstrapAuthStore {
|
||||
|
||||
struct UnavailableStore;
|
||||
|
||||
struct UnavailableAccountGroupStore;
|
||||
|
||||
#[async_trait]
|
||||
impl AccountGroupStore for UnavailableAccountGroupStore {
|
||||
async fn list_account_groups(
|
||||
&self,
|
||||
_: AccountGroupListQuery,
|
||||
) -> AdminStoreResult<AccountGroupPage> {
|
||||
Err(unavailable("account groups"))
|
||||
}
|
||||
|
||||
async fn account_group_members(
|
||||
&self,
|
||||
_: &gateway_core::routing::AccountGroupId,
|
||||
) -> AdminStoreResult<AccountGroupMembers> {
|
||||
Err(unavailable("account group members"))
|
||||
}
|
||||
|
||||
async fn create_account_group(
|
||||
&self,
|
||||
_: NewAccountGroup,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
Err(unavailable("account group create"))
|
||||
}
|
||||
|
||||
async fn update_account_group(
|
||||
&self,
|
||||
_: UpdateAccountGroup,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
Err(unavailable("account group update"))
|
||||
}
|
||||
|
||||
async fn set_account_group_enabled(
|
||||
&self,
|
||||
_: SetAccountGroupEnabled,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
Err(unavailable("account group state"))
|
||||
}
|
||||
|
||||
async fn delete_account_group(
|
||||
&self,
|
||||
_: DeleteAccountGroup,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
Err(unavailable("account group delete"))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AccountStore for UnavailableStore {
|
||||
async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult<AccountPage> {
|
||||
@@ -319,14 +376,22 @@ impl AccountStore for UnavailableStore {
|
||||
Err(unavailable("credential refresh"))
|
||||
}
|
||||
|
||||
async fn set_account_enabled(
|
||||
async fn update_account(
|
||||
&self,
|
||||
_: SetAccountEnabled,
|
||||
_: UpdateAccount,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
) -> AdminStoreResult<AccountUpdateResult> {
|
||||
Err(unavailable("account enabled"))
|
||||
}
|
||||
|
||||
async fn batch_update_accounts(
|
||||
&self,
|
||||
_: BatchUpdateAccounts,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountsUpdateResult> {
|
||||
Err(unavailable("account batch update"))
|
||||
}
|
||||
|
||||
async fn delete_accounts(
|
||||
&self,
|
||||
_: DeleteAccounts,
|
||||
|
||||
@@ -183,7 +183,8 @@ async fn dashboard_summary_should_project_rebuildable_runtime_slots() {
|
||||
let now = Utc::now();
|
||||
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
|
||||
store.replace_runtime_slots(Some(DashboardRuntimeSlots {
|
||||
active_accounts: 3,
|
||||
inherited_accounts: 2,
|
||||
overridden_slots: 5,
|
||||
used_slots: Some(2),
|
||||
}));
|
||||
let services = observability_services(store).await;
|
||||
@@ -195,9 +196,9 @@ async fn dashboard_summary_should_project_rebuildable_runtime_slots() {
|
||||
.expect("dashboard summary")
|
||||
.capacity;
|
||||
|
||||
assert_eq!(capacity.total_slots, 3);
|
||||
assert_eq!(capacity.total_slots, 7);
|
||||
assert_eq!(capacity.used_slots, Some(2));
|
||||
assert_eq!(capacity.available_slots, Some(1));
|
||||
assert_eq!(capacity.available_slots, Some(5));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -654,6 +655,9 @@ fn total_record(
|
||||
id: id.to_owned(),
|
||||
client_api_key_ref: "key_billing".to_owned(),
|
||||
config_revision: 1,
|
||||
routing_scope: "all".to_owned(),
|
||||
routing_group_refs: Vec::new(),
|
||||
routing_group_names_snapshot: Vec::new(),
|
||||
protocol: "openai".to_owned(),
|
||||
operation: "generate".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
|
||||
@@ -5,8 +5,8 @@ use gateway_core::engine::credential::ProviderAccountId;
|
||||
use gateway_admin::{
|
||||
AdminServices,
|
||||
model::provider_credentials::{
|
||||
AuthorizationMutationTarget, CompleteAuthorization, CredentialDeletion, CredentialMutation,
|
||||
ImportCredentials, ProviderQuotaRequest, StartAuthorization,
|
||||
AuthorizationMutationTarget, CompleteAuthorization, CredentialDeletion, ImportCredentials,
|
||||
ProviderQuotaRequest, StartAuthorization,
|
||||
},
|
||||
ports::provider::ProviderAdminErrorKind,
|
||||
};
|
||||
@@ -15,55 +15,6 @@ use super::accounts::{
|
||||
FakeAccountStore, FakeProviderAdmin, context, document, events, recorded, revision,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_enable_should_use_scoped_store_without_provider_mutation() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("openai", events.clone());
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
let services = service(provider, store.clone()).await;
|
||||
|
||||
services
|
||||
.openai()
|
||||
.enable(mutation("acct_test"))
|
||||
.await
|
||||
.expect("enable credential");
|
||||
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
[
|
||||
"store.credential_details",
|
||||
"store.set_enabled",
|
||||
"provider.account_facts_changed",
|
||||
]
|
||||
);
|
||||
assert_eq!(store.audit_requests(), ["request-openai"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_disable_should_commit_then_release_provider_resources() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("openai", events.clone());
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
let services = service(provider, store.clone()).await;
|
||||
|
||||
services
|
||||
.openai()
|
||||
.disable(mutation("acct_test"))
|
||||
.await
|
||||
.expect("disable credential");
|
||||
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
[
|
||||
"store.credential_details",
|
||||
"store.set_enabled",
|
||||
"provider.account_unavailable",
|
||||
"provider.account_facts_changed",
|
||||
]
|
||||
);
|
||||
assert_eq!(store.audit_requests(), ["request-openai"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_delete_should_commit_then_release_provider_resources() {
|
||||
let events = events();
|
||||
@@ -89,26 +40,6 @@ async fn openai_delete_should_commit_then_release_provider_resources() {
|
||||
assert_eq!(store.audit_requests(), ["request-openai"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_disable_store_failure_should_not_release_provider_resources() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("openai", events.clone());
|
||||
let store = FakeAccountStore::new("openai", events.clone());
|
||||
store.fail_next_commit();
|
||||
let services = service(provider, store).await;
|
||||
|
||||
services
|
||||
.openai()
|
||||
.disable(mutation("acct_test"))
|
||||
.await
|
||||
.expect_err("failed Store transaction");
|
||||
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
["store.credential_details", "store.set_enabled"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn openai_import_should_prepare_before_atomic_store_commit() {
|
||||
let events = events();
|
||||
@@ -450,13 +381,6 @@ async fn await_quota_requests(provider: &FakeProviderAdmin, expected: usize) {
|
||||
assert_eq!(provider.quota_requests().len(), expected);
|
||||
}
|
||||
|
||||
fn mutation(account_id: &str) -> CredentialMutation {
|
||||
CredentialMutation {
|
||||
context: context("request-openai"),
|
||||
account_id: ProviderAccountId::new(account_id).expect("account ID"),
|
||||
}
|
||||
}
|
||||
|
||||
fn deletion(account_id: &str) -> CredentialDeletion {
|
||||
CredentialDeletion {
|
||||
context: context("request-openai"),
|
||||
|
||||
@@ -4,37 +4,12 @@ use gateway_core::engine::credential::ProviderAccountId;
|
||||
|
||||
use gateway_admin::{
|
||||
AdminServices,
|
||||
model::provider_credentials::{CredentialMutation, ImportCredentials, ProviderQuotaRequest},
|
||||
model::provider_credentials::{ImportCredentials, ProviderQuotaRequest},
|
||||
ports::provider::ProviderAdminErrorKind,
|
||||
};
|
||||
|
||||
use super::accounts::{FakeAccountStore, FakeProviderAdmin, context, document, events, recorded};
|
||||
|
||||
#[tokio::test]
|
||||
async fn xai_disable_should_commit_then_notify_the_stateless_provider() {
|
||||
let events = events();
|
||||
let provider = FakeProviderAdmin::new("xai", events.clone());
|
||||
let store = FakeAccountStore::new("xai", events.clone());
|
||||
let services = service(provider.clone(), store.clone()).await;
|
||||
|
||||
services
|
||||
.xai()
|
||||
.disable(mutation("acct_test"))
|
||||
.await
|
||||
.expect("disable credential");
|
||||
|
||||
assert_eq!(
|
||||
recorded(&events),
|
||||
[
|
||||
"store.credential_details",
|
||||
"store.set_enabled",
|
||||
"provider.account_unavailable",
|
||||
"provider.account_facts_changed",
|
||||
]
|
||||
);
|
||||
assert_eq!(store.audit_requests(), ["request-xai"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn xai_import_should_prepare_before_atomic_store_commit() {
|
||||
let events = events();
|
||||
@@ -132,10 +107,3 @@ async fn service(provider: Arc<FakeProviderAdmin>, store: Arc<FakeAccountStore>)
|
||||
.build()
|
||||
.await
|
||||
}
|
||||
|
||||
fn mutation(account_id: &str) -> CredentialMutation {
|
||||
CredentialMutation {
|
||||
context: context("request-xai"),
|
||||
account_id: ProviderAccountId::new(account_id).expect("account ID"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,488 @@
|
||||
//! Account group HTTP wire and fixed routes.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Query, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
};
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_admin::model::{
|
||||
PageSize,
|
||||
account_groups::{
|
||||
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor, AccountGroupListQuery,
|
||||
AccountGroupMember, AccountGroupMembers, AccountGroupMutation, AccountGroupPage,
|
||||
AccountGroupRecord, AccountGroupUsage, CreateAccountGroup, DeleteAccountGroup,
|
||||
SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
},
|
||||
};
|
||||
use gateway_core::routing::AccountGroupId;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{
|
||||
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, PageMeta,
|
||||
WireValidationError, wire::map_admin_service_error,
|
||||
};
|
||||
|
||||
const DEFAULT_PAGE_SIZE: u32 = 50;
|
||||
const MAX_PAGE_SIZE: u32 = 200;
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct ListAccountGroupsQuery {
|
||||
page: Option<u32>,
|
||||
page_size: Option<u32>,
|
||||
search: Option<String>,
|
||||
enabled: Option<bool>,
|
||||
}
|
||||
|
||||
impl ListAccountGroupsQuery {
|
||||
fn into_command(self) -> Result<AccountGroupListQuery, WireValidationError> {
|
||||
let page = self.page.unwrap_or(1);
|
||||
let page_size = self.page_size.unwrap_or(DEFAULT_PAGE_SIZE);
|
||||
if page == 0 || page_size == 0 || page_size > MAX_PAGE_SIZE {
|
||||
return Err(WireValidationError::new("page"));
|
||||
}
|
||||
let search = self
|
||||
.search
|
||||
.map(|value| value.trim().to_owned())
|
||||
.filter(|value| !value.is_empty());
|
||||
if search
|
||||
.as_deref()
|
||||
.is_some_and(|value| value.len() > 256 || value.chars().any(char::is_control))
|
||||
{
|
||||
return Err(WireValidationError::new("search"));
|
||||
}
|
||||
Ok(AccountGroupListQuery {
|
||||
page,
|
||||
page_size: PageSize::new(
|
||||
u16::try_from(page_size).map_err(|_| WireValidationError::new("pageSize"))?,
|
||||
)
|
||||
.map_err(|_| WireValidationError::new("pageSize"))?,
|
||||
search,
|
||||
enabled: self.enabled,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct CreateAccountGroupRequest {
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
color: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct UpdateAccountGroupRequest {
|
||||
id: String,
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
color: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct AccountGroupIdRequest {
|
||||
id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
struct AccountGroupIdQuery {
|
||||
id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupView {
|
||||
id: String,
|
||||
name: String,
|
||||
description: Option<String>,
|
||||
color: String,
|
||||
enabled: bool,
|
||||
member_count: u64,
|
||||
provider_counts: BTreeMap<String, u64>,
|
||||
client_key_count: u64,
|
||||
account_summary: AccountGroupAccountSummaryView,
|
||||
capacity: AccountGroupCapacityView,
|
||||
usage: AccountGroupUsageView,
|
||||
created_at: DateTime<Utc>,
|
||||
updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupAccountSummaryView {
|
||||
available: u64,
|
||||
limited: u64,
|
||||
total: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupCapacityView {
|
||||
used_slots: Option<u64>,
|
||||
total_slots: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupUsageView {
|
||||
today_usd: String,
|
||||
total_usd: String,
|
||||
}
|
||||
|
||||
impl From<AccountGroupRecord> for AccountGroupView {
|
||||
fn from(record: AccountGroupRecord) -> Self {
|
||||
Self {
|
||||
id: record.id.to_string(),
|
||||
name: record.name,
|
||||
description: record.description,
|
||||
color: record.color.as_str().to_owned(),
|
||||
enabled: record.enabled,
|
||||
member_count: record.member_count,
|
||||
provider_counts: record.provider_counts,
|
||||
client_key_count: record.client_key_count,
|
||||
account_summary: account_summary_view(record.account_summary),
|
||||
capacity: capacity_view(record.capacity),
|
||||
usage: usage_view(record.usage),
|
||||
created_at: record.created_at,
|
||||
updated_at: record.updated_at,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn account_summary_view(summary: AccountGroupAccountSummary) -> AccountGroupAccountSummaryView {
|
||||
AccountGroupAccountSummaryView {
|
||||
available: summary.available,
|
||||
limited: summary.limited,
|
||||
total: summary.total,
|
||||
}
|
||||
}
|
||||
|
||||
fn capacity_view(capacity: AccountGroupCapacity) -> AccountGroupCapacityView {
|
||||
AccountGroupCapacityView {
|
||||
used_slots: capacity.used_slots,
|
||||
total_slots: capacity.total_slots,
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_view(usage: AccountGroupUsage) -> AccountGroupUsageView {
|
||||
AccountGroupUsageView {
|
||||
today_usd: usage.today_usd.to_string(),
|
||||
total_usd: usage.total_usd.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupPageData {
|
||||
items: Vec<AccountGroupView>,
|
||||
page: PageMeta,
|
||||
config_revision: u64,
|
||||
}
|
||||
|
||||
impl From<AccountGroupPage> for AccountGroupPageData {
|
||||
fn from(page: AccountGroupPage) -> Self {
|
||||
let total_pages = if page.total == 0 {
|
||||
0
|
||||
} else {
|
||||
page.total.div_ceil(u64::from(page.page_size))
|
||||
};
|
||||
Self {
|
||||
items: page.items.into_iter().map(Into::into).collect(),
|
||||
page: PageMeta::new(
|
||||
page.page,
|
||||
u32::from(page.page_size),
|
||||
page.total,
|
||||
u32::try_from(total_pages).unwrap_or(u32::MAX),
|
||||
),
|
||||
config_revision: page.config_revision.get(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupMemberView {
|
||||
id: String,
|
||||
name: String,
|
||||
provider_kind: String,
|
||||
email: Option<String>,
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl From<AccountGroupMember> for AccountGroupMemberView {
|
||||
fn from(member: AccountGroupMember) -> Self {
|
||||
Self {
|
||||
id: member.id,
|
||||
name: member.name,
|
||||
provider_kind: member.provider_kind.to_string(),
|
||||
email: member.email,
|
||||
enabled: member.enabled,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupMembersData {
|
||||
id: String,
|
||||
items: Vec<AccountGroupMemberView>,
|
||||
total: u64,
|
||||
config_revision: u64,
|
||||
}
|
||||
|
||||
impl From<AccountGroupMembers> for AccountGroupMembersData {
|
||||
fn from(members: AccountGroupMembers) -> Self {
|
||||
Self {
|
||||
id: members.id.to_string(),
|
||||
items: members.items.into_iter().map(Into::into).collect(),
|
||||
total: members.total,
|
||||
config_revision: members.config_revision.get(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AccountGroupMutationData {
|
||||
id: String,
|
||||
record: Option<AccountGroupView>,
|
||||
config_revision: u64,
|
||||
}
|
||||
|
||||
impl From<AccountGroupMutation> for AccountGroupMutationData {
|
||||
fn from(mutation: AccountGroupMutation) -> Self {
|
||||
Self {
|
||||
id: mutation.id.to_string(),
|
||||
record: mutation.record.map(Into::into),
|
||||
config_revision: mutation.config_revision.get(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct all fixed account-group management routes.
|
||||
pub fn router<S>() -> Router<S>
|
||||
where
|
||||
S: AdminSessionState + Clone + Send + Sync + 'static,
|
||||
{
|
||||
Router::new()
|
||||
.route("/api/admin/account-groups", get(list::<S>))
|
||||
.route("/api/admin/account-groups/create", post(create::<S>))
|
||||
.route("/api/admin/account-groups/update", post(update::<S>))
|
||||
.route("/api/admin/account-groups/enable", post(enable::<S>))
|
||||
.route("/api/admin/account-groups/disable", post(disable::<S>))
|
||||
.route("/api/admin/account-groups/delete", post(delete::<S>))
|
||||
.route("/api/admin/account-groups/members", get(members::<S>))
|
||||
}
|
||||
|
||||
async fn list<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Query(query): Query<ListAccountGroupsQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let result = state
|
||||
.admin_services()
|
||||
.account_groups()
|
||||
.list(query.into_command().map_err(map_wire_error)?)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AccountGroupPageData::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn members<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Query(query): Query<AccountGroupIdQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let id = group_id(query.id)?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.account_groups()
|
||||
.members(&id)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AccountGroupMembersData::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn create<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Json(request): Json<CreateAccountGroupRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
validate_group_fields(&request.name, request.description.as_deref())?;
|
||||
mutation_response(
|
||||
StatusCode::CREATED,
|
||||
state
|
||||
.admin_services()
|
||||
.account_groups()
|
||||
.create(
|
||||
&auth.context().mutation_context(),
|
||||
CreateAccountGroup {
|
||||
name: request.name,
|
||||
description: request.description,
|
||||
color: group_color(&request.color)?,
|
||||
},
|
||||
)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
|
||||
async fn update<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Json(request): Json<UpdateAccountGroupRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
validate_group_fields(&request.name, request.description.as_deref())?;
|
||||
mutation_response(
|
||||
StatusCode::OK,
|
||||
state
|
||||
.admin_services()
|
||||
.account_groups()
|
||||
.update(
|
||||
&auth.context().mutation_context(),
|
||||
UpdateAccountGroup {
|
||||
id: group_id(request.id)?,
|
||||
name: request.name,
|
||||
description: request.description,
|
||||
color: group_color(&request.color)?,
|
||||
},
|
||||
)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
|
||||
async fn enable<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Json(request): Json<AccountGroupIdRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
set_enabled(auth, state, request, true).await
|
||||
}
|
||||
|
||||
async fn disable<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Json(request): Json<AccountGroupIdRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
set_enabled(auth, state, request, false).await
|
||||
}
|
||||
|
||||
async fn set_enabled<S>(
|
||||
auth: AdminAuth,
|
||||
state: S,
|
||||
request: AccountGroupIdRequest,
|
||||
enabled: bool,
|
||||
) -> Result<AdminResponse<AdminEnvelope<AccountGroupMutationData>>, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
mutation_response(
|
||||
StatusCode::OK,
|
||||
state
|
||||
.admin_services()
|
||||
.account_groups()
|
||||
.set_enabled(
|
||||
&auth.context().mutation_context(),
|
||||
SetAccountGroupEnabled {
|
||||
id: group_id(request.id)?,
|
||||
enabled,
|
||||
},
|
||||
)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
|
||||
async fn delete<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Json(request): Json<AccountGroupIdRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
mutation_response(
|
||||
StatusCode::OK,
|
||||
state
|
||||
.admin_services()
|
||||
.account_groups()
|
||||
.delete(
|
||||
&auth.context().mutation_context(),
|
||||
DeleteAccountGroup {
|
||||
id: group_id(request.id)?,
|
||||
},
|
||||
)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
|
||||
fn mutation_response(
|
||||
status: StatusCode,
|
||||
result: Result<AccountGroupMutation, gateway_admin::model::AdminError>,
|
||||
) -> Result<AdminResponse<AdminEnvelope<AccountGroupMutationData>>, AdminError> {
|
||||
let result = result.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
status,
|
||||
AdminEnvelope::ok(AccountGroupMutationData::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
fn group_id(value: String) -> Result<AccountGroupId, AdminError> {
|
||||
AccountGroupId::new(value).map_err(|_| AdminError::bad_request("Invalid account group ID"))
|
||||
}
|
||||
|
||||
fn group_color(value: &str) -> Result<AccountGroupColor, AdminError> {
|
||||
AccountGroupColor::parse(value)
|
||||
.ok_or_else(|| AdminError::bad_request("Invalid account group color"))
|
||||
}
|
||||
|
||||
fn validate_group_fields(name: &str, description: Option<&str>) -> Result<(), AdminError> {
|
||||
if name.trim() != name
|
||||
|| name.is_empty()
|
||||
|| name.chars().count() > 100
|
||||
|| name.chars().any(char::is_control)
|
||||
|| description
|
||||
.is_some_and(|value| value.len() > 4096 || value.chars().any(char::is_control))
|
||||
{
|
||||
return Err(AdminError::bad_request("Invalid account group request"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn map_wire_error(_: WireValidationError) -> AdminError {
|
||||
AdminError::bad_request("Invalid account group query")
|
||||
}
|
||||
|
||||
fn map_service_error(error: gateway_admin::model::AdminError) -> AdminError {
|
||||
map_admin_service_error(error, "Account group repository unavailable")
|
||||
}
|
||||
@@ -18,9 +18,10 @@ use futures::{Stream, StreamExt as _};
|
||||
use gateway_admin::model::{
|
||||
AdminError as AdminServiceError, PageSize,
|
||||
accounts::{
|
||||
AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost, AccountListQuery,
|
||||
AccountModelUsage, AccountSort, AccountSortField, AccountStatus as DomainAccountStatus,
|
||||
AccountUsage, SortDirection,
|
||||
AccountConcurrencyLimit, AccountConnectionTestEvent as DomainConnectionTestEvent,
|
||||
AccountCost, AccountGroupFilter, AccountListQuery, AccountModelUsage, AccountSort,
|
||||
AccountSortField, AccountStatus as DomainAccountStatus, AccountUpdateResult, AccountUsage,
|
||||
AccountWeight, AccountsUpdateResult, BatchUpdateAccounts, SortDirection, UpdateAccount,
|
||||
},
|
||||
provider_credentials::{
|
||||
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
|
||||
@@ -32,9 +33,9 @@ use gateway_admin::model::{
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::credential::{OpaqueProviderData, ProviderAccountId},
|
||||
routing::{ProviderKind, UpstreamModelId},
|
||||
routing::{AccountGroupId, ProviderKind, UpstreamModelId},
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde::{Deserialize, Deserializer, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
|
||||
@@ -54,6 +55,7 @@ const MAX_REFRESH_TOKEN_BYTES: usize = 64 * 1024;
|
||||
const MAX_ID_TOKEN_BYTES: usize = 16 * 1024;
|
||||
const MAX_CALLBACK_URL_BYTES: usize = 64 * 1024;
|
||||
const MAX_ACCOUNT_DELETE_BATCH: usize = 200;
|
||||
const MAX_ACCOUNT_GROUP_BATCH: usize = 1000;
|
||||
|
||||
/// 账号列表查询参数。
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
@@ -62,12 +64,74 @@ pub struct ListQuery {
|
||||
pub page: Option<u32>,
|
||||
pub page_size: Option<u32>,
|
||||
pub provider: Option<String>,
|
||||
pub group_id: Option<String>,
|
||||
pub search: Option<String>,
|
||||
pub status: Option<String>,
|
||||
pub sort_by: Option<String>,
|
||||
pub sort_direction: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct BatchUpdateAccountsRequest {
|
||||
pub account_ids: Vec<String>,
|
||||
pub enabled: bool,
|
||||
#[serde(deserialize_with = "deserialize_required_nullable")]
|
||||
pub concurrency_limit: Option<u64>,
|
||||
pub weight: u64,
|
||||
pub group_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct BatchUpdatedAccountsData {
|
||||
account_ids: Vec<String>,
|
||||
config_revision: u64,
|
||||
}
|
||||
|
||||
impl From<AccountsUpdateResult> for BatchUpdatedAccountsData {
|
||||
fn from(result: AccountsUpdateResult) -> Self {
|
||||
Self {
|
||||
account_ids: result
|
||||
.account_ids
|
||||
.into_iter()
|
||||
.map(|id| id.to_string())
|
||||
.collect(),
|
||||
config_revision: result.config_revision.get(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl BatchUpdateAccountsRequest {
|
||||
pub fn validate(&self) -> Result<(), WireValidationError> {
|
||||
if self.account_ids.is_empty()
|
||||
|| self.account_ids.len() > MAX_ACCOUNT_GROUP_BATCH
|
||||
|| self
|
||||
.account_ids
|
||||
.iter()
|
||||
.any(|id| require_account_id(id, "accountIds").is_err())
|
||||
|| self.account_ids.iter().collect::<BTreeSet<_>>().len() != self.account_ids.len()
|
||||
{
|
||||
return Err(WireValidationError::new("accountIds"));
|
||||
}
|
||||
validate_wire_group_ids(&self.group_ids)?;
|
||||
parse_concurrency_limit(self.concurrency_limit)?;
|
||||
parse_account_weight(self.weight)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn into_command(self) -> Result<BatchUpdateAccounts, WireValidationError> {
|
||||
self.validate()?;
|
||||
Ok(BatchUpdateAccounts {
|
||||
account_ids: self.account_ids,
|
||||
enabled: self.enabled,
|
||||
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
|
||||
weight: parse_account_weight(self.weight)?,
|
||||
group_ids: validate_wire_group_ids(&self.group_ids)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ListQuery {
|
||||
/// 解析并校验全部 wire 字段,生成 Admin 查询命令。
|
||||
pub fn validate(self) -> Result<AccountListQuery, WireValidationError> {
|
||||
@@ -80,6 +144,14 @@ impl ListQuery {
|
||||
return Err(WireValidationError::new("pageSize"));
|
||||
}
|
||||
let provider_kind = parse_provider(self.provider.as_deref().unwrap_or("all"))?;
|
||||
let group_filter = match self.group_id.as_deref().map(str::trim) {
|
||||
None | Some("") => None,
|
||||
Some("ungrouped") => Some(AccountGroupFilter::Ungrouped),
|
||||
Some(value) => Some(AccountGroupFilter::Group(
|
||||
AccountGroupId::new(value.to_owned())
|
||||
.map_err(|_| WireValidationError::new("groupId"))?,
|
||||
)),
|
||||
};
|
||||
let search = self
|
||||
.search
|
||||
.map(|value| value.trim().to_owned())
|
||||
@@ -112,6 +184,7 @@ impl ListQuery {
|
||||
)
|
||||
.map_err(|_| WireValidationError::new("pageSize"))?,
|
||||
provider_kind,
|
||||
group_filter,
|
||||
search,
|
||||
status,
|
||||
sort,
|
||||
@@ -177,6 +250,7 @@ pub struct AccountView {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub provider: String,
|
||||
pub groups: Vec<AccountGroupRefView>,
|
||||
pub resource_ref: String,
|
||||
pub email: Option<String>,
|
||||
pub account_id: Option<String>,
|
||||
@@ -191,6 +265,8 @@ pub struct AccountView {
|
||||
/// 最近一次失败的上游错误描述;仅错误状态存在。
|
||||
pub error_message: Option<String>,
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<u32>,
|
||||
pub weight: u16,
|
||||
pub access_token_expires_at: Option<String>,
|
||||
pub access_token_expires_at_display: Option<String>,
|
||||
pub refresh_token_expires_at: Option<String>,
|
||||
@@ -203,6 +279,15 @@ pub struct AccountView {
|
||||
pub usage: AccountUsageView,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountGroupRefView {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub color: String,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
/// Provider quota 安全视图。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
@@ -682,31 +767,49 @@ impl CompleteAccountAuthorizationRequest {
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
pub struct AccountMutationRequest {
|
||||
pub provider: String,
|
||||
pub struct UpdateAccountRequest {
|
||||
pub account_id: String,
|
||||
pub enabled: bool,
|
||||
#[serde(deserialize_with = "deserialize_required_nullable")]
|
||||
pub concurrency_limit: Option<u64>,
|
||||
pub weight: u64,
|
||||
pub group_ids: Vec<String>,
|
||||
}
|
||||
|
||||
impl AccountMutationRequest {
|
||||
impl UpdateAccountRequest {
|
||||
pub fn validate(&self) -> Result<(), WireValidationError> {
|
||||
AccountProvider::parse(&self.provider)?;
|
||||
require_account_id(&self.account_id, "accountId")
|
||||
require_account_id(&self.account_id, "accountId")?;
|
||||
parse_concurrency_limit(self.concurrency_limit)?;
|
||||
parse_account_weight(self.weight)?;
|
||||
validate_wire_group_ids(&self.group_ids)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn into_command(
|
||||
self,
|
||||
context: gateway_admin::model::MutationContext,
|
||||
) -> Result<(AccountProvider, CredentialMutation), WireValidationError> {
|
||||
fn into_command(self) -> Result<UpdateAccount, WireValidationError> {
|
||||
self.validate()?;
|
||||
let provider = AccountProvider::parse(&self.provider)?;
|
||||
Ok((
|
||||
provider,
|
||||
CredentialMutation {
|
||||
context,
|
||||
account_id: ProviderAccountId::new(self.account_id)
|
||||
.map_err(|_| WireValidationError::new("accountId"))?,
|
||||
},
|
||||
))
|
||||
Ok(UpdateAccount {
|
||||
account_id: self.account_id,
|
||||
enabled: self.enabled,
|
||||
concurrency_limit: parse_concurrency_limit(self.concurrency_limit)?,
|
||||
weight: parse_account_weight(self.weight)?,
|
||||
group_ids: validate_wire_group_ids(&self.group_ids)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UpdatedAccountData {
|
||||
pub account_id: String,
|
||||
pub config_revision: u64,
|
||||
}
|
||||
|
||||
impl From<AccountUpdateResult> for UpdatedAccountData {
|
||||
fn from(result: AccountUpdateResult) -> Self {
|
||||
Self {
|
||||
account_id: result.account_id.to_string(),
|
||||
config_revision: result.config_revision.get(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -810,8 +913,8 @@ pub struct AccountImportData {
|
||||
pub account_ids: Vec<String>,
|
||||
}
|
||||
|
||||
impl From<CredentialImportResult> for AccountImportData {
|
||||
fn from(result: CredentialImportResult) -> Self {
|
||||
impl AccountImportData {
|
||||
pub fn from_result(result: CredentialImportResult) -> Self {
|
||||
let account_ids = result
|
||||
.credential_ids
|
||||
.into_iter()
|
||||
@@ -888,6 +991,46 @@ fn require_account_id(value: &str, field: &'static str) -> Result<(), WireValida
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_concurrency_limit(
|
||||
value: Option<u64>,
|
||||
) -> Result<Option<AccountConcurrencyLimit>, WireValidationError> {
|
||||
value
|
||||
.map(|value| {
|
||||
u32::try_from(value)
|
||||
.ok()
|
||||
.and_then(AccountConcurrencyLimit::new)
|
||||
.ok_or_else(|| WireValidationError::new("concurrencyLimit"))
|
||||
})
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn deserialize_required_nullable<'de, D>(deserializer: D) -> Result<Option<u64>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
Option::<u64>::deserialize(deserializer)
|
||||
}
|
||||
|
||||
fn parse_account_weight(value: u64) -> Result<AccountWeight, WireValidationError> {
|
||||
u16::try_from(value)
|
||||
.ok()
|
||||
.and_then(AccountWeight::new)
|
||||
.ok_or_else(|| WireValidationError::new("weight"))
|
||||
}
|
||||
|
||||
fn validate_wire_group_ids(values: &[String]) -> Result<Vec<AccountGroupId>, WireValidationError> {
|
||||
if values.len() > MAX_ACCOUNT_GROUP_BATCH
|
||||
|| values.iter().collect::<BTreeSet<_>>().len() != values.len()
|
||||
{
|
||||
return Err(WireValidationError::new("groupIds"));
|
||||
}
|
||||
values
|
||||
.iter()
|
||||
.cloned()
|
||||
.map(|value| AccountGroupId::new(value).map_err(|_| WireValidationError::new("groupIds")))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 构造统一账号管理路由。
|
||||
pub fn router<S>() -> Router<S>
|
||||
where
|
||||
@@ -900,9 +1043,12 @@ where
|
||||
.route("/api/admin/accounts/import", post(import_accounts::<S>))
|
||||
.route("/api/admin/accounts/refresh", post(refresh_account::<S>))
|
||||
.route("/api/admin/accounts/rotate", post(rotate_account::<S>))
|
||||
.route("/api/admin/accounts/enable", post(enable_account::<S>))
|
||||
.route("/api/admin/accounts/disable", post(disable_account::<S>))
|
||||
.route("/api/admin/accounts/update", post(update_account::<S>))
|
||||
.route("/api/admin/accounts/delete", post(delete_accounts::<S>))
|
||||
.route(
|
||||
"/api/admin/accounts/batch-update",
|
||||
post(batch_update_accounts::<S>),
|
||||
)
|
||||
.route("/api/admin/accounts/quota", get(account_quota::<S>))
|
||||
.route(
|
||||
"/api/admin/accounts/quota/refresh",
|
||||
@@ -927,6 +1073,27 @@ where
|
||||
)
|
||||
}
|
||||
|
||||
async fn batch_update_accounts<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Json(request): Json<BatchUpdateAccountsRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let command = request.into_command().map_err(map_wire_error)?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.accounts()
|
||||
.batch_update(&auth.context().mutation_context(), command)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(BatchUpdatedAccountsData::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
async fn list_accounts<S>(
|
||||
_auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
@@ -1016,7 +1183,7 @@ where
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::CREATED,
|
||||
AdminEnvelope::ok(AccountImportData::from(result)),
|
||||
AdminEnvelope::ok(AccountImportData::from_result(result)),
|
||||
))
|
||||
}
|
||||
|
||||
@@ -1111,64 +1278,24 @@ where
|
||||
))
|
||||
}
|
||||
|
||||
async fn enable_account<S>(
|
||||
async fn update_account<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Json(request): Json<AccountMutationRequest>,
|
||||
Json(request): Json<UpdateAccountRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
mutate_account(auth, state, request, AccountMutationAction::Enable).await
|
||||
}
|
||||
|
||||
async fn disable_account<S>(
|
||||
auth: AdminAuth,
|
||||
State(state): State<S>,
|
||||
Json(request): Json<AccountMutationRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
mutate_account(auth, state, request, AccountMutationAction::Disable).await
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum AccountMutationAction {
|
||||
Enable,
|
||||
Disable,
|
||||
}
|
||||
|
||||
async fn mutate_account<S>(
|
||||
auth: AdminAuth,
|
||||
state: S,
|
||||
request: AccountMutationRequest,
|
||||
action: AccountMutationAction,
|
||||
) -> Result<impl IntoResponse, AdminError>
|
||||
where
|
||||
S: AdminSessionState + Send + Sync,
|
||||
{
|
||||
let (provider, command) = request
|
||||
.into_command(auth.context().mutation_context())
|
||||
.map_err(map_wire_error)?;
|
||||
let result = match (provider, action) {
|
||||
(AccountProvider::OpenAi, AccountMutationAction::Enable) => {
|
||||
state.admin_services().openai().enable(command).await
|
||||
}
|
||||
(AccountProvider::OpenAi, AccountMutationAction::Disable) => {
|
||||
state.admin_services().openai().disable(command).await
|
||||
}
|
||||
(AccountProvider::Xai, AccountMutationAction::Enable) => {
|
||||
state.admin_services().xai().enable(command).await
|
||||
}
|
||||
(AccountProvider::Xai, AccountMutationAction::Disable) => {
|
||||
state.admin_services().xai().disable(command).await
|
||||
}
|
||||
}
|
||||
.map_err(map_service_error)?;
|
||||
let command = request.into_command().map_err(map_wire_error)?;
|
||||
let result = state
|
||||
.admin_services()
|
||||
.accounts()
|
||||
.update(&auth.context().mutation_context(), command)
|
||||
.await
|
||||
.map_err(map_service_error)?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AccountMutationData::from(result)),
|
||||
AdminEnvelope::ok(UpdatedAccountData::from(result)),
|
||||
))
|
||||
}
|
||||
|
||||
@@ -1384,6 +1511,16 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
|
||||
id: account.id.clone(),
|
||||
name: account.name,
|
||||
provider: account.provider_kind.to_string(),
|
||||
groups: account
|
||||
.groups
|
||||
.into_iter()
|
||||
.map(|group| AccountGroupRefView {
|
||||
id: group.id.to_string(),
|
||||
name: group.name,
|
||||
color: group.color.as_str().to_owned(),
|
||||
enabled: group.enabled,
|
||||
})
|
||||
.collect(),
|
||||
resource_ref: account.id,
|
||||
email: account.email,
|
||||
account_id: account.upstream_account_id,
|
||||
@@ -1398,6 +1535,8 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
|
||||
.map(|reason| reason.as_str().to_owned()),
|
||||
error_message: projection.error_message,
|
||||
enabled: account.enabled,
|
||||
concurrency_limit: account.concurrency_limit.map(|limit| limit.get()),
|
||||
weight: account.weight.get(),
|
||||
access_token_expires_at: expires_at,
|
||||
access_token_expires_at_display: account
|
||||
.access_token_expires_at
|
||||
|
||||
@@ -12,7 +12,7 @@ use gateway_admin::model::client_keys::{
|
||||
};
|
||||
use gateway_core::{
|
||||
policy::{ClientApiKeyId, RateLimits},
|
||||
routing::ProviderKind,
|
||||
routing::AccountGroupId,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
@@ -137,7 +137,7 @@ impl ClientKeySort {
|
||||
pub struct CreateClientKeyRequest {
|
||||
name: String,
|
||||
label: Option<String>,
|
||||
provider_kind: String,
|
||||
group_ids: Vec<String>,
|
||||
max_concurrency: u64,
|
||||
requests_per_minute: u64,
|
||||
}
|
||||
@@ -147,14 +147,13 @@ impl CreateClientKeyRequest {
|
||||
pub fn into_command(self) -> Result<CreateClientKey, WireValidationError> {
|
||||
validate_required_text(&self.name, "name")?;
|
||||
validate_optional_text(self.label.as_deref(), "label")?;
|
||||
validate_provider_kind(&self.provider_kind)?;
|
||||
let group_ids = validate_group_ids(self.group_ids)?;
|
||||
validate_limit(self.max_concurrency, "maxConcurrency")?;
|
||||
validate_limit(self.requests_per_minute, "requestsPerMinute")?;
|
||||
Ok(CreateClientKey {
|
||||
name: self.name,
|
||||
label: self.label,
|
||||
provider_kind: ProviderKind::new(self.provider_kind)
|
||||
.map_err(|_| WireValidationError::new("providerKind"))?,
|
||||
group_ids,
|
||||
limits: RateLimits {
|
||||
max_concurrency: self.max_concurrency,
|
||||
requests_per_minute: self.requests_per_minute,
|
||||
@@ -170,7 +169,7 @@ pub struct UpdateClientKeyRequest {
|
||||
id: String,
|
||||
name: String,
|
||||
label: Option<String>,
|
||||
provider_kind: String,
|
||||
group_ids: Vec<String>,
|
||||
max_concurrency: u64,
|
||||
requests_per_minute: u64,
|
||||
}
|
||||
@@ -181,15 +180,14 @@ impl UpdateClientKeyRequest {
|
||||
validate_required_text(&self.id, "id")?;
|
||||
validate_required_text(&self.name, "name")?;
|
||||
validate_optional_text(self.label.as_deref(), "label")?;
|
||||
validate_provider_kind(&self.provider_kind)?;
|
||||
let group_ids = validate_group_ids(self.group_ids)?;
|
||||
validate_limit(self.max_concurrency, "maxConcurrency")?;
|
||||
validate_limit(self.requests_per_minute, "requestsPerMinute")?;
|
||||
Ok(UpdateClientKey {
|
||||
id: client_key_id(self.id, "clientKeyMutationNotFound")?,
|
||||
name: self.name,
|
||||
label: self.label,
|
||||
provider_kind: ProviderKind::new(self.provider_kind)
|
||||
.map_err(|_| WireValidationError::new("providerKind"))?,
|
||||
group_ids,
|
||||
limits: RateLimits {
|
||||
max_concurrency: self.max_concurrency,
|
||||
requests_per_minute: self.requests_per_minute,
|
||||
@@ -242,7 +240,9 @@ pub struct ClientKeyView {
|
||||
id: String,
|
||||
name: String,
|
||||
label: Option<String>,
|
||||
provider_kind: String,
|
||||
routing_scope: &'static str,
|
||||
groups: Vec<ClientKeyGroupView>,
|
||||
provider_kinds: Vec<String>,
|
||||
prefix: String,
|
||||
enabled: bool,
|
||||
max_concurrency: u64,
|
||||
@@ -252,13 +252,42 @@ pub struct ClientKeyView {
|
||||
last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ClientKeyGroupView {
|
||||
id: String,
|
||||
name: String,
|
||||
color: String,
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl From<ClientKeyRecord> for ClientKeyView {
|
||||
fn from(record: ClientKeyRecord) -> Self {
|
||||
let routing_scope = if record.groups.is_empty() {
|
||||
"all"
|
||||
} else {
|
||||
"groups"
|
||||
};
|
||||
Self {
|
||||
id: record.id.to_string(),
|
||||
name: record.name,
|
||||
label: record.label,
|
||||
provider_kind: record.provider_kind.to_string(),
|
||||
routing_scope,
|
||||
groups: record
|
||||
.groups
|
||||
.into_iter()
|
||||
.map(|group| ClientKeyGroupView {
|
||||
id: group.id.to_string(),
|
||||
name: group.name,
|
||||
color: group.color.as_str().to_owned(),
|
||||
enabled: group.enabled,
|
||||
})
|
||||
.collect(),
|
||||
provider_kinds: record
|
||||
.provider_kinds
|
||||
.into_iter()
|
||||
.map(|provider| provider.to_string())
|
||||
.collect(),
|
||||
prefix: record.prefix,
|
||||
enabled: record.enabled,
|
||||
max_concurrency: record.limits.max_concurrency,
|
||||
@@ -572,21 +601,20 @@ fn validate_required_text(value: &str, field: &'static str) -> Result<(), WireVa
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_provider_kind(value: &str) -> Result<(), WireValidationError> {
|
||||
let valid = !value.is_empty()
|
||||
&& value.len() <= 64
|
||||
&& value.bytes().enumerate().all(|(index, byte)| {
|
||||
byte.is_ascii_lowercase()
|
||||
|| byte.is_ascii_digit()
|
||||
|| (byte == b'-' && index > 0 && index + 1 < value.len())
|
||||
})
|
||||
&& !value.starts_with('-')
|
||||
&& !value.ends_with('-');
|
||||
if valid {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(WireValidationError::new("providerKind"))
|
||||
fn validate_group_ids(values: Vec<String>) -> Result<Vec<AccountGroupId>, WireValidationError> {
|
||||
if values.len() > 1000
|
||||
|| values
|
||||
.iter()
|
||||
.collect::<std::collections::BTreeSet<_>>()
|
||||
.len()
|
||||
!= values.len()
|
||||
{
|
||||
return Err(WireValidationError::new("groupIds"));
|
||||
}
|
||||
values
|
||||
.into_iter()
|
||||
.map(|value| AccountGroupId::new(value).map_err(|_| WireValidationError::new("groupIds")))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn validate_optional_text(
|
||||
|
||||
@@ -7,6 +7,7 @@ use axum::{
|
||||
response::Response,
|
||||
};
|
||||
|
||||
pub mod account_groups;
|
||||
pub mod accounts;
|
||||
pub mod auth;
|
||||
pub mod backups;
|
||||
@@ -29,6 +30,7 @@ where
|
||||
S: AdminSessionState + Clone + Send + Sync + 'static,
|
||||
{
|
||||
Router::new()
|
||||
.merge(account_groups::router::<S>())
|
||||
.merge(accounts::router::<S>())
|
||||
.merge(auth::router::<S>())
|
||||
.merge(backups::router::<S>())
|
||||
|
||||
@@ -232,6 +232,9 @@ pub(crate) fn usage_record_view(record: domain::UsageRecord) -> UsageRecordView
|
||||
id: record.id.clone(),
|
||||
request_id: record.id,
|
||||
client_api_key_id: Some(record.client_api_key_ref),
|
||||
routing_scope: record.routing_scope,
|
||||
routing_group_refs: record.routing_group_refs,
|
||||
routing_group_names_snapshot: record.routing_group_names_snapshot,
|
||||
kind: record.operation,
|
||||
provider: record.provider_kind,
|
||||
authentication_kind: record.provider_account_authentication_kind,
|
||||
|
||||
@@ -84,6 +84,9 @@ pub struct UsageRecordView {
|
||||
pub id: String,
|
||||
pub request_id: String,
|
||||
pub client_api_key_id: Option<String>,
|
||||
pub routing_scope: String,
|
||||
pub routing_group_refs: Vec<String>,
|
||||
pub routing_group_names_snapshot: Vec<String>,
|
||||
pub kind: String,
|
||||
pub provider: Option<String>,
|
||||
pub authentication_kind: Option<String>,
|
||||
|
||||
@@ -89,6 +89,7 @@ pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
|
||||
| EngineError::ProviderMetadataMismatch
|
||||
| EngineError::ContinuationPinMismatch
|
||||
| EngineError::RequiredAccountMismatch
|
||||
| EngineError::AccountOutsideClientScope
|
||||
| EngineError::DownstreamCommitRequired
|
||||
| EngineError::NoActiveAttempt
|
||||
| EngineError::InvalidDeliveryState => {
|
||||
|
||||
@@ -0,0 +1,292 @@
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::{Method, Request, StatusCode, header},
|
||||
};
|
||||
use gateway_api::admin::account_groups;
|
||||
use serde_json::{Value, json};
|
||||
use tower::ServiceExt as _;
|
||||
|
||||
use super::{AdminTestFixture, AdminTestState, PRIMARY_GROUP_ID};
|
||||
|
||||
const SESSION_COOKIE: &str = "cpr_admin_session=valid-session";
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_route_should_keep_camel_case_group_and_page_wire() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
let response = request(
|
||||
router(&fixture),
|
||||
Method::GET,
|
||||
"/api/admin/account-groups?page=1&pageSize=1&search=alpha&enabled=true",
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let value = response_json(response).await;
|
||||
assert_eq!(value["data"]["page"]["page"], 1);
|
||||
assert_eq!(value["data"]["page"]["pageSize"], 1);
|
||||
assert_eq!(value["data"]["page"]["total"], 1);
|
||||
assert_eq!(value["data"]["page"]["totalPages"], 1);
|
||||
assert_eq!(value["data"]["configRevision"], 7);
|
||||
assert_eq!(value["data"]["items"][0]["id"], PRIMARY_GROUP_ID);
|
||||
assert_eq!(value["data"]["items"][0]["memberCount"], 2);
|
||||
assert_eq!(value["data"]["items"][0]["providerCounts"]["openai"], 1);
|
||||
assert_eq!(value["data"]["items"][0]["clientKeyCount"], 2);
|
||||
assert_eq!(value["data"]["items"][0]["accountSummary"]["available"], 1);
|
||||
assert_eq!(value["data"]["items"][0]["accountSummary"]["limited"], 1);
|
||||
assert_eq!(value["data"]["items"][0]["accountSummary"]["total"], 2);
|
||||
assert_eq!(value["data"]["items"][0]["capacity"]["usedSlots"], 0);
|
||||
assert_eq!(value["data"]["items"][0]["capacity"]["totalSlots"], 1);
|
||||
assert_eq!(value["data"]["items"][0]["usage"]["todayUsd"], "1.25");
|
||||
assert_eq!(value["data"]["items"][0]["usage"]["totalUsd"], "5.5");
|
||||
assert!(value["data"]["items"][0].get("member_count").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn members_route_should_keep_provider_kind_and_member_wire() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
let response = request(
|
||||
router(&fixture),
|
||||
Method::GET,
|
||||
&format!("/api/admin/account-groups/members?id={PRIMARY_GROUP_ID}"),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let value = response_json(response).await;
|
||||
assert_eq!(value["data"]["id"], PRIMARY_GROUP_ID);
|
||||
assert_eq!(value["data"]["total"], 2);
|
||||
assert_eq!(value["data"]["items"][0]["providerKind"], "openai");
|
||||
assert_eq!(value["data"]["items"][0]["email"], "openai@example.invalid");
|
||||
assert!(value["data"]["items"][0].get("provider_kind").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn create_route_should_create_an_empty_group() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
let response = request(
|
||||
router(&fixture),
|
||||
Method::POST,
|
||||
"/api/admin/account-groups/create",
|
||||
Some(json!({
|
||||
"name": "Gamma routing",
|
||||
"description": "New traffic pool",
|
||||
"color": "#a855f780"
|
||||
})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::CREATED);
|
||||
let value = response_json(response).await;
|
||||
assert!(
|
||||
value["data"]["id"]
|
||||
.as_str()
|
||||
.is_some_and(|id| id.starts_with("grp_") && id.len() == 36)
|
||||
);
|
||||
assert_eq!(value["data"]["record"]["name"], "Gamma routing");
|
||||
assert_eq!(value["data"]["record"]["memberCount"], 0);
|
||||
assert_eq!(value["data"]["record"]["color"], "#A855F780");
|
||||
assert_eq!(value["data"]["configRevision"], 8);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn update_route_should_replace_group_fields() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
let response = request(
|
||||
router(&fixture),
|
||||
Method::POST,
|
||||
"/api/admin/account-groups/update",
|
||||
Some(json!({
|
||||
"id": PRIMARY_GROUP_ID,
|
||||
"name": "Renamed routing",
|
||||
"description": null,
|
||||
"color": "#F43F5ECC"
|
||||
})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let value = response_json(response).await;
|
||||
assert_eq!(value["data"]["record"]["name"], "Renamed routing");
|
||||
assert!(value["data"]["record"]["description"].is_null());
|
||||
assert_eq!(value["data"]["record"]["color"], "#F43F5ECC");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn enable_and_disable_routes_should_set_explicit_group_state() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
let disabled = request(
|
||||
router(&fixture),
|
||||
Method::POST,
|
||||
"/api/admin/account-groups/disable",
|
||||
Some(json!({ "id": PRIMARY_GROUP_ID })),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
let enabled = request(
|
||||
router(&fixture),
|
||||
Method::POST,
|
||||
"/api/admin/account-groups/enable",
|
||||
Some(json!({ "id": PRIMARY_GROUP_ID })),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(disabled.status(), StatusCode::OK);
|
||||
assert!(
|
||||
!response_json(disabled).await["data"]["record"]["enabled"]
|
||||
.as_bool()
|
||||
.expect("disabled flag")
|
||||
);
|
||||
assert_eq!(enabled.status(), StatusCode::OK);
|
||||
assert!(
|
||||
response_json(enabled).await["data"]["record"]["enabled"]
|
||||
.as_bool()
|
||||
.expect("enabled flag")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_route_should_return_deleted_id_without_a_live_record() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
let response = request(
|
||||
router(&fixture),
|
||||
Method::POST,
|
||||
"/api/admin/account-groups/delete",
|
||||
Some(json!({ "id": PRIMARY_GROUP_ID })),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let value = response_json(response).await;
|
||||
assert_eq!(value["data"]["id"], PRIMARY_GROUP_ID);
|
||||
assert!(value["data"]["record"].is_null());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_route_should_reject_unknown_and_invalid_pagination_fields() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
for uri in [
|
||||
"/api/admin/account-groups?other=true",
|
||||
"/api/admin/account-groups?page=0",
|
||||
"/api/admin/account-groups?pageSize=0",
|
||||
"/api/admin/account-groups?pageSize=201",
|
||||
] {
|
||||
let response = request(router(&fixture), Method::GET, uri, None, true).await;
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{uri}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mutation_route_should_reject_unknown_json_fields() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
let response = request(
|
||||
router(&fixture),
|
||||
Method::POST,
|
||||
"/api/admin/account-groups/create",
|
||||
Some(json!({
|
||||
"name": "Unknown field",
|
||||
"description": null,
|
||||
"color": "#2563EBFF",
|
||||
"providerKind": "openai"
|
||||
})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::UNPROCESSABLE_ENTITY);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn create_route_should_reject_six_digit_color_without_alpha() {
|
||||
let fixture = authenticated_fixture().await;
|
||||
let response = request(
|
||||
router(&fixture),
|
||||
Method::POST,
|
||||
"/api/admin/account-groups/create",
|
||||
Some(json!({
|
||||
"name": "Opaque legacy color",
|
||||
"description": null,
|
||||
"color": "#2563EB"
|
||||
})),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_group_routes_should_require_admin_authentication() {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
let get_response = request(
|
||||
router(&fixture),
|
||||
Method::GET,
|
||||
"/api/admin/account-groups",
|
||||
None,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
let post_response = request(
|
||||
router(&fixture),
|
||||
Method::POST,
|
||||
"/api/admin/account-groups/create",
|
||||
Some(json!({ "name": "Unauthorized", "description": null, "color": "#2563EBFF" })),
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(get_response.status(), StatusCode::UNAUTHORIZED);
|
||||
assert_eq!(post_response.status(), StatusCode::UNAUTHORIZED);
|
||||
}
|
||||
|
||||
async fn authenticated_fixture() -> AdminTestFixture {
|
||||
let fixture = AdminTestFixture::new().await;
|
||||
fixture.auth.insert_session("valid-session");
|
||||
fixture
|
||||
}
|
||||
|
||||
fn router(fixture: &AdminTestFixture) -> Router {
|
||||
account_groups::router::<AdminTestState>().with_state(fixture.state())
|
||||
}
|
||||
|
||||
async fn request(
|
||||
router: Router,
|
||||
method: Method,
|
||||
uri: &str,
|
||||
body: Option<Value>,
|
||||
authenticated: bool,
|
||||
) -> axum::response::Response {
|
||||
let mut builder = Request::builder()
|
||||
.method(method)
|
||||
.uri(uri)
|
||||
.header("x-request-id", "req_account_groups");
|
||||
if authenticated {
|
||||
builder = builder.header(header::COOKIE, SESSION_COOKIE);
|
||||
}
|
||||
let body = match body {
|
||||
Some(value) => {
|
||||
builder = builder.header(header::CONTENT_TYPE, "application/json");
|
||||
Body::from(serde_json::to_vec(&value).expect("serialize request body"))
|
||||
}
|
||||
None => Body::empty(),
|
||||
};
|
||||
router
|
||||
.oneshot(builder.body(body).expect("account group request"))
|
||||
.await
|
||||
.expect("account group response")
|
||||
}
|
||||
|
||||
async fn response_json(response: axum::response::Response) -> Value {
|
||||
let body = to_bytes(response.into_body(), 64 * 1024)
|
||||
.await
|
||||
.expect("account group response body");
|
||||
serde_json::from_slice(&body).expect("account group response JSON")
|
||||
}
|
||||
@@ -70,6 +70,135 @@ mod query {
|
||||
}
|
||||
}
|
||||
|
||||
mod batch_update {
|
||||
use gateway_api::admin::accounts::{BatchUpdateAccountsRequest, UpdateAccountRequest};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn batch_update_should_accept_complete_atomic_payload() {
|
||||
let request: BatchUpdateAccountsRequest = serde_json::from_value(json!({
|
||||
"accountIds": ["acct_openai", "acct_xai"],
|
||||
"enabled": false,
|
||||
"concurrencyLimit": null,
|
||||
"weight": 1,
|
||||
"groupIds": ["grp_00000000000000000000000000000001"]
|
||||
}))
|
||||
.expect("deserialize batch update");
|
||||
|
||||
request.validate().expect("validate batch update");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_update_should_reject_duplicate_or_invalid_account_ids() {
|
||||
for account_ids in [
|
||||
json!(["acct_same", "acct_same"]),
|
||||
json!(["not-an-account"]),
|
||||
json!([]),
|
||||
] {
|
||||
let request: BatchUpdateAccountsRequest = serde_json::from_value(json!({
|
||||
"accountIds": account_ids,
|
||||
"enabled": true,
|
||||
"concurrencyLimit": 8,
|
||||
"weight": 100,
|
||||
"groupIds": []
|
||||
}))
|
||||
.expect("deserialize invalid batch update");
|
||||
assert_eq!(
|
||||
request.validate().expect_err("reject account IDs").field(),
|
||||
"accountIds"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_update_should_require_enabled_and_reject_unknown_fields() {
|
||||
assert!(
|
||||
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
|
||||
"accountIds": ["acct_test"],
|
||||
"concurrencyLimit": null,
|
||||
"weight": 1,
|
||||
"groupIds": []
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
|
||||
"accountIds": ["acct_test"],
|
||||
"enabled": true,
|
||||
"concurrencyLimit": null,
|
||||
"weight": 1,
|
||||
"groupIds": [],
|
||||
"legacy": true
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_update_should_reject_invalid_scheduling_bounds_and_missing_fields() {
|
||||
for (concurrency_limit, weight, field) in [
|
||||
(json!(0), json!(1), "concurrencyLimit"),
|
||||
(json!(4294967296_u64), json!(1), "concurrencyLimit"),
|
||||
(json!(null), json!(0), "weight"),
|
||||
(json!(null), json!(101), "weight"),
|
||||
] {
|
||||
let request: BatchUpdateAccountsRequest = serde_json::from_value(json!({
|
||||
"accountIds": ["acct_test"],
|
||||
"enabled": true,
|
||||
"concurrencyLimit": concurrency_limit,
|
||||
"weight": weight,
|
||||
"groupIds": []
|
||||
}))
|
||||
.expect("deserialize invalid scheduling");
|
||||
assert_eq!(
|
||||
request.validate().expect_err("reject scheduling").field(),
|
||||
field
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
|
||||
"accountIds": ["acct_test"],
|
||||
"enabled": true,
|
||||
"weight": 1,
|
||||
"groupIds": []
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
serde_json::from_value::<BatchUpdateAccountsRequest>(json!({
|
||||
"accountIds": ["acct_test"],
|
||||
"enabled": true,
|
||||
"concurrencyLimit": null,
|
||||
"groupIds": []
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_update_should_require_the_same_complete_scheduling_contract() {
|
||||
let request: UpdateAccountRequest = serde_json::from_value(json!({
|
||||
"accountId": "acct_test",
|
||||
"enabled": true,
|
||||
"concurrencyLimit": 4294967295_u64,
|
||||
"weight": 100,
|
||||
"groupIds": []
|
||||
}))
|
||||
.expect("deserialize single update");
|
||||
request.validate().expect("validate single update");
|
||||
|
||||
assert!(
|
||||
serde_json::from_value::<UpdateAccountRequest>(json!({
|
||||
"accountId": "acct_test",
|
||||
"enabled": true,
|
||||
"weight": 1,
|
||||
"groupIds": []
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
mod response {
|
||||
use gateway_api::admin::accounts::AccountUsageView;
|
||||
use gateway_api::admin::presenter::format_decimal_currency;
|
||||
@@ -125,6 +254,7 @@ mod response {
|
||||
}
|
||||
|
||||
mod actions {
|
||||
use base64::Engine as _;
|
||||
use gateway_admin::model::{
|
||||
Revision,
|
||||
accounts::AccountConnectionTestEvent as DomainConnectionTestEvent,
|
||||
@@ -136,7 +266,8 @@ mod actions {
|
||||
AccountActionRequest, AccountConnectionTestEvent, AccountDeletionData,
|
||||
AccountDeletionRequest, AccountExportData, AccountExportQuery, AccountIdQuery,
|
||||
AccountImportData, AccountImportRequest, AccountMutationData, AccountRefreshRequest,
|
||||
AccountTestQuery, RotateAccountRequest, StartAccountAuthorizationRequest,
|
||||
AccountTestQuery, CompleteAccountAuthorizationRequest, RotateAccountRequest,
|
||||
StartAccountAuthorizationRequest,
|
||||
};
|
||||
use gateway_core::engine::credential::ProviderAccountId;
|
||||
use serde_json::json;
|
||||
@@ -280,6 +411,35 @@ mod actions {
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
serde_json::from_value::<AccountImportRequest>(json!({
|
||||
"provider": "openai",
|
||||
"data": {},
|
||||
"groupIds": []
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oauth_complete_should_not_accept_group_assignment() {
|
||||
let flow_id = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([7_u8; 32]);
|
||||
let request: CompleteAccountAuthorizationRequest = serde_json::from_value(json!({
|
||||
"provider": "openai",
|
||||
"flowId": flow_id,
|
||||
"callbackUrl": "http://localhost/callback?code=ok"
|
||||
}))
|
||||
.expect("decode OAuth completion");
|
||||
assert!(request.validate().is_ok());
|
||||
assert!(
|
||||
serde_json::from_value::<CompleteAccountAuthorizationRequest>(json!({
|
||||
"provider": "xai",
|
||||
"flowId": "flow_test",
|
||||
"callbackUrl": "http://localhost/callback?code=ok",
|
||||
"groupIds": []
|
||||
}))
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -316,7 +476,7 @@ mod actions {
|
||||
|
||||
#[test]
|
||||
fn account_import_response_should_emit_account_ids() {
|
||||
let response = AccountImportData::from(CredentialImportResult {
|
||||
let response = AccountImportData::from_result(CredentialImportResult {
|
||||
config_revision: Revision::new(8).expect("revision"),
|
||||
credential_ids: vec![ProviderAccountId::new("acct_imported").expect("account ID")],
|
||||
});
|
||||
|
||||
@@ -105,7 +105,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
|
||||
let valid = serde_json::from_value::<CreateClientKeyRequest>(json!({
|
||||
"name": "terminal key",
|
||||
"label": "production",
|
||||
"providerKind": "openai",
|
||||
"groupIds": [],
|
||||
"maxConcurrency": 2,
|
||||
"requestsPerMinute": 60
|
||||
}))
|
||||
@@ -118,7 +118,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
|
||||
(
|
||||
json!({
|
||||
"name": " ",
|
||||
"providerKind": "openai",
|
||||
"groupIds": [],
|
||||
"maxConcurrency": 0,
|
||||
"requestsPerMinute": 0
|
||||
}),
|
||||
@@ -127,7 +127,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
|
||||
(
|
||||
json!({
|
||||
"name": "key",
|
||||
"providerKind": "openai",
|
||||
"groupIds": [],
|
||||
"maxConcurrency": u64::MAX,
|
||||
"requestsPerMinute": 0
|
||||
}),
|
||||
@@ -150,7 +150,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
|
||||
"id": "key_1",
|
||||
"expectedConfigRevision": 1,
|
||||
"name": "key",
|
||||
"providerKind": "openai",
|
||||
"groupIds": [],
|
||||
"maxConcurrency": 0,
|
||||
"requestsPerMinute": 0,
|
||||
"tokensPerMinute": 0
|
||||
@@ -161,7 +161,7 @@ fn client_key_mutations_should_validate_text_limits_and_unknown_fields() {
|
||||
serde_json::from_value::<CreateClientKeyRequest>(json!({
|
||||
"expectedConfigRevision": 7,
|
||||
"name": "terminal key",
|
||||
"providerKind": "openai",
|
||||
"groupIds": [],
|
||||
"maxConcurrency": 2,
|
||||
"requestsPerMinute": 60
|
||||
}))
|
||||
@@ -210,7 +210,10 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
|
||||
id: gateway_core::policy::ClientApiKeyId::new("key_visible").expect("Client Key ID"),
|
||||
name: "visible".to_owned(),
|
||||
label: None,
|
||||
provider_kind: gateway_core::routing::ProviderKind::new("openai").expect("Provider kind"),
|
||||
groups: Vec::new(),
|
||||
provider_kinds: vec![
|
||||
gateway_core::routing::ProviderKind::new("openai").expect("Provider kind"),
|
||||
],
|
||||
prefix: "sk_visible12".to_owned(),
|
||||
enabled: true,
|
||||
limits: gateway_core::policy::RateLimits {
|
||||
@@ -226,7 +229,11 @@ fn client_key_responses_should_keep_shape_and_redact_creation_debug() {
|
||||
assert!(list.get("configRevision").is_none());
|
||||
assert_eq!(list["total"], 1);
|
||||
assert_eq!(list["items"][0]["id"], "key_visible");
|
||||
assert_eq!(list["items"][0]["providerKind"], "openai");
|
||||
assert_eq!(
|
||||
list["items"][0]["providerKinds"],
|
||||
serde_json::json!(["openai"])
|
||||
);
|
||||
assert_eq!(list["items"][0]["routingScope"], "all");
|
||||
assert_eq!(list["items"][0]["maxConcurrency"], 2);
|
||||
assert_eq!(list["items"][0]["requestsPerMinute"], 60);
|
||||
assert!(list["items"][0].get("tokensPerMinute").is_none());
|
||||
|
||||
@@ -13,9 +13,16 @@ use gateway_admin::{
|
||||
AdminConfig, AdminServices, InitialAdminPassword,
|
||||
model::{
|
||||
MutationContext, Revision,
|
||||
account_groups::{
|
||||
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
|
||||
AccountGroupListQuery, AccountGroupMember, AccountGroupMembers, AccountGroupMutation,
|
||||
AccountGroupPage, AccountGroupRecord, AccountGroupUsage, DeleteAccountGroup,
|
||||
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
},
|
||||
accounts::{
|
||||
AccountListQuery, AccountPage, AccountPageItem, AccountUsage, AccountUsageWindowQuery,
|
||||
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
|
||||
AccountListQuery, AccountPage, AccountPageItem, AccountUpdateResult, AccountUsage,
|
||||
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
|
||||
},
|
||||
auth::{AdminAuditEvent, AdminSession},
|
||||
client_keys::{
|
||||
@@ -23,9 +30,9 @@ use gateway_admin::{
|
||||
NewClientKey, SetClientKeyEnabled, UpdateClientKey,
|
||||
},
|
||||
observability::{
|
||||
DashboardObservation, DiagnosticDimension, DiagnosticObservation, OpsError,
|
||||
OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail, UsageFilter,
|
||||
UsageOverview, UsagePage, UsageQuery, UsageRecord,
|
||||
DashboardObservation, DecimalAmount, DiagnosticDimension, DiagnosticObservation,
|
||||
OpsError, OpsErrorPage, OpsErrorQuery, RequestMetricPoint, TimeRange, UsageDetail,
|
||||
UsageFilter, UsageOverview, UsagePage, UsageQuery, UsageRecord,
|
||||
},
|
||||
provider_credentials::{
|
||||
AuthorizationCommit, AuthorizationStarted, CompleteAuthorization, CredentialDetails,
|
||||
@@ -45,8 +52,8 @@ use gateway_admin::{
|
||||
ports::{
|
||||
provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind},
|
||||
store::{
|
||||
AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts, AdminStoreResult,
|
||||
AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
|
||||
AccountGroupStore, AccountStore, AdminStoreError, AdminStoreErrorKind, AdminStorePorts,
|
||||
AdminStoreResult, AuthStore, ClientKeyStore, ObservabilityStore, SettingsStore,
|
||||
},
|
||||
system::{
|
||||
SystemOperationError, SystemOperationErrorKind, SystemOperations,
|
||||
@@ -66,6 +73,7 @@ use gateway_core::{
|
||||
},
|
||||
};
|
||||
|
||||
mod account_groups;
|
||||
mod accounts;
|
||||
mod auth;
|
||||
mod client_keys;
|
||||
@@ -96,6 +104,7 @@ impl AdminTestFixture {
|
||||
let auth = Arc::new(MemoryAuthStore::new(api_key.clone()));
|
||||
let settings = Arc::new(MemorySettingsStore::new(api_key));
|
||||
let client_keys = Arc::new(MemoryClientKeyStore);
|
||||
let account_groups = Arc::new(MemoryAccountGroupStore::new());
|
||||
let usage_records = Arc::new(Mutex::new(Vec::new()));
|
||||
let usage_detail = Arc::new(Mutex::new(None));
|
||||
let diagnostics = Arc::new(Mutex::new(Vec::new()));
|
||||
@@ -112,6 +121,7 @@ impl AdminTestFixture {
|
||||
});
|
||||
let stores = AdminStorePorts::new(
|
||||
unused.clone(),
|
||||
account_groups.clone(),
|
||||
auth.clone(),
|
||||
client_keys.clone(),
|
||||
unused,
|
||||
@@ -345,6 +355,258 @@ impl SettingsStore for MemorySettingsStore {
|
||||
|
||||
pub(super) struct MemoryClientKeyStore;
|
||||
|
||||
pub(super) const PRIMARY_GROUP_ID: &str = "grp_11111111111111111111111111111111";
|
||||
pub(super) const SECONDARY_GROUP_ID: &str = "grp_22222222222222222222222222222222";
|
||||
|
||||
struct MemoryAccountGroupState {
|
||||
revision: Revision,
|
||||
groups: BTreeMap<gateway_core::routing::AccountGroupId, AccountGroupRecord>,
|
||||
members: BTreeMap<gateway_core::routing::AccountGroupId, Vec<AccountGroupMember>>,
|
||||
}
|
||||
|
||||
pub(super) struct MemoryAccountGroupStore {
|
||||
state: Mutex<MemoryAccountGroupState>,
|
||||
}
|
||||
|
||||
impl MemoryAccountGroupStore {
|
||||
fn new() -> Self {
|
||||
let now = Utc::now();
|
||||
let primary_id = group_id(PRIMARY_GROUP_ID);
|
||||
let secondary_id = group_id(SECONDARY_GROUP_ID);
|
||||
let openai = AccountGroupMember {
|
||||
id: "acct_openai".to_owned(),
|
||||
name: "OpenAI primary".to_owned(),
|
||||
provider_kind: ProviderKind::new("openai").expect("provider kind"),
|
||||
email: Some("openai@example.invalid".to_owned()),
|
||||
enabled: true,
|
||||
};
|
||||
let xai = AccountGroupMember {
|
||||
id: "acct_xai".to_owned(),
|
||||
name: "xAI primary".to_owned(),
|
||||
provider_kind: ProviderKind::new("xai").expect("provider kind"),
|
||||
email: None,
|
||||
enabled: false,
|
||||
};
|
||||
let groups = BTreeMap::from([
|
||||
(
|
||||
primary_id.clone(),
|
||||
AccountGroupRecord {
|
||||
id: primary_id.clone(),
|
||||
name: "Alpha routing".to_owned(),
|
||||
description: Some("Primary traffic".to_owned()),
|
||||
color: group_color("#2563ebff"),
|
||||
enabled: true,
|
||||
member_count: 2,
|
||||
provider_counts: BTreeMap::from([
|
||||
("openai".to_owned(), 1),
|
||||
("xai".to_owned(), 1),
|
||||
]),
|
||||
client_key_count: 2,
|
||||
account_summary: account_summary(1, 1, 2),
|
||||
capacity: capacity(Some(0), 1),
|
||||
usage: usage("1.25", "5.5"),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
},
|
||||
),
|
||||
(
|
||||
secondary_id.clone(),
|
||||
AccountGroupRecord {
|
||||
id: secondary_id,
|
||||
name: "Beta routing".to_owned(),
|
||||
description: None,
|
||||
color: group_color("#64748B80"),
|
||||
enabled: false,
|
||||
member_count: 0,
|
||||
provider_counts: BTreeMap::new(),
|
||||
client_key_count: 0,
|
||||
account_summary: account_summary(0, 0, 0),
|
||||
capacity: capacity(Some(0), 0),
|
||||
usage: usage("0", "0"),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
},
|
||||
),
|
||||
]);
|
||||
Self {
|
||||
state: Mutex::new(MemoryAccountGroupState {
|
||||
revision: Revision::new(7).expect("revision"),
|
||||
groups,
|
||||
members: BTreeMap::from([(primary_id, vec![openai, xai])]),
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AccountGroupStore for MemoryAccountGroupStore {
|
||||
async fn list_account_groups(
|
||||
&self,
|
||||
query: AccountGroupListQuery,
|
||||
) -> AdminStoreResult<AccountGroupPage> {
|
||||
let state = self.state.lock().expect("account groups");
|
||||
let search = query.search.as_ref().map(|value| value.to_lowercase());
|
||||
let mut matching = state
|
||||
.groups
|
||||
.values()
|
||||
.filter(|record| {
|
||||
query
|
||||
.enabled
|
||||
.is_none_or(|enabled| record.enabled == enabled)
|
||||
})
|
||||
.filter(|record| {
|
||||
search.as_ref().is_none_or(|search| {
|
||||
record.name.to_lowercase().contains(search)
|
||||
|| record
|
||||
.description
|
||||
.as_ref()
|
||||
.is_some_and(|description| description.to_lowercase().contains(search))
|
||||
})
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
matching.sort_by(|left, right| left.name.cmp(&right.name));
|
||||
let total = matching.len() as u64;
|
||||
let page_size = usize::from(query.page_size.get());
|
||||
let offset = usize::try_from(query.page.saturating_sub(1))
|
||||
.unwrap_or(usize::MAX)
|
||||
.saturating_mul(page_size);
|
||||
let items = matching.into_iter().skip(offset).take(page_size).collect();
|
||||
Ok(AccountGroupPage {
|
||||
config_revision: state.revision,
|
||||
items,
|
||||
total,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn account_group_members(
|
||||
&self,
|
||||
id: &gateway_core::routing::AccountGroupId,
|
||||
) -> AdminStoreResult<AccountGroupMembers> {
|
||||
let state = self.state.lock().expect("account groups");
|
||||
if !state.groups.contains_key(id) {
|
||||
return Err(not_found("account group"));
|
||||
}
|
||||
let items = state.members.get(id).cloned().unwrap_or_default();
|
||||
Ok(AccountGroupMembers {
|
||||
config_revision: state.revision,
|
||||
id: id.clone(),
|
||||
total: items.len() as u64,
|
||||
items,
|
||||
})
|
||||
}
|
||||
|
||||
async fn create_account_group(
|
||||
&self,
|
||||
command: NewAccountGroup,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
let mut state = self.state.lock().expect("account groups");
|
||||
let now = Utc::now();
|
||||
let record = AccountGroupRecord {
|
||||
id: command.id.clone(),
|
||||
name: command.name,
|
||||
description: command.description,
|
||||
color: command.color,
|
||||
enabled: true,
|
||||
member_count: 0,
|
||||
provider_counts: BTreeMap::new(),
|
||||
client_key_count: 0,
|
||||
account_summary: account_summary(0, 0, 0),
|
||||
capacity: capacity(Some(0), 0),
|
||||
usage: usage("0", "0"),
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
};
|
||||
state.groups.insert(command.id.clone(), record);
|
||||
state.members.insert(command.id.clone(), Vec::new());
|
||||
refresh_group_counts(&mut state, &command.id);
|
||||
mutation(&mut state, command.id, true)
|
||||
}
|
||||
|
||||
async fn update_account_group(
|
||||
&self,
|
||||
command: UpdateAccountGroup,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
let mut state = self.state.lock().expect("account groups");
|
||||
let record = state
|
||||
.groups
|
||||
.get_mut(&command.id)
|
||||
.ok_or_else(|| not_found("account group"))?;
|
||||
record.name = command.name;
|
||||
record.description = command.description;
|
||||
record.color = command.color;
|
||||
record.updated_at = Utc::now();
|
||||
mutation(&mut state, command.id, true)
|
||||
}
|
||||
|
||||
async fn set_account_group_enabled(
|
||||
&self,
|
||||
command: SetAccountGroupEnabled,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
let mut state = self.state.lock().expect("account groups");
|
||||
let record = state
|
||||
.groups
|
||||
.get_mut(&command.id)
|
||||
.ok_or_else(|| not_found("account group"))?;
|
||||
record.enabled = command.enabled;
|
||||
record.updated_at = Utc::now();
|
||||
mutation(&mut state, command.id, true)
|
||||
}
|
||||
|
||||
async fn delete_account_group(
|
||||
&self,
|
||||
command: DeleteAccountGroup,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
let mut state = self.state.lock().expect("account groups");
|
||||
if state.groups.remove(&command.id).is_none() {
|
||||
return Err(not_found("account group"));
|
||||
}
|
||||
state.members.remove(&command.id);
|
||||
mutation(&mut state, command.id, false)
|
||||
}
|
||||
}
|
||||
|
||||
fn group_id(value: &str) -> gateway_core::routing::AccountGroupId {
|
||||
gateway_core::routing::AccountGroupId::new(value).expect("account group ID")
|
||||
}
|
||||
|
||||
fn refresh_group_counts(
|
||||
state: &mut MemoryAccountGroupState,
|
||||
id: &gateway_core::routing::AccountGroupId,
|
||||
) {
|
||||
let members = state.members.get(id).map(Vec::as_slice).unwrap_or_default();
|
||||
let provider_counts = members.iter().fold(BTreeMap::new(), |mut counts, member| {
|
||||
*counts
|
||||
.entry(member.provider_kind.as_str().to_owned())
|
||||
.or_insert(0) += 1;
|
||||
counts
|
||||
});
|
||||
if let Some(record) = state.groups.get_mut(id) {
|
||||
record.member_count = members.len() as u64;
|
||||
record.provider_counts = provider_counts;
|
||||
record.updated_at = Utc::now();
|
||||
}
|
||||
}
|
||||
|
||||
fn mutation(
|
||||
state: &mut MemoryAccountGroupState,
|
||||
id: gateway_core::routing::AccountGroupId,
|
||||
include_record: bool,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
state.revision = next_revision(state.revision);
|
||||
Ok(AccountGroupMutation {
|
||||
config_revision: state.revision,
|
||||
record: include_record.then(|| state.groups.get(&id).expect("group record").clone()),
|
||||
id,
|
||||
})
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ClientKeyStore for MemoryClientKeyStore {
|
||||
async fn list_client_keys(&self, _: ClientKeyListQuery) -> AdminStoreResult<ClientKeyPage> {
|
||||
@@ -366,7 +628,8 @@ impl ClientKeyStore for MemoryClientKeyStore {
|
||||
id: id.clone(),
|
||||
name: "revealed".to_owned(),
|
||||
label: None,
|
||||
provider_kind: ProviderKind::new("openai").expect("provider kind"),
|
||||
groups: Vec::new(),
|
||||
provider_kinds: vec![ProviderKind::new("openai").expect("provider kind")],
|
||||
prefix: "sk_aaaaaaaaa".to_owned(),
|
||||
enabled: true,
|
||||
limits: RateLimits::unlimited(),
|
||||
@@ -501,14 +764,22 @@ impl AccountStore for UnusedStore {
|
||||
Err(unavailable("credential refresh"))
|
||||
}
|
||||
|
||||
async fn set_account_enabled(
|
||||
async fn update_account(
|
||||
&self,
|
||||
_: SetAccountEnabled,
|
||||
_: UpdateAccount,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<Revision> {
|
||||
) -> AdminStoreResult<AccountUpdateResult> {
|
||||
Err(unavailable("account enabled"))
|
||||
}
|
||||
|
||||
async fn batch_update_accounts(
|
||||
&self,
|
||||
_: BatchUpdateAccounts,
|
||||
_: &MutationContext,
|
||||
) -> AdminStoreResult<AccountsUpdateResult> {
|
||||
Err(unavailable("account batch update"))
|
||||
}
|
||||
|
||||
async fn delete_accounts(
|
||||
&self,
|
||||
_: DeleteAccounts,
|
||||
@@ -792,6 +1063,32 @@ fn next_revision(revision: Revision) -> Revision {
|
||||
Revision::new(revision.get().saturating_add(1)).expect("next revision")
|
||||
}
|
||||
|
||||
fn account_summary(available: u64, limited: u64, total: u64) -> AccountGroupAccountSummary {
|
||||
AccountGroupAccountSummary {
|
||||
available,
|
||||
limited,
|
||||
total,
|
||||
}
|
||||
}
|
||||
|
||||
fn group_color(value: &str) -> AccountGroupColor {
|
||||
AccountGroupColor::parse(value).expect("group color")
|
||||
}
|
||||
|
||||
fn capacity(used_slots: Option<u64>, total_slots: u64) -> AccountGroupCapacity {
|
||||
AccountGroupCapacity {
|
||||
used_slots,
|
||||
total_slots,
|
||||
}
|
||||
}
|
||||
|
||||
fn usage(today_usd: &str, total_usd: &str) -> AccountGroupUsage {
|
||||
AccountGroupUsage {
|
||||
today_usd: today_usd.parse::<DecimalAmount>().expect("today cost"),
|
||||
total_usd: total_usd.parse::<DecimalAmount>().expect("total cost"),
|
||||
}
|
||||
}
|
||||
|
||||
fn unavailable(resource: &'static str) -> AdminStoreError {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Unavailable,
|
||||
@@ -800,6 +1097,10 @@ fn unavailable(resource: &'static str) -> AdminStoreError {
|
||||
)
|
||||
}
|
||||
|
||||
fn not_found(resource: &'static str) -> AdminStoreError {
|
||||
AdminStoreError::new(AdminStoreErrorKind::NotFound, resource, "not found")
|
||||
}
|
||||
|
||||
fn unsupported_provider() -> ProviderAdminError {
|
||||
ProviderAdminError::new(ProviderAdminErrorKind::Unsupported)
|
||||
}
|
||||
|
||||
@@ -674,6 +674,9 @@ fn usage_record_with_account(
|
||||
id: id.to_owned(),
|
||||
client_api_key_ref: "key_detail".to_owned(),
|
||||
config_revision: 1,
|
||||
routing_scope: "all".to_owned(),
|
||||
routing_group_refs: Vec::new(),
|
||||
routing_group_names_snapshot: Vec::new(),
|
||||
protocol: "openai".to_owned(),
|
||||
operation: "responses".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
@@ -800,6 +803,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
id: "request_endpoint".to_owned(),
|
||||
client_api_key_ref: "key_endpoint".to_owned(),
|
||||
config_revision: 1,
|
||||
routing_scope: "groups".to_owned(),
|
||||
routing_group_refs: vec!["grp_usage".to_owned()],
|
||||
routing_group_names_snapshot: vec!["Usage group".to_owned()],
|
||||
protocol: "openai".to_owned(),
|
||||
operation: "generate".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
@@ -919,6 +925,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
assert_eq!(
|
||||
serde_json::json!({
|
||||
"route": value["data"]["items"][0]["route"],
|
||||
"routingScope": value["data"]["items"][0]["routingScope"],
|
||||
"routingGroupRefs": value["data"]["items"][0]["routingGroupRefs"],
|
||||
"routingGroupNamesSnapshot": value["data"]["items"][0]["routingGroupNamesSnapshot"],
|
||||
"serviceTier": value["data"]["items"][0]["serviceTier"],
|
||||
"accountId": value["data"]["items"][0]["accountId"],
|
||||
"accountName": value["data"]["items"][0]["accountName"],
|
||||
@@ -949,6 +958,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
|
||||
}),
|
||||
serde_json::json!({
|
||||
"route": "/v1/responses",
|
||||
"routingScope": "groups",
|
||||
"routingGroupRefs": ["grp_usage"],
|
||||
"routingGroupNamesSnapshot": ["Usage group"],
|
||||
"serviceTier": "priority",
|
||||
"accountId": "acct_snapshot",
|
||||
"accountName": "Snapshot Alpha",
|
||||
|
||||
@@ -33,6 +33,7 @@ fn source_tree_should_match_frozen_machine_manifest() {
|
||||
let mut actual = rust_files(&root.join("src"));
|
||||
actual.sort();
|
||||
let mut expected = vec![
|
||||
"src/admin/account_groups.rs",
|
||||
"src/admin/accounts.rs",
|
||||
"src/admin/auth.rs",
|
||||
"src/admin/backups.rs",
|
||||
@@ -73,6 +74,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
|
||||
let mut actual = rust_files(&root.join("tests"));
|
||||
actual.sort();
|
||||
let mut expected = vec![
|
||||
"tests/admin/account_groups.rs",
|
||||
"tests/admin/accounts.rs",
|
||||
"tests/admin/auth.rs",
|
||||
"tests/admin/client_keys.rs",
|
||||
|
||||
@@ -4,7 +4,7 @@ mod models;
|
||||
mod responses;
|
||||
mod router;
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
@@ -18,7 +18,9 @@ use gateway_core::engine::admission::{
|
||||
use gateway_core::engine::continuation::{
|
||||
NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId,
|
||||
};
|
||||
use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
|
||||
};
|
||||
use gateway_core::engine::execution::{
|
||||
AuthenticatedClient, ClientApiKeyUsageSink, DefaultExecutionService, ExecutionService,
|
||||
ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
|
||||
@@ -34,8 +36,8 @@ use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLif
|
||||
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
|
||||
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
|
||||
use gateway_core::routing::{
|
||||
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, RuntimeSnapshot,
|
||||
UpstreamModelId,
|
||||
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
|
||||
ProviderModel, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
|
||||
};
|
||||
|
||||
pub(super) async fn api_router(execution: Arc<dyn ExecutionService>) -> axum::Router {
|
||||
@@ -94,6 +96,10 @@ impl ClientApiKeyUsageSink for IgnoredClientApiKeyUsage {
|
||||
|
||||
fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
|
||||
let provider = ProviderKind::new(provider_name).expect("provider");
|
||||
let account_directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
|
||||
ProviderAccountId::new("acct_api_test").expect("account ID"),
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
|
||||
)])));
|
||||
let capabilities = ModelCapabilities::new(
|
||||
BTreeSet::from([gateway_core::operation::OperationKind::Generate]),
|
||||
Some(16_000),
|
||||
@@ -119,12 +125,16 @@ fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
|
||||
vec![ClientPolicy::new(
|
||||
ClientApiKeyId::new("key_api_test").expect("key ID"),
|
||||
PlaintextClientApiKey::new(plaintext).expect("plaintext key"),
|
||||
provider,
|
||||
Arc::new(FrozenAccountScope::new(
|
||||
Arc::clone(&account_directory),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
)),
|
||||
true,
|
||||
RateLimits::unlimited(),
|
||||
)],
|
||||
)
|
||||
.expect("runtime snapshot")
|
||||
.with_account_directory(account_directory)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -260,6 +270,7 @@ struct UnusedContinuation;
|
||||
impl NativeContinuationPort for UnusedContinuation {
|
||||
fn resolve<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientApiKeyId,
|
||||
_: &'a PreviousResponseId,
|
||||
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
|
||||
Box::pin(async { unreachable!("authentication fixture does not execute") })
|
||||
|
||||
@@ -451,10 +451,12 @@ fn decoder_should_not_reject_large_bodies_using_catalog_context_limits() {
|
||||
.to_string();
|
||||
let decoded = decode_request(body.as_bytes()).expect("large request should decode");
|
||||
|
||||
let plan = super::snapshot("sk_context_test", "openai")
|
||||
let snapshot = super::snapshot("sk_context_test", "openai");
|
||||
let plan = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new(decoded.metadata().public_model()).expect("public model"),
|
||||
decoded.operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext::default(),
|
||||
)
|
||||
.expect("large body must not be locally context-gated");
|
||||
|
||||
@@ -9,6 +9,7 @@ use futures::future::BoxFuture;
|
||||
|
||||
use crate::engine::credential::ProviderAccountId;
|
||||
use crate::operation::ProviderSessionState;
|
||||
use crate::policy::ClientApiKeyId;
|
||||
use crate::routing::ProviderKind;
|
||||
|
||||
/// 客户端或 Provider 传递的 opaque response handle。
|
||||
@@ -48,14 +49,15 @@ pub enum NativeContinuationScope {
|
||||
|
||||
/// 从可丢失会话亲和存储恢复的 native previous-response pin。
|
||||
///
|
||||
/// 该值不依赖调用方 API Key;它只阻止 native handle 与不透明 Provider state
|
||||
/// 被发送到错误的 Provider 或账号。
|
||||
/// 该值同时绑定调用方 API Key、Provider 与账号,阻止 native handle 和不透明
|
||||
/// Provider state 被跨租户或跨上游目标复用。
|
||||
#[derive(Clone, PartialEq)]
|
||||
pub struct NativeContinuationPin {
|
||||
/// 客户端提交、用于查找可丢失会话亲和记录的 response ID。
|
||||
previous_response_id: PreviousResponseId,
|
||||
/// Provider 原生 response handle。
|
||||
upstream_response_id: PreviousResponseId,
|
||||
client_api_key_id: ClientApiKeyId,
|
||||
provider: ProviderKind,
|
||||
account: ProviderAccountId,
|
||||
scope: NativeContinuationScope,
|
||||
@@ -67,12 +69,14 @@ impl NativeContinuationPin {
|
||||
pub const fn new(
|
||||
previous_response_id: PreviousResponseId,
|
||||
upstream_response_id: PreviousResponseId,
|
||||
client_api_key_id: ClientApiKeyId,
|
||||
provider: ProviderKind,
|
||||
account: ProviderAccountId,
|
||||
) -> Self {
|
||||
Self {
|
||||
previous_response_id,
|
||||
upstream_response_id,
|
||||
client_api_key_id,
|
||||
provider,
|
||||
account,
|
||||
scope: NativeContinuationScope::ConnectionLocal,
|
||||
@@ -105,6 +109,16 @@ impl NativeContinuationPin {
|
||||
&self.upstream_response_id
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn client_api_key_id(&self) -> &ClientApiKeyId {
|
||||
&self.client_api_key_id
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn matches_client(&self, client_api_key_id: &ClientApiKeyId) -> bool {
|
||||
self.client_api_key_id == *client_api_key_id
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn provider(&self) -> &ProviderKind {
|
||||
&self.provider
|
||||
@@ -139,6 +153,7 @@ impl fmt::Debug for NativeContinuationPin {
|
||||
.debug_struct("NativeContinuationPin")
|
||||
.field("previous_response_id", &"<redacted>")
|
||||
.field("upstream_response_id", &"<redacted>")
|
||||
.field("client_api_key_id", &self.client_api_key_id)
|
||||
.field("provider", &self.provider)
|
||||
.field("account", &self.account)
|
||||
.field("scope", &self.scope)
|
||||
@@ -190,21 +205,57 @@ impl fmt::Debug for ContinuationBinding {
|
||||
}
|
||||
}
|
||||
|
||||
/// Native continuation lookup 的稳定失败分类。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum NativeContinuationStoreErrorKind {
|
||||
/// Redis 等可丢失协调存储暂不可用;调用方可以按外部 continuation fail-open。
|
||||
Unavailable,
|
||||
/// 存储中的记录无法还原为可信 pin;不得将原始 handle 继续透传上游。
|
||||
InvalidData,
|
||||
/// 已存在的记录属于另一个客户端 API Key;必须 fail closed。
|
||||
OwnershipMismatch,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
|
||||
#[error("native continuation store is unavailable: {detail}")]
|
||||
#[error("native continuation store failed: {kind:?}: {detail}")]
|
||||
pub struct NativeContinuationStoreError {
|
||||
kind: NativeContinuationStoreErrorKind,
|
||||
detail: String,
|
||||
}
|
||||
|
||||
impl NativeContinuationStoreError {
|
||||
/// 构造带可排查细节的亲和存储错误;`detail` 不应包含 response handle。
|
||||
/// 构造可 fail-open 的协调存储不可用错误;`detail` 不应包含 response handle。
|
||||
#[must_use]
|
||||
pub fn new(detail: impl Into<String>) -> Self {
|
||||
pub fn unavailable(detail: impl Into<String>) -> Self {
|
||||
Self {
|
||||
kind: NativeContinuationStoreErrorKind::Unavailable,
|
||||
detail: detail.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造必须 fail closed 的无效记录错误;`detail` 不应包含 response handle。
|
||||
#[must_use]
|
||||
pub fn invalid_data(detail: impl Into<String>) -> Self {
|
||||
Self {
|
||||
kind: NativeContinuationStoreErrorKind::InvalidData,
|
||||
detail: detail.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造不泄露记录归属细节的跨调用方错误。
|
||||
#[must_use]
|
||||
pub fn ownership_mismatch() -> Self {
|
||||
Self {
|
||||
kind: NativeContinuationStoreErrorKind::OwnershipMismatch,
|
||||
detail: "continuation belongs to a different client API key".to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn kind(&self) -> NativeContinuationStoreErrorKind {
|
||||
self.kind
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn detail(&self) -> &str {
|
||||
&self.detail
|
||||
@@ -213,11 +264,12 @@ impl NativeContinuationStoreError {
|
||||
|
||||
/// 可丢失的 previous-response 亲和存储端口。
|
||||
///
|
||||
/// 亲和记录不按调用方 API Key 隔离:与旧版一样,response ID 是会话级 opaque
|
||||
/// handle;Redis 不可用或超时必须由调用方 fail-open。
|
||||
/// 亲和记录按调用方 API Key 隔离;Redis 不可用或超时必须由调用方 fail-open,
|
||||
/// 但命中其他 Key 的记录必须 fail closed。
|
||||
pub trait NativeContinuationPort: Send + Sync {
|
||||
fn resolve<'a>(
|
||||
&'a self,
|
||||
client_api_key_id: &'a ClientApiKeyId,
|
||||
previous_response_id: &'a PreviousResponseId,
|
||||
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>>;
|
||||
|
||||
|
||||
@@ -46,6 +46,10 @@ impl AccountSelection {
|
||||
Self::Diagnostic(account) => Some(account),
|
||||
}
|
||||
}
|
||||
|
||||
const fn is_diagnostic(&self) -> bool {
|
||||
matches!(self, Self::Diagnostic(_))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: ?Sized> AttemptCoordinator<S>
|
||||
@@ -126,6 +130,15 @@ where
|
||||
.map(ProviderAccountStateOwner::from_continuation);
|
||||
let continuation_attempt =
|
||||
initial_continuation_attempt(&operation, &plan, continuation.as_ref());
|
||||
let candidate_index = continuation
|
||||
.as_ref()
|
||||
.and_then(ContinuationBinding::pinned)
|
||||
.map_or(Ok(0), |pin| {
|
||||
plan.candidates()
|
||||
.iter()
|
||||
.position(|candidate| candidate.provider() == pin.provider())
|
||||
.ok_or(EngineError::ContinuationPinMismatch)
|
||||
})?;
|
||||
let image_generation_requested = operation.image_generation_requested();
|
||||
let mut session = ResponseExecutionSession {
|
||||
engine: Arc::clone(&self.engine),
|
||||
@@ -149,6 +162,7 @@ where
|
||||
account_state_owner,
|
||||
cancellation,
|
||||
attempts: 0,
|
||||
candidate_index,
|
||||
excluded_accounts: BTreeSet::new(),
|
||||
credential_recovery_attempted_accounts: BTreeSet::new(),
|
||||
recovery_account: None,
|
||||
@@ -225,6 +239,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
|
||||
account_state_owner: Option<ProviderAccountStateOwner>,
|
||||
cancellation: CancellationToken,
|
||||
attempts: u32,
|
||||
candidate_index: usize,
|
||||
excluded_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
|
||||
credential_recovery_attempted_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
|
||||
/// 凭据恢复后的一次性同账号钉选;只约束紧随其后的 replay attempt,
|
||||
@@ -461,6 +476,7 @@ where
|
||||
NativeContinuationPin::new(
|
||||
previous_response_id,
|
||||
upstream_response_id,
|
||||
self.client_api_key_ref.clone(),
|
||||
provider,
|
||||
account,
|
||||
)
|
||||
@@ -565,7 +581,7 @@ where
|
||||
if self.attempts >= self.plan.max_attempts().get() {
|
||||
return Err(EngineError::EmptyRoutingPlan);
|
||||
}
|
||||
let Some(candidate) = self.plan.candidates().first().cloned() else {
|
||||
let Some(candidate) = self.plan.candidates().get(self.candidate_index).cloned() else {
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::NoAvailableProvider,
|
||||
"no upstream Provider is available",
|
||||
@@ -606,7 +622,8 @@ where
|
||||
self.excluded_accounts.clone(),
|
||||
pinned_account.clone(),
|
||||
self.account_state_owner.clone(),
|
||||
),
|
||||
)
|
||||
.with_account_scope(Arc::clone(self.plan.account_scope())),
|
||||
}
|
||||
.with_credential_recovery_attempted(pinned_account.as_ref().is_some_and(|account| {
|
||||
self.credential_recovery_attempted_accounts
|
||||
@@ -663,6 +680,20 @@ where
|
||||
if self.prepare_unavailable_native_continuation_replay(&error) {
|
||||
return Ok(Some(PullOutcome::AttemptDiscarded));
|
||||
}
|
||||
if matches!(
|
||||
error.kind(),
|
||||
ProviderErrorKind::AccountCapacityUnavailable
|
||||
| ProviderErrorKind::NoEligibleAccount
|
||||
| ProviderErrorKind::ProviderInfrastructureUnavailable
|
||||
) && error.send_state() == UpstreamSendState::NotSent
|
||||
&& matches!(
|
||||
self.continuation_attempt,
|
||||
ContinuationAttempt::None | ContinuationAttempt::ReplayAny
|
||||
)
|
||||
&& self.advance_provider_candidate()
|
||||
{
|
||||
return Ok(Some(PullOutcome::AttemptDiscarded));
|
||||
}
|
||||
if matches!(
|
||||
error.kind(),
|
||||
ProviderErrorKind::AccountCapacityUnavailable
|
||||
@@ -716,6 +747,27 @@ where
|
||||
}
|
||||
|
||||
let metadata = stream.metadata().clone();
|
||||
if !self.account_selection.is_diagnostic()
|
||||
&& metadata
|
||||
.provider_account_id()
|
||||
.is_none_or(|account| !candidate.account_scope().allows(account))
|
||||
{
|
||||
let error = GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"provider selected an account outside the frozen client scope",
|
||||
);
|
||||
self.finish_failure(FailureFinalization {
|
||||
outcome: ExecutionOutcome::Failed,
|
||||
send_state: self.current_send_state(),
|
||||
error,
|
||||
upstream_status_code: None,
|
||||
provider_error_code: None,
|
||||
retry_after_ms: None,
|
||||
provider_response_id: None,
|
||||
})
|
||||
.await?;
|
||||
return Err(EngineError::AccountOutsideClientScope);
|
||||
}
|
||||
if pinned_account
|
||||
.as_ref()
|
||||
.is_some_and(|required| metadata.provider_account_id() != Some(required))
|
||||
@@ -819,6 +871,17 @@ where
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn advance_provider_candidate(&mut self) -> bool {
|
||||
let Some(next) = self.candidate_index.checked_add(1) else {
|
||||
return false;
|
||||
};
|
||||
if next >= self.plan.candidates().len() {
|
||||
return false;
|
||||
}
|
||||
self.candidate_index = next;
|
||||
true
|
||||
}
|
||||
|
||||
async fn observe_event(&mut self, event: &GatewayEvent) {
|
||||
self.observe_event_facts(event);
|
||||
self.mark_send_observed().await;
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
use std::fmt;
|
||||
use std::num::{NonZeroU32, NonZeroU64};
|
||||
use std::num::{NonZeroU16, NonZeroU32, NonZeroU64};
|
||||
use std::sync::RwLock;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{Duration, SystemTime};
|
||||
@@ -38,6 +38,58 @@ impl ProviderAccountId {
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-account scheduling concurrency override.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct AccountConcurrencyLimit(NonZeroU32);
|
||||
|
||||
impl AccountConcurrencyLimit {
|
||||
#[must_use]
|
||||
pub const fn new(value: u32) -> Option<Self> {
|
||||
match NonZeroU32::new(value) {
|
||||
Some(value) => Some(Self(value)),
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn get(self) -> u32 {
|
||||
self.0.get()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn into_non_zero(self) -> NonZeroU32 {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
/// Relative scheduling priority for an account.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct AccountWeight(NonZeroU16);
|
||||
|
||||
impl AccountWeight {
|
||||
pub const DEFAULT: Self = Self(NonZeroU16::MIN);
|
||||
pub const MAX: u16 = 100;
|
||||
|
||||
#[must_use]
|
||||
pub const fn new(value: u16) -> Option<Self> {
|
||||
match NonZeroU16::new(value) {
|
||||
Some(value) if value.get() <= Self::MAX => Some(Self(value)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn get(self) -> u16 {
|
||||
self.0.get()
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for AccountWeight {
|
||||
fn default() -> Self {
|
||||
Self::DEFAULT
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for ProviderAccountId {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str(&self.0)
|
||||
@@ -619,6 +671,8 @@ pub struct ProviderAccount {
|
||||
authentication_kind: String,
|
||||
revision: CredentialRevision,
|
||||
enabled: bool,
|
||||
concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
weight: AccountWeight,
|
||||
credential_state: CredentialState,
|
||||
quota: QuotaState,
|
||||
last_error_reason: Option<AccountErrorReason>,
|
||||
@@ -651,6 +705,8 @@ impl ProviderAccount {
|
||||
authentication_kind,
|
||||
revision,
|
||||
enabled: true,
|
||||
concurrency_limit: None,
|
||||
weight: AccountWeight::DEFAULT,
|
||||
credential_state: CredentialState::Unknown,
|
||||
quota: QuotaState::unknown(),
|
||||
last_error_reason: None,
|
||||
@@ -695,6 +751,17 @@ impl ProviderAccount {
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn with_scheduling(
|
||||
mut self,
|
||||
concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
weight: AccountWeight,
|
||||
) -> Self {
|
||||
self.concurrency_limit = concurrency_limit;
|
||||
self.weight = weight;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置 RT 存在性与失败后的最早重试时刻。
|
||||
///
|
||||
/// 正常 OAuth 预刷新由 worker 使用 AT 原始过期时间与当前运行时策略动态判断,
|
||||
@@ -780,6 +847,24 @@ impl ProviderAccount {
|
||||
self.enabled
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn concurrency_limit(&self) -> Option<AccountConcurrencyLimit> {
|
||||
self.concurrency_limit
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn weight(&self) -> AccountWeight {
|
||||
self.weight
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn effective_concurrency(&self, default: NonZeroU32) -> NonZeroU32 {
|
||||
match self.concurrency_limit {
|
||||
Some(limit) => limit.into_non_zero(),
|
||||
None => default,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn access_token_expires_at(&self) -> Option<SystemTime> {
|
||||
self.access_token_expires_at
|
||||
@@ -1393,6 +1478,7 @@ pub struct AccountSelectionContext {
|
||||
pub preferred_account: Option<ProviderAccountId>,
|
||||
pub round_robin_cursor: u64,
|
||||
pub eligibility: AccountEligibilityPolicy,
|
||||
pub account_scope: Option<std::sync::Arc<crate::routing::FrozenAccountScope>>,
|
||||
}
|
||||
|
||||
/// 选择账号时是否执行本地调度资格投影。
|
||||
@@ -1417,10 +1503,12 @@ impl AccountEligibilityPolicy {
|
||||
/// 候选账号未进入调度池的约束。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AccountSchedulingBlocker {
|
||||
OutsideClientScope,
|
||||
LocalAvailability,
|
||||
Excluded,
|
||||
ConcurrencyLimit,
|
||||
RequestInterval,
|
||||
LowerWeight,
|
||||
}
|
||||
|
||||
/// 优先账号在本次选择中的处理结果。
|
||||
@@ -1463,26 +1551,6 @@ impl AccountSelector {
|
||||
candidates: &'a [AccountCandidate],
|
||||
context: &AccountSelectionContext,
|
||||
) -> Option<AccountSelection<'a>> {
|
||||
let preferred = if let Some(preferred) = context.preferred_account.as_ref() {
|
||||
match candidates
|
||||
.iter()
|
||||
.find(|candidate| candidate.account.id() == preferred)
|
||||
{
|
||||
Some(candidate) => match self.scheduling_blocker(candidate, context) {
|
||||
None => {
|
||||
return Some(AccountSelection {
|
||||
candidate,
|
||||
preferred: PreferredAccountSelection::Hit,
|
||||
});
|
||||
}
|
||||
Some(blocker) => PreferredAccountSelection::Blocked(blocker),
|
||||
},
|
||||
None => PreferredAccountSelection::Missing,
|
||||
}
|
||||
} else {
|
||||
PreferredAccountSelection::NotRequested
|
||||
};
|
||||
|
||||
let mut eligible = candidates
|
||||
.iter()
|
||||
.filter(|candidate| self.scheduling_blocker(candidate, context).is_none())
|
||||
@@ -1490,6 +1558,33 @@ impl AccountSelector {
|
||||
if eligible.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let highest_weight = eligible
|
||||
.iter()
|
||||
.map(|candidate| candidate.account.weight())
|
||||
.max()?;
|
||||
let preferred = if let Some(preferred) = context.preferred_account.as_ref() {
|
||||
match candidates
|
||||
.iter()
|
||||
.find(|candidate| candidate.account.id() == preferred)
|
||||
{
|
||||
Some(candidate) => match self.scheduling_blocker(candidate, context) {
|
||||
Some(blocker) => PreferredAccountSelection::Blocked(blocker),
|
||||
None if candidate.account.weight() < highest_weight => {
|
||||
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight)
|
||||
}
|
||||
None => {
|
||||
return Some(AccountSelection {
|
||||
candidate,
|
||||
preferred: PreferredAccountSelection::Hit,
|
||||
});
|
||||
}
|
||||
},
|
||||
None => PreferredAccountSelection::Missing,
|
||||
}
|
||||
} else {
|
||||
PreferredAccountSelection::NotRequested
|
||||
};
|
||||
eligible.retain(|candidate| candidate.account.weight() == highest_weight);
|
||||
|
||||
let candidate = match context.policy.strategy() {
|
||||
RotationStrategy::QuotaResetPriority => {
|
||||
@@ -1533,6 +1628,13 @@ impl AccountSelector {
|
||||
candidate: &AccountCandidate,
|
||||
context: &AccountSelectionContext,
|
||||
) -> Option<AccountSchedulingBlocker> {
|
||||
if context
|
||||
.account_scope
|
||||
.as_ref()
|
||||
.is_some_and(|scope| !scope.allows(candidate.account.id()))
|
||||
{
|
||||
return Some(AccountSchedulingBlocker::OutsideClientScope);
|
||||
}
|
||||
if !context.eligibility.bypasses_local_eligibility() {
|
||||
let status = candidate
|
||||
.account
|
||||
@@ -1545,7 +1647,12 @@ impl AccountSelector {
|
||||
if context.excluded_accounts.contains(candidate.account.id()) {
|
||||
return Some(AccountSchedulingBlocker::Excluded);
|
||||
}
|
||||
if candidate.signals.in_flight >= context.policy.max_concurrent_per_account().get() {
|
||||
if candidate.signals.in_flight
|
||||
>= candidate
|
||||
.account
|
||||
.effective_concurrency(context.policy.max_concurrent_per_account())
|
||||
.get()
|
||||
{
|
||||
return Some(AccountSchedulingBlocker::ConcurrencyLimit);
|
||||
}
|
||||
if candidate
|
||||
|
||||
@@ -15,7 +15,8 @@ use crate::engine::admission::{
|
||||
ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest,
|
||||
};
|
||||
use crate::engine::continuation::{
|
||||
ContinuationBinding, NativeContinuationPin, NativeContinuationPort, PreviousResponseId,
|
||||
ContinuationBinding, NativeContinuationPin, NativeContinuationPort,
|
||||
NativeContinuationStoreErrorKind, PreviousResponseId,
|
||||
};
|
||||
use crate::engine::coordinator::ResponseExecutionSession;
|
||||
use crate::engine::probe::{
|
||||
@@ -261,22 +262,32 @@ impl DefaultExecutionService {
|
||||
})?;
|
||||
let request_id = new_request_id()?;
|
||||
let routing_context = self
|
||||
.route_context(request.client.policy.provider_kind())
|
||||
.route_context(request.client.policy.account_scope().provider_kinds())
|
||||
.await?;
|
||||
let plan = request
|
||||
.client
|
||||
.snapshot
|
||||
.plan(&request.public_model, &request.operation, &routing_context)
|
||||
.plan(
|
||||
&request.public_model,
|
||||
&request.operation,
|
||||
Arc::clone(request.client.policy.account_scope()),
|
||||
&routing_context,
|
||||
)
|
||||
.map_err(map_routing_error)?;
|
||||
let continuation = match request.metadata.previous_response_id.as_ref() {
|
||||
Some(previous) => {
|
||||
let resolve = self.continuation.resolve(previous).fuse();
|
||||
let resolve = self
|
||||
.continuation
|
||||
.resolve(request.client.policy.key_id(), previous)
|
||||
.fuse();
|
||||
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
|
||||
pin_mut!(resolve, timeout);
|
||||
let pin = select_biased! {
|
||||
result = resolve => match result {
|
||||
Ok(pin) => pin,
|
||||
Err(error) => {
|
||||
Err(error)
|
||||
if error.kind() == NativeContinuationStoreErrorKind::Unavailable =>
|
||||
{
|
||||
tracing::debug!(
|
||||
request_id = request_id.as_str(),
|
||||
%error,
|
||||
@@ -284,10 +295,53 @@ impl DefaultExecutionService {
|
||||
);
|
||||
None
|
||||
}
|
||||
Err(error)
|
||||
if error.kind()
|
||||
== NativeContinuationStoreErrorKind::OwnershipMismatch =>
|
||||
{
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::PolicyDenied,
|
||||
"continuation does not belong to this client API key",
|
||||
));
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
request_id = request_id.as_str(),
|
||||
%error,
|
||||
"Continuation affinity 记录无效,已拒绝续接"
|
||||
);
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::Internal,
|
||||
"continuation state is invalid",
|
||||
));
|
||||
}
|
||||
},
|
||||
_ = timeout => None,
|
||||
};
|
||||
match pin {
|
||||
Some(pin) if !pin.matches_client(request.client.policy.key_id()) => {
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::PolicyDenied,
|
||||
"continuation does not belong to this client API key",
|
||||
));
|
||||
}
|
||||
Some(pin) if !request.client.policy.account_scope().allows(pin.account()) => {
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::PolicyDenied,
|
||||
"continuation account is outside the client account scope",
|
||||
));
|
||||
}
|
||||
Some(pin)
|
||||
if !plan
|
||||
.candidates()
|
||||
.iter()
|
||||
.any(|candidate| candidate.provider() == pin.provider()) =>
|
||||
{
|
||||
return Err(GatewayError::new(
|
||||
GatewayErrorKind::NoAvailableProvider,
|
||||
"continuation provider is not available",
|
||||
));
|
||||
}
|
||||
Some(pin) => {
|
||||
attach_continuation_session_state(&mut request.operation, &pin);
|
||||
ContinuationBinding::Pinned(pin)
|
||||
@@ -360,14 +414,19 @@ impl DefaultExecutionService {
|
||||
client_api_key_id: request.client.policy.key_id().clone(),
|
||||
model_request_id: request_id.clone(),
|
||||
};
|
||||
let observation = self
|
||||
.providers
|
||||
.request_observation(request.client.policy.provider_kind(), &request.operation);
|
||||
let observation = plan
|
||||
.candidates()
|
||||
.first()
|
||||
.map_or_else(Default::default, |candidate| {
|
||||
self.providers
|
||||
.request_observation(candidate.provider(), &request.operation)
|
||||
});
|
||||
let new_request = NewModelRequest {
|
||||
id: request_id.clone(),
|
||||
client_api_key_id: Some(request.client.policy.key_id().clone()),
|
||||
client_api_key_ref: request.client.policy.key_id().clone(),
|
||||
config_revision: plan.config_revision(),
|
||||
routing: request.client.policy.account_scope().routing_snapshot(),
|
||||
protocol: request.metadata.protocol,
|
||||
operation: request.operation.kind(),
|
||||
endpoint: request.metadata.endpoint,
|
||||
@@ -417,38 +476,42 @@ impl DefaultExecutionService {
|
||||
|
||||
async fn route_context(
|
||||
&self,
|
||||
provider_kind: &ProviderKind,
|
||||
provider_kinds: &BTreeSet<ProviderKind>,
|
||||
) -> Result<RoutingContext, GatewayError> {
|
||||
let decision = self.circuits.decision(provider_kind).fuse();
|
||||
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
|
||||
pin_mut!(decision, timeout);
|
||||
let decision = select_biased! {
|
||||
result = decision => Some(result),
|
||||
_ = timeout => None,
|
||||
};
|
||||
let blocked_providers = match decision {
|
||||
Some(Ok(ProviderCircuitDecision::Allow)) => BTreeSet::new(),
|
||||
Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => {
|
||||
BTreeSet::from([provider_kind.clone()])
|
||||
let decisions = futures::future::join_all(provider_kinds.iter().map(|provider_kind| {
|
||||
let circuits = Arc::clone(&self.circuits);
|
||||
async move {
|
||||
let decision = circuits.decision(provider_kind).fuse();
|
||||
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
|
||||
pin_mut!(decision, timeout);
|
||||
let decision = select_biased! {
|
||||
result = decision => Some(result),
|
||||
_ = timeout => None,
|
||||
};
|
||||
(provider_kind, decision)
|
||||
}
|
||||
Some(Err(error)) => {
|
||||
tracing::warn!(
|
||||
}))
|
||||
.await;
|
||||
let mut blocked_providers = BTreeSet::new();
|
||||
for (provider_kind, decision) in decisions {
|
||||
match decision {
|
||||
Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => {
|
||||
blocked_providers.insert(provider_kind.clone());
|
||||
}
|
||||
Some(Err(error)) => tracing::warn!(
|
||||
provider = provider_kind.as_str(),
|
||||
%error,
|
||||
"Provider circuit 读取失败,按可重建协调状态 fail-open"
|
||||
);
|
||||
BTreeSet::new()
|
||||
}
|
||||
None => {
|
||||
tracing::warn!(
|
||||
),
|
||||
None => tracing::warn!(
|
||||
provider = provider_kind.as_str(),
|
||||
"Provider circuit 读取超时,按可重建协调状态 fail-open"
|
||||
);
|
||||
BTreeSet::new()
|
||||
),
|
||||
Some(Ok(ProviderCircuitDecision::Allow)) => {}
|
||||
}
|
||||
};
|
||||
}
|
||||
Ok(RoutingContext {
|
||||
provider_kind: Some(provider_kind.clone()),
|
||||
required_provider: None,
|
||||
blocked_providers,
|
||||
})
|
||||
}
|
||||
@@ -479,11 +542,16 @@ impl DefaultExecutionService {
|
||||
GatewayError::new(GatewayErrorKind::Unsupported, "requested model is invalid")
|
||||
})?;
|
||||
let routing_context = RoutingContext {
|
||||
provider_kind: Some(provider_kind),
|
||||
required_provider: Some(provider_kind),
|
||||
..RoutingContext::default()
|
||||
};
|
||||
let plan = snapshot
|
||||
.plan(&public_model, &operation, &routing_context)
|
||||
.plan(
|
||||
&public_model,
|
||||
&operation,
|
||||
snapshot.all_account_scope(),
|
||||
&routing_context,
|
||||
)
|
||||
.map_err(map_routing_error)?;
|
||||
let started_at = SystemTime::now();
|
||||
let deadline_at = started_at
|
||||
@@ -499,6 +567,7 @@ impl DefaultExecutionService {
|
||||
client_api_key_id: None,
|
||||
client_api_key_ref: actor,
|
||||
config_revision: plan.config_revision(),
|
||||
routing: crate::routing::AccountRoutingSnapshot::all(),
|
||||
protocol: "admin_connection_test".to_owned(),
|
||||
operation: operation.kind(),
|
||||
endpoint: "/api/admin/accounts/connection-test".to_owned(),
|
||||
@@ -693,19 +762,19 @@ impl ExecutionService for DefaultExecutionService {
|
||||
fn public_models(&self, client: &AuthenticatedClient) -> Vec<PublicModelId> {
|
||||
client
|
||||
.snapshot
|
||||
.public_models_for_provider(client.policy.provider_kind())
|
||||
.public_models_for_scope(client.policy.account_scope())
|
||||
}
|
||||
|
||||
fn public_model_profiles(&self, client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
|
||||
client
|
||||
.snapshot
|
||||
.public_model_profiles_for_provider(client.policy.provider_kind())
|
||||
.public_model_profiles_for_scope(client.policy.account_scope())
|
||||
}
|
||||
|
||||
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool {
|
||||
client
|
||||
.snapshot
|
||||
.contains_public_model_for_provider(model, client.policy.provider_kind())
|
||||
.contains_public_model_for_scope(model, client.policy.account_scope())
|
||||
}
|
||||
|
||||
fn start(
|
||||
@@ -1006,7 +1075,8 @@ fn new_request_id() -> Result<ModelRequestId, GatewayError> {
|
||||
|
||||
fn map_routing_error(error: crate::error::RoutingError) -> GatewayError {
|
||||
match error {
|
||||
crate::error::RoutingError::NoCapableProvider { .. } => GatewayError::new(
|
||||
crate::error::RoutingError::NoCapableProvider { .. }
|
||||
| crate::error::RoutingError::EmptyAccountScope => GatewayError::new(
|
||||
GatewayErrorKind::NoAvailableProvider,
|
||||
"no provider can execute this request",
|
||||
),
|
||||
|
||||
@@ -233,6 +233,7 @@ pub struct AccountAttemptContext {
|
||||
state_owner: Option<ProviderAccountStateOwner>,
|
||||
credential_recovery_attempted: bool,
|
||||
diagnostic_required_account: bool,
|
||||
account_scope: Option<Arc<crate::routing::FrozenAccountScope>>,
|
||||
}
|
||||
|
||||
impl AccountAttemptContext {
|
||||
@@ -248,6 +249,7 @@ impl AccountAttemptContext {
|
||||
state_owner,
|
||||
credential_recovery_attempted: false,
|
||||
diagnostic_required_account: false,
|
||||
account_scope: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -266,9 +268,17 @@ impl AccountAttemptContext {
|
||||
state_owner,
|
||||
credential_recovery_attempted: false,
|
||||
diagnostic_required_account: true,
|
||||
account_scope: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 附着普通请求认证时冻结的账号范围。
|
||||
#[must_use]
|
||||
pub fn with_account_scope(mut self, scope: Arc<crate::routing::FrozenAccountScope>) -> Self {
|
||||
self.account_scope = Some(scope);
|
||||
self
|
||||
}
|
||||
|
||||
/// 标记本请求已对即将选择的固定账号执行过一次凭据恢复。
|
||||
#[must_use]
|
||||
pub const fn with_credential_recovery_attempted(mut self, attempted: bool) -> Self {
|
||||
@@ -301,6 +311,11 @@ impl AccountAttemptContext {
|
||||
pub const fn is_diagnostic_required_account(&self) -> bool {
|
||||
self.diagnostic_required_account
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
|
||||
self.account_scope.as_ref()
|
||||
}
|
||||
}
|
||||
|
||||
/// 请求中 Provider 账号绑定状态的唯一归属。
|
||||
@@ -506,6 +521,12 @@ impl AttemptContext {
|
||||
self.account.is_diagnostic_required_account()
|
||||
}
|
||||
|
||||
/// 普通请求认证时冻结的账号范围;管理端诊断为 `None`。
|
||||
#[must_use]
|
||||
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
|
||||
self.account.account_scope()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn continuation(&self) -> Option<&ContinuationBinding> {
|
||||
self.continuation.as_ref()
|
||||
@@ -529,6 +550,7 @@ pub struct NewModelRequest {
|
||||
pub client_api_key_id: Option<ClientApiKeyId>,
|
||||
pub client_api_key_ref: ClientApiKeyId,
|
||||
pub config_revision: ConfigRevision,
|
||||
pub routing: crate::routing::AccountRoutingSnapshot,
|
||||
pub protocol: String,
|
||||
pub operation: OperationKind,
|
||||
pub endpoint: String,
|
||||
@@ -749,6 +771,8 @@ pub enum EngineError {
|
||||
ContinuationPinMismatch,
|
||||
#[error("provider did not use the account required by this execution")]
|
||||
RequiredAccountMismatch,
|
||||
#[error("provider selected an account outside the frozen client scope")]
|
||||
AccountOutsideClientScope,
|
||||
#[error("provider execution failed")]
|
||||
Provider(ProviderError),
|
||||
#[error("request was cancelled")]
|
||||
|
||||
@@ -33,6 +33,9 @@ pub enum IdentifierError {
|
||||
/// 规定前缀。
|
||||
expected: &'static str,
|
||||
},
|
||||
/// 标识不满足该领域值对象的固定格式。
|
||||
#[error("identifier has an invalid format")]
|
||||
InvalidFormat,
|
||||
}
|
||||
|
||||
/// Operation 构造或校验失败。
|
||||
@@ -67,6 +70,12 @@ pub enum RoutingError {
|
||||
/// 配置 revision 必须为正数。
|
||||
#[error("config revision must be greater than zero")]
|
||||
InvalidRevision,
|
||||
/// Restricted scope 必须保留至少一个持久化分组 binding。
|
||||
#[error("account routing scope is invalid")]
|
||||
InvalidAccountScope,
|
||||
/// Key 的账号范围中当前没有任何可参与路由的账号。
|
||||
#[error("account routing scope is empty")]
|
||||
EmptyAccountScope,
|
||||
/// 快照中存在重复实体。
|
||||
#[error("duplicate {entity} `{id}`")]
|
||||
DuplicateEntity {
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
//! 下游 Client API Key 的准入策略。
|
||||
//!
|
||||
//! Client API Key 绑定一个 Provider 平台;模型名称不参与改写或 allowlist 判断。
|
||||
//! Client API Key 冻结账号分组权限;模型名称不参与权限判断。
|
||||
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::error::{IdentifierError, PolicyError, validate_text};
|
||||
use crate::routing::ProviderKind;
|
||||
use crate::routing::FrozenAccountScope;
|
||||
|
||||
/// `client_api_keys.id` 的核心值对象。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
@@ -88,7 +89,7 @@ impl RateLimits {
|
||||
pub struct ClientPolicy {
|
||||
key_id: ClientApiKeyId,
|
||||
plaintext_key: PlaintextClientApiKey,
|
||||
provider_kind: ProviderKind,
|
||||
account_scope: Arc<FrozenAccountScope>,
|
||||
enabled: bool,
|
||||
limits: RateLimits,
|
||||
}
|
||||
@@ -98,14 +99,14 @@ impl ClientPolicy {
|
||||
pub const fn new(
|
||||
key_id: ClientApiKeyId,
|
||||
plaintext_key: PlaintextClientApiKey,
|
||||
provider_kind: ProviderKind,
|
||||
account_scope: Arc<FrozenAccountScope>,
|
||||
enabled: bool,
|
||||
limits: RateLimits,
|
||||
) -> Self {
|
||||
Self {
|
||||
key_id,
|
||||
plaintext_key,
|
||||
provider_kind,
|
||||
account_scope,
|
||||
enabled,
|
||||
limits,
|
||||
}
|
||||
@@ -122,8 +123,8 @@ impl ClientPolicy {
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn provider_kind(&self) -> &ProviderKind {
|
||||
&self.provider_kind
|
||||
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
|
||||
&self.account_scope
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
|
||||
@@ -13,6 +13,7 @@ use crate::engine::credential::{
|
||||
OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
|
||||
};
|
||||
use crate::error::{IdentifierError, validate_text};
|
||||
use crate::policy::ClientApiKeyId;
|
||||
use crate::routing::{ProviderKind, UpstreamModelId};
|
||||
|
||||
const MAX_PENDING_FLOW_TTL: Duration = Duration::from_secs(30 * 60);
|
||||
@@ -170,6 +171,7 @@ pub enum ProviderLeaseRequest {
|
||||
pub trait ProviderLeasePort: Send + Sync {
|
||||
fn load_state<'a>(
|
||||
&'a self,
|
||||
client_api_key_id: &'a ClientApiKeyId,
|
||||
provider_kind: &'a ProviderKind,
|
||||
accounts: &'a [ProviderAccountId],
|
||||
) -> BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>>;
|
||||
|
||||
@@ -9,12 +9,293 @@ use std::fmt;
|
||||
use std::num::{NonZeroU32, NonZeroU64};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::engine::credential::AccountSelectionPolicy;
|
||||
use crate::engine::credential::{AccountSelectionPolicy, ProviderAccountId};
|
||||
use crate::error::{IdentifierError, RoutingError, validate_text};
|
||||
use crate::operation::{CapabilityRequirements, Feature, OperationKind};
|
||||
|
||||
const MAX_REQUEST_ATTEMPTS: u32 = 32;
|
||||
|
||||
/// `account_groups.id` 的核心值对象。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct AccountGroupId(String);
|
||||
|
||||
impl AccountGroupId {
|
||||
/// 校验并创建账号分组 ID。
|
||||
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
|
||||
let value = value.into();
|
||||
let Some(suffix) = value.strip_prefix("grp_") else {
|
||||
return Err(IdentifierError::MissingPrefix { expected: "grp_" });
|
||||
};
|
||||
if suffix.len() != 32
|
||||
|| !suffix
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
|
||||
{
|
||||
return Err(IdentifierError::InvalidFormat);
|
||||
}
|
||||
Ok(Self(value))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for AccountGroupId {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
formatter.write_str(&self.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// 快照中一个账号的 Provider 与分组归属。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RuntimeAccount {
|
||||
provider_kind: ProviderKind,
|
||||
group_ids: Arc<BTreeSet<AccountGroupId>>,
|
||||
}
|
||||
|
||||
impl RuntimeAccount {
|
||||
#[must_use]
|
||||
pub fn new(provider_kind: ProviderKind, group_ids: BTreeSet<AccountGroupId>) -> Self {
|
||||
Self {
|
||||
provider_kind,
|
||||
group_ids: Arc::new(group_ids),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn provider_kind(&self) -> &ProviderKind {
|
||||
&self.provider_kind
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn group_ids(&self) -> &BTreeSet<AccountGroupId> {
|
||||
&self.group_ids
|
||||
}
|
||||
}
|
||||
|
||||
/// 全快照共享的账号、Provider 与分组反向索引。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct RuntimeAccountDirectory {
|
||||
accounts: BTreeMap<ProviderAccountId, RuntimeAccount>,
|
||||
providers_with_accounts: BTreeSet<ProviderKind>,
|
||||
providers_by_group: BTreeMap<AccountGroupId, BTreeSet<ProviderKind>>,
|
||||
}
|
||||
|
||||
impl RuntimeAccountDirectory {
|
||||
#[must_use]
|
||||
pub fn new(accounts: BTreeMap<ProviderAccountId, RuntimeAccount>) -> Self {
|
||||
let mut providers_with_accounts = BTreeSet::new();
|
||||
let mut providers_by_group = BTreeMap::<AccountGroupId, BTreeSet<ProviderKind>>::new();
|
||||
for account in accounts.values() {
|
||||
providers_with_accounts.insert(account.provider_kind.clone());
|
||||
for group_id in account.group_ids.iter() {
|
||||
providers_by_group
|
||||
.entry(group_id.clone())
|
||||
.or_default()
|
||||
.insert(account.provider_kind.clone());
|
||||
}
|
||||
}
|
||||
Self {
|
||||
accounts,
|
||||
providers_with_accounts,
|
||||
providers_by_group,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn account(&self, account_id: &ProviderAccountId) -> Option<&RuntimeAccount> {
|
||||
self.accounts.get(account_id)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn providers_with_accounts(&self) -> &BTreeSet<ProviderKind> {
|
||||
&self.providers_with_accounts
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn providers_for_groups<'a>(
|
||||
&self,
|
||||
group_ids: impl IntoIterator<Item = &'a AccountGroupId>,
|
||||
) -> BTreeSet<ProviderKind> {
|
||||
group_ids
|
||||
.into_iter()
|
||||
.filter_map(|group_id| self.providers_by_group.get(group_id))
|
||||
.flatten()
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// 历史请求保存的账号范围种类。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AccountRoutingScopeKind {
|
||||
All,
|
||||
Groups,
|
||||
}
|
||||
|
||||
impl AccountRoutingScopeKind {
|
||||
#[must_use]
|
||||
pub const fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::All => "all",
|
||||
Self::Groups => "groups",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 请求开始时冻结的分组 ID 与名称。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RoutingGroupSnapshot {
|
||||
id: AccountGroupId,
|
||||
name: String,
|
||||
}
|
||||
|
||||
impl RoutingGroupSnapshot {
|
||||
#[must_use]
|
||||
pub fn new(id: AccountGroupId, name: String) -> Self {
|
||||
Self { id, name }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn id(&self) -> &AccountGroupId {
|
||||
&self.id
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
}
|
||||
|
||||
/// 请求历史所需的完整、稳定账号范围快照。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AccountRoutingSnapshot {
|
||||
kind: AccountRoutingScopeKind,
|
||||
groups: Arc<[RoutingGroupSnapshot]>,
|
||||
}
|
||||
|
||||
impl AccountRoutingSnapshot {
|
||||
#[must_use]
|
||||
pub fn all() -> Self {
|
||||
Self {
|
||||
kind: AccountRoutingScopeKind::All,
|
||||
groups: Arc::from([]),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn groups(groups: Vec<RoutingGroupSnapshot>) -> Self {
|
||||
Self {
|
||||
kind: AccountRoutingScopeKind::Groups,
|
||||
groups: Arc::from(groups),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn kind(&self) -> AccountRoutingScopeKind {
|
||||
self.kind
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn groups_snapshot(&self) -> &[RoutingGroupSnapshot] {
|
||||
&self.groups
|
||||
}
|
||||
}
|
||||
|
||||
/// Key 持久 binding 编译出的账号权限。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ClientRoutingScope {
|
||||
AllAccounts,
|
||||
Restricted {
|
||||
bound_groups: Arc<[RoutingGroupSnapshot]>,
|
||||
enabled_group_ids: Arc<BTreeSet<AccountGroupId>>,
|
||||
provider_kinds: Arc<BTreeSet<ProviderKind>>,
|
||||
},
|
||||
}
|
||||
|
||||
impl ClientRoutingScope {
|
||||
#[must_use]
|
||||
pub fn all_accounts() -> Self {
|
||||
Self::AllAccounts
|
||||
}
|
||||
|
||||
pub fn restricted(
|
||||
bound_groups: Vec<RoutingGroupSnapshot>,
|
||||
enabled_group_ids: BTreeSet<AccountGroupId>,
|
||||
provider_kinds: BTreeSet<ProviderKind>,
|
||||
) -> Result<Self, RoutingError> {
|
||||
if bound_groups.is_empty() {
|
||||
return Err(RoutingError::InvalidAccountScope);
|
||||
}
|
||||
Ok(Self::Restricted {
|
||||
bound_groups: Arc::from(bound_groups),
|
||||
enabled_group_ids: Arc::new(enabled_group_ids),
|
||||
provider_kinds: Arc::new(provider_kinds),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// 一次认证随 RuntimeSnapshot 冻结的账号目录与 Key 权限。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct FrozenAccountScope {
|
||||
directory: Arc<RuntimeAccountDirectory>,
|
||||
client_scope: ClientRoutingScope,
|
||||
}
|
||||
|
||||
impl FrozenAccountScope {
|
||||
#[must_use]
|
||||
pub const fn new(
|
||||
directory: Arc<RuntimeAccountDirectory>,
|
||||
client_scope: ClientRoutingScope,
|
||||
) -> Self {
|
||||
Self {
|
||||
directory,
|
||||
client_scope,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn allows(&self, account_id: &ProviderAccountId) -> bool {
|
||||
let Some(account) = self.directory.account(account_id) else {
|
||||
return false;
|
||||
};
|
||||
match &self.client_scope {
|
||||
ClientRoutingScope::AllAccounts => true,
|
||||
ClientRoutingScope::Restricted {
|
||||
enabled_group_ids, ..
|
||||
} => account
|
||||
.group_ids()
|
||||
.iter()
|
||||
.any(|group_id| enabled_group_ids.contains(group_id)),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn provider_kinds(&self) -> &BTreeSet<ProviderKind> {
|
||||
match &self.client_scope {
|
||||
ClientRoutingScope::AllAccounts => self.directory.providers_with_accounts(),
|
||||
ClientRoutingScope::Restricted { provider_kinds, .. } => provider_kinds,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn routing_snapshot(&self) -> AccountRoutingSnapshot {
|
||||
match &self.client_scope {
|
||||
ClientRoutingScope::AllAccounts => AccountRoutingSnapshot::all(),
|
||||
ClientRoutingScope::Restricted { bound_groups, .. } => {
|
||||
AccountRoutingSnapshot::groups(bound_groups.to_vec())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn directory(&self) -> &Arc<RuntimeAccountDirectory> {
|
||||
&self.directory
|
||||
}
|
||||
}
|
||||
|
||||
/// 编译进二进制的 Provider adapter slug。
|
||||
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct ProviderKind(String);
|
||||
@@ -494,8 +775,8 @@ impl ProviderModel {
|
||||
/// 本次请求选择 Provider 时使用的动态过滤事实。
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct RoutingContext {
|
||||
/// 已认证 Client API Key 绑定的平台;模型名称不参与平台猜测。
|
||||
pub provider_kind: Option<ProviderKind>,
|
||||
/// 管理端 connection test 显式限制的 Provider;普通请求留空。
|
||||
pub required_provider: Option<ProviderKind>,
|
||||
pub blocked_providers: BTreeSet<ProviderKind>,
|
||||
}
|
||||
|
||||
@@ -505,6 +786,7 @@ pub struct ProviderCandidate {
|
||||
provider: ProviderKind,
|
||||
upstream_model: UpstreamModelId,
|
||||
emulated_features: BTreeSet<Feature>,
|
||||
account_scope: Arc<FrozenAccountScope>,
|
||||
}
|
||||
|
||||
impl ProviderCandidate {
|
||||
@@ -522,6 +804,11 @@ impl ProviderCandidate {
|
||||
pub const fn emulated_features(&self) -> &BTreeSet<Feature> {
|
||||
&self.emulated_features
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
|
||||
&self.account_scope
|
||||
}
|
||||
}
|
||||
|
||||
/// 一次请求冻结的 Provider 尝试顺序。
|
||||
@@ -532,6 +819,7 @@ pub struct RoutingPlan {
|
||||
public_model: PublicModelId,
|
||||
operation: OperationKind,
|
||||
max_attempts: NonZeroU32,
|
||||
account_scope: Arc<FrozenAccountScope>,
|
||||
candidates: Arc<[ProviderCandidate]>,
|
||||
}
|
||||
|
||||
@@ -561,6 +849,11 @@ impl RoutingPlan {
|
||||
self.max_attempts
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
|
||||
&self.account_scope
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn candidates(&self) -> &[ProviderCandidate] {
|
||||
&self.candidates
|
||||
|
||||
@@ -11,7 +11,7 @@ use futures::{FutureExt as _, Stream, StreamExt as _, pin_mut, select_biased};
|
||||
use futures_timer::Delay;
|
||||
|
||||
use crate::engine::CancellationToken;
|
||||
use crate::engine::credential::{AccountSelectionPolicy, RotationStrategy};
|
||||
use crate::engine::credential::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
|
||||
use crate::engine::provider::{ProviderCatalogGeneration, ProviderRegistry};
|
||||
use crate::error::RoutingError;
|
||||
use crate::health::{HealthProbe, HealthState};
|
||||
@@ -24,8 +24,9 @@ use crate::task::{
|
||||
};
|
||||
|
||||
use super::{
|
||||
ConfigRevision, ModelCapabilities, ProviderCandidate, ProviderKind, ProviderModel,
|
||||
PublicModelId, RoutingContext, RoutingPlan, UpstreamModelId,
|
||||
AccountGroupId, ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
|
||||
ProviderCandidate, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
|
||||
RoutingGroupSnapshot, RoutingPlan, RuntimeAccount, RuntimeAccountDirectory, UpstreamModelId,
|
||||
};
|
||||
|
||||
const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5);
|
||||
@@ -66,7 +67,7 @@ impl SnapshotSettingsFacts {
|
||||
pub struct SnapshotClientPolicyFacts {
|
||||
key_id: ClientApiKeyId,
|
||||
plaintext_key: PlaintextClientApiKey,
|
||||
provider_kind: String,
|
||||
group_ids: Vec<AccountGroupId>,
|
||||
limits: RateLimits,
|
||||
}
|
||||
|
||||
@@ -75,18 +76,67 @@ impl SnapshotClientPolicyFacts {
|
||||
pub fn new(
|
||||
key_id: ClientApiKeyId,
|
||||
plaintext_key: PlaintextClientApiKey,
|
||||
provider_kind: impl Into<String>,
|
||||
group_ids: Vec<AccountGroupId>,
|
||||
limits: RateLimits,
|
||||
) -> Self {
|
||||
Self {
|
||||
key_id,
|
||||
plaintext_key,
|
||||
provider_kind: provider_kind.into(),
|
||||
group_ids,
|
||||
limits,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Store 读取到的账号分组事实。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotAccountGroupFacts {
|
||||
id: AccountGroupId,
|
||||
name: String,
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl SnapshotAccountGroupFacts {
|
||||
#[must_use]
|
||||
pub fn new(id: AccountGroupId, name: String, enabled: bool) -> Self {
|
||||
Self { id, name, enabled }
|
||||
}
|
||||
}
|
||||
|
||||
/// Store 读取到的账号及其固有 Provider 事实。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotProviderAccountFacts {
|
||||
account_id: ProviderAccountId,
|
||||
provider_kind: String,
|
||||
}
|
||||
|
||||
impl SnapshotProviderAccountFacts {
|
||||
#[must_use]
|
||||
pub fn new(account_id: ProviderAccountId, provider_kind: impl Into<String>) -> Self {
|
||||
Self {
|
||||
account_id,
|
||||
provider_kind: provider_kind.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Store 读取到的一条分组成员关系。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotAccountGroupMemberFacts {
|
||||
group_id: AccountGroupId,
|
||||
account_id: ProviderAccountId,
|
||||
}
|
||||
|
||||
impl SnapshotAccountGroupMemberFacts {
|
||||
#[must_use]
|
||||
pub const fn new(group_id: AccountGroupId, account_id: ProviderAccountId) -> Self {
|
||||
Self {
|
||||
group_id,
|
||||
account_id,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 一次一致性读取产生的全部 RuntimeSnapshot 持久事实。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotFacts {
|
||||
@@ -94,6 +144,9 @@ pub struct SnapshotFacts {
|
||||
observed_current_revision: ConfigRevision,
|
||||
settings: SnapshotSettingsFacts,
|
||||
client_policies: Vec<SnapshotClientPolicyFacts>,
|
||||
account_groups: Vec<SnapshotAccountGroupFacts>,
|
||||
provider_accounts: Vec<SnapshotProviderAccountFacts>,
|
||||
group_memberships: Vec<SnapshotAccountGroupMemberFacts>,
|
||||
}
|
||||
|
||||
impl SnapshotFacts {
|
||||
@@ -103,12 +156,18 @@ impl SnapshotFacts {
|
||||
observed_current_revision: ConfigRevision,
|
||||
settings: SnapshotSettingsFacts,
|
||||
client_policies: Vec<SnapshotClientPolicyFacts>,
|
||||
account_groups: Vec<SnapshotAccountGroupFacts>,
|
||||
provider_accounts: Vec<SnapshotProviderAccountFacts>,
|
||||
group_memberships: Vec<SnapshotAccountGroupMemberFacts>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config_revision,
|
||||
observed_current_revision,
|
||||
settings,
|
||||
client_policies,
|
||||
account_groups,
|
||||
provider_accounts,
|
||||
group_memberships,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -223,13 +282,16 @@ async fn compile_runtime_snapshot(
|
||||
providers: &ProviderRegistry,
|
||||
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
|
||||
let provider_kinds = providers.provider_kinds().cloned().collect::<Vec<_>>();
|
||||
let registered_providers = provider_kinds.iter().cloned().collect::<BTreeSet<_>>();
|
||||
|
||||
// 实时目录只提供公开模型与能力提示;查询失败时保留 Provider 透传语义。
|
||||
// 目录查询失败表示未知;查询成功后,即使为空,也必须与“已知缺少模型”区分。
|
||||
let mut provider_models = Vec::new();
|
||||
let mut known_provider_catalogs = BTreeSet::new();
|
||||
for provider in &provider_kinds {
|
||||
let Ok(models) = providers.query_model_capabilities(provider).await else {
|
||||
continue;
|
||||
};
|
||||
known_provider_catalogs.insert(provider.clone());
|
||||
provider_models.extend(models.into_iter().map(|model| {
|
||||
let compiled = ProviderModel::new(
|
||||
provider.clone(),
|
||||
@@ -243,6 +305,56 @@ async fn compile_runtime_snapshot(
|
||||
}));
|
||||
}
|
||||
|
||||
let mut groups = BTreeMap::new();
|
||||
for group in facts.account_groups {
|
||||
if group.name.trim() != group.name
|
||||
|| group.name.is_empty()
|
||||
|| group.name.chars().count() > 100
|
||||
|| group.name.chars().any(char::is_control)
|
||||
|| groups.insert(group.id.clone(), group).is_some()
|
||||
{
|
||||
return Err(RuntimeSnapshotCompileError::InvalidData);
|
||||
}
|
||||
}
|
||||
|
||||
let mut account_groups = BTreeMap::<ProviderAccountId, BTreeSet<AccountGroupId>>::new();
|
||||
let mut accounts = BTreeMap::new();
|
||||
for account in facts.provider_accounts {
|
||||
let provider_kind = ProviderKind::new(account.provider_kind)
|
||||
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
|
||||
if !registered_providers.contains(&provider_kind)
|
||||
|| accounts
|
||||
.insert(
|
||||
account.account_id.clone(),
|
||||
RuntimeAccount::new(provider_kind, BTreeSet::new()),
|
||||
)
|
||||
.is_some()
|
||||
{
|
||||
return Err(RuntimeSnapshotCompileError::InvalidData);
|
||||
}
|
||||
account_groups.insert(account.account_id, BTreeSet::new());
|
||||
}
|
||||
let mut memberships = BTreeSet::new();
|
||||
for membership in facts.group_memberships {
|
||||
if !groups.contains_key(&membership.group_id)
|
||||
|| !accounts.contains_key(&membership.account_id)
|
||||
|| !memberships.insert((membership.group_id.clone(), membership.account_id.clone()))
|
||||
{
|
||||
return Err(RuntimeSnapshotCompileError::InvalidData);
|
||||
}
|
||||
account_groups
|
||||
.get_mut(&membership.account_id)
|
||||
.ok_or(RuntimeSnapshotCompileError::InvalidData)?
|
||||
.insert(membership.group_id);
|
||||
}
|
||||
for (account_id, group_ids) in account_groups {
|
||||
let account = accounts
|
||||
.get_mut(&account_id)
|
||||
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
|
||||
*account = RuntimeAccount::new(account.provider_kind().clone(), group_ids);
|
||||
}
|
||||
let account_directory = Arc::new(RuntimeAccountDirectory::new(accounts));
|
||||
|
||||
let model_mappings = facts.settings.model_mappings;
|
||||
let rotation_strategy = RotationStrategy::parse(facts.settings.rotation_strategy.as_str())
|
||||
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
|
||||
@@ -252,21 +364,48 @@ async fn compile_runtime_snapshot(
|
||||
.ok_or(RuntimeSnapshotCompileError::InvalidData)?,
|
||||
Duration::from_millis(facts.settings.request_interval_ms),
|
||||
);
|
||||
let client_policies = facts
|
||||
.client_policies
|
||||
.into_iter()
|
||||
.map(|policy| {
|
||||
let provider = ProviderKind::new(policy.provider_kind)
|
||||
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
|
||||
Ok(ClientPolicy::new(
|
||||
policy.key_id,
|
||||
policy.plaintext_key,
|
||||
provider,
|
||||
true,
|
||||
policy.limits,
|
||||
))
|
||||
})
|
||||
.collect::<Result<Vec<_>, RuntimeSnapshotCompileError>>()?;
|
||||
let mut client_policies = Vec::with_capacity(facts.client_policies.len());
|
||||
for policy in facts.client_policies {
|
||||
let account_scope = if policy.group_ids.is_empty() {
|
||||
FrozenAccountScope::new(
|
||||
Arc::clone(&account_directory),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
)
|
||||
} else {
|
||||
let mut seen = BTreeSet::new();
|
||||
let mut bound_groups = Vec::with_capacity(policy.group_ids.len());
|
||||
let mut enabled_group_ids = BTreeSet::new();
|
||||
for group_id in policy.group_ids {
|
||||
if !seen.insert(group_id.clone()) {
|
||||
return Err(RuntimeSnapshotCompileError::InvalidData);
|
||||
}
|
||||
let group = groups
|
||||
.get(&group_id)
|
||||
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
|
||||
bound_groups.push(RoutingGroupSnapshot::new(
|
||||
group.id.clone(),
|
||||
group.name.clone(),
|
||||
));
|
||||
if group.enabled {
|
||||
enabled_group_ids.insert(group_id);
|
||||
}
|
||||
}
|
||||
bound_groups.sort_by(|left, right| left.id().cmp(right.id()));
|
||||
let provider_kinds = account_directory.providers_for_groups(&enabled_group_ids);
|
||||
FrozenAccountScope::new(
|
||||
Arc::clone(&account_directory),
|
||||
ClientRoutingScope::restricted(bound_groups, enabled_group_ids, provider_kinds)
|
||||
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?,
|
||||
)
|
||||
};
|
||||
client_policies.push(ClientPolicy::new(
|
||||
policy.key_id,
|
||||
policy.plaintext_key,
|
||||
Arc::new(account_scope),
|
||||
true,
|
||||
policy.limits,
|
||||
));
|
||||
}
|
||||
|
||||
RuntimeSnapshot::new(
|
||||
facts.config_revision,
|
||||
@@ -276,7 +415,12 @@ async fn compile_runtime_snapshot(
|
||||
client_policies,
|
||||
)
|
||||
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)
|
||||
.map(|snapshot| snapshot.with_model_mappings(model_mappings))
|
||||
.map(|snapshot| {
|
||||
snapshot
|
||||
.with_model_mappings(model_mappings)
|
||||
.with_account_directory(account_directory)
|
||||
.with_known_provider_catalogs(known_provider_catalogs)
|
||||
})
|
||||
}
|
||||
|
||||
/// 数据面使用的不可变配置快照。
|
||||
@@ -290,6 +434,8 @@ pub struct RuntimeSnapshot {
|
||||
Arc<BTreeMap<ProviderKind, BTreeMap<UpstreamModelId, super::ModelPresentation>>>,
|
||||
model_mappings: Arc<BTreeMap<String, String>>,
|
||||
provider_catalog_generations: Arc<BTreeMap<ProviderKind, ProviderCatalogGeneration>>,
|
||||
known_provider_catalogs: Arc<BTreeSet<ProviderKind>>,
|
||||
account_directory: Arc<RuntimeAccountDirectory>,
|
||||
client_policies: Arc<BTreeMap<ClientApiKeyId, ClientPolicy>>,
|
||||
}
|
||||
|
||||
@@ -312,6 +458,7 @@ impl RuntimeSnapshot {
|
||||
}
|
||||
}
|
||||
|
||||
let mut known_provider_catalogs = BTreeSet::new();
|
||||
let mut model_map =
|
||||
BTreeMap::<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>::new();
|
||||
let mut presentation_map =
|
||||
@@ -323,6 +470,7 @@ impl RuntimeSnapshot {
|
||||
capabilities,
|
||||
presentation,
|
||||
} = model;
|
||||
known_provider_catalogs.insert(provider.clone());
|
||||
if !provider_set.contains(&provider) {
|
||||
return Err(RoutingError::NotFound {
|
||||
entity: "provider",
|
||||
@@ -367,6 +515,8 @@ impl RuntimeSnapshot {
|
||||
provider_model_presentations: Arc::new(presentation_map),
|
||||
model_mappings: Arc::new(BTreeMap::new()),
|
||||
provider_catalog_generations: Arc::new(BTreeMap::new()),
|
||||
known_provider_catalogs: Arc::new(known_provider_catalogs),
|
||||
account_directory: Arc::new(RuntimeAccountDirectory::default()),
|
||||
client_policies: Arc::new(client_policy_map),
|
||||
})
|
||||
}
|
||||
@@ -377,6 +527,26 @@ impl RuntimeSnapshot {
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_account_directory(mut self, directory: Arc<RuntimeAccountDirectory>) -> Self {
|
||||
self.account_directory = directory;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
fn with_known_provider_catalogs(mut self, providers: BTreeSet<ProviderKind>) -> Self {
|
||||
self.known_provider_catalogs = Arc::new(providers);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn all_account_scope(&self) -> Arc<FrozenAccountScope> {
|
||||
Arc::new(FrozenAccountScope::new(
|
||||
Arc::clone(&self.account_directory),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
fn with_provider_catalog_generations(
|
||||
mut self,
|
||||
@@ -459,8 +629,39 @@ impl RuntimeSnapshot {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 已取得目录时按目录与全局别名判断模型可用性;目录暂不可用时沿用旧版
|
||||
/// fail-open 语义,把具体 wire 可接受性留给上游。
|
||||
/// 合并冻结账号范围内实际存在 Provider 的公开模型。
|
||||
#[must_use]
|
||||
pub fn public_models_for_scope(&self, scope: &FrozenAccountScope) -> Vec<PublicModelId> {
|
||||
scope
|
||||
.provider_kinds()
|
||||
.iter()
|
||||
.flat_map(|provider| self.public_models_for_provider(provider))
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 合并冻结账号范围内实际存在 Provider 的公开模型画像。
|
||||
#[must_use]
|
||||
pub fn public_model_profiles_for_scope(
|
||||
&self,
|
||||
scope: &FrozenAccountScope,
|
||||
) -> Vec<super::PublicModelProfile> {
|
||||
let mut profiles = BTreeMap::new();
|
||||
for provider in scope.provider_kinds() {
|
||||
for profile in self.public_model_profiles_for_provider(provider) {
|
||||
profiles
|
||||
.entry(profile.model().clone())
|
||||
.or_insert_with(|| profile.presentation().clone());
|
||||
}
|
||||
}
|
||||
profiles
|
||||
.into_iter()
|
||||
.map(|(model, presentation)| super::PublicModelProfile::new(model, presentation))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 已取得目录时以映射后的上游模型为准;目录不可用时由 Provider 在发送前判定。
|
||||
#[must_use]
|
||||
pub fn contains_public_model_for_provider(
|
||||
&self,
|
||||
@@ -470,17 +671,25 @@ impl RuntimeSnapshot {
|
||||
if !self.providers.contains(provider) {
|
||||
return false;
|
||||
}
|
||||
if self.model_mappings.contains_key(public_model.as_str()) {
|
||||
return true;
|
||||
}
|
||||
let upstream_model = self.mapped_model(public_model.as_str());
|
||||
match self.provider_models.get(provider) {
|
||||
Some(models) if !models.is_empty() => models
|
||||
.keys()
|
||||
.any(|model| model.as_str() == public_model.as_str()),
|
||||
Some(_) | None => true,
|
||||
Some(models) => models.keys().any(|model| model.as_str() == upstream_model),
|
||||
None => !self.known_provider_catalogs.contains(provider),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn contains_public_model_for_scope(
|
||||
&self,
|
||||
public_model: &PublicModelId,
|
||||
scope: &FrozenAccountScope,
|
||||
) -> bool {
|
||||
scope
|
||||
.provider_kinds()
|
||||
.iter()
|
||||
.any(|provider| self.contains_public_model_for_provider(public_model, provider))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn mapped_model(&self, requested: &str) -> String {
|
||||
let original = requested;
|
||||
@@ -506,14 +715,26 @@ impl RuntimeSnapshot {
|
||||
&self,
|
||||
public_model: &PublicModelId,
|
||||
operation: &Operation,
|
||||
account_scope: Arc<FrozenAccountScope>,
|
||||
context: &RoutingContext,
|
||||
) -> Result<RoutingPlan, RoutingError> {
|
||||
let requirements = operation.capability_requirements();
|
||||
let mut candidates = Vec::new();
|
||||
|
||||
for provider in self.providers.iter() {
|
||||
if context.required_provider.is_none() && account_scope.provider_kinds().is_empty() {
|
||||
return Err(RoutingError::EmptyAccountScope);
|
||||
}
|
||||
|
||||
let providers = context.required_provider.as_ref().map_or_else(
|
||||
|| account_scope.provider_kinds().clone(),
|
||||
|provider| BTreeSet::from([provider.clone()]),
|
||||
);
|
||||
for provider in &providers {
|
||||
if !self.providers.contains(provider) {
|
||||
continue;
|
||||
}
|
||||
if context
|
||||
.provider_kind
|
||||
.required_provider
|
||||
.as_ref()
|
||||
.is_some_and(|expected| expected != provider)
|
||||
|| context.blocked_providers.contains(provider)
|
||||
@@ -539,12 +760,14 @@ impl RuntimeSnapshot {
|
||||
};
|
||||
emulated
|
||||
}
|
||||
None if self.known_provider_catalogs.contains(provider) => continue,
|
||||
None => BTreeSet::new(),
|
||||
};
|
||||
candidates.push(ProviderCandidate {
|
||||
provider: provider.clone(),
|
||||
upstream_model,
|
||||
emulated_features,
|
||||
account_scope: Arc::clone(&account_scope),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -561,6 +784,7 @@ impl RuntimeSnapshot {
|
||||
operation: operation.kind(),
|
||||
max_attempts: NonZeroU32::new(super::MAX_REQUEST_ATTEMPTS)
|
||||
.expect("constant request attempt limit is non-zero"),
|
||||
account_scope,
|
||||
candidates: Arc::from(candidates),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2,12 +2,14 @@ use gateway_core::engine::continuation::{
|
||||
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
|
||||
};
|
||||
use gateway_core::engine::credential::ProviderAccountId;
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::ProviderKind;
|
||||
|
||||
fn pin() -> NativeContinuationPin {
|
||||
NativeContinuationPin::new(
|
||||
PreviousResponseId::new("response-private"),
|
||||
PreviousResponseId::new("upstream-private"),
|
||||
ClientApiKeyId::new("key_client").expect("valid client key"),
|
||||
ProviderKind::new("openai").expect("valid provider"),
|
||||
ProviderAccountId::new("acct_codex").expect("valid account"),
|
||||
)
|
||||
@@ -43,3 +45,8 @@ fn native_pin_should_reject_different_account() {
|
||||
&ProviderAccountId::new("acct_other").expect("valid account"),
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn native_pin_should_reject_different_client_api_key() {
|
||||
assert!(!pin().matches_client(&ClientApiKeyId::new("key_other").expect("valid client key")));
|
||||
}
|
||||
|
||||
@@ -43,8 +43,9 @@ use gateway_core::operation::{
|
||||
};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::{
|
||||
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
|
||||
RoutingPlan, RuntimeSnapshot, UpstreamModelId,
|
||||
AccountRoutingSnapshot, ClientRoutingScope, ConfigRevision, FrozenAccountScope,
|
||||
ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan,
|
||||
RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
|
||||
};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
@@ -531,7 +532,32 @@ fn plan_with_policy(
|
||||
let public_model = PublicModelId::new("gpt-5").expect("public model");
|
||||
let capabilities = ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000))
|
||||
.with_upstream_feature_validation();
|
||||
RuntimeSnapshot::new(
|
||||
let directory = Arc::new(RuntimeAccountDirectory::new(
|
||||
[
|
||||
"acct_failed",
|
||||
"acct_first",
|
||||
"acct_image",
|
||||
"acct_observation_mismatch",
|
||||
"acct_observed",
|
||||
"acct_one",
|
||||
"acct_only",
|
||||
"acct_other",
|
||||
"acct_required",
|
||||
"acct_second",
|
||||
"acct_tool",
|
||||
"acct_two",
|
||||
"acct_wrong",
|
||||
]
|
||||
.into_iter()
|
||||
.map(|id| {
|
||||
(
|
||||
ProviderAccountId::new(id).expect("account"),
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
));
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("config revision"),
|
||||
account_selection_policy,
|
||||
vec![provider.clone()],
|
||||
@@ -543,15 +569,21 @@ fn plan_with_policy(
|
||||
Vec::new(),
|
||||
)
|
||||
.expect("snapshot")
|
||||
.plan(
|
||||
&public_model,
|
||||
operation,
|
||||
&RoutingContext {
|
||||
provider_kind: Some(provider),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
.expect("routing plan")
|
||||
.with_account_directory(Arc::clone(&directory));
|
||||
snapshot
|
||||
.plan(
|
||||
&public_model,
|
||||
operation,
|
||||
Arc::new(FrozenAccountScope::new(
|
||||
directory,
|
||||
ClientRoutingScope::all_accounts(),
|
||||
)),
|
||||
&RoutingContext {
|
||||
required_provider: Some(provider),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
.expect("routing plan")
|
||||
}
|
||||
|
||||
fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest {
|
||||
@@ -561,6 +593,7 @@ fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest
|
||||
client_api_key_id: Some(client_key.clone()),
|
||||
client_api_key_ref: client_key,
|
||||
config_revision: ConfigRevision::new(1).expect("config revision"),
|
||||
routing: AccountRoutingSnapshot::all(),
|
||||
protocol: "openai".to_owned(),
|
||||
operation: operation.kind(),
|
||||
endpoint: "responses".to_owned(),
|
||||
@@ -1354,6 +1387,7 @@ fn authenticated_native_continuation_reaches_every_attempt_context() {
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-secret-id"),
|
||||
PreviousResponseId::new("provider-native-id"),
|
||||
ClientApiKeyId::new("key_client_1").expect("client key"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_one").expect("account"),
|
||||
);
|
||||
@@ -1411,6 +1445,7 @@ fn streaming_native_continuation_delivers_started_event_before_later_output() {
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-response"),
|
||||
PreviousResponseId::new("upstream-response"),
|
||||
ClientApiKeyId::new("key_client_1").expect("client key"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_one").expect("account"),
|
||||
);
|
||||
@@ -1487,6 +1522,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-secret-id"),
|
||||
PreviousResponseId::new("provider-native-id"),
|
||||
ClientApiKeyId::new("key_client_1").expect("client key"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_one").expect("account"),
|
||||
);
|
||||
@@ -1552,6 +1588,7 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-response"),
|
||||
PreviousResponseId::new("upstream-response"),
|
||||
ClientApiKeyId::new("key_client_1").expect("client key"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
original.clone(),
|
||||
);
|
||||
@@ -2498,6 +2535,7 @@ fn native_continuation_explicit_429_is_not_retried() {
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-secret-id"),
|
||||
PreviousResponseId::new("provider-native-id"),
|
||||
ClientApiKeyId::new("key_client_1").expect("client key"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_first").expect("account"),
|
||||
);
|
||||
|
||||
@@ -1,18 +1,21 @@
|
||||
use std::collections::BTreeSet;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use gateway_core::engine::credential::{
|
||||
AccountAttemptFeedback, AccountCandidate, AccountEligibilityPolicy, AccountFeedbackStats,
|
||||
AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker, AccountSelectionContext,
|
||||
AccountSelectionPolicy, AccountSelector, AccountStatus, CredentialCasUpdate,
|
||||
CredentialRevision, CredentialState, OpaqueProviderData, PlaintextCredential,
|
||||
PreferredAccountSelection, ProviderAccount, ProviderAccountId, ProviderAccountIdentity,
|
||||
ProviderAccountUpdate, QuotaEvidence, QuotaState, RotationStrategy,
|
||||
AccountAttemptFeedback, AccountCandidate, AccountConcurrencyLimit, AccountEligibilityPolicy,
|
||||
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker,
|
||||
AccountSelectionContext, AccountSelectionPolicy, AccountSelector, AccountStatus, AccountWeight,
|
||||
CredentialCasUpdate, CredentialRevision, CredentialState, OpaqueProviderData,
|
||||
PlaintextCredential, PreferredAccountSelection, ProviderAccount, ProviderAccountId,
|
||||
ProviderAccountIdentity, ProviderAccountUpdate, QuotaEvidence, QuotaState, RotationStrategy,
|
||||
};
|
||||
use gateway_core::routing::{
|
||||
ClientRoutingScope, FrozenAccountScope, ProviderKind, RuntimeAccount, RuntimeAccountDirectory,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
|
||||
fn account(id: &str) -> ProviderAccount {
|
||||
ProviderAccount::new(
|
||||
@@ -48,6 +51,34 @@ fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandida
|
||||
}
|
||||
}
|
||||
|
||||
fn weighted_candidate(id: &str, weight: u16, in_flight: u32) -> AccountCandidate {
|
||||
let mut candidate = candidate(id, in_flight, None);
|
||||
candidate.account = candidate.account.with_scheduling(
|
||||
None,
|
||||
AccountWeight::new(weight).expect("valid account weight"),
|
||||
);
|
||||
candidate
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_scheduling_should_default_and_resolve_concurrency_override() {
|
||||
let default = NonZeroU32::new(12).expect("default concurrency");
|
||||
let account = account("acct_scheduling_default");
|
||||
assert_eq!(account.concurrency_limit(), None);
|
||||
assert_eq!(account.weight(), AccountWeight::DEFAULT);
|
||||
assert_eq!(account.effective_concurrency(default), default);
|
||||
|
||||
let limit = AccountConcurrencyLimit::new(7).expect("concurrency override");
|
||||
let weight = AccountWeight::new(80).expect("weight");
|
||||
let account = account.with_scheduling(Some(limit), weight);
|
||||
assert_eq!(account.concurrency_limit(), Some(limit));
|
||||
assert_eq!(account.weight(), weight);
|
||||
assert_eq!(account.effective_concurrency(default).get(), 7);
|
||||
assert!(AccountConcurrencyLimit::new(0).is_none());
|
||||
assert!(AccountWeight::new(0).is_none());
|
||||
assert!(AccountWeight::new(101).is_none());
|
||||
}
|
||||
|
||||
fn context(strategy: RotationStrategy) -> AccountSelectionContext {
|
||||
AccountSelectionContext {
|
||||
policy: AccountSelectionPolicy::new(
|
||||
@@ -60,6 +91,7 @@ fn context(strategy: RotationStrategy) -> AccountSelectionContext {
|
||||
preferred_account: None,
|
||||
round_robin_cursor: 0,
|
||||
eligibility: AccountEligibilityPolicy::Enforce,
|
||||
account_scope: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -215,6 +247,35 @@ fn diagnostic_selection_bypasses_all_local_account_eligibility() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_should_reject_a_preferred_account_outside_the_frozen_client_scope() {
|
||||
let allowed = candidate("acct_allowed", 0, None);
|
||||
let outside = candidate("acct_outside", 0, None);
|
||||
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
|
||||
allowed.account.id().clone(),
|
||||
RuntimeAccount::new(
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
BTreeSet::new(),
|
||||
),
|
||||
)])));
|
||||
let mut context = context(RotationStrategy::Sticky);
|
||||
context.preferred_account = Some(outside.account.id().clone());
|
||||
context.account_scope = Some(Arc::new(FrozenAccountScope::new(
|
||||
directory,
|
||||
ClientRoutingScope::all_accounts(),
|
||||
)));
|
||||
|
||||
let candidates = [outside, allowed];
|
||||
let selection = AccountSelector
|
||||
.select(&candidates, &context)
|
||||
.expect("allowed account should remain eligible");
|
||||
assert_eq!(selection.candidate().account.id().as_str(), "acct_allowed");
|
||||
assert_eq!(
|
||||
selection.preferred(),
|
||||
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::OutsideClientScope)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_account_should_not_be_schedulable() {
|
||||
let account = account("acct_disabled").with_account_facts(
|
||||
@@ -287,6 +348,70 @@ fn smart_selector_should_prefer_lower_inflight_count() {
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_idle");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_should_prioritize_the_highest_weight_for_every_rotation_strategy() {
|
||||
let candidates = vec![
|
||||
weighted_candidate("acct_low_weight", 1, 0),
|
||||
weighted_candidate("acct_high_weight", 100, 2),
|
||||
];
|
||||
|
||||
for strategy in [
|
||||
RotationStrategy::Smart,
|
||||
RotationStrategy::QuotaResetPriority,
|
||||
RotationStrategy::RoundRobin,
|
||||
RotationStrategy::Sticky,
|
||||
] {
|
||||
let selected = AccountSelector
|
||||
.select(&candidates, &context(strategy))
|
||||
.expect("highest weight account should be eligible");
|
||||
assert_eq!(
|
||||
selected.candidate().account.id().as_str(),
|
||||
"acct_high_weight",
|
||||
"strategy {strategy:?} must select inside the highest weight tier"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_should_not_let_soft_affinity_bypass_a_higher_weight_account() {
|
||||
let preferred = weighted_candidate("acct_preferred", 1, 0);
|
||||
let preferred_id = preferred.account.id().clone();
|
||||
let higher = weighted_candidate("acct_higher", 10, 0);
|
||||
let mut selection = context(RotationStrategy::Smart);
|
||||
selection.preferred_account = Some(preferred_id);
|
||||
let candidates = [preferred, higher];
|
||||
|
||||
let selected = AccountSelector
|
||||
.select(&candidates, &selection)
|
||||
.expect("higher weight account should be eligible");
|
||||
|
||||
assert_eq!(selected.candidate().account.id().as_str(), "acct_higher");
|
||||
assert_eq!(
|
||||
selected.preferred(),
|
||||
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_should_fall_back_when_a_higher_weight_account_reaches_its_override() {
|
||||
let mut saturated = weighted_candidate("acct_high_weight", 100, 1);
|
||||
saturated.account = saturated.account.with_scheduling(
|
||||
Some(AccountConcurrencyLimit::new(1).expect("concurrency override")),
|
||||
AccountWeight::new(100).expect("weight"),
|
||||
);
|
||||
let fallback = weighted_candidate("acct_low_weight", 1, 0);
|
||||
let candidates = [saturated, fallback];
|
||||
|
||||
let selected = AccountSelector
|
||||
.select(&candidates, &context(RotationStrategy::Smart))
|
||||
.expect("lower weight account should be used after saturation");
|
||||
|
||||
assert_eq!(
|
||||
selected.candidate().account.id().as_str(),
|
||||
"acct_low_weight"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_selector_should_balance_signals_instead_of_using_lexicographic_load() {
|
||||
let mut healthy = candidate("acct_healthy", 1, Some(100));
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
use std::{
|
||||
collections::BTreeSet,
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
@@ -42,7 +42,8 @@ use gateway_core::operation::{GenerateRequest, Operation, OperationKind, Protoco
|
||||
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
|
||||
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
|
||||
use gateway_core::routing::{
|
||||
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RuntimeSnapshot,
|
||||
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
|
||||
ProviderModel, PublicModelId, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot,
|
||||
UpstreamModelId,
|
||||
};
|
||||
use serde_json::json;
|
||||
@@ -283,7 +284,7 @@ fn slow_circuit_store_should_time_out_and_fail_open_during_request_start() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
|
||||
fn known_catalog_should_reject_a_model_that_the_provider_did_not_publish() {
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(start_snapshot()),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
@@ -298,7 +299,7 @@ fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
|
||||
.expect("authenticated client");
|
||||
let model = "gpt-future-not-in-catalog";
|
||||
|
||||
let started = block_on(service.start(StartExecution {
|
||||
let result = block_on(service.start(StartExecution {
|
||||
client,
|
||||
public_model: PublicModelId::from_client_wire(model).expect("client model"),
|
||||
operation: start_operation_for_model(model),
|
||||
@@ -311,10 +312,80 @@ fn stale_nonempty_catalog_should_not_block_a_new_model_during_request_start() {
|
||||
user_agent: None,
|
||||
previous_response_id: None,
|
||||
},
|
||||
}))
|
||||
.expect("stale model catalog must not reject the request");
|
||||
}));
|
||||
let Err(error) = result else {
|
||||
panic!("a known provider catalog must be authoritative for model availability");
|
||||
};
|
||||
|
||||
assert!(!started.session.is_finalized());
|
||||
assert_eq!(error.kind(), GatewayErrorKind::NoAvailableProvider);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn continuation_owned_by_another_client_api_key_should_fail_closed() {
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(start_snapshot()),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::default(),
|
||||
Arc::new(UnusedAdmissions),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(RejectedContinuation::OwnershipMismatch),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
let client = service
|
||||
.authenticate("sk_start_test")
|
||||
.expect("authenticated client");
|
||||
|
||||
let result = block_on(service.start(StartExecution {
|
||||
client,
|
||||
public_model: PublicModelId::new("gpt-start").expect("public model"),
|
||||
operation: start_operation(),
|
||||
metadata: execution_metadata_with_continuation(),
|
||||
}));
|
||||
let Err(error) = result else {
|
||||
panic!("cross-client continuation must be rejected");
|
||||
};
|
||||
|
||||
assert_eq!(error.kind(), GatewayErrorKind::PolicyDenied);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_continuation_record_should_not_be_forwarded_as_an_external_handle() {
|
||||
let service = DefaultExecutionService::new(
|
||||
RuntimeSnapshotHandle::new(start_snapshot()),
|
||||
Arc::new(TrackingExecutionStore::default()),
|
||||
ProviderRegistry::default(),
|
||||
Arc::new(UnusedAdmissions),
|
||||
Arc::new(UnusedCircuits),
|
||||
Arc::new(RejectedContinuation::InvalidData),
|
||||
Arc::new(RecordingClientApiKeyUsage::default()),
|
||||
);
|
||||
let client = service
|
||||
.authenticate("sk_start_test")
|
||||
.expect("authenticated client");
|
||||
|
||||
let result = block_on(service.start(StartExecution {
|
||||
client,
|
||||
public_model: PublicModelId::new("gpt-start").expect("public model"),
|
||||
operation: start_operation(),
|
||||
metadata: execution_metadata_with_continuation(),
|
||||
}));
|
||||
let Err(error) = result else {
|
||||
panic!("invalid continuation state must be rejected");
|
||||
};
|
||||
|
||||
assert_eq!(error.kind(), GatewayErrorKind::Internal);
|
||||
}
|
||||
|
||||
fn execution_metadata_with_continuation() -> ExecutionRequestMetadata {
|
||||
ExecutionRequestMetadata {
|
||||
protocol: "openai".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
transport: ClientTransport::HttpJson,
|
||||
stream: false,
|
||||
client_ip: None,
|
||||
user_agent: None,
|
||||
previous_response_id: Some(PreviousResponseId::new("response-private")),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
@@ -526,6 +597,7 @@ struct UnusedContinuation;
|
||||
impl NativeContinuationPort for UnusedContinuation {
|
||||
fn resolve<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientApiKeyId,
|
||||
_: &'a PreviousResponseId,
|
||||
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
|
||||
Box::pin(async { Ok(None) })
|
||||
@@ -539,6 +611,43 @@ impl NativeContinuationPort for UnusedContinuation {
|
||||
}
|
||||
}
|
||||
|
||||
enum RejectedContinuation {
|
||||
OwnershipMismatch,
|
||||
InvalidData,
|
||||
}
|
||||
|
||||
impl NativeContinuationPort for RejectedContinuation {
|
||||
fn resolve<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientApiKeyId,
|
||||
_: &'a PreviousResponseId,
|
||||
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
|
||||
Box::pin(async move {
|
||||
Err(match self {
|
||||
Self::OwnershipMismatch => NativeContinuationStoreError::ownership_mismatch(),
|
||||
Self::InvalidData => NativeContinuationStoreError::invalid_data("invalid record"),
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn record<'a>(
|
||||
&'a self,
|
||||
_: NativeContinuationPin,
|
||||
) -> BoxFuture<'a, Result<(), NativeContinuationStoreError>> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
}
|
||||
|
||||
fn account_scope(provider: &ProviderKind, account_id: &str) -> Arc<FrozenAccountScope> {
|
||||
Arc::new(FrozenAccountScope::new(
|
||||
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
|
||||
ProviderAccountId::new(account_id).expect("account ID"),
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
|
||||
)]))),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
))
|
||||
}
|
||||
|
||||
fn probe_snapshot() -> RuntimeSnapshot {
|
||||
let provider = ProviderKind::new("openai").expect("provider kind");
|
||||
let capabilities =
|
||||
@@ -575,7 +684,7 @@ fn client_snapshot() -> RuntimeSnapshot {
|
||||
vec![ClientPolicy::new(
|
||||
ClientApiKeyId::new("key_usage_test").expect("client API key ID"),
|
||||
PlaintextClientApiKey::new("sk_usage_test").expect("plaintext client API key"),
|
||||
provider,
|
||||
account_scope(&provider, "acct_usage"),
|
||||
true,
|
||||
RateLimits::unlimited(),
|
||||
)],
|
||||
@@ -603,7 +712,7 @@ fn start_snapshot() -> RuntimeSnapshot {
|
||||
vec![ClientPolicy::new(
|
||||
ClientApiKeyId::new("key_start_test").expect("client API key ID"),
|
||||
PlaintextClientApiKey::new("sk_start_test").expect("plaintext client API key"),
|
||||
provider,
|
||||
account_scope(&provider, "acct_start"),
|
||||
true,
|
||||
RateLimits::unlimited(),
|
||||
)],
|
||||
|
||||
@@ -1,16 +1,25 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use gateway_core::routing::{ClientRoutingScope, FrozenAccountScope, RuntimeAccountDirectory};
|
||||
|
||||
fn plaintext(value: &str) -> PlaintextClientApiKey {
|
||||
PlaintextClientApiKey::new(value).expect("valid plaintext client key")
|
||||
}
|
||||
|
||||
fn account_scope() -> Arc<FrozenAccountScope> {
|
||||
Arc::new(FrozenAccountScope::new(
|
||||
Arc::new(RuntimeAccountDirectory::default()),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_client_key_should_be_denied() {
|
||||
let policy = ClientPolicy::new(
|
||||
ClientApiKeyId::new("key_disabled").expect("valid key ID"),
|
||||
plaintext("sk_disabled_secret"),
|
||||
ProviderKind::new("openai").expect("valid provider"),
|
||||
account_scope(),
|
||||
false,
|
||||
RateLimits::unlimited(),
|
||||
);
|
||||
@@ -23,7 +32,7 @@ fn enabled_client_key_should_be_authorized() {
|
||||
let policy = ClientPolicy::new(
|
||||
ClientApiKeyId::new("key_enabled").expect("valid key ID"),
|
||||
plaintext("sk_enabled_secret"),
|
||||
ProviderKind::new("openai").expect("valid provider"),
|
||||
account_scope(),
|
||||
true,
|
||||
RateLimits::unlimited(),
|
||||
);
|
||||
|
||||
@@ -1,15 +1,20 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
|
||||
};
|
||||
use gateway_core::operation::{
|
||||
CapabilityRequirements, Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload,
|
||||
};
|
||||
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
|
||||
use gateway_core::routing::{
|
||||
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
|
||||
RuntimeSnapshot, UpstreamModelId,
|
||||
AccountGroupId, AccountRoutingScopeKind, ClientRoutingScope, ConfigRevision,
|
||||
FrozenAccountScope, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId,
|
||||
RoutingContext, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot,
|
||||
UpstreamModelId,
|
||||
};
|
||||
|
||||
mod snapshot;
|
||||
@@ -45,11 +50,34 @@ fn operation() -> Operation {
|
||||
))
|
||||
}
|
||||
|
||||
fn account_directory() -> Arc<RuntimeAccountDirectory> {
|
||||
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
|
||||
(
|
||||
ProviderAccountId::new("acct_openai").expect("account"),
|
||||
RuntimeAccount::new(
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
BTreeSet::new(),
|
||||
),
|
||||
),
|
||||
(
|
||||
ProviderAccountId::new("acct_xai").expect("account"),
|
||||
RuntimeAccount::new(ProviderKind::new("xai").expect("provider"), BTreeSet::new()),
|
||||
),
|
||||
])))
|
||||
}
|
||||
|
||||
fn account_scope() -> Arc<FrozenAccountScope> {
|
||||
Arc::new(FrozenAccountScope::new(
|
||||
account_directory(),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
))
|
||||
}
|
||||
|
||||
fn client_policy(id: &str, plaintext: &str, enabled: bool) -> ClientPolicy {
|
||||
ClientPolicy::new(
|
||||
ClientApiKeyId::new(id).expect("client key ID"),
|
||||
PlaintextClientApiKey::new(plaintext).expect("plaintext client key"),
|
||||
ProviderKind::new("openai").expect("valid provider"),
|
||||
account_scope(),
|
||||
enabled,
|
||||
RateLimits::unlimited(),
|
||||
)
|
||||
@@ -70,6 +98,7 @@ fn snapshot() -> RuntimeSnapshot {
|
||||
Vec::new(),
|
||||
)
|
||||
.expect("snapshot")
|
||||
.with_account_directory(account_directory())
|
||||
.with_model_mappings(BTreeMap::from([
|
||||
("gpt-5.4".to_owned(), "gpt-5.5".to_owned()),
|
||||
("grok-latest".to_owned(), "grok-4.5".to_owned()),
|
||||
@@ -81,6 +110,92 @@ fn config_revision_should_reject_zero() {
|
||||
assert!(ConfigRevision::new(0).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn account_group_id_should_match_the_database_contract_exactly() {
|
||||
let valid = "grp_0123456789abcdef0123456789abcdef";
|
||||
assert_eq!(
|
||||
AccountGroupId::new(valid).expect("valid group ID").as_str(),
|
||||
valid
|
||||
);
|
||||
for invalid in [
|
||||
"0123456789abcdef0123456789abcdef",
|
||||
"grp_0123456789abcdef0123456789abcde",
|
||||
"grp_0123456789abcdef0123456789abcdef0",
|
||||
"grp_0123456789abcdef0123456789abcdeF",
|
||||
"grp_0123456789abcdef0123456789abcdeg",
|
||||
] {
|
||||
assert!(
|
||||
AccountGroupId::new(invalid).is_err(),
|
||||
"invalid group ID was accepted: {invalid}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restricted_scope_should_use_enabled_group_union_and_preserve_all_bindings_for_history() {
|
||||
let enabled =
|
||||
AccountGroupId::new("grp_11111111111111111111111111111111").expect("enabled group");
|
||||
let disabled =
|
||||
AccountGroupId::new("grp_22222222222222222222222222222222").expect("disabled group");
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let allowed_account = ProviderAccountId::new("acct_allowed").expect("account");
|
||||
let disabled_only_account = ProviderAccountId::new("acct_disabled").expect("account");
|
||||
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
|
||||
(
|
||||
allowed_account.clone(),
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::from([enabled.clone()])),
|
||||
),
|
||||
(
|
||||
disabled_only_account.clone(),
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::from([disabled.clone()])),
|
||||
),
|
||||
])));
|
||||
let scope = FrozenAccountScope::new(
|
||||
directory,
|
||||
ClientRoutingScope::restricted(
|
||||
vec![
|
||||
RoutingGroupSnapshot::new(enabled.clone(), "Enabled".to_owned()),
|
||||
RoutingGroupSnapshot::new(disabled, "Disabled".to_owned()),
|
||||
],
|
||||
BTreeSet::from([enabled]),
|
||||
BTreeSet::from([provider]),
|
||||
)
|
||||
.expect("restricted scope"),
|
||||
);
|
||||
|
||||
assert!(scope.allows(&allowed_account));
|
||||
assert!(!scope.allows(&disabled_only_account));
|
||||
let history = scope.routing_snapshot();
|
||||
assert_eq!(history.kind(), AccountRoutingScopeKind::Groups);
|
||||
assert_eq!(history.groups_snapshot().len(), 2);
|
||||
assert_eq!(history.groups_snapshot()[1].name(), "Disabled");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn restricted_scope_with_no_enabled_group_should_fail_closed() {
|
||||
let group = AccountGroupId::new("grp_33333333333333333333333333333333").expect("group");
|
||||
let account = ProviderAccountId::new("acct_group_member").expect("account");
|
||||
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
|
||||
account.clone(),
|
||||
RuntimeAccount::new(
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
BTreeSet::from([group.clone()]),
|
||||
),
|
||||
)])));
|
||||
let scope = FrozenAccountScope::new(
|
||||
directory,
|
||||
ClientRoutingScope::restricted(
|
||||
vec![RoutingGroupSnapshot::new(group, "Disabled".to_owned())],
|
||||
BTreeSet::new(),
|
||||
BTreeSet::new(),
|
||||
)
|
||||
.expect("restricted scope"),
|
||||
);
|
||||
|
||||
assert!(scope.provider_kinds().is_empty());
|
||||
assert!(!scope.allows(&account));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
@@ -98,7 +213,12 @@ fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
|
||||
let policies = snapshot.client_policies().collect::<Vec<_>>();
|
||||
assert_eq!(policies.len(), 1);
|
||||
assert_eq!(policies[0].key_id().as_str(), "key_enabled");
|
||||
assert_eq!(policies[0].provider_kind().as_str(), "openai");
|
||||
assert!(
|
||||
policies[0]
|
||||
.account_scope()
|
||||
.provider_kinds()
|
||||
.contains(&ProviderKind::new("openai").expect("provider"))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -137,8 +257,9 @@ fn selected_provider_should_use_global_model_mapping() {
|
||||
.plan(
|
||||
&PublicModelId::new("gpt-5.4").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext {
|
||||
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
|
||||
required_provider: Some(ProviderKind::new("openai").expect("provider")),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
@@ -170,26 +291,37 @@ fn cyclic_model_mapping_should_fall_back_to_the_original_name() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn global_mapping_should_apply_to_every_provider() {
|
||||
fn known_provider_catalog_missing_mapped_model_should_be_filtered() {
|
||||
let snapshot = snapshot();
|
||||
let plan = snapshot
|
||||
let error = snapshot
|
||||
.plan(
|
||||
&PublicModelId::new("gpt-5.4").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext {
|
||||
provider_kind: Some(ProviderKind::new("xai").expect("provider")),
|
||||
required_provider: Some(ProviderKind::new("xai").expect("provider")),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
.expect("transparent plan");
|
||||
.expect_err("known xAI catalog does not contain the mapped OpenAI model");
|
||||
|
||||
assert_eq!(plan.candidates()[0].provider().as_str(), "xai");
|
||||
assert_eq!(plan.candidates()[0].upstream_model().as_str(), "gpt-5.5");
|
||||
assert!(matches!(
|
||||
error,
|
||||
gateway_core::error::RoutingError::NoCapableProvider { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unmapped_model_should_pass_through_unchanged() {
|
||||
let snapshot = snapshot();
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("revision"),
|
||||
scheduling(),
|
||||
vec![ProviderKind::new("openai").expect("provider")],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
)
|
||||
.expect("snapshot")
|
||||
.with_account_directory(account_directory());
|
||||
for requested in [
|
||||
"future-openai-model".to_owned(),
|
||||
format!("future-{}", "x".repeat(512)),
|
||||
@@ -200,8 +332,9 @@ fn unmapped_model_should_pass_through_unchanged() {
|
||||
.plan(
|
||||
&PublicModelId::from_client_wire(requested.clone()).expect("client model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext {
|
||||
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
|
||||
required_provider: Some(ProviderKind::new("openai").expect("provider")),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
@@ -229,8 +362,9 @@ fn blocked_provider_should_be_filtered() {
|
||||
.plan(
|
||||
&PublicModelId::new("gpt-5.5").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext {
|
||||
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
|
||||
required_provider: Some(ProviderKind::new("openai").expect("provider")),
|
||||
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
|
||||
},
|
||||
)
|
||||
@@ -262,8 +396,9 @@ fn known_unsupported_operation_should_not_be_bypassed() {
|
||||
.plan(
|
||||
&PublicModelId::new("gpt-known-unsupported").expect("model"),
|
||||
&operation(),
|
||||
snapshot.all_account_scope(),
|
||||
&RoutingContext {
|
||||
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
|
||||
required_provider: Some(ProviderKind::new("openai").expect("provider")),
|
||||
..RoutingContext::default()
|
||||
},
|
||||
)
|
||||
@@ -316,3 +451,20 @@ fn unknown_model_should_be_hidden_when_provider_catalog_is_available() {
|
||||
&ProviderKind::new("openai").expect("provider"),
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn alias_should_only_be_available_from_a_provider_with_its_mapped_model() {
|
||||
let snapshot = snapshot();
|
||||
let alias = PublicModelId::new("grok-latest").expect("model");
|
||||
|
||||
assert!(!snapshot.contains_public_model_for_provider(
|
||||
&alias,
|
||||
&ProviderKind::new("openai").expect("provider"),
|
||||
));
|
||||
assert!(
|
||||
snapshot.contains_public_model_for_provider(
|
||||
&alias,
|
||||
&ProviderKind::new("xai").expect("provider"),
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
@@ -19,7 +19,8 @@ use gateway_core::operation::OperationKind;
|
||||
use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
|
||||
use gateway_core::routing::snapshot::{
|
||||
RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, RuntimeSnapshotHandle,
|
||||
RuntimeSnapshotPublisher, SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts,
|
||||
RuntimeSnapshotPublisher, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts,
|
||||
SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts, SnapshotProviderAccountFacts,
|
||||
SnapshotRevisionStream, SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort,
|
||||
SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh,
|
||||
};
|
||||
@@ -348,9 +349,12 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
|
||||
vec![SnapshotClientPolicyFacts::new(
|
||||
ClientApiKeyId::new("key_one").expect("key ID"),
|
||||
PlaintextClientApiKey::new("sk_test").expect("plaintext key"),
|
||||
"alpha",
|
||||
Vec::new(),
|
||||
RateLimits::unlimited(),
|
||||
)],
|
||||
Vec::<SnapshotAccountGroupFacts>::new(),
|
||||
Vec::<SnapshotProviderAccountFacts>::new(),
|
||||
Vec::<SnapshotAccountGroupMemberFacts>::new(),
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -90,6 +90,11 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
|
||||
pool.clone(),
|
||||
Some(Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>),
|
||||
)),
|
||||
Arc::new(postgres::PgAccountGroupRepository::with_runtime_state(
|
||||
pool.clone(),
|
||||
credential_leases.clone(),
|
||||
Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>,
|
||||
)),
|
||||
Arc::new(AdminAuthStoreAdapter {
|
||||
security: postgres::PgAdminSecurityAuditRepository::new(pool.clone()),
|
||||
settings: postgres::PgRuntimeSettingsRepository::new(pool.clone()),
|
||||
|
||||
@@ -0,0 +1,783 @@
|
||||
//! PostgreSQL owner for provider-neutral account groups and memberships.
|
||||
|
||||
use std::{
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
str::FromStr as _,
|
||||
sync::Arc,
|
||||
};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationContext,
|
||||
account_groups::{
|
||||
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
|
||||
AccountGroupListQuery, AccountGroupMember, AccountGroupMembers, AccountGroupMutation,
|
||||
AccountGroupPage, AccountGroupRecord, AccountGroupUsage, DeleteAccountGroup,
|
||||
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
|
||||
},
|
||||
observability::DecimalAmount,
|
||||
},
|
||||
ports::store::{AccountGroupStore, AdminStoreError, AdminStoreResult},
|
||||
};
|
||||
use gateway_core::{
|
||||
engine::credential::AccountStatus,
|
||||
provider_ports::ProviderCooldownPort,
|
||||
routing::{AccountGroupId, ProviderKind},
|
||||
};
|
||||
use sqlx::{PgPool, Postgres, QueryBuilder, Row as _, Transaction};
|
||||
|
||||
use crate::redis::CredentialLeaseRepository as _;
|
||||
use crate::{
|
||||
ConflictKind, StoreError, StoreResult, admin_revision, admin_store_error, mutation_audit,
|
||||
postgres_unavailable,
|
||||
};
|
||||
|
||||
use super::{
|
||||
PgProviderAccountRepository, PgRuntimeSettingsRepository, ProviderAccountRepository,
|
||||
RuntimeSettingsRepository, account_status_projection, append_admin_audit_event_in_transaction,
|
||||
bump_config_revision_in_transaction, load_rate_limited_until,
|
||||
};
|
||||
|
||||
const ENTITY: &str = "account group";
|
||||
|
||||
/// Account group store with transactional revision and audit ownership.
|
||||
#[derive(Clone)]
|
||||
pub struct PgAccountGroupRepository {
|
||||
pool: PgPool,
|
||||
runtime_signals: Option<crate::redis::RedisCredentialLeaseRepository>,
|
||||
cooldowns: Option<Arc<dyn ProviderCooldownPort>>,
|
||||
}
|
||||
|
||||
impl PgAccountGroupRepository {
|
||||
#[must_use]
|
||||
pub const fn new(pool: PgPool) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
runtime_signals: None,
|
||||
cooldowns: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_runtime_state(
|
||||
pool: PgPool,
|
||||
runtime_signals: crate::redis::RedisCredentialLeaseRepository,
|
||||
cooldowns: Arc<dyn ProviderCooldownPort>,
|
||||
) -> Self {
|
||||
Self {
|
||||
pool,
|
||||
runtime_signals: Some(runtime_signals),
|
||||
cooldowns: Some(cooldowns),
|
||||
}
|
||||
}
|
||||
|
||||
async fn current_revision(&self) -> AdminStoreResult<gateway_admin::model::Revision> {
|
||||
RuntimeSettingsRepository::load_runtime_settings(&PgRuntimeSettingsRepository::new(
|
||||
self.pool.clone(),
|
||||
))
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))
|
||||
.and_then(|settings| admin_revision(settings.config_revision))
|
||||
}
|
||||
|
||||
async fn required_record(&self, id: &AccountGroupId) -> AdminStoreResult<AccountGroupRecord> {
|
||||
let record = load_record(&self.pool, id.as_str())
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?
|
||||
.ok_or_else(|| not_found(id.as_str()))?;
|
||||
let mut records = vec![record];
|
||||
self.enrich_records(&mut records).await?;
|
||||
records
|
||||
.pop()
|
||||
.ok_or_else(|| invalid_admin("account group record disappeared"))
|
||||
}
|
||||
|
||||
async fn enrich_records(&self, records: &mut [AccountGroupRecord]) -> AdminStoreResult<()> {
|
||||
enrich_group_records(
|
||||
&self.pool,
|
||||
self.runtime_signals.as_ref(),
|
||||
self.cooldowns.as_deref(),
|
||||
records,
|
||||
)
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))
|
||||
}
|
||||
|
||||
async fn mutate<F>(
|
||||
&self,
|
||||
audit: super::AdminAuditEvent,
|
||||
mutation: F,
|
||||
) -> AdminStoreResult<gateway_admin::model::Revision>
|
||||
where
|
||||
F: for<'a> FnOnce(&'a mut Transaction<'_, Postgres>) -> BoxFuture<'a, StoreResult<()>>,
|
||||
{
|
||||
let mut transaction =
|
||||
self.pool.begin().await.map_err(|_| {
|
||||
admin_store_error(ENTITY, unavailable("begin account group mutation"))
|
||||
})?;
|
||||
let result = async {
|
||||
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
|
||||
mutation(&mut transaction).await?;
|
||||
append_admin_audit_event_in_transaction(&mut transaction, audit, revision).await?;
|
||||
Ok(revision)
|
||||
}
|
||||
.await;
|
||||
match result {
|
||||
Ok(revision) => {
|
||||
transaction.commit().await.map_err(|_| {
|
||||
admin_store_error(ENTITY, unavailable("commit account group mutation"))
|
||||
})?;
|
||||
admin_revision(revision)
|
||||
}
|
||||
Err(error) => {
|
||||
transaction.rollback().await.map_err(|_| {
|
||||
admin_store_error(ENTITY, unavailable("rollback account group mutation"))
|
||||
})?;
|
||||
Err(admin_store_error(ENTITY, error))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl AccountGroupStore for PgAccountGroupRepository {
|
||||
async fn list_account_groups(
|
||||
&self,
|
||||
query: AccountGroupListQuery,
|
||||
) -> AdminStoreResult<AccountGroupPage> {
|
||||
validate_page_query(&query)?;
|
||||
let total = count_groups(&self.pool, &query)
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
let offset = u64::from(query.page.saturating_sub(1)) * u64::from(query.page_size.get());
|
||||
let mut statement = group_select();
|
||||
push_group_filter(&mut statement, &query);
|
||||
statement.push(" order by g.created_at desc, g.id desc limit ");
|
||||
statement.push_bind(i64::from(query.page_size.get()));
|
||||
statement.push(" offset ");
|
||||
statement.push_bind(i64::try_from(offset).map_err(|_| invalid_admin("page is too large"))?);
|
||||
let rows = statement
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| admin_store_error(ENTITY, unavailable("list account groups")))?;
|
||||
let mut items = rows
|
||||
.iter()
|
||||
.map(group_record)
|
||||
.collect::<StoreResult<Vec<_>>>()
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
self.enrich_records(&mut items).await?;
|
||||
Ok(AccountGroupPage {
|
||||
config_revision: self.current_revision().await?,
|
||||
items,
|
||||
total,
|
||||
page: query.page,
|
||||
page_size: query.page_size.get(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn account_group_members(
|
||||
&self,
|
||||
id: &AccountGroupId,
|
||||
) -> AdminStoreResult<AccountGroupMembers> {
|
||||
self.required_record(id).await?;
|
||||
let rows = sqlx::query(
|
||||
"select a.id, a.name, a.provider_kind, a.email, a.enabled
|
||||
from account_group_accounts m
|
||||
join provider_accounts a on a.id = m.provider_account_id
|
||||
where m.account_group_id = $1
|
||||
order by a.provider_kind, lower(a.name), a.id",
|
||||
)
|
||||
.bind(id.as_str())
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| admin_store_error(ENTITY, unavailable("load account group members")))?;
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(group_member)
|
||||
.collect::<StoreResult<Vec<_>>>()
|
||||
.map_err(|error| admin_store_error(ENTITY, error))?;
|
||||
Ok(AccountGroupMembers {
|
||||
config_revision: self.current_revision().await?,
|
||||
id: id.clone(),
|
||||
total: u64::try_from(items.len())
|
||||
.map_err(|_| invalid_admin("member count overflow"))?,
|
||||
items,
|
||||
})
|
||||
}
|
||||
|
||||
async fn create_account_group(
|
||||
&self,
|
||||
command: NewAccountGroup,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
validate_group_fields(&command.name, command.description.as_deref())?;
|
||||
let id = command.id.clone();
|
||||
let audit = mutation_audit(
|
||||
context,
|
||||
"create",
|
||||
"account_group",
|
||||
id.as_str(),
|
||||
vec!["name".to_owned(), "description".to_owned()],
|
||||
);
|
||||
let revision = self
|
||||
.mutate(audit, |transaction| {
|
||||
Box::pin(async move {
|
||||
sqlx::query(
|
||||
"insert into account_groups
|
||||
(id, name, description, color, enabled, created_at, updated_at)
|
||||
values ($1, $2, $3, $4, true, now(), now())",
|
||||
)
|
||||
.bind(command.id.as_str())
|
||||
.bind(command.name)
|
||||
.bind(command.description)
|
||||
.bind(command.color.as_str())
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
|
||||
Ok(())
|
||||
})
|
||||
})
|
||||
.await?;
|
||||
Ok(AccountGroupMutation {
|
||||
config_revision: revision,
|
||||
id: id.clone(),
|
||||
record: Some(self.required_record(&id).await?),
|
||||
})
|
||||
}
|
||||
|
||||
async fn update_account_group(
|
||||
&self,
|
||||
command: UpdateAccountGroup,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
validate_group_fields(&command.name, command.description.as_deref())?;
|
||||
let id = command.id.clone();
|
||||
let audit = mutation_audit(
|
||||
context,
|
||||
"update",
|
||||
"account_group",
|
||||
id.as_str(),
|
||||
vec!["name".to_owned(), "description".to_owned()],
|
||||
);
|
||||
let revision = self
|
||||
.mutate(audit, |transaction| {
|
||||
Box::pin(async move {
|
||||
let result = sqlx::query(
|
||||
"update account_groups
|
||||
set name = $2, description = $3, color = $4, updated_at = now()
|
||||
where id = $1",
|
||||
)
|
||||
.bind(command.id.as_str())
|
||||
.bind(command.name)
|
||||
.bind(command.description)
|
||||
.bind(command.color.as_str())
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|error| map_group_write_error(error, command.id.as_str()))?;
|
||||
require_one(result.rows_affected(), command.id.as_str())
|
||||
})
|
||||
})
|
||||
.await?;
|
||||
Ok(AccountGroupMutation {
|
||||
config_revision: revision,
|
||||
id: id.clone(),
|
||||
record: Some(self.required_record(&id).await?),
|
||||
})
|
||||
}
|
||||
|
||||
async fn set_account_group_enabled(
|
||||
&self,
|
||||
command: SetAccountGroupEnabled,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
let id = command.id.clone();
|
||||
let audit = mutation_audit(
|
||||
context,
|
||||
if command.enabled { "enable" } else { "disable" },
|
||||
"account_group",
|
||||
id.as_str(),
|
||||
vec!["enabled".to_owned()],
|
||||
);
|
||||
let revision = self
|
||||
.mutate(audit, |transaction| {
|
||||
Box::pin(async move {
|
||||
let result = sqlx::query(
|
||||
"update account_groups set enabled = $2, updated_at = now() where id = $1",
|
||||
)
|
||||
.bind(command.id.as_str())
|
||||
.bind(command.enabled)
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| unavailable("set account group state"))?;
|
||||
require_one(result.rows_affected(), command.id.as_str())
|
||||
})
|
||||
})
|
||||
.await?;
|
||||
Ok(AccountGroupMutation {
|
||||
config_revision: revision,
|
||||
id: id.clone(),
|
||||
record: Some(self.required_record(&id).await?),
|
||||
})
|
||||
}
|
||||
|
||||
async fn delete_account_group(
|
||||
&self,
|
||||
command: DeleteAccountGroup,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountGroupMutation> {
|
||||
let id = command.id.clone();
|
||||
let audit = mutation_audit(context, "delete", "account_group", id.as_str(), Vec::new());
|
||||
let revision = self
|
||||
.mutate(audit, |transaction| {
|
||||
Box::pin(async move {
|
||||
let result = sqlx::query(
|
||||
"delete from account_groups g
|
||||
where g.id = $1
|
||||
and not exists (
|
||||
select 1 from client_api_key_groups k where k.account_group_id = g.id
|
||||
)",
|
||||
)
|
||||
.bind(command.id.as_str())
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| unavailable("delete account group"))?;
|
||||
if result.rows_affected() == 1 {
|
||||
return Ok(());
|
||||
}
|
||||
let exists = sqlx::query_scalar::<_, bool>(
|
||||
"select exists(select 1 from account_groups where id = $1)",
|
||||
)
|
||||
.bind(command.id.as_str())
|
||||
.fetch_one(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| unavailable("check account group delete conflict"))?;
|
||||
if exists {
|
||||
Err(conflict(command.id.as_str()))
|
||||
} else {
|
||||
Err(not_found_store(command.id.as_str()))
|
||||
}
|
||||
})
|
||||
})
|
||||
.await?;
|
||||
Ok(AccountGroupMutation {
|
||||
config_revision: revision,
|
||||
id,
|
||||
record: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn group_select() -> QueryBuilder<Postgres> {
|
||||
QueryBuilder::new(
|
||||
"select g.id, g.name, g.description, g.color, g.enabled, g.created_at, g.updated_at,
|
||||
coalesce(members.member_count, 0)::bigint as member_count,
|
||||
coalesce(keys.client_key_count, 0)::bigint as client_key_count,
|
||||
coalesce(members.provider_counts, '{}'::jsonb) as provider_counts
|
||||
from account_groups g
|
||||
left join lateral (
|
||||
select coalesce(sum(counts.provider_count), 0)::bigint as member_count,
|
||||
jsonb_object_agg(counts.provider_kind, counts.provider_count)
|
||||
as provider_counts
|
||||
from (
|
||||
select a.provider_kind, count(*)::bigint as provider_count
|
||||
from account_group_accounts gm
|
||||
join provider_accounts a on a.id = gm.provider_account_id
|
||||
where gm.account_group_id = g.id
|
||||
group by a.provider_kind
|
||||
) counts
|
||||
) members on true
|
||||
left join lateral (
|
||||
select count(*)::bigint as client_key_count
|
||||
from client_api_key_groups kg
|
||||
where kg.account_group_id = g.id
|
||||
) keys on true
|
||||
where true",
|
||||
)
|
||||
}
|
||||
|
||||
fn push_group_filter(statement: &mut QueryBuilder<Postgres>, query: &AccountGroupListQuery) {
|
||||
if let Some(search) = &query.search {
|
||||
statement.push(" and lower(g.name) like ");
|
||||
statement.push_bind(format!("%{}%", search.to_lowercase()));
|
||||
}
|
||||
if let Some(enabled) = query.enabled {
|
||||
statement.push(" and g.enabled = ");
|
||||
statement.push_bind(enabled);
|
||||
}
|
||||
}
|
||||
|
||||
async fn count_groups(pool: &PgPool, query: &AccountGroupListQuery) -> StoreResult<u64> {
|
||||
let mut statement =
|
||||
QueryBuilder::<Postgres>::new("select count(*)::bigint from account_groups g where true");
|
||||
if let Some(search) = &query.search {
|
||||
statement.push(" and lower(g.name) like ");
|
||||
statement.push_bind(format!("%{}%", search.to_lowercase()));
|
||||
}
|
||||
if let Some(enabled) = query.enabled {
|
||||
statement.push(" and g.enabled = ");
|
||||
statement.push_bind(enabled);
|
||||
}
|
||||
let count = statement
|
||||
.build_query_scalar::<i64>()
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.map_err(|_| unavailable("count account groups"))?;
|
||||
u64::try_from(count).map_err(|_| invalid("negative account group count"))
|
||||
}
|
||||
|
||||
async fn load_record(pool: &PgPool, id: &str) -> StoreResult<Option<AccountGroupRecord>> {
|
||||
let mut statement = group_select();
|
||||
statement.push(" and g.id = ");
|
||||
statement.push_bind(id.to_owned());
|
||||
statement
|
||||
.build()
|
||||
.fetch_optional(pool)
|
||||
.await
|
||||
.map_err(|_| unavailable("load account group"))?
|
||||
.as_ref()
|
||||
.map(group_record)
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn group_record(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupRecord> {
|
||||
let provider_counts: serde_json::Value = row
|
||||
.try_get("provider_counts")
|
||||
.map_err(|_| invalid("invalid provider counts"))?;
|
||||
let provider_counts = serde_json::from_value::<BTreeMap<String, u64>>(provider_counts)
|
||||
.map_err(|_| invalid("invalid provider counts"))?;
|
||||
Ok(AccountGroupRecord {
|
||||
id: AccountGroupId::new(
|
||||
row.try_get::<String, _>("id")
|
||||
.map_err(|_| invalid("invalid id"))?,
|
||||
)
|
||||
.map_err(|_| invalid("invalid id"))?,
|
||||
name: row.try_get("name").map_err(|_| invalid("invalid name"))?,
|
||||
description: row
|
||||
.try_get("description")
|
||||
.map_err(|_| invalid("invalid description"))?,
|
||||
color: AccountGroupColor::parse(
|
||||
row.try_get::<String, _>("color")
|
||||
.map_err(|_| invalid("invalid color"))?
|
||||
.as_str(),
|
||||
)
|
||||
.ok_or_else(|| invalid("invalid color"))?,
|
||||
enabled: row
|
||||
.try_get("enabled")
|
||||
.map_err(|_| invalid("invalid enabled"))?,
|
||||
member_count: count_value(row, "member_count")?,
|
||||
provider_counts,
|
||||
client_key_count: count_value(row, "client_key_count")?,
|
||||
account_summary: AccountGroupAccountSummary {
|
||||
available: 0,
|
||||
limited: 0,
|
||||
total: 0,
|
||||
},
|
||||
capacity: AccountGroupCapacity {
|
||||
used_slots: None,
|
||||
total_slots: 0,
|
||||
},
|
||||
usage: AccountGroupUsage {
|
||||
today_usd: DecimalAmount::from_str("0").map_err(|_| invalid("invalid zero cost"))?,
|
||||
total_usd: DecimalAmount::from_str("0").map_err(|_| invalid("invalid zero cost"))?,
|
||||
},
|
||||
created_at: row
|
||||
.try_get("created_at")
|
||||
.map_err(|_| invalid("invalid created_at"))?,
|
||||
updated_at: row
|
||||
.try_get("updated_at")
|
||||
.map_err(|_| invalid("invalid updated_at"))?,
|
||||
})
|
||||
}
|
||||
|
||||
async fn enrich_group_records(
|
||||
pool: &PgPool,
|
||||
runtime_signals: Option<&crate::redis::RedisCredentialLeaseRepository>,
|
||||
cooldowns: Option<&dyn ProviderCooldownPort>,
|
||||
records: &mut [AccountGroupRecord],
|
||||
) -> StoreResult<()> {
|
||||
if records.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let settings = RuntimeSettingsRepository::load_runtime_settings(
|
||||
&PgRuntimeSettingsRepository::new(pool.clone()),
|
||||
)
|
||||
.await?;
|
||||
let accounts = PgProviderAccountRepository::new(pool.clone())
|
||||
.list_provider_accounts(None, true)
|
||||
.await?;
|
||||
let now = chrono::Utc::now();
|
||||
let cooldowns = load_rate_limited_until(cooldowns, &accounts, now.into()).await;
|
||||
let mut status_by_account = BTreeMap::new();
|
||||
let mut capacity_by_account = BTreeMap::new();
|
||||
for account in &accounts {
|
||||
let status =
|
||||
account_status_projection(account, now.into(), cooldowns.get(&account.id).copied())
|
||||
.status;
|
||||
status_by_account.insert(account.id.as_str(), status);
|
||||
capacity_by_account.insert(
|
||||
account.id.as_str(),
|
||||
u64::from(
|
||||
account
|
||||
.concurrency_limit
|
||||
.map_or(settings.max_concurrent_per_account, |limit| limit.get()),
|
||||
),
|
||||
);
|
||||
}
|
||||
let group_ids = records
|
||||
.iter()
|
||||
.map(|record| record.id.as_str().to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
let memberships = sqlx::query(
|
||||
"select account_group_id, provider_account_id
|
||||
from account_group_accounts
|
||||
where account_group_id = any($1::text[])",
|
||||
)
|
||||
.bind(&group_ids)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| unavailable("load account group availability"))?;
|
||||
let mut accounts_by_group = BTreeMap::<String, Vec<String>>::new();
|
||||
for row in memberships {
|
||||
accounts_by_group
|
||||
.entry(
|
||||
row.try_get("account_group_id")
|
||||
.map_err(|_| invalid("invalid group ID"))?,
|
||||
)
|
||||
.or_default()
|
||||
.push(
|
||||
row.try_get("provider_account_id")
|
||||
.map_err(|_| invalid("invalid account ID"))?,
|
||||
);
|
||||
}
|
||||
let normal_ids = accounts_by_group
|
||||
.values()
|
||||
.flatten()
|
||||
.filter(|account_id| {
|
||||
status_by_account.get(account_id.as_str()) == Some(&AccountStatus::Normal)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<BTreeSet<_>>()
|
||||
.into_iter()
|
||||
.collect::<Vec<_>>();
|
||||
let used_by_account = if normal_ids.is_empty() {
|
||||
Some(BTreeMap::new())
|
||||
} else if let Some(runtime_signals) = runtime_signals {
|
||||
runtime_signals
|
||||
.credential_runtime_signals(&normal_ids)
|
||||
.await
|
||||
.ok()
|
||||
.map(|signals| {
|
||||
signals
|
||||
.into_iter()
|
||||
.map(|signal| (signal.resource_id, u64::from(signal.in_flight)))
|
||||
.collect::<BTreeMap<_, _>>()
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let costs = group_costs(pool, &group_ids).await?;
|
||||
for record in records {
|
||||
let account_ids = accounts_by_group
|
||||
.get(record.id.as_str())
|
||||
.map_or(&[][..], Vec::as_slice);
|
||||
let available = account_ids
|
||||
.iter()
|
||||
.filter(|id| status_by_account.get(id.as_str()) == Some(&AccountStatus::Normal))
|
||||
.count() as u64;
|
||||
let total = u64::try_from(account_ids.len())
|
||||
.map_err(|_| invalid("group account count overflow"))?;
|
||||
record.account_summary = AccountGroupAccountSummary {
|
||||
available,
|
||||
limited: total.saturating_sub(available),
|
||||
total,
|
||||
};
|
||||
record.capacity = AccountGroupCapacity {
|
||||
used_slots: if available == 0 {
|
||||
Some(0)
|
||||
} else {
|
||||
used_by_account.as_ref().map(|used| {
|
||||
account_ids.iter().fold(0_u64, |sum, id| {
|
||||
sum.saturating_add(used.get(id).copied().unwrap_or(0))
|
||||
})
|
||||
})
|
||||
},
|
||||
total_slots: account_ids.iter().fold(0_u64, |total, account_id| {
|
||||
if status_by_account.get(account_id.as_str()) == Some(&AccountStatus::Normal) {
|
||||
total.saturating_add(
|
||||
capacity_by_account
|
||||
.get(account_id.as_str())
|
||||
.copied()
|
||||
.unwrap_or(0),
|
||||
)
|
||||
} else {
|
||||
total
|
||||
}
|
||||
}),
|
||||
};
|
||||
if let Some(usage) = costs.get(record.id.as_str()) {
|
||||
record.usage = usage.clone();
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn group_costs(
|
||||
pool: &PgPool,
|
||||
group_ids: &[String],
|
||||
) -> StoreResult<BTreeMap<String, AccountGroupUsage>> {
|
||||
let rows = sqlx::query(
|
||||
"select gm.account_group_id,
|
||||
coalesce(sum(mr.cost_amount) filter (
|
||||
where mr.started_at >= date_trunc('day', now() at time zone 'Asia/Shanghai')
|
||||
at time zone 'Asia/Shanghai'
|
||||
), 0)::text as today_usd,
|
||||
coalesce(sum(mr.cost_amount), 0)::text as total_usd
|
||||
from account_group_accounts gm
|
||||
left join model_requests mr
|
||||
on mr.provider_account_ref = gm.provider_account_id
|
||||
and mr.outcome = 'succeeded'
|
||||
and mr.downstream_committed_at is not null
|
||||
and mr.client_status_code between 200 and 399
|
||||
and mr.cost_currency = 'USD'
|
||||
and mr.cost_amount is not null
|
||||
where gm.account_group_id = any($1::text[])
|
||||
group by gm.account_group_id",
|
||||
)
|
||||
.bind(group_ids)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.map_err(|_| unavailable("load account group costs"))?;
|
||||
rows.into_iter()
|
||||
.map(|row| {
|
||||
let group_id = row
|
||||
.try_get("account_group_id")
|
||||
.map_err(|_| invalid("invalid group ID"))?;
|
||||
let today = row
|
||||
.try_get::<String, _>("today_usd")
|
||||
.map_err(|_| invalid("invalid today cost"))?;
|
||||
let total = row
|
||||
.try_get::<String, _>("total_usd")
|
||||
.map_err(|_| invalid("invalid total cost"))?;
|
||||
Ok((
|
||||
group_id,
|
||||
AccountGroupUsage {
|
||||
today_usd: DecimalAmount::from_str(&today)
|
||||
.map_err(|_| invalid("invalid today cost"))?,
|
||||
total_usd: DecimalAmount::from_str(&total)
|
||||
.map_err(|_| invalid("invalid total cost"))?,
|
||||
},
|
||||
))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn group_member(row: &sqlx::postgres::PgRow) -> StoreResult<AccountGroupMember> {
|
||||
Ok(AccountGroupMember {
|
||||
id: row
|
||||
.try_get("id")
|
||||
.map_err(|_| invalid("invalid member id"))?,
|
||||
name: row
|
||||
.try_get("name")
|
||||
.map_err(|_| invalid("invalid member name"))?,
|
||||
provider_kind: ProviderKind::new(
|
||||
row.try_get::<String, _>("provider_kind")
|
||||
.map_err(|_| invalid("invalid member provider"))?,
|
||||
)
|
||||
.map_err(|_| invalid("invalid member provider"))?,
|
||||
email: row
|
||||
.try_get("email")
|
||||
.map_err(|_| invalid("invalid member email"))?,
|
||||
enabled: row
|
||||
.try_get("enabled")
|
||||
.map_err(|_| invalid("invalid member enabled"))?,
|
||||
})
|
||||
}
|
||||
|
||||
fn count_value(row: &sqlx::postgres::PgRow, field: &str) -> StoreResult<u64> {
|
||||
u64::try_from(
|
||||
row.try_get::<i64, _>(field)
|
||||
.map_err(|_| invalid("invalid count"))?,
|
||||
)
|
||||
.map_err(|_| invalid("negative count"))
|
||||
}
|
||||
|
||||
fn validate_page_query(query: &AccountGroupListQuery) -> AdminStoreResult<()> {
|
||||
if query.page == 0 {
|
||||
return Err(invalid_admin("page must be positive"));
|
||||
}
|
||||
if query.search.as_deref().is_some_and(|search| {
|
||||
search.trim().is_empty() || search.len() > 256 || search.chars().any(char::is_control)
|
||||
}) {
|
||||
return Err(invalid_admin("invalid search"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_group_fields(name: &str, description: Option<&str>) -> AdminStoreResult<()> {
|
||||
if name.trim() != name
|
||||
|| name.is_empty()
|
||||
|| name.chars().count() > 100
|
||||
|| name.chars().any(char::is_control)
|
||||
|| description
|
||||
.is_some_and(|value| value.len() > 4096 || value.chars().any(char::is_control))
|
||||
{
|
||||
return Err(invalid_admin("invalid account group fields"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn require_one(rows: u64, id: &str) -> StoreResult<()> {
|
||||
if rows == 1 {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(not_found_store(id))
|
||||
}
|
||||
}
|
||||
|
||||
fn map_group_write_error(error: sqlx::Error, id: &str) -> StoreError {
|
||||
if error
|
||||
.as_database_error()
|
||||
.is_some_and(|database| database.is_unique_violation())
|
||||
{
|
||||
conflict(id)
|
||||
} else {
|
||||
unavailable("write account group")
|
||||
}
|
||||
}
|
||||
|
||||
fn not_found(id: &str) -> AdminStoreError {
|
||||
admin_store_error(ENTITY, not_found_store(id))
|
||||
}
|
||||
|
||||
fn not_found_store(id: &str) -> StoreError {
|
||||
StoreError::NotFound {
|
||||
entity: ENTITY,
|
||||
id: id.to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
fn conflict(id: &str) -> StoreError {
|
||||
StoreError::Conflict {
|
||||
entity: ENTITY,
|
||||
id: id.to_owned(),
|
||||
kind: ConflictKind::InvalidTransition,
|
||||
}
|
||||
}
|
||||
|
||||
fn invalid_admin(message: &str) -> AdminStoreError {
|
||||
admin_store_error(ENTITY, invalid(message))
|
||||
}
|
||||
|
||||
fn invalid(message: &str) -> StoreError {
|
||||
StoreError::InvalidData {
|
||||
entity: ENTITY,
|
||||
message: message.to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
fn unavailable(message: &'static str) -> StoreError {
|
||||
postgres_unavailable(message)
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
//! 明文 `client_api_keys` 的 PostgreSQL owner。
|
||||
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
collections::{BTreeMap, BTreeSet},
|
||||
fmt,
|
||||
sync::{
|
||||
Arc, Mutex,
|
||||
@@ -15,6 +15,7 @@ use chrono::{DateTime, Utc};
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationContext,
|
||||
account_groups::AccountGroupRef as AdminAccountGroupRef,
|
||||
client_keys::{
|
||||
ClientKeyCursor as AdminClientKeyCursor,
|
||||
ClientKeyCursorValue as AdminClientKeyCursorValue,
|
||||
@@ -27,11 +28,12 @@ use gateway_admin::{
|
||||
},
|
||||
ports::store::{AdminStoreResult, ClientKeyStore},
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use gateway_core::{
|
||||
engine::execution::ClientApiKeyUsageSink,
|
||||
policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits},
|
||||
routing::{AccountGroupId, ProviderKind},
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use sqlx::{PgPool, Postgres, QueryBuilder, Transaction};
|
||||
|
||||
use crate::{
|
||||
@@ -49,7 +51,7 @@ const CLIENT_API_KEY_LAST_USED_FLUSH_DELAY: Duration = Duration::from_secs(1);
|
||||
pub struct ClientApiKeySnapshot {
|
||||
pub id: ClientApiKeyId,
|
||||
pub plaintext_key: PlaintextClientApiKey,
|
||||
pub provider_kind: String,
|
||||
pub group_ids: Vec<AccountGroupId>,
|
||||
pub limits: RateLimits,
|
||||
}
|
||||
|
||||
@@ -57,7 +59,7 @@ impl ClientApiKeySnapshot {
|
||||
pub(crate) fn from_persisted(
|
||||
id: String,
|
||||
key: String,
|
||||
provider_kind: String,
|
||||
group_ids: Vec<String>,
|
||||
max_concurrency: i64,
|
||||
requests_per_minute: i64,
|
||||
) -> StoreResult<Self> {
|
||||
@@ -65,7 +67,12 @@ impl ClientApiKeySnapshot {
|
||||
id: ClientApiKeyId::new(id).map_err(|_| invalid("persisted key ID is invalid"))?,
|
||||
plaintext_key: PlaintextClientApiKey::new(key)
|
||||
.map_err(|_| invalid("persisted plaintext key is invalid"))?,
|
||||
provider_kind,
|
||||
group_ids: group_ids
|
||||
.into_iter()
|
||||
.map(|id| {
|
||||
AccountGroupId::new(id).map_err(|_| invalid("persisted group ID is invalid"))
|
||||
})
|
||||
.collect::<StoreResult<Vec<_>>>()?,
|
||||
limits: RateLimits {
|
||||
max_concurrency: to_u64(max_concurrency)?,
|
||||
requests_per_minute: to_u64(requests_per_minute)?,
|
||||
@@ -78,7 +85,6 @@ impl ClientApiKeySnapshot {
|
||||
pub struct ClientApiKeySecret {
|
||||
pub id: String,
|
||||
pub key: String,
|
||||
pub provider_kind: String,
|
||||
pub enabled: bool,
|
||||
pub max_concurrency: u64,
|
||||
pub requests_per_minute: u64,
|
||||
@@ -100,7 +106,8 @@ pub struct ClientApiKeyRecord {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub provider_kind: String,
|
||||
pub groups: Vec<ClientApiKeyGroupRecord>,
|
||||
pub provider_kinds: Vec<String>,
|
||||
pub prefix: String,
|
||||
pub enabled: bool,
|
||||
pub max_concurrency: u64,
|
||||
@@ -110,6 +117,14 @@ pub struct ClientApiKeyRecord {
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
|
||||
pub struct ClientApiKeyGroupRecord {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub color: String,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub enum ClientApiKeySortField {
|
||||
Name,
|
||||
@@ -250,7 +265,7 @@ pub struct NewClientApiKey {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub provider_kind: String,
|
||||
pub group_ids: Vec<String>,
|
||||
pub key: String,
|
||||
pub max_concurrency: u64,
|
||||
pub requests_per_minute: u64,
|
||||
@@ -271,28 +286,17 @@ impl NewClientApiKey {
|
||||
pub fn validate(&self) -> StoreResult<()> {
|
||||
require_nonempty(ENTITY, "id", &self.id)?;
|
||||
require_nonempty(ENTITY, "name", &self.name)?;
|
||||
require_nonempty(ENTITY, "provider_kind", &self.provider_kind)?;
|
||||
validate_group_ids(&self.group_ids)?;
|
||||
validate_key(&self.key)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpdateClientApiKey {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub provider_kind: String,
|
||||
pub enabled: bool,
|
||||
pub max_concurrency: u64,
|
||||
pub requests_per_minute: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpdateClientApiKeyDetails {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub provider_kind: String,
|
||||
pub group_ids: Vec<String>,
|
||||
pub max_concurrency: u64,
|
||||
pub requests_per_minute: u64,
|
||||
}
|
||||
@@ -301,7 +305,7 @@ impl UpdateClientApiKeyDetails {
|
||||
pub fn validate(&self) -> StoreResult<()> {
|
||||
require_nonempty(ENTITY, "id", &self.id)?;
|
||||
require_nonempty(ENTITY, "name", &self.name)?;
|
||||
require_nonempty(ENTITY, "provider_kind", &self.provider_kind)?;
|
||||
validate_group_ids(&self.group_ids)?;
|
||||
to_i64(self.max_concurrency)?;
|
||||
to_i64(self.requests_per_minute)?;
|
||||
Ok(())
|
||||
@@ -316,7 +320,6 @@ pub trait ClientApiKeyRepository: Send + Sync {
|
||||
) -> StoreResult<ClientApiKeyPage>;
|
||||
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>>;
|
||||
async fn reveal_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeySecret>>;
|
||||
async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult<bool>;
|
||||
async fn touch_client_api_keys(
|
||||
&self,
|
||||
touched_at: &BTreeMap<String, DateTime<Utc>>,
|
||||
@@ -344,9 +347,38 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
query.validate()?;
|
||||
let total = count_client_api_keys(&self.pool, query.search.as_deref()).await?;
|
||||
let mut statement = QueryBuilder::<Postgres>::new(
|
||||
"select id, name, label, provider_kind, left(key, 10) as prefix, enabled, max_concurrency,
|
||||
requests_per_minute, last_used_at, created_at, updated_at
|
||||
from client_api_keys where true",
|
||||
"select k.id, k.name, k.label, left(k.key, 10) as prefix, k.enabled,
|
||||
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
|
||||
k.updated_at,
|
||||
coalesce(groups.groups, '[]'::jsonb) as groups,
|
||||
case
|
||||
when groups.binding_count = 0 then coalesce(
|
||||
(select array_agg(distinct a.provider_kind order by a.provider_kind)
|
||||
from provider_accounts a),
|
||||
'{}'
|
||||
)
|
||||
else coalesce(groups.provider_kinds, '{}')
|
||||
end as provider_kinds
|
||||
from client_api_keys k
|
||||
left join lateral (
|
||||
select
|
||||
(select count(*)::bigint
|
||||
from client_api_key_groups kg
|
||||
where kg.client_api_key_id = k.id) as binding_count,
|
||||
(select jsonb_agg(jsonb_build_object(
|
||||
'id', g.id, 'name', g.name, 'color', g.color, 'enabled', g.enabled
|
||||
) order by g.id)
|
||||
from client_api_key_groups kg
|
||||
join account_groups g on g.id = kg.account_group_id
|
||||
where kg.client_api_key_id = k.id) as groups,
|
||||
(select array_agg(distinct a.provider_kind order by a.provider_kind)
|
||||
from client_api_key_groups kg
|
||||
join account_groups g on g.id = kg.account_group_id and g.enabled
|
||||
join account_group_accounts gm on gm.account_group_id = g.id
|
||||
join provider_accounts a on a.id = gm.provider_account_id
|
||||
where kg.client_api_key_id = k.id) as provider_kinds
|
||||
) groups on true
|
||||
where true",
|
||||
);
|
||||
push_client_key_search(&mut statement, query.search.as_deref());
|
||||
if let Some(cursor) = &query.cursor {
|
||||
@@ -356,12 +388,12 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
statement.push(" limit ");
|
||||
statement.push_bind(i64::from(query.page_size) + 1);
|
||||
let rows = statement
|
||||
.build_query_as::<ClientRecordRow>()
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("list client API keys"))?;
|
||||
let mut items = rows
|
||||
.into_iter()
|
||||
.iter()
|
||||
.map(client_record_from_row)
|
||||
.collect::<StoreResult<Vec<_>>>()?;
|
||||
let has_more = items.len() > usize::from(query.page_size);
|
||||
@@ -384,8 +416,8 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
|
||||
async fn reveal_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeySecret>> {
|
||||
require_nonempty(ENTITY, "id", id)?;
|
||||
sqlx::query_as::<_, (String, String, String, bool, i64, i64)>(
|
||||
"select id, key, provider_kind, enabled, max_concurrency, requests_per_minute
|
||||
sqlx::query_as::<_, (String, String, bool, i64, i64)>(
|
||||
"select id, key, enabled, max_concurrency, requests_per_minute
|
||||
from client_api_keys where id = $1",
|
||||
)
|
||||
.bind(id)
|
||||
@@ -398,41 +430,48 @@ impl ClientApiKeyRepository for PgClientApiKeyRepository {
|
||||
|
||||
async fn get_client_api_key(&self, id: &str) -> StoreResult<Option<ClientApiKeyRecord>> {
|
||||
require_nonempty(ENTITY, "id", id)?;
|
||||
sqlx::query_as::<_, ClientRecordRow>(
|
||||
"select id, name, label, provider_kind, left(key, 10) as prefix, enabled, max_concurrency,
|
||||
requests_per_minute, last_used_at, created_at, updated_at
|
||||
from client_api_keys where id = $1",
|
||||
sqlx::query(
|
||||
"select k.id, k.name, k.label, left(k.key, 10) as prefix, k.enabled,
|
||||
k.max_concurrency, k.requests_per_minute, k.last_used_at, k.created_at,
|
||||
k.updated_at, coalesce(groups.groups, '[]'::jsonb) as groups,
|
||||
case
|
||||
when groups.binding_count = 0 then coalesce(
|
||||
(select array_agg(distinct a.provider_kind order by a.provider_kind)
|
||||
from provider_accounts a),
|
||||
'{}'
|
||||
)
|
||||
else coalesce(groups.provider_kinds, '{}')
|
||||
end as provider_kinds
|
||||
from client_api_keys k
|
||||
left join lateral (
|
||||
select
|
||||
(select count(*)::bigint
|
||||
from client_api_key_groups kg
|
||||
where kg.client_api_key_id = k.id) as binding_count,
|
||||
(select jsonb_agg(jsonb_build_object(
|
||||
'id', g.id, 'name', g.name, 'color', g.color, 'enabled', g.enabled
|
||||
) order by g.id)
|
||||
from client_api_key_groups kg
|
||||
join account_groups g on g.id = kg.account_group_id
|
||||
where kg.client_api_key_id = k.id) as groups,
|
||||
(select array_agg(distinct a.provider_kind order by a.provider_kind)
|
||||
from client_api_key_groups kg
|
||||
join account_groups g on g.id = kg.account_group_id and g.enabled
|
||||
join account_group_accounts gm on gm.account_group_id = g.id
|
||||
join provider_accounts a on a.id = gm.provider_account_id
|
||||
where kg.client_api_key_id = k.id) as provider_kinds
|
||||
) groups on true
|
||||
where k.id = $1",
|
||||
)
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("get client API key"))?
|
||||
.as_ref()
|
||||
.map(client_record_from_row)
|
||||
.transpose()
|
||||
}
|
||||
|
||||
async fn update_client_api_key(&self, key: UpdateClientApiKey) -> StoreResult<bool> {
|
||||
require_nonempty(ENTITY, "id", &key.id)?;
|
||||
require_nonempty(ENTITY, "name", &key.name)?;
|
||||
let result = sqlx::query(
|
||||
"update client_api_keys
|
||||
set name = $2, label = $3, provider_kind = $4, enabled = $5, max_concurrency = $6,
|
||||
requests_per_minute = $7, updated_at = now()
|
||||
where id = $1",
|
||||
)
|
||||
.bind(key.id)
|
||||
.bind(key.name)
|
||||
.bind(key.label)
|
||||
.bind(key.provider_kind)
|
||||
.bind(key.enabled)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
.bind(to_i64(key.requests_per_minute)?)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("update client API key"))?;
|
||||
Ok(result.rows_affected() == 1)
|
||||
}
|
||||
|
||||
async fn touch_client_api_keys(
|
||||
&self,
|
||||
touched_at: &BTreeMap<String, DateTime<Utc>>,
|
||||
@@ -637,7 +676,11 @@ impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
id: id.as_str().to_owned(),
|
||||
name: command.name,
|
||||
label: command.label,
|
||||
provider_kind: command.provider_kind.as_str().to_owned(),
|
||||
group_ids: command
|
||||
.group_ids
|
||||
.iter()
|
||||
.map(|id| id.as_str().to_owned())
|
||||
.collect(),
|
||||
key: command.plaintext,
|
||||
max_concurrency: command.limits.max_concurrency,
|
||||
requests_per_minute: command.limits.requests_per_minute,
|
||||
@@ -650,7 +693,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
[
|
||||
"name",
|
||||
"label",
|
||||
"provider_kind",
|
||||
"group_ids",
|
||||
"key",
|
||||
"enabled",
|
||||
"max_concurrency",
|
||||
@@ -679,7 +722,11 @@ impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
id: id.as_str().to_owned(),
|
||||
name: command.name,
|
||||
label: command.label,
|
||||
provider_kind: command.provider_kind.as_str().to_owned(),
|
||||
group_ids: command
|
||||
.group_ids
|
||||
.iter()
|
||||
.map(|id| id.as_str().to_owned())
|
||||
.collect(),
|
||||
max_concurrency: command.limits.max_concurrency,
|
||||
requests_per_minute: command.limits.requests_per_minute,
|
||||
},
|
||||
@@ -691,7 +738,7 @@ impl ClientKeyStore for PgAdminClientKeyStore {
|
||||
[
|
||||
"name",
|
||||
"label",
|
||||
"provider_kind",
|
||||
"group_ids",
|
||||
"max_concurrency",
|
||||
"requests_per_minute",
|
||||
]
|
||||
@@ -828,8 +875,30 @@ fn admin_client_key_record(record: ClientApiKeyRecord) -> AdminStoreResult<Admin
|
||||
.map_err(|_| admin_store_error(ENTITY, invalid("invalid client key id")))?,
|
||||
name: record.name,
|
||||
label: record.label,
|
||||
provider_kind: ProviderKind::new(record.provider_kind)
|
||||
.map_err(|_| admin_store_error(ENTITY, invalid("invalid provider kind")))?,
|
||||
groups: record
|
||||
.groups
|
||||
.into_iter()
|
||||
.map(|group| {
|
||||
Ok(AdminAccountGroupRef {
|
||||
id: AccountGroupId::new(group.id)
|
||||
.map_err(|_| admin_store_error(ENTITY, invalid("invalid group id")))?,
|
||||
name: group.name,
|
||||
color: gateway_admin::model::account_groups::AccountGroupColor::parse(
|
||||
&group.color,
|
||||
)
|
||||
.ok_or_else(|| admin_store_error(ENTITY, invalid("invalid group color")))?,
|
||||
enabled: group.enabled,
|
||||
})
|
||||
})
|
||||
.collect::<AdminStoreResult<Vec<_>>>()?,
|
||||
provider_kinds: record
|
||||
.provider_kinds
|
||||
.into_iter()
|
||||
.map(|provider| {
|
||||
ProviderKind::new(provider)
|
||||
.map_err(|_| admin_store_error(ENTITY, invalid("invalid provider kind")))
|
||||
})
|
||||
.collect::<AdminStoreResult<Vec<_>>>()?,
|
||||
prefix: record.prefix,
|
||||
enabled: record.enabled,
|
||||
limits: RateLimits {
|
||||
@@ -849,20 +918,20 @@ pub(crate) async fn insert_client_api_key_in_transaction(
|
||||
key.validate()?;
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, label, provider_kind, key, enabled, max_concurrency, requests_per_minute,
|
||||
id, name, label, key, enabled, max_concurrency, requests_per_minute,
|
||||
last_used_at, created_at, updated_at
|
||||
) values ($1, $2, $3, $4, $5, true, $6, $7, null, now(), now())",
|
||||
) values ($1, $2, $3, $4, true, $5, $6, null, now(), now())",
|
||||
)
|
||||
.bind(&key.id)
|
||||
.bind(&key.name)
|
||||
.bind(&key.label)
|
||||
.bind(&key.provider_kind)
|
||||
.bind(&key.key)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
.bind(to_i64(key.requests_per_minute)?)
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("insert client API key in transaction"))?;
|
||||
replace_client_api_key_groups_in_transaction(transaction, &key.id, &key.group_ids).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -873,20 +942,20 @@ pub(crate) async fn update_client_api_key_in_transaction(
|
||||
key.validate()?;
|
||||
let result = sqlx::query(
|
||||
"update client_api_keys
|
||||
set name = $2, label = $3, provider_kind = $4, max_concurrency = $5,
|
||||
requests_per_minute = $6, updated_at = now()
|
||||
set name = $2, label = $3, max_concurrency = $4,
|
||||
requests_per_minute = $5, updated_at = now()
|
||||
where id = $1",
|
||||
)
|
||||
.bind(&key.id)
|
||||
.bind(&key.name)
|
||||
.bind(&key.label)
|
||||
.bind(&key.provider_kind)
|
||||
.bind(to_i64(key.max_concurrency)?)
|
||||
.bind(to_i64(key.requests_per_minute)?)
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("update client API key in transaction"))?;
|
||||
require_changed(result.rows_affected(), &key.id)
|
||||
require_changed(result.rows_affected(), &key.id)?;
|
||||
replace_client_api_key_groups_in_transaction(transaction, &key.id, &key.group_ids).await
|
||||
}
|
||||
|
||||
pub(crate) async fn set_client_api_key_enabled_in_transaction(
|
||||
@@ -929,6 +998,59 @@ fn require_changed(rows_affected: u64, id: &str) -> StoreResult<()> {
|
||||
}
|
||||
}
|
||||
|
||||
async fn replace_client_api_key_groups_in_transaction(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
key_id: &str,
|
||||
group_ids: &[String],
|
||||
) -> StoreResult<()> {
|
||||
validate_group_ids(group_ids)?;
|
||||
if !group_ids.is_empty() {
|
||||
let count = sqlx::query_scalar::<_, i64>(
|
||||
"select count(*)::bigint from account_groups where id = any($1::text[])",
|
||||
)
|
||||
.bind(group_ids)
|
||||
.fetch_one(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("validate client API key groups"))?;
|
||||
if usize::try_from(count).ok() != Some(group_ids.len()) {
|
||||
return Err(StoreError::NotFound {
|
||||
entity: "account group",
|
||||
id: "client API key group selection".to_owned(),
|
||||
});
|
||||
}
|
||||
}
|
||||
sqlx::query("delete from client_api_key_groups where client_api_key_id = $1")
|
||||
.bind(key_id)
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("delete client API key groups"))?;
|
||||
if !group_ids.is_empty() {
|
||||
sqlx::query(
|
||||
"insert into client_api_key_groups
|
||||
(client_api_key_id, account_group_id, created_at)
|
||||
select $1, group_id, now() from unnest($2::text[]) group_id",
|
||||
)
|
||||
.bind(key_id)
|
||||
.bind(group_ids)
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("insert client API key groups"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_group_ids(group_ids: &[String]) -> StoreResult<()> {
|
||||
if group_ids.len() > 1000
|
||||
|| group_ids.iter().collect::<BTreeSet<_>>().len() != group_ids.len()
|
||||
|| group_ids
|
||||
.iter()
|
||||
.any(|id| AccountGroupId::new(id.clone()).is_err())
|
||||
{
|
||||
return Err(invalid("group IDs are invalid or duplicated"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_key(key: &str) -> StoreResult<()> {
|
||||
let valid = key.len() == KEY_LENGTH
|
||||
&& key.starts_with("sk_")
|
||||
@@ -945,45 +1067,53 @@ fn validate_key(key: &str) -> StoreResult<()> {
|
||||
}
|
||||
|
||||
fn client_secret_from_row(
|
||||
row: (String, String, String, bool, i64, i64),
|
||||
row: (String, String, bool, i64, i64),
|
||||
) -> StoreResult<ClientApiKeySecret> {
|
||||
Ok(ClientApiKeySecret {
|
||||
id: row.0,
|
||||
key: row.1,
|
||||
provider_kind: row.2,
|
||||
enabled: row.3,
|
||||
max_concurrency: to_u64(row.4)?,
|
||||
requests_per_minute: to_u64(row.5)?,
|
||||
enabled: row.2,
|
||||
max_concurrency: to_u64(row.3)?,
|
||||
requests_per_minute: to_u64(row.4)?,
|
||||
})
|
||||
}
|
||||
|
||||
type ClientRecordRow = (
|
||||
String,
|
||||
String,
|
||||
Option<String>,
|
||||
String,
|
||||
String,
|
||||
bool,
|
||||
i64,
|
||||
i64,
|
||||
Option<DateTime<Utc>>,
|
||||
DateTime<Utc>,
|
||||
DateTime<Utc>,
|
||||
);
|
||||
|
||||
fn client_record_from_row(row: ClientRecordRow) -> StoreResult<ClientApiKeyRecord> {
|
||||
fn client_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<ClientApiKeyRecord> {
|
||||
use sqlx::Row as _;
|
||||
let groups: serde_json::Value = row
|
||||
.try_get("groups")
|
||||
.map_err(|_| invalid("invalid groups"))?;
|
||||
Ok(ClientApiKeyRecord {
|
||||
id: row.0,
|
||||
name: row.1,
|
||||
label: row.2,
|
||||
provider_kind: row.3,
|
||||
prefix: row.4,
|
||||
enabled: row.5,
|
||||
max_concurrency: to_u64(row.6)?,
|
||||
requests_per_minute: to_u64(row.7)?,
|
||||
last_used_at: row.8,
|
||||
created_at: row.9,
|
||||
updated_at: row.10,
|
||||
id: row.try_get("id").map_err(|_| invalid("invalid id"))?,
|
||||
name: row.try_get("name").map_err(|_| invalid("invalid name"))?,
|
||||
label: row.try_get("label").map_err(|_| invalid("invalid label"))?,
|
||||
groups: serde_json::from_value(groups).map_err(|_| invalid("invalid groups"))?,
|
||||
provider_kinds: row
|
||||
.try_get("provider_kinds")
|
||||
.map_err(|_| invalid("invalid provider kinds"))?,
|
||||
prefix: row
|
||||
.try_get("prefix")
|
||||
.map_err(|_| invalid("invalid prefix"))?,
|
||||
enabled: row
|
||||
.try_get("enabled")
|
||||
.map_err(|_| invalid("invalid enabled"))?,
|
||||
max_concurrency: to_u64(
|
||||
row.try_get("max_concurrency")
|
||||
.map_err(|_| invalid("invalid max concurrency"))?,
|
||||
)?,
|
||||
requests_per_minute: to_u64(
|
||||
row.try_get("requests_per_minute")
|
||||
.map_err(|_| invalid("invalid requests per minute"))?,
|
||||
)?,
|
||||
last_used_at: row
|
||||
.try_get("last_used_at")
|
||||
.map_err(|_| invalid("invalid last used at"))?,
|
||||
created_at: row
|
||||
.try_get("created_at")
|
||||
.map_err(|_| invalid("invalid created at"))?,
|
||||
updated_at: row
|
||||
.try_get("updated_at")
|
||||
.map_err(|_| invalid("invalid updated at"))?,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ use gateway_core::engine::{
|
||||
UpstreamSendState as CoreUpstreamSendState,
|
||||
};
|
||||
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
|
||||
use gateway_core::routing::{AccountRoutingScopeKind, AccountRoutingSnapshot};
|
||||
|
||||
use crate::{
|
||||
ConflictKind, DecimalAmount, StoreError, StoreResult, postgres_unavailable, require_nonempty,
|
||||
@@ -83,6 +84,9 @@ pub struct NewModelRequest {
|
||||
pub client_api_key_id: Option<String>,
|
||||
pub client_api_key_ref: String,
|
||||
pub config_revision: u64,
|
||||
pub routing_scope: String,
|
||||
pub routing_group_refs: Vec<String>,
|
||||
pub routing_group_names_snapshot: Value,
|
||||
pub protocol: String,
|
||||
pub operation: String,
|
||||
pub endpoint: String,
|
||||
@@ -109,6 +113,11 @@ impl NewModelRequest {
|
||||
require_nonempty(ENTITY, "endpoint", &self.endpoint)?;
|
||||
require_nonempty(ENTITY, "client_transport", &self.client_transport)?;
|
||||
require_nonempty(ENTITY, "requested_model_id", &self.requested_model_id)?;
|
||||
validate_routing_snapshot(
|
||||
&self.routing_scope,
|
||||
&self.routing_group_refs,
|
||||
&self.routing_group_names_snapshot,
|
||||
)?;
|
||||
if self.config_revision == 0 || self.started_at > self.deadline_at {
|
||||
return Err(invalid(
|
||||
"revision and deadline violate the frozen constraints",
|
||||
@@ -362,13 +371,15 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
sqlx::query(
|
||||
"insert into model_requests (
|
||||
id, client_api_key_id, client_api_key_ref, config_revision, protocol,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot,
|
||||
operation, endpoint, client_transport, requested_model_id,
|
||||
client_ip, user_agent, reasoning_effort,
|
||||
reasoning_preset, request_kind, subagent_kind, compact,
|
||||
image_generation_requested, started_at, deadline_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10::inet, $11, $12,
|
||||
$13, $14, $15, $16, $17, $18, $19
|
||||
$1, $2, $3, $4, $5, $6, $7, $8,
|
||||
$9, $10, $11, $12, $13::inet, $14, $15,
|
||||
$16, $17, $18, $19, $20, $21, $22
|
||||
)",
|
||||
)
|
||||
.bind(request.id)
|
||||
@@ -376,6 +387,9 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
.bind(request.client_api_key_ref)
|
||||
.bind(to_i64(request.config_revision, "config_revision")?)
|
||||
.bind(request.protocol)
|
||||
.bind(request.routing_scope)
|
||||
.bind(request.routing_group_refs)
|
||||
.bind(sqlx::types::Json(request.routing_group_names_snapshot))
|
||||
.bind(request.operation)
|
||||
.bind(request.endpoint)
|
||||
.bind(request.client_transport)
|
||||
@@ -409,6 +423,7 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
sqlx::query(
|
||||
"insert into model_requests (
|
||||
id, client_api_key_id, client_api_key_ref, config_revision, protocol,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot,
|
||||
operation, endpoint, client_transport, requested_model_id,
|
||||
client_ip, user_agent, reasoning_effort,
|
||||
reasoning_preset, request_kind, subagent_kind, compact,
|
||||
@@ -419,13 +434,14 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
upstream_model_id, upstream_transport, http_version,
|
||||
attempt_count, upstream_send_state
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, $10::inet, $11, $12,
|
||||
$13, $14, $15, $16, $17, $18, $19,
|
||||
$20, $21, $22,
|
||||
(select name from provider_accounts where id = $21),
|
||||
(select email from provider_accounts where id = $21),
|
||||
(select authentication_kind from provider_accounts where id = $21),
|
||||
$23, $24, $25, 1, 'not_sent'
|
||||
$1, $2, $3, $4, $5, $6, $7, $8,
|
||||
$9, $10, $11, $12, $13::inet, $14, $15,
|
||||
$16, $17, $18, $19, $20, $21, $22,
|
||||
$23, $24, $25,
|
||||
(select name from provider_accounts where id = $24),
|
||||
(select email from provider_accounts where id = $24),
|
||||
(select authentication_kind from provider_accounts where id = $24),
|
||||
$26, $27, $28, 1, 'not_sent'
|
||||
)",
|
||||
)
|
||||
.bind(request.id)
|
||||
@@ -433,6 +449,9 @@ impl ModelRequestRepository for PgExecutionStore {
|
||||
.bind(request.client_api_key_ref)
|
||||
.bind(to_i64(request.config_revision, "config_revision")?)
|
||||
.bind(request.protocol)
|
||||
.bind(request.routing_scope)
|
||||
.bind(request.routing_group_refs)
|
||||
.bind(sqlx::types::Json(request.routing_group_names_snapshot))
|
||||
.bind(request.operation)
|
||||
.bind(request.endpoint)
|
||||
.bind(request.client_transport)
|
||||
@@ -1045,6 +1064,8 @@ fn core_store_error(error: StoreError) -> CoreStoreError {
|
||||
}
|
||||
|
||||
fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
|
||||
let (routing_scope, routing_group_refs, routing_group_names_snapshot) =
|
||||
routing_snapshot_row(&request.routing);
|
||||
NewModelRequest {
|
||||
id: request.id.as_str().to_owned(),
|
||||
client_api_key_id: request
|
||||
@@ -1053,6 +1074,9 @@ fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
|
||||
.map(|id| id.as_str().to_owned()),
|
||||
client_api_key_ref: request.client_api_key_ref.as_str().to_owned(),
|
||||
config_revision: request.config_revision.get(),
|
||||
routing_scope,
|
||||
routing_group_refs,
|
||||
routing_group_names_snapshot,
|
||||
protocol: request.protocol,
|
||||
operation: request.operation.as_str().to_owned(),
|
||||
endpoint: request.endpoint,
|
||||
@@ -1071,6 +1095,48 @@ fn new_model_request_row(request: CoreNewModelRequest) -> NewModelRequest {
|
||||
}
|
||||
}
|
||||
|
||||
fn routing_snapshot_row(snapshot: &AccountRoutingSnapshot) -> (String, Vec<String>, Value) {
|
||||
let groups = snapshot.groups_snapshot();
|
||||
(
|
||||
snapshot.kind().as_str().to_owned(),
|
||||
groups
|
||||
.iter()
|
||||
.map(|group| group.id().as_str().to_owned())
|
||||
.collect(),
|
||||
Value::Array(
|
||||
groups
|
||||
.iter()
|
||||
.map(|group| Value::String(group.name().to_owned()))
|
||||
.collect(),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
fn validate_routing_snapshot(scope: &str, refs: &[String], names: &Value) -> StoreResult<()> {
|
||||
let Some(names) = names.as_array() else {
|
||||
return Err(invalid("routing group names snapshot must be an array"));
|
||||
};
|
||||
let valid = match scope {
|
||||
value if value == AccountRoutingScopeKind::All.as_str() => {
|
||||
refs.is_empty() && names.is_empty()
|
||||
}
|
||||
value if value == AccountRoutingScopeKind::Groups.as_str() => {
|
||||
!refs.is_empty()
|
||||
&& refs.len() == names.len()
|
||||
&& refs.iter().all(|id| !id.is_empty())
|
||||
&& names
|
||||
.iter()
|
||||
.all(|name| name.as_str().is_some_and(|name| !name.is_empty()))
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
if valid {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(invalid("invalid routing scope snapshot"))
|
||||
}
|
||||
}
|
||||
|
||||
fn attempt_start_row(attempt: CoreAttemptRecord) -> ModelRequestAttemptStart {
|
||||
ModelRequestAttemptStart {
|
||||
model_request_id: attempt.request_id.as_str().to_owned(),
|
||||
|
||||
@@ -7,6 +7,7 @@ use crate::{
|
||||
Revision, StoreBackend, StoreError, StorePoolConfig, StoreResult, postgres_unavailable,
|
||||
};
|
||||
|
||||
mod account_groups;
|
||||
mod admin_security_audit;
|
||||
mod admission_recovery;
|
||||
mod backup;
|
||||
@@ -20,6 +21,7 @@ mod retention;
|
||||
mod runtime_settings;
|
||||
mod snapshot;
|
||||
|
||||
pub use account_groups::*;
|
||||
pub use admin_security_audit::*;
|
||||
pub use admission_recovery::*;
|
||||
pub use backup::*;
|
||||
@@ -251,7 +253,7 @@ impl PgControlPlaneRepository {
|
||||
async fn apply_targeted_mutation(
|
||||
&self,
|
||||
mutation: ControlPlaneMutation,
|
||||
audit: AdminAuditEvent,
|
||||
mut audit: AdminAuditEvent,
|
||||
) -> StoreResult<Revision> {
|
||||
let mut transaction = self
|
||||
.pool
|
||||
@@ -262,9 +264,34 @@ impl PgControlPlaneRepository {
|
||||
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
|
||||
match mutation {
|
||||
ControlPlaneMutation::CreateClientApiKey(key) => {
|
||||
audit.changed_fields.push(if key.group_ids.is_empty() {
|
||||
"routing_scope:all".to_owned()
|
||||
} else {
|
||||
"routing_scope:groups".to_owned()
|
||||
});
|
||||
insert_client_api_key_in_transaction(&mut transaction, &key).await?;
|
||||
}
|
||||
ControlPlaneMutation::UpdateClientApiKey(key) => {
|
||||
let previously_restricted = sqlx::query_scalar::<_, bool>(
|
||||
"select exists(
|
||||
select 1 from client_api_key_groups where client_api_key_id = $1
|
||||
)",
|
||||
)
|
||||
.bind(&key.id)
|
||||
.fetch_one(&mut *transaction)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
postgres_unavailable("load client API key routing scope for audit")
|
||||
})?;
|
||||
if previously_restricted && key.group_ids.is_empty() {
|
||||
audit
|
||||
.changed_fields
|
||||
.push("routing_scope:groups->all".to_owned());
|
||||
} else if !previously_restricted && !key.group_ids.is_empty() {
|
||||
audit
|
||||
.changed_fields
|
||||
.push("routing_scope:all->groups".to_owned());
|
||||
}
|
||||
update_client_api_key_in_transaction(&mut transaction, &key).await?;
|
||||
}
|
||||
ControlPlaneMutation::SetClientApiKeyEnabled { id, enabled } => {
|
||||
|
||||
@@ -6,8 +6,8 @@ use std::sync::Arc;
|
||||
use super::*;
|
||||
|
||||
use crate::postgres::{
|
||||
PgProviderAccountRepository, ProviderAccountRepository, account_status_projection,
|
||||
load_rate_limited_until,
|
||||
PgProviderAccountRepository, ProviderAccountRepository, ProviderAccountSummary,
|
||||
account_status_projection, load_rate_limited_until,
|
||||
};
|
||||
|
||||
use crate::redis::{CredentialLeaseRepository as _, RedisCredentialLeaseRepository};
|
||||
@@ -28,7 +28,7 @@ impl PgObservabilityRepository {
|
||||
async fn account_status_snapshot(
|
||||
&self,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<(ProviderAccountMetrics, Vec<String>)> {
|
||||
) -> StoreResult<(ProviderAccountMetrics, Vec<ProviderAccountSummary>)> {
|
||||
let accounts = PgProviderAccountRepository::new(self.pool.clone())
|
||||
.list_provider_accounts(None, true)
|
||||
.await?;
|
||||
@@ -39,7 +39,7 @@ impl PgObservabilityRepository {
|
||||
total: u64::try_from(accounts.len()).unwrap_or(u64::MAX),
|
||||
..ProviderAccountMetrics::default()
|
||||
};
|
||||
let mut normal_account_ids = Vec::new();
|
||||
let mut normal_accounts = Vec::new();
|
||||
for account in &accounts {
|
||||
let projection = account_status_projection(
|
||||
account,
|
||||
@@ -49,7 +49,7 @@ impl PgObservabilityRepository {
|
||||
match projection.status {
|
||||
gateway_core::engine::credential::AccountStatus::Normal => {
|
||||
metrics.normal = metrics.normal.saturating_add(1);
|
||||
normal_account_ids.push(account.id.clone());
|
||||
normal_accounts.push(account.clone());
|
||||
}
|
||||
gateway_core::engine::credential::AccountStatus::QuotaExhausted => {
|
||||
metrics.quota_exhausted = metrics.quota_exhausted.saturating_add(1);
|
||||
@@ -65,7 +65,7 @@ impl PgObservabilityRepository {
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok((metrics, normal_account_ids))
|
||||
Ok((metrics, normal_accounts))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -230,21 +230,42 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
|
||||
&self,
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> AdminStoreResult<Option<admin_observability::DashboardRuntimeSlots>> {
|
||||
let (_, normal_account_ids) = self
|
||||
let (_, normal_accounts) = self
|
||||
.repository
|
||||
.account_status_snapshot(observed_at)
|
||||
.await
|
||||
.map_err(observability_error)?;
|
||||
let active_accounts = u64::try_from(normal_account_ids.len())
|
||||
.map_err(|_| observability_error(invalid("active account count overflows u64")))?;
|
||||
if normal_account_ids.is_empty() {
|
||||
let inherited_accounts = u64::try_from(
|
||||
normal_accounts
|
||||
.iter()
|
||||
.filter(|account| account.concurrency_limit.is_none())
|
||||
.count(),
|
||||
)
|
||||
.map_err(|_| observability_error(invalid("inherited account count overflows u64")))?;
|
||||
let overridden_slots = normal_accounts.iter().fold(0_u64, |total, account| {
|
||||
total.saturating_add(
|
||||
account
|
||||
.concurrency_limit
|
||||
.map_or(0, |limit| u64::from(limit.get())),
|
||||
)
|
||||
});
|
||||
if normal_accounts.is_empty() {
|
||||
return Ok(Some(admin_observability::DashboardRuntimeSlots {
|
||||
active_accounts,
|
||||
inherited_accounts,
|
||||
overridden_slots,
|
||||
used_slots: Some(0),
|
||||
}));
|
||||
}
|
||||
let normal_account_ids = normal_accounts
|
||||
.iter()
|
||||
.map(|account| account.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let Some(runtime_signals) = &self.runtime_signals else {
|
||||
return Ok(None);
|
||||
return Ok(Some(admin_observability::DashboardRuntimeSlots {
|
||||
inherited_accounts,
|
||||
overridden_slots,
|
||||
used_slots: None,
|
||||
}));
|
||||
};
|
||||
let signals = match runtime_signals
|
||||
.credential_runtime_signals(&normal_account_ids)
|
||||
@@ -253,7 +274,8 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
|
||||
Ok(signals) => signals,
|
||||
Err(_) => {
|
||||
return Ok(Some(admin_observability::DashboardRuntimeSlots {
|
||||
active_accounts,
|
||||
inherited_accounts,
|
||||
overridden_slots,
|
||||
used_slots: None,
|
||||
}));
|
||||
}
|
||||
@@ -262,7 +284,8 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
|
||||
total.saturating_add(u64::from(signal.in_flight))
|
||||
});
|
||||
Ok(Some(admin_observability::DashboardRuntimeSlots {
|
||||
active_accounts,
|
||||
inherited_accounts,
|
||||
overridden_slots,
|
||||
used_slots: Some(used_slots),
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -436,6 +436,9 @@ pub(crate) fn admin_usage_record(
|
||||
id: record.id,
|
||||
client_api_key_ref: record.client_api_key_ref,
|
||||
config_revision: record.config_revision,
|
||||
routing_scope: record.routing_scope,
|
||||
routing_group_refs: record.routing_group_refs,
|
||||
routing_group_names_snapshot: record.routing_group_names_snapshot,
|
||||
protocol: record.protocol,
|
||||
operation: record.operation,
|
||||
endpoint: record.endpoint,
|
||||
@@ -655,6 +658,13 @@ pub(crate) fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> StoreResult<
|
||||
id: get(row, "id")?,
|
||||
client_api_key_ref: get(row, "client_api_key_ref")?,
|
||||
config_revision: unsigned(row, "config_revision")?,
|
||||
routing_scope: get(row, "routing_scope")?,
|
||||
routing_group_refs: get(row, "routing_group_refs")?,
|
||||
routing_group_names_snapshot: get::<sqlx::types::Json<Vec<String>>>(
|
||||
row,
|
||||
"routing_group_names_snapshot",
|
||||
)?
|
||||
.0,
|
||||
protocol: get(row, "protocol")?,
|
||||
operation: get(row, "operation")?,
|
||||
endpoint: get(row, "endpoint")?,
|
||||
|
||||
@@ -497,6 +497,9 @@ pub struct UsageRecord {
|
||||
pub id: String,
|
||||
pub client_api_key_ref: String,
|
||||
pub config_revision: u64,
|
||||
pub routing_scope: String,
|
||||
pub routing_group_refs: Vec<String>,
|
||||
pub routing_group_names_snapshot: Vec<String>,
|
||||
pub protocol: String,
|
||||
pub operation: String,
|
||||
pub endpoint: String,
|
||||
|
||||
@@ -103,7 +103,9 @@ pub(crate) fn push_completed_usage_fact_filter(query: &mut QueryBuilder<Postgres
|
||||
}
|
||||
|
||||
pub(crate) const USAGE_RECORD_SELECT: &str =
|
||||
"select mr.id, mr.client_api_key_ref, mr.config_revision, mr.protocol, mr.operation,
|
||||
"select mr.id, mr.client_api_key_ref, mr.config_revision,
|
||||
mr.routing_scope, mr.routing_group_refs, mr.routing_group_names_snapshot,
|
||||
mr.protocol, mr.operation,
|
||||
mr.endpoint, mr.client_transport, mr.requested_model_id,
|
||||
mr.provider_kind, mr.provider_account_ref,
|
||||
mr.provider_account_name_snapshot as provider_account_name,
|
||||
|
||||
@@ -231,6 +231,111 @@ impl PgAdminAccountStore {
|
||||
credential_revision: Some(admin_revision(rotation.credential_revision)?),
|
||||
})
|
||||
}
|
||||
|
||||
async fn account_groups_by_account(
|
||||
&self,
|
||||
account_ids: &[String],
|
||||
) -> AdminStoreResult<BTreeMap<String, Vec<AccountGroupRef>>> {
|
||||
if account_ids.is_empty() {
|
||||
return Ok(BTreeMap::new());
|
||||
}
|
||||
let rows = sqlx::query_as::<_, (String, String, String, String, bool)>(
|
||||
"select m.provider_account_id, g.id, g.name, g.color, g.enabled
|
||||
from account_group_accounts m
|
||||
join account_groups g on g.id = m.account_group_id
|
||||
where m.provider_account_id = any($1::text[])
|
||||
order by m.provider_account_id, g.id",
|
||||
)
|
||||
.bind(account_ids)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
admin_store_error(
|
||||
ENTITY,
|
||||
postgres_unavailable("load account group references"),
|
||||
)
|
||||
})?;
|
||||
let mut groups = BTreeMap::<String, Vec<AccountGroupRef>>::new();
|
||||
for (account_id, group_id, name, color, enabled) in rows {
|
||||
let id = AccountGroupId::new(group_id).map_err(|_| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"persisted account group ID is invalid",
|
||||
)
|
||||
})?;
|
||||
groups.entry(account_id).or_default().push(AccountGroupRef {
|
||||
id,
|
||||
name,
|
||||
color: gateway_admin::model::account_groups::AccountGroupColor::parse(&color)
|
||||
.ok_or_else(|| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"persisted account group color is invalid",
|
||||
)
|
||||
})?,
|
||||
enabled,
|
||||
});
|
||||
}
|
||||
Ok(groups)
|
||||
}
|
||||
|
||||
async fn account_ids_matching_group(
|
||||
&self,
|
||||
filter: Option<&AccountGroupFilter>,
|
||||
) -> AdminStoreResult<Option<BTreeSet<String>>> {
|
||||
let Some(filter) = filter else {
|
||||
return Ok(None);
|
||||
};
|
||||
let rows = match filter {
|
||||
AccountGroupFilter::Group(group_id) => {
|
||||
let exists = sqlx::query_scalar::<_, bool>(
|
||||
"select exists(select 1 from account_groups where id = $1)",
|
||||
)
|
||||
.bind(group_id.as_str())
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_err(|_| {
|
||||
admin_store_error(
|
||||
ENTITY,
|
||||
postgres_unavailable("validate account group filter"),
|
||||
)
|
||||
})?;
|
||||
if !exists {
|
||||
return Err(admin_store_error(
|
||||
ENTITY,
|
||||
StoreError::NotFound {
|
||||
entity: "account group",
|
||||
id: group_id.as_str().to_owned(),
|
||||
},
|
||||
));
|
||||
}
|
||||
sqlx::query_scalar::<_, String>(
|
||||
"select provider_account_id from account_group_accounts
|
||||
where account_group_id = $1",
|
||||
)
|
||||
.bind(group_id.as_str())
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
}
|
||||
AccountGroupFilter::Ungrouped => {
|
||||
sqlx::query_scalar::<_, String>(
|
||||
"select a.id from provider_accounts a
|
||||
where not exists (
|
||||
select 1 from account_group_accounts m
|
||||
where m.provider_account_id = a.id
|
||||
)",
|
||||
)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
}
|
||||
}
|
||||
.map_err(|_| {
|
||||
admin_store_error(ENTITY, postgres_unavailable("load account group filter"))
|
||||
})?;
|
||||
Ok(Some(rows.into_iter().collect()))
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -252,18 +357,23 @@ impl AccountStore for PgAdminAccountStore {
|
||||
let rate_limited_until =
|
||||
load_rate_limited_until(self.cooldowns.as_deref(), &accounts, now.into()).await;
|
||||
let summary = admin_account_summary(&accounts, now, &rate_limited_until);
|
||||
let group_matches = self
|
||||
.account_ids_matching_group(query.group_filter.as_ref())
|
||||
.await?;
|
||||
let mut items = accounts
|
||||
.into_iter()
|
||||
.filter_map(|account| {
|
||||
let until = rate_limited_until.get(&account.id).copied();
|
||||
let projection = account_status_projection(&account, now.into(), until);
|
||||
account_matches_admin_query(&account, projection.status, &query).then_some(
|
||||
AdminAccountListItem {
|
||||
account,
|
||||
projection,
|
||||
usage: None,
|
||||
},
|
||||
)
|
||||
(account_matches_admin_query(&account, projection.status, &query)
|
||||
&& group_matches
|
||||
.as_ref()
|
||||
.is_none_or(|matches| matches.contains(&account.id)))
|
||||
.then_some(AdminAccountListItem {
|
||||
account,
|
||||
projection,
|
||||
usage: None,
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
@@ -298,9 +408,20 @@ impl AccountStore for PgAdminAccountStore {
|
||||
.into_iter()
|
||||
.skip(offset)
|
||||
.take(page_size)
|
||||
.collect::<Vec<_>>();
|
||||
let item_ids = items
|
||||
.iter()
|
||||
.map(|item| item.account.id.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let mut groups_by_account = self.account_groups_by_account(&item_ids).await?;
|
||||
let items = items
|
||||
.into_iter()
|
||||
.map(|item| {
|
||||
let account_id = item.account.id.clone();
|
||||
let mut account = admin_account_record(item.account)?;
|
||||
account.groups = groups_by_account.remove(&account_id).unwrap_or_default();
|
||||
Ok(AccountPageItem {
|
||||
account: admin_account_record(item.account)?,
|
||||
account,
|
||||
projection: item.projection,
|
||||
})
|
||||
})
|
||||
@@ -332,8 +453,14 @@ impl AccountStore for PgAdminAccountStore {
|
||||
.get(account_id)
|
||||
.copied();
|
||||
let projection = account_status_projection(&record.summary, now.into(), rate_limited_until);
|
||||
let account_id = record.summary.id.clone();
|
||||
let mut groups = self
|
||||
.account_groups_by_account(std::slice::from_ref(&account_id))
|
||||
.await?;
|
||||
let mut account = admin_account_record(record.summary)?;
|
||||
account.groups = groups.remove(&account_id).unwrap_or_default();
|
||||
Ok(Some(AccountPageItem {
|
||||
account: admin_account_record(record.summary)?,
|
||||
account,
|
||||
projection,
|
||||
}))
|
||||
}
|
||||
@@ -572,29 +699,99 @@ impl AccountStore for PgAdminAccountStore {
|
||||
.await
|
||||
}
|
||||
|
||||
async fn set_account_enabled(
|
||||
async fn update_account(
|
||||
&self,
|
||||
command: SetAccountEnabled,
|
||||
command: UpdateAccount,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AdminRevision> {
|
||||
let scope = self.required_scope(&command.account_id).await?;
|
||||
let action = if command.enabled { "enable" } else { "disable" };
|
||||
self.accounts
|
||||
.set_provider_account_enabled_admin(SetProviderAccountEnabled {
|
||||
scope,
|
||||
account_id: command.account_id.clone(),
|
||||
) -> AdminStoreResult<AccountUpdateResult> {
|
||||
let account_id = CoreProviderAccountId::new(command.account_id.clone()).map_err(|_| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"invalid provider account ID",
|
||||
)
|
||||
})?;
|
||||
let config_revision = self
|
||||
.accounts
|
||||
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
|
||||
account_ids: vec![command.account_id.clone()],
|
||||
enabled: command.enabled,
|
||||
concurrency_limit: command.concurrency_limit,
|
||||
weight: command.weight,
|
||||
group_ids: command.group_ids,
|
||||
audit: mutation_audit(
|
||||
context,
|
||||
action,
|
||||
"update",
|
||||
"provider_account",
|
||||
&command.account_id,
|
||||
vec!["enabled".to_owned()],
|
||||
vec![
|
||||
"enabled".to_owned(),
|
||||
"concurrency_limit".to_owned(),
|
||||
"weight".to_owned(),
|
||||
"groups".to_owned(),
|
||||
],
|
||||
),
|
||||
})
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))
|
||||
.and_then(admin_revision)
|
||||
.and_then(admin_revision)?;
|
||||
Ok(AccountUpdateResult {
|
||||
config_revision,
|
||||
account_id,
|
||||
})
|
||||
}
|
||||
|
||||
async fn batch_update_accounts(
|
||||
&self,
|
||||
command: BatchUpdateAccounts,
|
||||
context: &MutationContext,
|
||||
) -> AdminStoreResult<AccountsUpdateResult> {
|
||||
let account_ids = command
|
||||
.account_ids
|
||||
.iter()
|
||||
.map(|id| {
|
||||
CoreProviderAccountId::new(id.clone()).map_err(|_| {
|
||||
AdminStoreError::new(
|
||||
AdminStoreErrorKind::Invalid,
|
||||
ENTITY,
|
||||
"invalid provider account ID",
|
||||
)
|
||||
})
|
||||
})
|
||||
.collect::<AdminStoreResult<Vec<_>>>()?;
|
||||
let audit_target = if command.account_ids.len() == 1 {
|
||||
command.account_ids[0].clone()
|
||||
} else {
|
||||
"provider_accounts".to_owned()
|
||||
};
|
||||
let config_revision = self
|
||||
.accounts
|
||||
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
|
||||
account_ids: command.account_ids,
|
||||
enabled: command.enabled,
|
||||
concurrency_limit: command.concurrency_limit,
|
||||
weight: command.weight,
|
||||
group_ids: command.group_ids,
|
||||
audit: mutation_audit(
|
||||
context,
|
||||
"batch_update",
|
||||
"provider_account",
|
||||
&audit_target,
|
||||
vec![
|
||||
"enabled".to_owned(),
|
||||
"concurrency_limit".to_owned(),
|
||||
"weight".to_owned(),
|
||||
"groups".to_owned(),
|
||||
],
|
||||
),
|
||||
})
|
||||
.await
|
||||
.map_err(|error| admin_store_error(ENTITY, error))
|
||||
.and_then(admin_revision)?;
|
||||
Ok(AccountsUpdateResult {
|
||||
config_revision,
|
||||
account_ids,
|
||||
})
|
||||
}
|
||||
|
||||
async fn delete_accounts(
|
||||
|
||||
@@ -31,6 +31,8 @@ impl ProviderAccountStore for PgProviderAccountRepository {
|
||||
.map(DateTime::<Utc>::from),
|
||||
next_refresh_at: account.account.next_refresh_at().map(DateTime::<Utc>::from),
|
||||
enabled: account.account.enabled(),
|
||||
concurrency_limit: account.account.concurrency_limit(),
|
||||
weight: account.account.weight(),
|
||||
credential_state: account.account.credential_state(),
|
||||
credential_observed_at: Utc::now(),
|
||||
})
|
||||
|
||||
@@ -146,6 +146,7 @@ pub(crate) fn admin_account_record(
|
||||
"persisted Provider kind is invalid",
|
||||
)
|
||||
})?,
|
||||
groups: Vec::new(),
|
||||
name: summary.name,
|
||||
email: summary.email,
|
||||
upstream_user_id: summary.upstream_user_id,
|
||||
@@ -157,6 +158,8 @@ pub(crate) fn admin_account_record(
|
||||
access_token_expires_at: summary.access_token_expires_at,
|
||||
next_refresh_at: summary.next_refresh_at,
|
||||
enabled: summary.enabled,
|
||||
concurrency_limit: summary.concurrency_limit,
|
||||
weight: summary.weight,
|
||||
credential_state: summary.credential_state,
|
||||
credential_observed_at: summary.credential_observed_at,
|
||||
quota: summary.quota,
|
||||
@@ -184,6 +187,8 @@ pub(crate) fn prepared_account(
|
||||
access_token_expires_at: credential.access_token_expires_at,
|
||||
next_refresh_at: credential.next_refresh_at,
|
||||
enabled: credential.enabled,
|
||||
concurrency_limit: None,
|
||||
weight: AccountWeight::DEFAULT,
|
||||
credential_state: credential.credential_state,
|
||||
credential_observed_at: credential.credential_observed_at,
|
||||
})
|
||||
|
||||
@@ -11,12 +11,15 @@ use chrono::{DateTime, TimeDelta, Utc};
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationContext, Revision as AdminRevision,
|
||||
account_groups::AccountGroupRef,
|
||||
accounts::{
|
||||
AccountCost, AccountListQuery as AdminAccountListQuery, AccountModelUsage, AccountPage,
|
||||
AccountPageItem, AccountRecord, AccountRequestBucket, AccountSort as AdminAccountSort,
|
||||
AccountSortField as AdminAccountSortField, AccountStatus as AdminAccountStatus,
|
||||
AccountSummary, AccountUsage, AccountUsageWindowQuery, AccountUsageWindowResult,
|
||||
DeleteAccounts, SetAccountEnabled, SortDirection as AdminSortDirection,
|
||||
AccountCost, AccountGroupFilter, AccountListQuery as AdminAccountListQuery,
|
||||
AccountModelUsage, AccountPage, AccountPageItem, AccountRecord, AccountRequestBucket,
|
||||
AccountSort as AdminAccountSort, AccountSortField as AdminAccountSortField,
|
||||
AccountStatus as AdminAccountStatus, AccountSummary, AccountUpdateResult, AccountUsage,
|
||||
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
|
||||
BatchUpdateAccounts, DeleteAccounts, SortDirection as AdminSortDirection,
|
||||
UpdateAccount,
|
||||
},
|
||||
observability::{
|
||||
CostCoverage as AdminCostCoverage, DecimalAmount as AdminDecimalAmount, TimeRange,
|
||||
@@ -34,16 +37,16 @@ use gateway_admin::{
|
||||
use sqlx::{PgPool, Postgres, Row, Transaction};
|
||||
|
||||
use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
|
||||
CredentialRevision as CoreCredentialRevision, CredentialState, LoadedCredential,
|
||||
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
|
||||
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
|
||||
ProviderAccountIdentity, ProviderAccountStore,
|
||||
AccountConcurrencyLimit, AccountErrorReason, AccountStateChange, AccountWeight,
|
||||
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision as CoreCredentialRevision,
|
||||
CredentialState, LoadedCredential, NewProviderAccount as CoreNewProviderAccount,
|
||||
OpaqueProviderData, PlaintextCredential, ProviderAccount as CoreProviderAccount,
|
||||
ProviderAccountId as CoreProviderAccountId, ProviderAccountIdentity, ProviderAccountStore,
|
||||
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaAccessChange, QuotaAccessState,
|
||||
QuotaEvidence, QuotaObservation, QuotaState, QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use gateway_core::routing::{AccountGroupId, ProviderKind};
|
||||
|
||||
use crate::{
|
||||
ConflictKind, JsonObject, Revision, StoreError, StoreResult, admin_revision, admin_store_error,
|
||||
|
||||
@@ -56,9 +56,9 @@ pub trait ProviderAccountAdminRepository: Send + Sync {
|
||||
command: RotateProviderAccount,
|
||||
) -> StoreResult<ProviderAccountAdminRotation>;
|
||||
|
||||
async fn set_provider_account_enabled_admin(
|
||||
async fn batch_update_provider_accounts_admin(
|
||||
&self,
|
||||
command: SetProviderAccountEnabled,
|
||||
command: BatchUpdateProviderAccountsAdmin,
|
||||
) -> StoreResult<Revision>;
|
||||
|
||||
async fn delete_provider_accounts_admin(
|
||||
@@ -99,7 +99,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
let rows = sqlx::query(
|
||||
"select id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, credential_revision, has_refresh_token,
|
||||
access_token_expires_at, next_refresh_at, enabled, credential_state,
|
||||
access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
|
||||
credential_observed_at, quota_access_state, quota_evidence,
|
||||
quota_access_observed_at, quota_reset_at,
|
||||
quota_observed_at, last_error_reason, last_error_message, created_at, updated_at
|
||||
@@ -127,11 +127,11 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
|
||||
credential_state, provider_quota_json,
|
||||
concurrency_limit, weight, credential_state, provider_quota_json,
|
||||
credential_observed_at, quota_access_observed_at, quota_observed_at, created_at, updated_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
|
||||
$14, null, $15, null, null, now(), greatest(now(), $15)
|
||||
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
|
||||
)",
|
||||
)
|
||||
.bind(account.id)
|
||||
@@ -147,6 +147,8 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
|
||||
.bind(account.access_token_expires_at)
|
||||
.bind(account.next_refresh_at)
|
||||
.bind(account.enabled)
|
||||
.bind(account.concurrency_limit.map(|limit| i64::from(limit.get())))
|
||||
.bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?)
|
||||
.bind(credential_state.as_str())
|
||||
.bind(account.credential_observed_at)
|
||||
.execute(&self.pool)
|
||||
@@ -471,12 +473,12 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
|
||||
finish_admin_transaction(transaction, result, "provider account admin rotation").await
|
||||
}
|
||||
|
||||
async fn set_provider_account_enabled_admin(
|
||||
async fn batch_update_provider_accounts_admin(
|
||||
&self,
|
||||
command: SetProviderAccountEnabled,
|
||||
command: BatchUpdateProviderAccountsAdmin,
|
||||
) -> StoreResult<Revision> {
|
||||
command.scope.validate()?;
|
||||
require_nonempty(ENTITY, "account_id", &command.account_id)?;
|
||||
validate_batch_update_account_ids(&command.account_ids)?;
|
||||
validate_batch_update_group_ids(&command.group_ids)?;
|
||||
let mut transaction = self
|
||||
.pool
|
||||
.begin()
|
||||
@@ -484,11 +486,18 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
|
||||
.map_err(|_| postgres_unavailable("begin provider account admin state change"))?;
|
||||
let result = async {
|
||||
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
|
||||
set_provider_account_enabled_in_transaction(
|
||||
update_provider_accounts_scheduling_in_transaction(
|
||||
&mut transaction,
|
||||
&command.scope,
|
||||
&command.account_id,
|
||||
&command.account_ids,
|
||||
command.enabled,
|
||||
command.concurrency_limit,
|
||||
command.weight,
|
||||
)
|
||||
.await?;
|
||||
replace_account_group_assignments_in_transaction(
|
||||
&mut transaction,
|
||||
&command.account_ids,
|
||||
&command.group_ids,
|
||||
)
|
||||
.await?;
|
||||
append_admin_audit_event_in_transaction(&mut transaction, command.audit, revision)
|
||||
@@ -527,6 +536,53 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
|
||||
}
|
||||
}
|
||||
|
||||
async fn replace_account_group_assignments_in_transaction(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
account_ids: &[String],
|
||||
group_ids: &[AccountGroupId],
|
||||
) -> StoreResult<()> {
|
||||
let group_ids = group_ids
|
||||
.iter()
|
||||
.map(|group_id| group_id.as_str().to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
if !group_ids.is_empty() {
|
||||
let known_group_count = sqlx::query_scalar::<_, i64>(
|
||||
"select count(*)::bigint from account_groups where id = any($1::text[])",
|
||||
)
|
||||
.bind(&group_ids)
|
||||
.fetch_one(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("validate account group assignment"))?;
|
||||
if usize::try_from(known_group_count).ok() != Some(group_ids.len()) {
|
||||
return Err(StoreError::NotFound {
|
||||
entity: "account group",
|
||||
id: "one or more group IDs".to_owned(),
|
||||
});
|
||||
}
|
||||
}
|
||||
sqlx::query("delete from account_group_accounts where provider_account_id = any($1::text[])")
|
||||
.bind(account_ids)
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("clear account group assignments"))?;
|
||||
if group_ids.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
sqlx::query(
|
||||
"insert into account_group_accounts
|
||||
(account_group_id, provider_account_id, created_at)
|
||||
select group_id, account_id, now()
|
||||
from unnest($1::text[]) group_id
|
||||
cross join unnest($2::text[]) account_id",
|
||||
)
|
||||
.bind(group_ids)
|
||||
.bind(account_ids)
|
||||
.execute(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("assign accounts to groups"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn upsert_provider_account_in_transaction(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
account: &NewProviderAccount,
|
||||
@@ -542,11 +598,11 @@ pub(crate) async fn upsert_provider_account_in_transaction(
|
||||
id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
|
||||
credential_state, provider_quota_json,
|
||||
concurrency_limit, weight, credential_state, provider_quota_json,
|
||||
credential_observed_at, quota_access_observed_at, quota_observed_at, created_at, updated_at
|
||||
) values (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, 1, $10, $11, $12, $13,
|
||||
$14, null, $15, null, null, now(), greatest(now(), $15)
|
||||
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
|
||||
)
|
||||
on conflict (
|
||||
provider_kind,
|
||||
@@ -589,6 +645,8 @@ pub(crate) async fn upsert_provider_account_in_transaction(
|
||||
.bind(account.access_token_expires_at)
|
||||
.bind(account.next_refresh_at)
|
||||
.bind(account.enabled)
|
||||
.bind(account.concurrency_limit.map(|limit| i64::from(limit.get())))
|
||||
.bind(i16::try_from(account.weight.get()).map_err(|_| invalid("invalid weight"))?)
|
||||
.bind(credential_state.as_str())
|
||||
.bind(account.credential_observed_at)
|
||||
.fetch_optional(&mut **transaction)
|
||||
@@ -691,23 +749,37 @@ pub(crate) async fn rotate_provider_account_in_transaction(
|
||||
Revision::new(to_u64(next)?)
|
||||
}
|
||||
|
||||
pub(crate) async fn set_provider_account_enabled_in_transaction(
|
||||
pub(crate) async fn update_provider_accounts_scheduling_in_transaction(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
scope: &ProviderAccountAdminScope,
|
||||
account_id: &str,
|
||||
account_ids: &[String],
|
||||
enabled: bool,
|
||||
concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
weight: AccountWeight,
|
||||
) -> StoreResult<()> {
|
||||
let result = sqlx::query(
|
||||
"update provider_accounts set enabled = $3, updated_at = now()
|
||||
where id = $1 and provider_kind = $2",
|
||||
let updated = sqlx::query_scalar::<_, String>(
|
||||
"update provider_accounts
|
||||
set enabled = $2, concurrency_limit = $3, weight = $4, updated_at = now()
|
||||
where id = any($1::text[])
|
||||
returning id",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(&scope.provider_kind)
|
||||
.bind(account_ids)
|
||||
.bind(enabled)
|
||||
.execute(&mut **transaction)
|
||||
.bind(concurrency_limit.map(|limit| i64::from(limit.get())))
|
||||
.bind(i16::try_from(weight.get()).map_err(|_| invalid("invalid weight"))?)
|
||||
.fetch_all(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("set provider account state in admin transaction"))?;
|
||||
require_admin_account_changed(result.rows_affected(), account_id)
|
||||
.map_err(|_| postgres_unavailable("set provider accounts state in admin transaction"))?
|
||||
.into_iter()
|
||||
.collect::<BTreeSet<_>>();
|
||||
let expected = account_ids.iter().cloned().collect::<BTreeSet<_>>();
|
||||
if updated == expected {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(StoreError::NotFound {
|
||||
entity: ENTITY,
|
||||
id: "one or more provider account IDs".to_owned(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_provider_accounts_in_transaction(
|
||||
@@ -749,20 +821,6 @@ pub(crate) fn validate_credential_update(update: &ProviderCredentialUpdate) -> S
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn require_admin_account_changed(
|
||||
rows_affected: u64,
|
||||
account_id: &str,
|
||||
) -> StoreResult<()> {
|
||||
if rows_affected == 1 {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(StoreError::NotFound {
|
||||
entity: ENTITY,
|
||||
id: account_id.to_owned(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn validate_admin_account_ids(account_ids: &[String]) -> StoreResult<()> {
|
||||
if account_ids.is_empty() || account_ids.len() > MAX_ADMIN_IMPORT_BATCH {
|
||||
return Err(invalid(
|
||||
@@ -779,6 +837,38 @@ pub(crate) fn validate_admin_account_ids(account_ids: &[String]) -> StoreResult<
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_batch_update_account_ids(account_ids: &[String]) -> StoreResult<()> {
|
||||
const MAX_BATCH_UPDATE_ACCOUNTS: usize = 1000;
|
||||
if account_ids.is_empty() || account_ids.len() > MAX_BATCH_UPDATE_ACCOUNTS {
|
||||
return Err(invalid(
|
||||
"account batch update must contain between 1 and 1000 IDs",
|
||||
));
|
||||
}
|
||||
let mut unique = BTreeSet::new();
|
||||
for account_id in account_ids {
|
||||
require_nonempty(ENTITY, "account_id", account_id)?;
|
||||
if !unique.insert(account_id.as_str()) {
|
||||
return Err(invalid("account batch update contains duplicate IDs"));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_batch_update_group_ids(group_ids: &[AccountGroupId]) -> StoreResult<()> {
|
||||
const MAX_BATCH_UPDATE_GROUPS: usize = 1000;
|
||||
if group_ids.len() > MAX_BATCH_UPDATE_GROUPS {
|
||||
return Err(invalid("account batch update contains too many group IDs"));
|
||||
}
|
||||
let mut unique = BTreeSet::new();
|
||||
if group_ids
|
||||
.iter()
|
||||
.any(|group_id| !unique.insert(group_id.as_str()))
|
||||
{
|
||||
return Err(invalid("account batch update contains duplicate group IDs"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn finish_admin_transaction<T>(
|
||||
transaction: Transaction<'_, Postgres>,
|
||||
result: StoreResult<T>,
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
//! `provider_accounts` row 模型、解码与校验。
|
||||
|
||||
use super::*;
|
||||
use gateway_core::routing::AccountGroupId;
|
||||
|
||||
pub(crate) const ENTITY: &str = "provider account";
|
||||
pub(crate) const CREDENTIALS_MAX_BYTES: usize = 256 * 1024;
|
||||
@@ -105,6 +106,8 @@ pub struct ProviderAccountSummary {
|
||||
pub access_token_expires_at: Option<DateTime<Utc>>,
|
||||
pub next_refresh_at: Option<DateTime<Utc>>,
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub credential_state: CredentialState,
|
||||
pub credential_observed_at: DateTime<Utc>,
|
||||
pub quota: QuotaState,
|
||||
@@ -150,6 +153,8 @@ pub struct NewProviderAccount {
|
||||
pub access_token_expires_at: Option<DateTime<Utc>>,
|
||||
pub next_refresh_at: Option<DateTime<Utc>>,
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub credential_state: CredentialState,
|
||||
pub credential_observed_at: DateTime<Utc>,
|
||||
}
|
||||
@@ -274,10 +279,12 @@ impl fmt::Debug for RotateProviderAccount {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SetProviderAccountEnabled {
|
||||
pub scope: ProviderAccountAdminScope,
|
||||
pub account_id: String,
|
||||
pub struct BatchUpdateProviderAccountsAdmin {
|
||||
pub account_ids: Vec<String>,
|
||||
pub enabled: bool,
|
||||
pub concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
pub weight: AccountWeight,
|
||||
pub group_ids: Vec<AccountGroupId>,
|
||||
pub audit: AdminAuditEvent,
|
||||
}
|
||||
|
||||
@@ -334,7 +341,7 @@ impl ProviderAccountStateUpdate {
|
||||
|
||||
pub(crate) const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, credential_state,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
|
||||
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
|
||||
last_error_reason, last_error_message,
|
||||
credential_observed_at, quota_observed_at, created_at, updated_at
|
||||
@@ -342,7 +349,7 @@ pub(crate) const ACCOUNT_SELECT: &str = "select id, provider_kind, name, email,
|
||||
|
||||
pub(crate) const ACCOUNT_SELECT_BY_IDS: &str = "select id, provider_kind, name, email, upstream_user_id,
|
||||
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, credential_state,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled, concurrency_limit, weight, credential_state,
|
||||
provider_quota_json, quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at,
|
||||
last_error_reason, last_error_message,
|
||||
credential_observed_at, quota_observed_at, created_at, updated_at
|
||||
@@ -401,6 +408,7 @@ pub(crate) fn core_account_from_summary(
|
||||
summary.last_error_reason,
|
||||
summary.last_error_message,
|
||||
)
|
||||
.with_scheduling(summary.concurrency_limit, summary.weight)
|
||||
.with_refresh_schedule(
|
||||
summary.has_refresh_token,
|
||||
summary.next_refresh_at.map(Into::into),
|
||||
@@ -446,6 +454,18 @@ pub(crate) fn account_summary_from_row(
|
||||
quota_reset_at.map(Into::into),
|
||||
)
|
||||
.ok_or_else(|| invalid("invalid persisted quota fact"))?;
|
||||
let concurrency_limit = get::<Option<i64>>(&row, "concurrency_limit")?
|
||||
.map(|value| {
|
||||
u32::try_from(value)
|
||||
.ok()
|
||||
.and_then(AccountConcurrencyLimit::new)
|
||||
.ok_or_else(|| invalid("invalid concurrency_limit"))
|
||||
})
|
||||
.transpose()?;
|
||||
let weight = u16::try_from(get::<i16>(&row, "weight")?)
|
||||
.ok()
|
||||
.and_then(AccountWeight::new)
|
||||
.ok_or_else(|| invalid("invalid weight"))?;
|
||||
Ok(ProviderAccountSummary {
|
||||
id: get(&row, "id")?,
|
||||
provider_kind: get(&row, "provider_kind")?,
|
||||
@@ -460,6 +480,8 @@ pub(crate) fn account_summary_from_row(
|
||||
access_token_expires_at: get(&row, "access_token_expires_at")?,
|
||||
next_refresh_at: get(&row, "next_refresh_at")?,
|
||||
enabled: get(&row, "enabled")?,
|
||||
concurrency_limit,
|
||||
weight,
|
||||
credential_state: parse_credential_state(&credential_state)?,
|
||||
credential_observed_at: get(&row, "credential_observed_at")?,
|
||||
quota,
|
||||
|
||||
@@ -4,10 +4,12 @@ use std::collections::BTreeMap;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_core::engine::credential::ProviderAccountId;
|
||||
use gateway_core::routing::{
|
||||
ConfigRevision,
|
||||
AccountGroupId, ConfigRevision,
|
||||
snapshot::{
|
||||
SnapshotClientPolicyFacts, SnapshotFacts, SnapshotSettingsFacts, SnapshotStoreError,
|
||||
SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts, SnapshotClientPolicyFacts,
|
||||
SnapshotFacts, SnapshotProviderAccountFacts, SnapshotSettingsFacts, SnapshotStoreError,
|
||||
SnapshotStorePort,
|
||||
},
|
||||
};
|
||||
@@ -41,6 +43,28 @@ pub struct RuntimeSnapshotData {
|
||||
pub observed_current_revision: Revision,
|
||||
pub settings: SnapshotRuntimeSettings,
|
||||
pub client_api_keys: Vec<ClientApiKeySnapshot>,
|
||||
pub account_groups: Vec<SnapshotAccountGroupData>,
|
||||
pub provider_accounts: Vec<SnapshotProviderAccountData>,
|
||||
pub group_memberships: Vec<SnapshotGroupMembershipData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotAccountGroupData {
|
||||
pub id: AccountGroupId,
|
||||
pub name: String,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotProviderAccountData {
|
||||
pub id: String,
|
||||
pub provider_kind: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SnapshotGroupMembershipData {
|
||||
pub group_id: AccountGroupId,
|
||||
pub account_id: String,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -77,6 +101,9 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
|
||||
|
||||
let (config_revision, settings) = load_settings(&mut transaction).await?;
|
||||
let client_api_keys = load_client_keys(&mut transaction).await?;
|
||||
let account_groups = load_account_groups(&mut transaction).await?;
|
||||
let provider_accounts = load_provider_accounts(&mut transaction).await?;
|
||||
let group_memberships = load_group_memberships(&mut transaction).await?;
|
||||
transaction
|
||||
.commit()
|
||||
.await
|
||||
@@ -89,6 +116,9 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
|
||||
observed_current_revision,
|
||||
settings,
|
||||
client_api_keys,
|
||||
account_groups,
|
||||
provider_accounts,
|
||||
group_memberships,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -151,16 +181,44 @@ impl SnapshotStorePort for PgRuntimeSnapshotRepository {
|
||||
SnapshotClientPolicyFacts::new(
|
||||
key.id,
|
||||
key.plaintext_key,
|
||||
key.provider_kind,
|
||||
key.group_ids,
|
||||
key.limits,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let account_groups = data
|
||||
.account_groups
|
||||
.into_iter()
|
||||
.map(|group| SnapshotAccountGroupFacts::new(group.id, group.name, group.enabled))
|
||||
.collect();
|
||||
let provider_accounts = data
|
||||
.provider_accounts
|
||||
.into_iter()
|
||||
.map(|account| {
|
||||
ProviderAccountId::new(account.id)
|
||||
.map(|id| SnapshotProviderAccountFacts::new(id, account.provider_kind))
|
||||
.map_err(|_| SnapshotStoreError::unavailable())
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let group_memberships = data
|
||||
.group_memberships
|
||||
.into_iter()
|
||||
.map(|membership| {
|
||||
ProviderAccountId::new(membership.account_id)
|
||||
.map(|account_id| {
|
||||
SnapshotAccountGroupMemberFacts::new(membership.group_id, account_id)
|
||||
})
|
||||
.map_err(|_| SnapshotStoreError::unavailable())
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(SnapshotFacts::new(
|
||||
config_revision,
|
||||
observed_current_revision,
|
||||
settings,
|
||||
client_policies,
|
||||
account_groups,
|
||||
provider_accounts,
|
||||
group_memberships,
|
||||
))
|
||||
})
|
||||
}
|
||||
@@ -224,9 +282,16 @@ async fn load_settings(
|
||||
async fn load_client_keys(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
) -> StoreResult<Vec<ClientApiKeySnapshot>> {
|
||||
let rows = sqlx::query_as::<_, (String, String, String, i64, i64)>(
|
||||
"select id, key, provider_kind, max_concurrency, requests_per_minute
|
||||
from client_api_keys where enabled order by id",
|
||||
let rows = sqlx::query_as::<_, (String, String, Vec<String>, i64, i64)>(
|
||||
"select k.id, k.key,
|
||||
coalesce(array_agg(kg.account_group_id order by kg.account_group_id)
|
||||
filter (where kg.account_group_id is not null), '{}') as group_ids,
|
||||
k.max_concurrency, k.requests_per_minute
|
||||
from client_api_keys k
|
||||
left join client_api_key_groups kg on kg.client_api_key_id = k.id
|
||||
where k.enabled
|
||||
group by k.id
|
||||
order by k.id",
|
||||
)
|
||||
.fetch_all(&mut **transaction)
|
||||
.await
|
||||
@@ -236,6 +301,63 @@ async fn load_client_keys(
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn load_account_groups(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
) -> StoreResult<Vec<SnapshotAccountGroupData>> {
|
||||
let rows = sqlx::query_as::<_, (String, String, bool)>(
|
||||
"select id, name, enabled from account_groups order by id",
|
||||
)
|
||||
.fetch_all(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load snapshot account groups"))?;
|
||||
rows.into_iter()
|
||||
.map(|(id, name, enabled)| {
|
||||
Ok(SnapshotAccountGroupData {
|
||||
id: AccountGroupId::new(id).map_err(|_| invalid("invalid account group id"))?,
|
||||
name,
|
||||
enabled,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn load_provider_accounts(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
) -> StoreResult<Vec<SnapshotProviderAccountData>> {
|
||||
sqlx::query_as::<_, (String, String)>(
|
||||
"select id, provider_kind from provider_accounts order by id",
|
||||
)
|
||||
.fetch_all(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load snapshot provider accounts"))
|
||||
.map(|rows| {
|
||||
rows.into_iter()
|
||||
.map(|(id, provider_kind)| SnapshotProviderAccountData { id, provider_kind })
|
||||
.collect()
|
||||
})
|
||||
}
|
||||
|
||||
async fn load_group_memberships(
|
||||
transaction: &mut Transaction<'_, Postgres>,
|
||||
) -> StoreResult<Vec<SnapshotGroupMembershipData>> {
|
||||
let rows = sqlx::query_as::<_, (String, String)>(
|
||||
"select account_group_id, provider_account_id
|
||||
from account_group_accounts order by account_group_id, provider_account_id",
|
||||
)
|
||||
.fetch_all(&mut **transaction)
|
||||
.await
|
||||
.map_err(|_| postgres_unavailable("load snapshot group memberships"))?;
|
||||
rows.into_iter()
|
||||
.map(|(group_id, account_id)| {
|
||||
Ok(SnapshotGroupMembershipData {
|
||||
group_id: AccountGroupId::new(group_id)
|
||||
.map_err(|_| invalid("invalid membership group id"))?,
|
||||
account_id,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn revision_from_i64(value: i64) -> StoreResult<Revision> {
|
||||
Revision::new(to_u64(value)?)
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::future::join_all;
|
||||
use gateway_core::engine::credential::{AccountRuntimeSignals, ProviderAccountId};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
|
||||
ProviderRefreshCapacityRequest, ProviderSchedulingLeaseRequest, ProviderSchedulingState,
|
||||
@@ -329,10 +330,13 @@ impl RedisCredentialLeaseRepository {
|
||||
}))
|
||||
}
|
||||
|
||||
/// 原子推进跨进程共享的 Provider 调度游标。
|
||||
pub async fn advance_scheduling_cursor(&self, provider_kind: &str) -> StoreResult<u64> {
|
||||
require_nonempty("provider scheduling cursor", "provider_kind", provider_kind)?;
|
||||
let key = self.scheduling_cursor_key(provider_kind)?;
|
||||
/// 原子推进跨进程共享、按 Client Key 与 Provider 隔离的调度游标。
|
||||
pub async fn advance_scheduling_cursor(
|
||||
&self,
|
||||
client_api_key_id: &ClientApiKeyId,
|
||||
provider_kind: &ProviderKind,
|
||||
) -> StoreResult<u64> {
|
||||
let key = self.scheduling_cursor_key(client_api_key_id, provider_kind)?;
|
||||
let mut connection = self.connection.clone();
|
||||
let cursor = Script::new(ADVANCE_SCHEDULING_CURSOR_SCRIPT)
|
||||
.key(key)
|
||||
@@ -369,10 +373,17 @@ impl RedisCredentialLeaseRepository {
|
||||
))
|
||||
}
|
||||
|
||||
fn scheduling_cursor_key(&self, provider_kind: &str) -> StoreResult<String> {
|
||||
let fingerprint = resource_fingerprint("provider scheduling cursor", provider_kind)?;
|
||||
fn scheduling_cursor_key(
|
||||
&self,
|
||||
client_api_key_id: &ClientApiKeyId,
|
||||
provider_kind: &ProviderKind,
|
||||
) -> StoreResult<String> {
|
||||
let fingerprint = resource_fingerprint(
|
||||
"provider scheduling cursor",
|
||||
&format!("{}\0{}", client_api_key_id.as_str(), provider_kind.as_str()),
|
||||
)?;
|
||||
Ok(format!(
|
||||
"{}:scheduler:provider:{{{fingerprint}}}:cursor",
|
||||
"{}:scheduler:client-provider:{{{fingerprint}}}:cursor",
|
||||
self.namespace
|
||||
))
|
||||
}
|
||||
@@ -475,10 +486,11 @@ impl RedisProviderLeaseCoordinator {
|
||||
|
||||
async fn next_scheduling_cursor(
|
||||
&self,
|
||||
client_api_key_id: &ClientApiKeyId,
|
||||
provider_kind: &ProviderKind,
|
||||
) -> Result<u64, ProviderStoreError> {
|
||||
self.repository
|
||||
.advance_scheduling_cursor(provider_kind.as_str())
|
||||
.advance_scheduling_cursor(client_api_key_id, provider_kind)
|
||||
.await
|
||||
.map_err(|_| provider_unavailable("advance scheduling cursor"))
|
||||
}
|
||||
@@ -589,12 +601,15 @@ impl RedisProviderLeaseCoordinator {
|
||||
impl ProviderLeasePort for RedisProviderLeaseCoordinator {
|
||||
fn load_state<'a>(
|
||||
&'a self,
|
||||
client_api_key_id: &'a ClientApiKeyId,
|
||||
provider_kind: &'a ProviderKind,
|
||||
accounts: &'a [ProviderAccountId],
|
||||
) -> futures::future::BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
|
||||
Box::pin(async move {
|
||||
let signals = self.load_signals(accounts).await?;
|
||||
let round_robin_cursor = self.next_scheduling_cursor(provider_kind).await?;
|
||||
let round_robin_cursor = self
|
||||
.next_scheduling_cursor(client_api_key_id, provider_kind)
|
||||
.await?;
|
||||
Ok(ProviderSchedulingState::new(signals, round_robin_cursor))
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
//! 旧版 Responses continuation 的 Redis 会话亲和记录。
|
||||
//! Responses continuation 的 Redis 会话亲和记录。
|
||||
//!
|
||||
//! 该存储只保存可丢失的 response → Provider/account pin;PostgreSQL 审计记录
|
||||
//! 不是 continuation 的可用性前提。键只含 response ID 的不可逆指纹,Provider
|
||||
//! 私有状态作为不透明 JSON 保存,Store 不解释其内容。
|
||||
//! 不是 continuation 的可用性前提。键只含 response ID 的不可逆指纹;Client Key
|
||||
//! 所有权写入 payload 并在读取时严格校验,使跨 Key 复用能够 fail closed。
|
||||
//! Provider 私有状态作为不透明 JSON 保存,Store 不解释其内容。
|
||||
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
@@ -13,6 +14,7 @@ use gateway_core::engine::continuation::{
|
||||
};
|
||||
use gateway_core::engine::credential::ProviderAccountId;
|
||||
use gateway_core::operation::ProviderSessionState;
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use redis::aio::ConnectionManager;
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -72,8 +74,9 @@ pub struct RedisNativeContinuationRepository {
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
#[serde(rename_all = "camelCase", deny_unknown_fields)]
|
||||
struct ContinuationWire {
|
||||
client_api_key_id: String,
|
||||
upstream_response_id: String,
|
||||
provider: String,
|
||||
account: String,
|
||||
@@ -86,13 +89,13 @@ impl RedisNativeContinuationRepository {
|
||||
pub fn new(connection: ConnectionManager, key_namespace: &str) -> StoreResult<Self> {
|
||||
Ok(Self {
|
||||
connection,
|
||||
namespace: format!("{}:response-affinity:v1", namespace(key_namespace)?),
|
||||
namespace: format!("{}:response-affinity:v2", namespace(key_namespace)?),
|
||||
})
|
||||
}
|
||||
|
||||
fn entry_fingerprint(&self, response_id: &str) -> Result<String, NativeContinuationStoreError> {
|
||||
resource_fingerprint("native continuation", response_id)
|
||||
.map_err(|error| NativeContinuationStoreError::new(error.to_string()))
|
||||
.map_err(|error| NativeContinuationStoreError::invalid_data(error.to_string()))
|
||||
}
|
||||
|
||||
fn entry_key(&self, fingerprint: &str) -> String {
|
||||
@@ -111,6 +114,7 @@ impl RedisNativeContinuationRepository {
|
||||
impl NativeContinuationPort for RedisNativeContinuationRepository {
|
||||
fn resolve<'a>(
|
||||
&'a self,
|
||||
client_api_key_id: &'a ClientApiKeyId,
|
||||
previous_response_id: &'a PreviousResponseId,
|
||||
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
|
||||
Box::pin(async move {
|
||||
@@ -124,10 +128,10 @@ impl NativeContinuationPort for RedisNativeContinuationRepository {
|
||||
.query_async::<Option<String>>(&mut connection)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
NativeContinuationStoreError::new(format!("redis GET failed: {error}"))
|
||||
NativeContinuationStoreError::unavailable(format!("redis GET failed: {error}"))
|
||||
})?;
|
||||
payload
|
||||
.map(|payload| decode_pin(previous_response_id, &payload))
|
||||
.map(|payload| decode_pin(client_api_key_id, previous_response_id, &payload))
|
||||
.transpose()
|
||||
})
|
||||
}
|
||||
@@ -143,20 +147,24 @@ impl NativeContinuationPort for RedisNativeContinuationRepository {
|
||||
let fingerprint = self.entry_fingerprint(pin.previous_response_id().as_str())?;
|
||||
let payload = encode_pin(&pin)?;
|
||||
if payload.len() > MAX_RECORD_BYTES {
|
||||
return Err(NativeContinuationStoreError::new(format!(
|
||||
return Err(NativeContinuationStoreError::invalid_data(format!(
|
||||
"pin payload is too large ({} bytes, limit {MAX_RECORD_BYTES})",
|
||||
payload.len()
|
||||
)));
|
||||
}
|
||||
let now_millis = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_err(|_| NativeContinuationStoreError::new("system clock predates Unix epoch"))?
|
||||
.map_err(|_| {
|
||||
NativeContinuationStoreError::unavailable("system clock predates Unix epoch")
|
||||
})?
|
||||
.as_millis();
|
||||
let now_millis = u64::try_from(now_millis).map_err(|_| {
|
||||
NativeContinuationStoreError::new("system time is outside supported range")
|
||||
NativeContinuationStoreError::unavailable("system time is outside supported range")
|
||||
})?;
|
||||
let ttl_millis = u64::try_from(CONTINUATION_TTL.as_millis()).map_err(|_| {
|
||||
NativeContinuationStoreError::new("continuation TTL is outside supported range")
|
||||
NativeContinuationStoreError::invalid_data(
|
||||
"continuation TTL is outside supported range",
|
||||
)
|
||||
})?;
|
||||
let cutoff = now_millis.saturating_sub(ttl_millis);
|
||||
let ttl_seconds = CONTINUATION_TTL.as_secs();
|
||||
@@ -175,7 +183,7 @@ impl NativeContinuationPort for RedisNativeContinuationRepository {
|
||||
.invoke_async::<i64>(&mut connection)
|
||||
.await
|
||||
.map_err(|error| {
|
||||
NativeContinuationStoreError::new(format!(
|
||||
NativeContinuationStoreError::unavailable(format!(
|
||||
"redis affinity script failed: {error}"
|
||||
))
|
||||
})?;
|
||||
@@ -189,11 +197,12 @@ fn encode_pin(pin: &NativeContinuationPin) -> Result<String, NativeContinuationS
|
||||
.session_state()
|
||||
.is_some_and(|state| state.provider() != pin.provider().as_str())
|
||||
{
|
||||
return Err(NativeContinuationStoreError::new(
|
||||
return Err(NativeContinuationStoreError::invalid_data(
|
||||
"session state provider does not match pin provider",
|
||||
));
|
||||
}
|
||||
serde_json::to_string(&ContinuationWire {
|
||||
client_api_key_id: pin.client_api_key_id().as_str().to_owned(),
|
||||
upstream_response_id: pin.upstream_response_id().as_str().to_owned(),
|
||||
provider: pin.provider().as_str().to_owned(),
|
||||
account: pin.account().as_str().to_owned(),
|
||||
@@ -201,34 +210,44 @@ fn encode_pin(pin: &NativeContinuationPin) -> Result<String, NativeContinuationS
|
||||
session_state: pin.session_state().cloned(),
|
||||
})
|
||||
.map_err(|error| {
|
||||
NativeContinuationStoreError::new(format!("pin serialization failed: {error}"))
|
||||
NativeContinuationStoreError::invalid_data(format!("pin serialization failed: {error}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn decode_pin(
|
||||
expected_client_api_key_id: &ClientApiKeyId,
|
||||
previous_response_id: &PreviousResponseId,
|
||||
payload: &str,
|
||||
) -> Result<NativeContinuationPin, NativeContinuationStoreError> {
|
||||
let wire: ContinuationWire = serde_json::from_str(payload)
|
||||
.map_err(|error| NativeContinuationStoreError::new(format!("{error}")))?;
|
||||
.map_err(|error| NativeContinuationStoreError::invalid_data(format!("{error}")))?;
|
||||
let provider = ProviderKind::new(wire.provider).map_err(|_| {
|
||||
NativeContinuationStoreError::new("pin payload contains an invalid provider")
|
||||
NativeContinuationStoreError::invalid_data("pin payload contains an invalid provider")
|
||||
})?;
|
||||
let account = ProviderAccountId::new(wire.account).map_err(|_| {
|
||||
NativeContinuationStoreError::new("pin payload contains an invalid account id")
|
||||
NativeContinuationStoreError::invalid_data("pin payload contains an invalid account id")
|
||||
})?;
|
||||
let client_api_key_id = ClientApiKeyId::new(wire.client_api_key_id).map_err(|_| {
|
||||
NativeContinuationStoreError::invalid_data(
|
||||
"pin payload contains an invalid client API key id",
|
||||
)
|
||||
})?;
|
||||
if client_api_key_id != *expected_client_api_key_id {
|
||||
return Err(NativeContinuationStoreError::ownership_mismatch());
|
||||
}
|
||||
let upstream_response_id = PreviousResponseId::new(wire.upstream_response_id);
|
||||
let scope = parse_scope(&wire.scope)?;
|
||||
let mut pin = NativeContinuationPin::new(
|
||||
previous_response_id.clone(),
|
||||
upstream_response_id,
|
||||
client_api_key_id,
|
||||
provider.clone(),
|
||||
account,
|
||||
)
|
||||
.with_scope(scope);
|
||||
if let Some(state) = wire.session_state {
|
||||
if state.provider() != provider.as_str() {
|
||||
return Err(NativeContinuationStoreError::new(
|
||||
return Err(NativeContinuationStoreError::invalid_data(
|
||||
"pin payload session state provider does not match pin provider",
|
||||
));
|
||||
}
|
||||
@@ -248,7 +267,7 @@ fn parse_scope(scope: &str) -> Result<NativeContinuationScope, NativeContinuatio
|
||||
match scope {
|
||||
"persisted" => Ok(NativeContinuationScope::Persisted),
|
||||
"connection_local" => Ok(NativeContinuationScope::ConnectionLocal),
|
||||
_ => Err(NativeContinuationStoreError::new(
|
||||
_ => Err(NativeContinuationStoreError::invalid_data(
|
||||
"pin payload contains an unsupported continuation scope",
|
||||
)),
|
||||
}
|
||||
|
||||
@@ -0,0 +1,313 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use gateway_admin::{
|
||||
model::{
|
||||
MutationActor, MutationContext, PageSize,
|
||||
account_groups::{
|
||||
AccountGroupColor, AccountGroupListQuery, DeleteAccountGroup, NewAccountGroup,
|
||||
},
|
||||
client_keys::NewClientKey,
|
||||
client_keys::UpdateClientKey,
|
||||
},
|
||||
ports::store::{AccountGroupStore, AdminStoreErrorKind, ClientKeyStore},
|
||||
};
|
||||
use gateway_core::{
|
||||
policy::{ClientApiKeyId, RateLimits},
|
||||
routing::AccountGroupId,
|
||||
};
|
||||
use gateway_store::postgres::{PgAccountGroupRepository, PgAdminClientKeyStore};
|
||||
|
||||
use super::TestDatabase;
|
||||
|
||||
const MIXED_GROUP: &str = "grp_00000000000000000000000000000001";
|
||||
const EMPTY_GROUP: &str = "grp_00000000000000000000000000000002";
|
||||
|
||||
#[tokio::test]
|
||||
async fn groups_aggregate_cross_provider_members_and_key_bindings_without_multiplication() {
|
||||
let Some(database) = TestDatabase::create("account_group_aggregate").await else {
|
||||
return;
|
||||
};
|
||||
seed_account(
|
||||
&database.pool,
|
||||
"acct_group_openai",
|
||||
"openai",
|
||||
"OpenAI Account",
|
||||
)
|
||||
.await;
|
||||
seed_account(&database.pool, "acct_group_xai", "xai", "xAI Account").await;
|
||||
sqlx::query(
|
||||
"update provider_accounts set concurrency_limit = 4 where id = 'acct_group_openai'",
|
||||
)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("set account concurrency override");
|
||||
|
||||
let groups = PgAccountGroupRepository::new(database.pool.clone());
|
||||
let keys = PgAdminClientKeyStore::new(database.pool.clone());
|
||||
let mixed_group = group_id(MIXED_GROUP);
|
||||
let empty_group = group_id(EMPTY_GROUP);
|
||||
groups
|
||||
.create_account_group(
|
||||
NewAccountGroup {
|
||||
id: mixed_group.clone(),
|
||||
name: "Mixed Production".to_owned(),
|
||||
description: Some("cross-provider".to_owned()),
|
||||
color: group_color("#2563EBFF"),
|
||||
},
|
||||
&context("create-mixed"),
|
||||
)
|
||||
.await
|
||||
.expect("create mixed account group");
|
||||
groups
|
||||
.create_account_group(
|
||||
NewAccountGroup {
|
||||
id: empty_group.clone(),
|
||||
name: "Empty Pool".to_owned(),
|
||||
description: None,
|
||||
color: group_color("#06B6D4CC"),
|
||||
},
|
||||
&context("create-empty"),
|
||||
)
|
||||
.await
|
||||
.expect("create empty account group");
|
||||
assign_accounts(
|
||||
&database.pool,
|
||||
MIXED_GROUP,
|
||||
&["acct_group_openai", "acct_group_xai"],
|
||||
)
|
||||
.await;
|
||||
|
||||
for (id, group_ids) in [
|
||||
("key_group_one", vec![mixed_group.clone()]),
|
||||
("key_group_two", vec![mixed_group.clone()]),
|
||||
("key_empty_pool", vec![empty_group.clone()]),
|
||||
("key_all_accounts", Vec::new()),
|
||||
] {
|
||||
keys.create_client_key(new_key(id, group_ids), &context(id))
|
||||
.await
|
||||
.expect("create scoped client key");
|
||||
}
|
||||
|
||||
let page = groups
|
||||
.list_account_groups(AccountGroupListQuery {
|
||||
page: 1,
|
||||
page_size: PageSize::new(20).expect("page size"),
|
||||
search: None,
|
||||
enabled: None,
|
||||
})
|
||||
.await
|
||||
.expect("list account groups");
|
||||
assert_eq!(page.total, 2);
|
||||
let by_id = page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|group| (group.id.to_string(), group))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let mixed = by_id.get(MIXED_GROUP).expect("mixed group");
|
||||
assert_eq!(mixed.member_count, 2);
|
||||
assert_eq!(
|
||||
mixed.provider_counts,
|
||||
BTreeMap::from([("openai".to_owned(), 1), ("xai".to_owned(), 1)])
|
||||
);
|
||||
assert_eq!(mixed.client_key_count, 2);
|
||||
assert_eq!(mixed.account_summary.available, 2);
|
||||
assert_eq!(mixed.account_summary.limited, 0);
|
||||
assert_eq!(mixed.account_summary.total, 2);
|
||||
assert_eq!(mixed.capacity.used_slots, None);
|
||||
assert_eq!(mixed.capacity.total_slots, 7);
|
||||
assert_eq!(mixed.usage.today_usd.as_str(), "0");
|
||||
assert_eq!(mixed.usage.total_usd.as_str(), "0");
|
||||
let empty = by_id.get(EMPTY_GROUP).expect("empty group");
|
||||
assert_eq!(empty.member_count, 0);
|
||||
assert!(empty.provider_counts.is_empty());
|
||||
assert_eq!(empty.client_key_count, 1);
|
||||
|
||||
let members = groups
|
||||
.account_group_members(&mixed_group)
|
||||
.await
|
||||
.expect("load mixed group members");
|
||||
assert_eq!(members.total, 2);
|
||||
assert_eq!(
|
||||
members
|
||||
.items
|
||||
.iter()
|
||||
.map(|member| member.provider_kind.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
["openai", "xai"]
|
||||
);
|
||||
|
||||
let all_key = keys
|
||||
.reveal_client_key(&client_key_id("key_all_accounts"))
|
||||
.await
|
||||
.expect("reveal all-accounts key")
|
||||
.expect("all-accounts key exists");
|
||||
assert!(all_key.record.groups.is_empty());
|
||||
assert_eq!(
|
||||
all_key
|
||||
.record
|
||||
.provider_kinds
|
||||
.iter()
|
||||
.map(|kind| kind.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
["openai", "xai"]
|
||||
);
|
||||
let empty_pool_key = keys
|
||||
.reveal_client_key(&client_key_id("key_empty_pool"))
|
||||
.await
|
||||
.expect("reveal empty-pool key")
|
||||
.expect("empty-pool key exists");
|
||||
assert_eq!(empty_pool_key.record.groups.len(), 1);
|
||||
assert!(empty_pool_key.record.provider_kinds.is_empty());
|
||||
|
||||
let (scope_revision, widened) = keys
|
||||
.update_client_key(
|
||||
UpdateClientKey {
|
||||
id: client_key_id("key_group_one"),
|
||||
name: "key_group_one".to_owned(),
|
||||
label: None,
|
||||
group_ids: Vec::new(),
|
||||
limits: RateLimits::unlimited(),
|
||||
},
|
||||
&context("widen-group-key"),
|
||||
)
|
||||
.await
|
||||
.expect("widen restricted key to all accounts");
|
||||
assert!(widened.groups.is_empty());
|
||||
let scope_audit: Vec<String> = sqlx::query_scalar(
|
||||
"select changed_fields from admin_audit_events
|
||||
where admin_request_id = 'widen-group-key'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load scope widening audit");
|
||||
assert!(scope_audit.contains(&"routing_scope:groups->all".to_owned()));
|
||||
assert_eq!(current_revision(&database.pool).await, scope_revision.get());
|
||||
let (restricted_revision, restricted) = keys
|
||||
.update_client_key(
|
||||
UpdateClientKey {
|
||||
id: client_key_id("key_group_one"),
|
||||
name: "key_group_one".to_owned(),
|
||||
label: None,
|
||||
group_ids: vec![empty_group],
|
||||
limits: RateLimits::unlimited(),
|
||||
},
|
||||
&context("restrict-all-key"),
|
||||
)
|
||||
.await
|
||||
.expect("restrict all-accounts key to groups");
|
||||
assert_eq!(restricted.groups.len(), 1);
|
||||
let restricted_audit: Vec<String> = sqlx::query_scalar(
|
||||
"select changed_fields from admin_audit_events
|
||||
where admin_request_id = 'restrict-all-key'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load scope restriction audit");
|
||||
assert!(restricted_audit.contains(&"routing_scope:all->groups".to_owned()));
|
||||
assert_eq!(
|
||||
current_revision(&database.pool).await,
|
||||
restricted_revision.get()
|
||||
);
|
||||
|
||||
let revision_before_delete = current_revision(&database.pool).await;
|
||||
let audit_before_delete = audit_count(&database.pool).await;
|
||||
let error = groups
|
||||
.delete_account_group(
|
||||
DeleteAccountGroup { id: mixed_group },
|
||||
&context("delete-referenced"),
|
||||
)
|
||||
.await
|
||||
.expect_err("referenced group must not be deleted");
|
||||
assert_eq!(error.kind(), AdminStoreErrorKind::Conflict);
|
||||
assert_eq!(
|
||||
current_revision(&database.pool).await,
|
||||
revision_before_delete
|
||||
);
|
||||
assert_eq!(audit_count(&database.pool).await, audit_before_delete);
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
fn new_key(id: &str, group_ids: Vec<AccountGroupId>) -> NewClientKey {
|
||||
let marker = char::from(id.as_bytes().last().copied().unwrap_or(b'k'));
|
||||
NewClientKey {
|
||||
id: client_key_id(id),
|
||||
name: id.to_owned(),
|
||||
label: None,
|
||||
group_ids,
|
||||
limits: RateLimits::unlimited(),
|
||||
plaintext: format!("sk_{}", marker.to_string().repeat(43)),
|
||||
}
|
||||
}
|
||||
|
||||
fn group_id(value: &str) -> AccountGroupId {
|
||||
AccountGroupId::new(value).expect("valid account group ID")
|
||||
}
|
||||
|
||||
fn group_color(value: &str) -> AccountGroupColor {
|
||||
AccountGroupColor::parse(value).expect("valid account group color")
|
||||
}
|
||||
|
||||
fn client_key_id(value: &str) -> ClientApiKeyId {
|
||||
ClientApiKeyId::new(value).expect("valid client key ID")
|
||||
}
|
||||
|
||||
fn context(request_id: &str) -> MutationContext {
|
||||
MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: request_id.to_owned(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn seed_account(pool: &sqlx::PgPool, id: &str, provider: &str, name: &str) {
|
||||
sqlx::query(
|
||||
"insert into provider_accounts (
|
||||
id, provider_kind, name, email, upstream_user_id, upstream_account_id,
|
||||
plan_type, authentication_kind, provider_credentials_json, credential_revision,
|
||||
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
|
||||
credential_state, credential_observed_at, created_at, updated_at
|
||||
) values (
|
||||
$1, $2, $3, null, $1 || '-user', null, null, 'oauth', '{}'::jsonb, 1,
|
||||
false, null, null, true, 'ready', now(), now(), now()
|
||||
)",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(provider)
|
||||
.bind(name)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("seed provider account");
|
||||
}
|
||||
|
||||
async fn assign_accounts(pool: &sqlx::PgPool, group_id: &str, account_ids: &[&str]) {
|
||||
for account_id in account_ids {
|
||||
sqlx::query(
|
||||
"insert into account_group_accounts (
|
||||
account_group_id, provider_account_id, created_at
|
||||
)
|
||||
values ($1, $2, now())",
|
||||
)
|
||||
.bind(group_id)
|
||||
.bind(account_id)
|
||||
.execute(pool)
|
||||
.await
|
||||
.expect("seed account group membership");
|
||||
}
|
||||
}
|
||||
|
||||
async fn current_revision(pool: &sqlx::PgPool) -> u64 {
|
||||
let value =
|
||||
sqlx::query_scalar::<_, i64>("select config_revision from runtime_settings where id = 1")
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.expect("load config revision");
|
||||
u64::try_from(value).expect("positive config revision")
|
||||
}
|
||||
|
||||
async fn audit_count(pool: &sqlx::PgPool) -> u64 {
|
||||
let value = sqlx::query_scalar::<_, i64>("select count(*) from admin_audit_events")
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.expect("count audit rows");
|
||||
u64::try_from(value).expect("non-negative audit count")
|
||||
}
|
||||
@@ -93,10 +93,12 @@ async fn seed_request(
|
||||
"insert into model_requests (
|
||||
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
|
||||
client_transport, requested_model_id, outcome,
|
||||
started_at, deadline_at, completed_at
|
||||
started_at, deadline_at, completed_at,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
) values (
|
||||
$1, 'key-recovery', 1, 'openai', 'responses', '/v1/responses',
|
||||
'http_sse', 'coding', $2, $3, $4, $5
|
||||
'http_sse', 'coding', $2, $3, $4, $5,
|
||||
'all', '{}'::text[], '[]'::jsonb
|
||||
)",
|
||||
)
|
||||
.bind(id)
|
||||
|
||||
@@ -23,7 +23,7 @@ fn client_key_requires_the_frozen_plaintext_format() {
|
||||
id: "key-1".to_owned(),
|
||||
name: "default".to_owned(),
|
||||
label: None,
|
||||
provider_kind: "openai".to_owned(),
|
||||
group_ids: Vec::new(),
|
||||
key: format!("sk_{}", "a".repeat(43)),
|
||||
max_concurrency: 0,
|
||||
requests_per_minute: 0,
|
||||
@@ -44,9 +44,9 @@ async fn client_key_list_uses_safe_keyset_search_and_filtered_total() {
|
||||
);
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, label, provider_kind, key, enabled, max_concurrency, requests_per_minute,
|
||||
id, name, label, key, enabled, max_concurrency, requests_per_minute,
|
||||
created_at, updated_at
|
||||
) values ($1, $2, $2, 'openai', $3, true, 0, 0,
|
||||
) values ($1, $2, $2, $3, true, 0, 0,
|
||||
now() - ($4::bigint * interval '1 minute'), now())",
|
||||
)
|
||||
.bind(id)
|
||||
@@ -186,9 +186,9 @@ async fn client_key_database_sort_is_stable_and_keeps_null_last_used_at_last() {
|
||||
{
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
|
||||
id, name, key, enabled, max_concurrency, requests_per_minute,
|
||||
last_used_at, created_at, updated_at
|
||||
) values ($1, $2, 'openai', $3, $4, 0, 0, $5::timestamptz, $6::timestamptz,
|
||||
) values ($1, $2, $3, $4, 0, 0, $5::timestamptz, $6::timestamptz,
|
||||
$6::timestamptz)",
|
||||
)
|
||||
.bind(id)
|
||||
@@ -272,9 +272,9 @@ async fn client_key_usage_touch_should_preserve_the_latest_observed_timestamp()
|
||||
};
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
|
||||
id, name, key, enabled, max_concurrency, requests_per_minute,
|
||||
created_at, updated_at
|
||||
) values ('key_usage_touch', 'usage', 'openai', $1, true, 0, 0, now(), now())",
|
||||
) values ('key_usage_touch', 'usage', $1, true, 0, 0, now(), now())",
|
||||
)
|
||||
.bind(format!("sk_{}", "t".repeat(43)))
|
||||
.execute(&database.pool)
|
||||
@@ -315,9 +315,9 @@ async fn dedicated_reveal_returns_plaintext_without_debug_exposure() {
|
||||
let plaintext = format!("sk_{}", "r".repeat(43));
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
|
||||
id, name, key, enabled, max_concurrency, requests_per_minute,
|
||||
created_at, updated_at
|
||||
) values ('key_reveal', 'reveal', 'openai', $1, true, 1, 2, now(), now())",
|
||||
) values ('key_reveal', 'reveal', $1, true, 1, 2, now(), now())",
|
||||
)
|
||||
.bind(&plaintext)
|
||||
.execute(&database.pool)
|
||||
@@ -340,7 +340,7 @@ fn client_key_debug_redacts_plaintext() {
|
||||
id: "key-1".to_owned(),
|
||||
name: "default".to_owned(),
|
||||
label: None,
|
||||
provider_kind: "openai".to_owned(),
|
||||
group_ids: Vec::new(),
|
||||
key: secret.clone(),
|
||||
max_concurrency: 0,
|
||||
requests_per_minute: 0,
|
||||
|
||||
@@ -29,6 +29,9 @@ fn model_request_rejects_mismatched_client_key_live_id() {
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: "coding".to_owned(),
|
||||
routing_scope: "all".to_owned(),
|
||||
routing_group_refs: Vec::new(),
|
||||
routing_group_names_snapshot: serde_json::json!([]),
|
||||
client_ip: None,
|
||||
user_agent: None,
|
||||
reasoning_effort: None,
|
||||
@@ -76,6 +79,9 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: "coding".to_owned(),
|
||||
routing_scope: "all".to_owned(),
|
||||
routing_group_refs: Vec::new(),
|
||||
routing_group_names_snapshot: serde_json::json!([]),
|
||||
client_ip: None,
|
||||
user_agent: None,
|
||||
reasoning_effort: None,
|
||||
@@ -153,6 +159,66 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn model_request_persists_group_routing_snapshot_without_live_group_foreign_keys() {
|
||||
let Some(database) = TestDatabase::create("execution_group_routing_history").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgExecutionStore::new(database.pool.clone());
|
||||
let started_at = Utc::now();
|
||||
repository
|
||||
.insert_model_request(NewModelRequest {
|
||||
id: "req_group_history".to_owned(),
|
||||
client_api_key_id: None,
|
||||
client_api_key_ref: "key_group_history".to_owned(),
|
||||
config_revision: 7,
|
||||
routing_scope: "groups".to_owned(),
|
||||
routing_group_refs: vec![
|
||||
"grp_00000000000000000000000000000001".to_owned(),
|
||||
"grp_00000000000000000000000000000002".to_owned(),
|
||||
],
|
||||
routing_group_names_snapshot: serde_json::json!(["Production Pool", "Overflow Pool"]),
|
||||
protocol: "openai".to_owned(),
|
||||
operation: "responses".to_owned(),
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: "coding".to_owned(),
|
||||
client_ip: None,
|
||||
user_agent: None,
|
||||
reasoning_effort: None,
|
||||
reasoning_preset: None,
|
||||
request_kind: None,
|
||||
subagent_kind: None,
|
||||
compact: false,
|
||||
image_generation_requested: false,
|
||||
started_at,
|
||||
deadline_at: started_at + Duration::seconds(30),
|
||||
})
|
||||
.await
|
||||
.expect("insert grouped request history");
|
||||
|
||||
let stored: (String, Vec<String>, serde_json::Value) = sqlx::query_as(
|
||||
"select routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
from model_requests where id = 'req_group_history'",
|
||||
)
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load grouped request history");
|
||||
assert_eq!(stored.0, "groups");
|
||||
assert_eq!(
|
||||
stored.1,
|
||||
[
|
||||
"grp_00000000000000000000000000000001",
|
||||
"grp_00000000000000000000000000000002",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
stored.2,
|
||||
serde_json::json!(["Production Pool", "Overflow Pool"])
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn downstream_commit_should_atomically_record_http_status_once() {
|
||||
let Some(database) = TestDatabase::create("execution_commit_status").await else {
|
||||
@@ -483,9 +549,11 @@ async fn seed_running_request(pool: &sqlx::PgPool, id: &str) -> Result<(), sqlx:
|
||||
"insert into model_requests (
|
||||
id, client_api_key_ref, config_revision, protocol, operation, endpoint,
|
||||
client_transport, requested_model_id, provider_kind, provider_account_ref, cost_source,
|
||||
started_at, deadline_at
|
||||
started_at, deadline_at,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
) values ($1, 'key_status', 1, 'openai_responses', 'generate', '/v1/responses',
|
||||
'http_json', 'status-model', 'openai', 'acct_status', 'unavailable', now(), now() + interval '1 minute')",
|
||||
'http_json', 'status-model', 'openai', 'acct_status', 'unavailable', now(), now() + interval '1 minute',
|
||||
'all', '{}'::text[], '[]'::jsonb)",
|
||||
)
|
||||
.bind(id)
|
||||
.execute(pool)
|
||||
|
||||
@@ -7,6 +7,7 @@ use sqlx::{
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
mod account_groups;
|
||||
mod admin_security_audit;
|
||||
mod admission_recovery;
|
||||
mod backup;
|
||||
@@ -146,6 +147,29 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
|
||||
.fetch_one(&second)
|
||||
.await
|
||||
.expect("check removed raw response ID index");
|
||||
let legacy_key_provider_column_exists = sqlx::query_scalar::<_, bool>(
|
||||
"select exists (
|
||||
select 1 from information_schema.columns
|
||||
where table_schema = 'public'
|
||||
and table_name = 'client_api_keys'
|
||||
and column_name = 'provider_kind'
|
||||
)",
|
||||
)
|
||||
.fetch_one(&second)
|
||||
.await
|
||||
.expect("check removed client key provider column");
|
||||
let routing_history_columns = sqlx::query_scalar::<_, String>(
|
||||
"select column_name from information_schema.columns
|
||||
where table_schema = 'public'
|
||||
and table_name = 'model_requests'
|
||||
and column_name in (
|
||||
'routing_scope', 'routing_group_refs', 'routing_group_names_snapshot'
|
||||
)
|
||||
order by column_name",
|
||||
)
|
||||
.fetch_all(&second)
|
||||
.await
|
||||
.expect("load routing history columns");
|
||||
second.close().await;
|
||||
|
||||
sqlx::raw_sql(sqlx::AssertSqlSafe(format!(
|
||||
@@ -156,7 +180,7 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
|
||||
.expect("drop migration test database");
|
||||
admin.close().await;
|
||||
|
||||
assert_eq!(first_table_count, 10);
|
||||
assert_eq!(first_table_count, 13);
|
||||
assert_eq!(
|
||||
migration_count,
|
||||
i64::try_from(TEST_MIGRATOR.iter().count())
|
||||
@@ -164,6 +188,15 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
|
||||
);
|
||||
assert_eq!(response_id_types, ["bytea", "bytea"]);
|
||||
assert!(!raw_response_id_index_exists);
|
||||
assert!(!legacy_key_provider_column_exists);
|
||||
assert_eq!(
|
||||
routing_history_columns,
|
||||
[
|
||||
"routing_group_names_snapshot",
|
||||
"routing_group_refs",
|
||||
"routing_scope",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -550,6 +550,12 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
.await
|
||||
.expect("fully filtered usage page");
|
||||
assert_eq!(filtered.total, 1);
|
||||
assert_eq!(filtered.items[0].routing_scope, "groups");
|
||||
assert_eq!(filtered.items[0].routing_group_refs, ["grp_history"]);
|
||||
assert_eq!(
|
||||
filtered.items[0].routing_group_names_snapshot,
|
||||
["Historical group"]
|
||||
);
|
||||
assert_eq!(filtered.items[0].service_tier.as_deref(), Some("priority"));
|
||||
assert_eq!(
|
||||
filtered.items[0].upstream_request_id.as_deref(),
|
||||
@@ -584,6 +590,12 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
|
||||
admin_observability::RequestOutcome::Succeeded
|
||||
);
|
||||
assert_eq!(detail.request.service_tier.as_deref(), Some("priority"));
|
||||
assert_eq!(detail.request.routing_scope, "groups");
|
||||
assert_eq!(detail.request.routing_group_refs, ["grp_history"]);
|
||||
assert_eq!(
|
||||
detail.request.routing_group_names_snapshot,
|
||||
["Historical group"]
|
||||
);
|
||||
assert_eq!(detail.attempts.len(), 1);
|
||||
assert_eq!(
|
||||
detail.attempts[0].outcome,
|
||||
@@ -1021,7 +1033,8 @@ async fn seed_observability_facts(
|
||||
provider_account_authentication_kind_snapshot,
|
||||
attempt_count, upstream_send_state, outcome, client_status_code,
|
||||
input_tokens, output_tokens, total_tokens, cost_source, latency_ms,
|
||||
started_at, deadline_at, completed_at
|
||||
started_at, deadline_at, completed_at,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
) values (
|
||||
'req_observe_uncommitted', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
|
||||
'http_sse', 'public-model', 'openai', 'acct_observe',
|
||||
@@ -1029,7 +1042,8 @@ async fn seed_observability_facts(
|
||||
'primary', 'account@example.invalid', 'oauth',
|
||||
1, 'sent', 'succeeded', 200,
|
||||
900, 900, 1800, 'unavailable', 650,
|
||||
$1 - interval '15 minutes', $1 + interval '10 minutes', $1 - interval '14 minutes'
|
||||
$1 - interval '15 minutes', $1 + interval '10 minutes', $1 - interval '14 minutes',
|
||||
'all', '{}'::text[], '[]'::jsonb
|
||||
)",
|
||||
)
|
||||
.bind(now)
|
||||
@@ -1051,7 +1065,8 @@ async fn seed_observability_facts(
|
||||
image_input_tokens, image_output_tokens, total_tokens,
|
||||
image_generation_requested, image_generation_succeeded,
|
||||
cost_source, cost_amount, cost_currency, first_token_ms, latency_ms,
|
||||
started_at, deadline_at, completed_at
|
||||
started_at, deadline_at, completed_at,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
) values (
|
||||
'req_observe_success', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
|
||||
'http_sse', 'public-model', 'openai', 'acct_observe',
|
||||
@@ -1062,7 +1077,8 @@ async fn seed_observability_facts(
|
||||
'resp_observe_success', 'upstream_req_success', 'upstream_resp_success',
|
||||
100, 20, 40, 3, 5, 31, 9, 120, true, true,
|
||||
'provider_reported', 1.25, 'USD', 120, 900,
|
||||
$1 - interval '20 minutes', $1 + interval '10 minutes', $1 - interval '19 minutes'
|
||||
$1 - interval '20 minutes', $1 + interval '10 minutes', $1 - interval '19 minutes',
|
||||
'groups', array['grp_history'], jsonb_build_array('Historical group')
|
||||
)",
|
||||
)
|
||||
.bind(now)
|
||||
@@ -1080,7 +1096,8 @@ async fn seed_observability_facts(
|
||||
error_kind, provider_error_code, error_message, retry_after_ms,
|
||||
input_tokens, cached_tokens, image_generation_requested,
|
||||
image_generation_succeeded, cost_source, latency_ms,
|
||||
started_at, deadline_at, completed_at
|
||||
started_at, deadline_at, completed_at,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
) values (
|
||||
'req_observe_failed', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
|
||||
'http_sse', 'public-model', 'openai', 'acct_observe',
|
||||
@@ -1089,7 +1106,8 @@ async fn seed_observability_facts(
|
||||
'primary', 'account@example.invalid', 'oauth',
|
||||
'sent', 'failed', 502, 429, 'rate_limited', 'rate_limit',
|
||||
'upstream limited', 1000, 0, 0, true, false, 'unavailable', 700,
|
||||
$1 - interval '10 minutes', $1 + interval '20 minutes', $1 - interval '9 minutes'
|
||||
$1 - interval '10 minutes', $1 + interval '20 minutes', $1 - interval '9 minutes',
|
||||
'all', '{}'::text[], '[]'::jsonb
|
||||
)",
|
||||
)
|
||||
.bind(now)
|
||||
@@ -1131,7 +1149,8 @@ async fn seed_calculated_billing_facts(
|
||||
provider_account_authentication_kind_snapshot,
|
||||
upstream_send_state, downstream_committed_at, outcome, client_status_code,
|
||||
input_tokens, output_tokens, cached_tokens, cache_write_tokens, total_tokens, service_tier,
|
||||
cost_source, cost_amount, cost_currency, started_at, deadline_at, completed_at
|
||||
cost_source, cost_amount, cost_currency, started_at, deadline_at, completed_at,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
) values (
|
||||
'req_observe_calculated', 'key_observe', 1, 'openai', 'responses', '/v1/responses',
|
||||
'http_sse', 'public-model', 'openai', 'acct_observe', 'acct_observe', 'gpt-5.5',
|
||||
@@ -1139,7 +1158,8 @@ async fn seed_calculated_billing_facts(
|
||||
'primary', 'account@example.invalid', 'oauth',
|
||||
'sent', $1 - interval '29 minutes', 'succeeded', 200,
|
||||
800, 200, 0, 0, 1000, 'priority', 'calculated', 1.25, 'USD',
|
||||
$1 - interval '30 minutes', $1 + interval '10 minutes', $1 - interval '29 minutes'
|
||||
$1 - interval '30 minutes', $1 + interval '10 minutes', $1 - interval '29 minutes',
|
||||
'all', '{}'::text[], '[]'::jsonb
|
||||
)",
|
||||
)
|
||||
.bind(now)
|
||||
@@ -1154,7 +1174,7 @@ async fn seed_calculated_billing_facts(
|
||||
provider_account_authentication_kind_snapshot,
|
||||
upstream_send_state, outcome, client_status_code, input_tokens, output_tokens,
|
||||
total_tokens, cost_source, cost_amount, cost_currency, started_at, deadline_at,
|
||||
completed_at
|
||||
completed_at, routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
) values (
|
||||
'req_observe_calculated_uncommitted', 'key_observe', 1, 'openai', 'responses',
|
||||
'/v1/responses', 'http_sse', 'public-model', 'openai', 'acct_observe',
|
||||
@@ -1162,7 +1182,8 @@ async fn seed_calculated_billing_facts(
|
||||
'primary', 'account@example.invalid', 'oauth',
|
||||
'sent', 'succeeded', 200, 800, 200,
|
||||
1000, 'calculated', 1.25, 'USD',
|
||||
$1 - interval '40 minutes', $1 + interval '10 minutes', $1 - interval '39 minutes'
|
||||
$1 - interval '40 minutes', $1 + interval '10 minutes', $1 - interval '39 minutes',
|
||||
'all', '{}'::text[], '[]'::jsonb
|
||||
)",
|
||||
)
|
||||
.bind(now)
|
||||
|
||||
@@ -6,8 +6,8 @@ use gateway_admin::{
|
||||
MutationActor, MutationContext, PageSize,
|
||||
accounts::{
|
||||
AccountListQuery, AccountSort, AccountSortField, AccountStatus,
|
||||
AccountUsageWindowQuery, DeleteAccounts, SetAccountEnabled as AdminSetAccountEnabled,
|
||||
SortDirection,
|
||||
AccountUsageWindowQuery, BatchUpdateAccounts, DeleteAccounts, SortDirection,
|
||||
UpdateAccount,
|
||||
},
|
||||
observability::TimeRange,
|
||||
provider_credentials::{CredentialListQuery, CredentialListWindow, CredentialStateFilter},
|
||||
@@ -21,14 +21,14 @@ use gateway_core::engine::credential::{
|
||||
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState,
|
||||
QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use gateway_core::routing::{AccountGroupId, ProviderKind};
|
||||
use gateway_store::{
|
||||
ConflictKind, JsonObject, Revision, StoreError,
|
||||
postgres::{
|
||||
AdminAuditActorKind, AdminAuditEvent, DeleteProviderAccounts, ImportProviderAccounts,
|
||||
NewProviderAccount, PgAdminAccountStore, PgProviderAccountRepository,
|
||||
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountRepository,
|
||||
ProviderCredentialUpdate, RotateProviderAccount, SetProviderAccountEnabled,
|
||||
AdminAuditActorKind, AdminAuditEvent, BatchUpdateProviderAccountsAdmin,
|
||||
DeleteProviderAccounts, ImportProviderAccounts, NewProviderAccount, PgAdminAccountStore,
|
||||
PgProviderAccountRepository, ProviderAccountAdminRepository, ProviderAccountAdminScope,
|
||||
ProviderAccountRepository, ProviderCredentialUpdate, RotateProviderAccount,
|
||||
UpdateProviderAccount,
|
||||
},
|
||||
};
|
||||
@@ -392,6 +392,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
page: 1,
|
||||
page_size: PageSize::new(2).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: Some(AccountSort {
|
||||
@@ -431,6 +432,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
page: 1,
|
||||
page_size: PageSize::new(2).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: None,
|
||||
sort: Some(AccountSort {
|
||||
@@ -454,6 +456,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: Some(ProviderKind::new("openai").expect("Provider kind")),
|
||||
group_filter: None,
|
||||
search: Some("ALPHA@EXAMPLE".to_owned()),
|
||||
status: Some(AccountStatus::Normal),
|
||||
sort: None,
|
||||
@@ -470,6 +473,7 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
|
||||
page: 1,
|
||||
page_size: PageSize::new(10).expect("page size"),
|
||||
provider_kind: None,
|
||||
group_filter: None,
|
||||
search: None,
|
||||
status: Some(AccountStatus::Error),
|
||||
sort: None,
|
||||
@@ -725,17 +729,20 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
|
||||
request_id: "request_terminal_mutation".to_owned(),
|
||||
};
|
||||
|
||||
let revision = store
|
||||
.set_account_enabled(
|
||||
AdminSetAccountEnabled {
|
||||
let result = store
|
||||
.update_account(
|
||||
UpdateAccount {
|
||||
account_id: "acct_terminal_mutation".to_owned(),
|
||||
enabled: false,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
|
||||
group_ids: Vec::new(),
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.expect("disable account atomically");
|
||||
assert_eq!(revision.get(), 2);
|
||||
assert_eq!(result.config_revision.get(), 2);
|
||||
let enabled: bool = sqlx::query_scalar(
|
||||
"select enabled from provider_accounts where id = 'acct_terminal_mutation'",
|
||||
)
|
||||
@@ -768,7 +775,16 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
|
||||
assert_eq!(
|
||||
audit_rows,
|
||||
vec![
|
||||
("disable".to_owned(), 2, vec!["enabled".to_owned()]),
|
||||
(
|
||||
"update".to_owned(),
|
||||
2,
|
||||
vec![
|
||||
"enabled".to_owned(),
|
||||
"concurrency_limit".to_owned(),
|
||||
"weight".to_owned(),
|
||||
"groups".to_owned(),
|
||||
],
|
||||
),
|
||||
("delete".to_owned(), 3, Vec::new()),
|
||||
]
|
||||
);
|
||||
@@ -776,6 +792,118 @@ async fn terminal_admin_mutations_keep_revision_account_and_audit_atomic() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_batch_update_replaces_state_and_groups_once_or_rolls_back_everything() {
|
||||
const GROUP_ID: &str = "grp_00000000000000000000000000000071";
|
||||
let Some(database) = TestDatabase::create("provider_account_batch_update").await else {
|
||||
return;
|
||||
};
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
for (account_id, upstream_user_id) in [
|
||||
("acct_batch_a", "user-batch-a"),
|
||||
("acct_batch_b", "user-batch-b"),
|
||||
] {
|
||||
repository
|
||||
.insert_provider_account(account(account_id, upstream_user_id))
|
||||
.await
|
||||
.expect("insert batch account");
|
||||
}
|
||||
sqlx::query(
|
||||
"insert into account_groups (id, name, description, color, enabled, created_at, updated_at)
|
||||
values ($1, 'Batch group', null, '#2563EBFF', true, now(), now())",
|
||||
)
|
||||
.bind(GROUP_ID)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("insert batch account group");
|
||||
let store = PgAdminAccountStore::new(database.pool.clone(), None);
|
||||
let account_ids = vec!["acct_batch_a".to_owned(), "acct_batch_b".to_owned()];
|
||||
let context = MutationContext {
|
||||
actor: MutationActor::System,
|
||||
request_id: "request_batch_update".to_owned(),
|
||||
};
|
||||
|
||||
let result = store
|
||||
.batch_update_accounts(
|
||||
BatchUpdateAccounts {
|
||||
account_ids: account_ids.clone(),
|
||||
enabled: false,
|
||||
concurrency_limit: gateway_core::engine::credential::AccountConcurrencyLimit::new(
|
||||
7,
|
||||
),
|
||||
weight: gateway_core::engine::credential::AccountWeight::new(25).expect("weight"),
|
||||
group_ids: vec![AccountGroupId::new(GROUP_ID).expect("group ID")],
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.expect("batch update accounts");
|
||||
|
||||
assert_eq!(result.config_revision.get(), 2);
|
||||
assert_eq!(result.account_ids.len(), 2);
|
||||
let scheduling: Vec<(bool, Option<i64>, i16)> = sqlx::query_as(
|
||||
"select enabled, concurrency_limit, weight
|
||||
from provider_accounts where id = any($1::text[]) order by id",
|
||||
)
|
||||
.bind(&account_ids)
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.expect("load batch account state");
|
||||
assert_eq!(scheduling, [(false, Some(7), 25), (false, Some(7), 25)]);
|
||||
for account_id in &account_ids {
|
||||
assert_eq!(
|
||||
account_group_ids(&database.pool, account_id).await,
|
||||
[GROUP_ID]
|
||||
);
|
||||
}
|
||||
assert_eq!(audit_count(&database.pool, &context.request_id).await, 1);
|
||||
|
||||
let revision_before_failure = current_revision(&database.pool).await;
|
||||
let audit_before_failure = audit_count(&database.pool, &context.request_id).await;
|
||||
store
|
||||
.batch_update_accounts(
|
||||
BatchUpdateAccounts {
|
||||
account_ids: account_ids.clone(),
|
||||
enabled: true,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
|
||||
group_ids: vec![
|
||||
AccountGroupId::new("grp_00000000000000000000000000000072")
|
||||
.expect("missing group ID"),
|
||||
],
|
||||
},
|
||||
&context,
|
||||
)
|
||||
.await
|
||||
.expect_err("missing group must roll back the batch");
|
||||
|
||||
assert_eq!(
|
||||
current_revision(&database.pool).await,
|
||||
revision_before_failure
|
||||
);
|
||||
assert_eq!(
|
||||
audit_count(&database.pool, &context.request_id).await,
|
||||
audit_before_failure
|
||||
);
|
||||
let scheduling: Vec<(bool, Option<i64>, i16)> = sqlx::query_as(
|
||||
"select enabled, concurrency_limit, weight
|
||||
from provider_accounts where id = any($1::text[]) order by id",
|
||||
)
|
||||
.bind(&account_ids)
|
||||
.fetch_all(&database.pool)
|
||||
.await
|
||||
.expect("load rolled back account state");
|
||||
assert_eq!(scheduling, [(false, Some(7), 25), (false, Some(7), 25)]);
|
||||
for account_id in &account_ids {
|
||||
assert_eq!(
|
||||
account_group_ids(&database.pool, account_id).await,
|
||||
[GROUP_ID]
|
||||
);
|
||||
}
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_admin_delete_removes_enabled_accounts_in_one_transaction() {
|
||||
let Some(database) = TestDatabase::create("provider_account_enabled_delete").await else {
|
||||
@@ -916,6 +1044,8 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
|
||||
access_token_expires_at: Some(Utc::now() + TimeDelta::minutes(5)),
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
|
||||
credential_state: CredentialState::Ready,
|
||||
credential_observed_at: Utc::now(),
|
||||
})
|
||||
@@ -1190,10 +1320,12 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
assert_eq!(unchanged.4.as_deref(), Some("workspace-admin-rebound"));
|
||||
|
||||
let revision = repository
|
||||
.set_provider_account_enabled_admin(SetProviderAccountEnabled {
|
||||
scope: scope.clone(),
|
||||
account_id: "acct_admin_a".to_owned(),
|
||||
.batch_update_provider_accounts_admin(BatchUpdateProviderAccountsAdmin {
|
||||
account_ids: vec!["acct_admin_a".to_owned()],
|
||||
enabled: false,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
|
||||
group_ids: Vec::new(),
|
||||
audit: audit("audit_account_disable", "disable", "acct_admin_a"),
|
||||
})
|
||||
.await
|
||||
@@ -1230,6 +1362,116 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_account_import_and_reauthorization_preserve_existing_memberships() {
|
||||
const GROUP_ID: &str = "grp_00000000000000000000000000000091";
|
||||
let Some(database) = TestDatabase::create("provider_account_group_assignment").await else {
|
||||
return;
|
||||
};
|
||||
sqlx::query(
|
||||
"insert into account_groups
|
||||
(id, name, description, color, enabled, created_at, updated_at)
|
||||
values ($1, 'Credential group', null, '#2563EBFF', true, now(), now())",
|
||||
)
|
||||
.bind(GROUP_ID)
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("seed account group");
|
||||
let repository = PgProviderAccountRepository::new(database.pool.clone());
|
||||
let scope = ProviderAccountAdminScope {
|
||||
provider_kind: "openai".to_owned(),
|
||||
};
|
||||
let imported = repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
scope: scope.clone(),
|
||||
accounts: vec![account("acct_grouped_import", "user-grouped-import")],
|
||||
audit: audit("audit_grouped_import", "import", "acct_grouped_import"),
|
||||
})
|
||||
.await
|
||||
.expect("import account without a group");
|
||||
assert_eq!(imported.config_revision.get(), 2);
|
||||
assert!(
|
||||
account_group_ids(&database.pool, "acct_grouped_import")
|
||||
.await
|
||||
.is_empty()
|
||||
);
|
||||
assert_eq!(
|
||||
audit_count(&database.pool, "request:audit_grouped_import").await,
|
||||
1
|
||||
);
|
||||
|
||||
sqlx::query(
|
||||
"insert into account_group_accounts (account_group_id, provider_account_id, created_at)
|
||||
values ($1, $2, now())",
|
||||
)
|
||||
.bind(GROUP_ID)
|
||||
.bind("acct_grouped_import")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("assign existing account group");
|
||||
sqlx::query("update provider_accounts set concurrency_limit = 9, weight = 40 where id = $1")
|
||||
.bind("acct_grouped_import")
|
||||
.execute(&database.pool)
|
||||
.await
|
||||
.expect("set account scheduling before reimport");
|
||||
repository
|
||||
.import_provider_accounts(ImportProviderAccounts {
|
||||
scope: scope.clone(),
|
||||
accounts: vec![account("acct_reimport_candidate", "user-grouped-import")],
|
||||
audit: audit("audit_grouped_reimport", "import", "acct_grouped_import"),
|
||||
})
|
||||
.await
|
||||
.expect("reimport existing identity");
|
||||
let scheduling: (Option<i64>, i16) =
|
||||
sqlx::query_as("select concurrency_limit, weight from provider_accounts where id = $1")
|
||||
.bind("acct_grouped_import")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load preserved scheduling");
|
||||
assert_eq!(scheduling, (Some(9), 40));
|
||||
assert_eq!(
|
||||
account_group_ids(&database.pool, "acct_grouped_import").await,
|
||||
[GROUP_ID]
|
||||
);
|
||||
|
||||
let revision_before_reauthorization = current_revision(&database.pool).await;
|
||||
let reauthorized = repository
|
||||
.rotate_provider_account(RotateProviderAccount {
|
||||
scope,
|
||||
profile: profile("acct_grouped_import", "reauthorized"),
|
||||
replacement_identity: None,
|
||||
credential: credential_update("acct_grouped_import", 2, "reauthorized-secret"),
|
||||
audit: audit(
|
||||
"audit_clear_groups_reauthorize",
|
||||
"reauthorize",
|
||||
"acct_grouped_import",
|
||||
),
|
||||
})
|
||||
.await
|
||||
.expect("reauthorization preserves groups atomically");
|
||||
assert_eq!(
|
||||
reauthorized.config_revision.get(),
|
||||
u64::try_from(revision_before_reauthorization + 1).expect("revision")
|
||||
);
|
||||
assert_eq!(
|
||||
account_group_ids(&database.pool, "acct_grouped_import").await,
|
||||
[GROUP_ID]
|
||||
);
|
||||
assert_eq!(
|
||||
audit_count(&database.pool, "request:audit_clear_groups_reauthorize").await,
|
||||
1
|
||||
);
|
||||
let scheduling: (Option<i64>, i16) =
|
||||
sqlx::query_as("select concurrency_limit, weight from provider_accounts where id = $1")
|
||||
.bind("acct_grouped_import")
|
||||
.fetch_one(&database.pool)
|
||||
.await
|
||||
.expect("load scheduling after reauthorization");
|
||||
assert_eq!(scheduling, (Some(9), 40));
|
||||
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn verified_credential_rotation_preserves_quota_exhaustion() {
|
||||
let Some(database) = TestDatabase::create("provider_account_rotation_quota").await else {
|
||||
@@ -1450,6 +1692,8 @@ fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
|
||||
access_token_expires_at: Some(Utc::now() + TimeDelta::hours(1)),
|
||||
next_refresh_at: None,
|
||||
enabled: true,
|
||||
concurrency_limit: None,
|
||||
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
|
||||
credential_state: CredentialState::Ready,
|
||||
credential_observed_at: Utc::now(),
|
||||
}
|
||||
@@ -1532,13 +1776,15 @@ async fn seed_model_request(
|
||||
upstream_status_code,
|
||||
input_tokens, output_tokens, cached_tokens, cache_write_tokens, reasoning_tokens,
|
||||
total_tokens, cost_source, cost_amount, cost_currency,
|
||||
started_at, deadline_at, completed_at
|
||||
started_at, deadline_at, completed_at,
|
||||
routing_scope, routing_group_refs, routing_group_names_snapshot
|
||||
) values (
|
||||
$1, 'key-provider-account-test', 1, 'openai', 'responses', '/v1/responses',
|
||||
'http_sse', $4, $3, $2, $2, $4, 'http_sse', 1,
|
||||
'sent', $7 + interval '1 second', 'succeeded', 200, 200, $5, 0, 0, 0, 0,
|
||||
$5, 'provider_reported', $6::numeric, 'USD', $7,
|
||||
$7 + interval '5 minutes', $7 + interval '1 second'
|
||||
$7 + interval '5 minutes', $7 + interval '1 second',
|
||||
'all', '{}'::text[], '[]'::jsonb
|
||||
)",
|
||||
)
|
||||
.bind(seed.request_id)
|
||||
@@ -1568,6 +1814,25 @@ async fn account_count(pool: &sqlx::PgPool, account_id: &str) -> i64 {
|
||||
.expect("count provider account")
|
||||
}
|
||||
|
||||
async fn account_group_ids(pool: &sqlx::PgPool, account_id: &str) -> Vec<String> {
|
||||
sqlx::query_scalar(
|
||||
"select account_group_id from account_group_accounts
|
||||
where provider_account_id = $1 order by account_group_id",
|
||||
)
|
||||
.bind(account_id)
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.expect("load account groups")
|
||||
}
|
||||
|
||||
async fn audit_count(pool: &sqlx::PgPool, request_id: &str) -> i64 {
|
||||
sqlx::query_scalar("select count(*) from admin_audit_events where admin_request_id = $1")
|
||||
.bind(request_id)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.expect("count audit events")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_credentials_are_redacted_from_debug() {
|
||||
let secret = "secret-access-token";
|
||||
|
||||
@@ -10,14 +10,14 @@ fn snapshot_client_policy_contains_only_common_limits() {
|
||||
let policy = ClientApiKeySnapshot {
|
||||
id: ClientApiKeyId::new("key-1").expect("client key ID"),
|
||||
plaintext_key: PlaintextClientApiKey::new("sk_snapshot_secret").expect("plaintext key"),
|
||||
provider_kind: "openai".to_owned(),
|
||||
group_ids: Vec::new(),
|
||||
limits: RateLimits {
|
||||
max_concurrency: 3,
|
||||
requests_per_minute: 60,
|
||||
},
|
||||
};
|
||||
assert_eq!(policy.limits.max_concurrency, 3);
|
||||
assert_eq!(policy.provider_kind, "openai");
|
||||
assert!(policy.group_ids.is_empty());
|
||||
assert!(!format!("{policy:?}").contains("sk_snapshot_secret"));
|
||||
}
|
||||
|
||||
@@ -29,9 +29,8 @@ async fn runtime_snapshot_loads_enabled_plaintext_key_without_debug_exposure() {
|
||||
let plaintext = format!("sk_{}", "s".repeat(43));
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (
|
||||
id, name, provider_kind, key, enabled, max_concurrency, requests_per_minute,
|
||||
created_at, updated_at
|
||||
) values ('key_snapshot', 'snapshot', 'xai', $1, true, 2, 60, now(), now())",
|
||||
id, name, key, enabled, max_concurrency, requests_per_minute, created_at, updated_at
|
||||
) values ('key_snapshot', 'snapshot', $1, true, 2, 60, now(), now())",
|
||||
)
|
||||
.bind(&plaintext)
|
||||
.execute(&database.pool)
|
||||
@@ -42,7 +41,7 @@ async fn runtime_snapshot_loads_enabled_plaintext_key_without_debug_exposure() {
|
||||
.await
|
||||
.expect("load runtime snapshot");
|
||||
assert_eq!(snapshot.client_api_keys.len(), 1);
|
||||
assert_eq!(snapshot.client_api_keys[0].provider_kind, "xai");
|
||||
assert!(snapshot.client_api_keys[0].group_ids.is_empty());
|
||||
assert_eq!(
|
||||
snapshot.client_api_keys[0].plaintext_key.expose_for_auth(),
|
||||
plaintext
|
||||
|
||||
@@ -665,6 +665,9 @@ fn new_request(id: &str, started_at: DateTime<Utc>) -> NewModelRequest {
|
||||
endpoint: "/v1/responses".to_owned(),
|
||||
client_transport: "http_sse".to_owned(),
|
||||
requested_model_id: "coding".to_owned(),
|
||||
routing_scope: "all".to_owned(),
|
||||
routing_group_refs: Vec::new(),
|
||||
routing_group_names_snapshot: serde_json::json!([]),
|
||||
client_ip: None,
|
||||
user_agent: None,
|
||||
reasoning_effort: None,
|
||||
@@ -756,8 +759,8 @@ async fn finalize_request(pool: &PgPool, id: &str, started_at: DateTime<Utc>) {
|
||||
|
||||
async fn seed_api_key(pool: &PgPool, id: &str, name: &str, now: DateTime<Utc>) {
|
||||
sqlx::query(
|
||||
"insert into client_api_keys (id, name, provider_kind, key, enabled, created_at, updated_at)
|
||||
values ($1, $2, 'openai', 'sk_' || $3, true, $4, $4)",
|
||||
"insert into client_api_keys (id, name, key, enabled, created_at, updated_at)
|
||||
values ($1, $2, 'sk_' || $3, true, $4, $4)",
|
||||
)
|
||||
.bind(id)
|
||||
.bind(name)
|
||||
|
||||
@@ -211,30 +211,41 @@ async fn scheduling_lease_ttl_and_cache_loss_rebuild_empty_runtime_signal() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn scheduling_cursor_is_shared_across_gateway_processes() {
|
||||
async fn scheduling_cursor_is_shared_across_processes_and_isolated_by_client_and_provider() {
|
||||
let Some((repository, mut connection, namespace)) = repository().await else {
|
||||
return;
|
||||
};
|
||||
let peer = RedisCredentialLeaseRepository::new(connection.clone(), &namespace)
|
||||
.expect("peer repository");
|
||||
|
||||
let primary = gateway_core::policy::ClientApiKeyId::new("key_primary").expect("client key");
|
||||
let secondary = gateway_core::policy::ClientApiKeyId::new("key_secondary").expect("client key");
|
||||
let openai = gateway_core::routing::ProviderKind::new("openai").expect("provider");
|
||||
let xai = gateway_core::routing::ProviderKind::new("xai").expect("provider");
|
||||
|
||||
assert_eq!(
|
||||
repository
|
||||
.advance_scheduling_cursor("route_openai_real")
|
||||
.advance_scheduling_cursor(&primary, &openai)
|
||||
.await
|
||||
.expect("first cursor"),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
peer.advance_scheduling_cursor("route_openai_real")
|
||||
peer.advance_scheduling_cursor(&primary, &openai)
|
||||
.await
|
||||
.expect("shared cursor"),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
peer.advance_scheduling_cursor("route_openai_backup")
|
||||
peer.advance_scheduling_cursor(&secondary, &openai)
|
||||
.await
|
||||
.expect("independent cursor"),
|
||||
.expect("client-isolated cursor"),
|
||||
0
|
||||
);
|
||||
assert_eq!(
|
||||
peer.advance_scheduling_cursor(&primary, &xai)
|
||||
.await
|
||||
.expect("provider-isolated cursor"),
|
||||
0
|
||||
);
|
||||
|
||||
@@ -243,12 +254,13 @@ async fn scheduling_cursor_is_shared_across_gateway_processes() {
|
||||
.query_async::<Vec<String>>(&mut connection)
|
||||
.await
|
||||
.expect("list cursor keys");
|
||||
assert_eq!(keys.len(), 2);
|
||||
assert!(
|
||||
keys.iter().all(|key| {
|
||||
!key.contains("route_openai_real") && !key.contains("route_openai_backup")
|
||||
})
|
||||
);
|
||||
assert_eq!(keys.len(), 3);
|
||||
assert!(keys.iter().all(|key| {
|
||||
!key.contains("key_primary")
|
||||
&& !key.contains("key_secondary")
|
||||
&& !key.contains("openai")
|
||||
&& !key.contains("xai")
|
||||
}));
|
||||
redis::cmd("DEL")
|
||||
.arg(keys)
|
||||
.query_async::<i64>(&mut connection)
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
use gateway_core::engine::continuation::{
|
||||
NativeContinuationPin, NativeContinuationPort, NativeContinuationScope, PreviousResponseId,
|
||||
NativeContinuationPin, NativeContinuationPort, NativeContinuationScope,
|
||||
NativeContinuationStoreErrorKind, PreviousResponseId,
|
||||
};
|
||||
use gateway_core::engine::credential::ProviderAccountId;
|
||||
use gateway_core::operation::ProviderSessionState;
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use gateway_store::redis::RedisNativeContinuationRepository;
|
||||
use redis::aio::ConnectionManager;
|
||||
@@ -30,6 +32,7 @@ async fn native_continuation_round_trips_opaque_state_without_exposing_response_
|
||||
NativeContinuationPin::new(
|
||||
previous_response_id.clone(),
|
||||
previous_response_id.clone(),
|
||||
client_key("key_primary"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_primary").expect("account"),
|
||||
)
|
||||
@@ -40,7 +43,7 @@ async fn native_continuation_round_trips_opaque_state_without_exposing_response_
|
||||
.expect("record continuation");
|
||||
|
||||
let resolved = repository
|
||||
.resolve(&previous_response_id)
|
||||
.resolve(&client_key("key_primary"), &previous_response_id)
|
||||
.await
|
||||
.expect("resolve continuation")
|
||||
.expect("stored continuation");
|
||||
@@ -79,7 +82,7 @@ async fn empty_opaque_continuation_is_a_store_miss_without_a_redis_record() {
|
||||
|
||||
assert!(
|
||||
repository
|
||||
.resolve(&response_id)
|
||||
.resolve(&client_key("key_primary"), &response_id)
|
||||
.await
|
||||
.expect("resolve")
|
||||
.is_none()
|
||||
@@ -88,6 +91,7 @@ async fn empty_opaque_continuation_is_a_store_miss_without_a_redis_record() {
|
||||
.record(NativeContinuationPin::new(
|
||||
response_id.clone(),
|
||||
response_id,
|
||||
client_key("key_primary"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_primary").expect("account"),
|
||||
))
|
||||
@@ -104,6 +108,7 @@ async fn native_continuation_rejects_provider_state_for_a_different_provider() {
|
||||
let pin = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("resp-mismatched-state"),
|
||||
PreviousResponseId::new("resp-upstream-handle"),
|
||||
client_key("key_primary"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_primary").expect("account"),
|
||||
)
|
||||
@@ -115,6 +120,35 @@ async fn native_continuation_rejects_provider_state_for_a_different_provider() {
|
||||
assert!(namespace_keys(&mut connection, &namespace).await.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_continuation_rejects_cross_client_reuse() {
|
||||
let Some((repository, mut connection, namespace)) = repository().await else {
|
||||
return;
|
||||
};
|
||||
let response_id = PreviousResponseId::new("resp-client-owned");
|
||||
repository
|
||||
.record(NativeContinuationPin::new(
|
||||
response_id.clone(),
|
||||
response_id.clone(),
|
||||
client_key("key_primary"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_primary").expect("account"),
|
||||
))
|
||||
.await
|
||||
.expect("record client-owned continuation");
|
||||
|
||||
let error = repository
|
||||
.resolve(&client_key("key_secondary"), &response_id)
|
||||
.await
|
||||
.expect_err("another client must not reuse the continuation");
|
||||
assert_eq!(
|
||||
error.kind(),
|
||||
NativeContinuationStoreErrorKind::OwnershipMismatch
|
||||
);
|
||||
|
||||
delete_namespace_keys(&mut connection, &namespace).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_continuation_record_should_bound_index_cleanup_work() {
|
||||
const STALE_ENTRY_COUNT: usize = 100;
|
||||
@@ -177,11 +211,16 @@ fn pin(response_id: &str) -> NativeContinuationPin {
|
||||
NativeContinuationPin::new(
|
||||
PreviousResponseId::new(response_id),
|
||||
PreviousResponseId::new(response_id),
|
||||
client_key("key_primary"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
ProviderAccountId::new("acct_primary").expect("account"),
|
||||
)
|
||||
}
|
||||
|
||||
fn client_key(id: &str) -> ClientApiKeyId {
|
||||
ClientApiKeyId::new(id).expect("client API key id")
|
||||
}
|
||||
|
||||
async fn namespace_keys(connection: &mut ConnectionManager, namespace: &str) -> Vec<String> {
|
||||
redis::cmd("KEYS")
|
||||
.arg(format!("{namespace}:*"))
|
||||
|
||||
@@ -9,6 +9,7 @@ use crate::transport::request::derive_conversation_anchor;
|
||||
|
||||
pub(crate) fn derive_codex_session_affinity_key(
|
||||
request: &CodexResponsesRequest,
|
||||
client_api_key_id: &ClientApiKeyId,
|
||||
) -> Option<ProviderSessionAffinityKey> {
|
||||
let session_key = if let Some(conversation_id) = request
|
||||
.local_conversation_id
|
||||
@@ -27,16 +28,25 @@ pub(crate) fn derive_codex_session_affinity_key(
|
||||
};
|
||||
|
||||
let session_discriminator = request.subagent_kind();
|
||||
session_key.and_then(|session_key| match session_discriminator.as_deref() {
|
||||
None => Some(session_key),
|
||||
Some(discriminator) => {
|
||||
// 仅扩展原会话锚点的键空间;后续查找、绑定和调度仍使用同一套亲和逻辑。
|
||||
opaque_affinity_key(
|
||||
"subagent-session",
|
||||
&format!("{discriminator}\0{}", session_key.expose_to_store()),
|
||||
)
|
||||
}
|
||||
})
|
||||
let session_key =
|
||||
session_key.and_then(|session_key| match session_discriminator.as_deref() {
|
||||
None => Some(session_key),
|
||||
Some(discriminator) => {
|
||||
// 仅扩展原会话锚点的键空间;后续查找、绑定和调度仍使用同一套亲和逻辑。
|
||||
opaque_affinity_key(
|
||||
"subagent-session",
|
||||
&format!("{discriminator}\0{}", session_key.expose_to_store()),
|
||||
)
|
||||
}
|
||||
})?;
|
||||
opaque_affinity_key(
|
||||
"client-session",
|
||||
&format!(
|
||||
"{}\0{}",
|
||||
client_api_key_id.as_str(),
|
||||
session_key.expose_to_store()
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
fn opaque_affinity_key(domain: &str, value: &str) -> Option<ProviderSessionAffinityKey> {
|
||||
|
||||
@@ -127,6 +127,7 @@ enum AffinityEscapeReason {
|
||||
QuotaExhausted,
|
||||
Cooldown,
|
||||
LeaseSaturated,
|
||||
HigherPriority,
|
||||
PinnedAccount,
|
||||
SelectionInvariant,
|
||||
}
|
||||
@@ -138,6 +139,7 @@ impl AffinityEscapeReason {
|
||||
Self::QuotaExhausted => "quota_exhausted",
|
||||
Self::Cooldown => "cooldown",
|
||||
Self::LeaseSaturated => "lease_saturated",
|
||||
Self::HigherPriority => "higher_priority",
|
||||
Self::PinnedAccount => "pinned_account",
|
||||
Self::SelectionInvariant => "selection_invariant",
|
||||
}
|
||||
@@ -195,8 +197,13 @@ impl AffinitySelection {
|
||||
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::RequestInterval) => {
|
||||
self.escape(AffinityEscapeReason::Cooldown);
|
||||
}
|
||||
PreferredAccountSelection::Blocked(AccountSchedulingBlocker::LowerWeight) => {
|
||||
self.escape(AffinityEscapeReason::HigherPriority);
|
||||
}
|
||||
PreferredAccountSelection::Blocked(
|
||||
AccountSchedulingBlocker::LocalAvailability | AccountSchedulingBlocker::Excluded,
|
||||
AccountSchedulingBlocker::LocalAvailability
|
||||
| AccountSchedulingBlocker::Excluded
|
||||
| AccountSchedulingBlocker::OutsideClientScope,
|
||||
)
|
||||
| PreferredAccountSelection::Missing => {
|
||||
self.escape(AffinityEscapeReason::HardUnavailable);
|
||||
@@ -282,6 +289,11 @@ impl CodexCredentialSelector {
|
||||
.into_iter()
|
||||
.filter(|account| {
|
||||
account.provider() == &self.provider_kind
|
||||
&& (diagnostic
|
||||
|| request
|
||||
.attempt
|
||||
.account_scope()
|
||||
.is_some_and(|scope| scope.allows(account.id())))
|
||||
&& (diagnostic || {
|
||||
let observed_support = self
|
||||
.catalog
|
||||
@@ -310,7 +322,11 @@ impl CodexCredentialSelector {
|
||||
.collect::<Vec<_>>();
|
||||
let scheduling = self
|
||||
.leases
|
||||
.load_state(&self.provider_kind, &account_ids)
|
||||
.load_state(
|
||||
request.attempt.client_api_key_ref(),
|
||||
&self.provider_kind,
|
||||
&account_ids,
|
||||
)
|
||||
.await?;
|
||||
let round_robin_cursor = scheduling.round_robin_cursor();
|
||||
let candidates = accounts
|
||||
@@ -411,6 +427,7 @@ impl CodexCredentialSelector {
|
||||
} else {
|
||||
AccountEligibilityPolicy::Enforce
|
||||
},
|
||||
account_scope: request.attempt.account_scope().cloned(),
|
||||
};
|
||||
let Some(selection) = AccountSelector.select(&candidates, &context) else {
|
||||
return match shortest_retry {
|
||||
@@ -436,7 +453,7 @@ impl CodexCredentialSelector {
|
||||
self.provider_kind.clone(),
|
||||
account.id().clone(),
|
||||
account.revision(),
|
||||
policy.max_concurrent_per_account(),
|
||||
account.effective_concurrency(policy.max_concurrent_per_account()),
|
||||
policy.request_interval(),
|
||||
request.attempt.deadline(),
|
||||
),
|
||||
|
||||
@@ -253,7 +253,8 @@ impl Provider for CodexProvider {
|
||||
.as_ref()
|
||||
.and_then(|state| state.turn_state.clone());
|
||||
}
|
||||
let session_affinity_key = derive_codex_session_affinity_key(&upstream_request);
|
||||
let session_affinity_key =
|
||||
derive_codex_session_affinity_key(&upstream_request, context.client_api_key_ref());
|
||||
let cyber_policy_session_key =
|
||||
derive_codex_cyber_policy_session_key(&upstream_request, context.client_api_key_ref());
|
||||
let transport = selected_transport(&upstream_request);
|
||||
|
||||
@@ -797,6 +797,7 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
|
||||
AccountRecord {
|
||||
id: account.id().to_string(),
|
||||
provider_kind: account.provider().clone(),
|
||||
groups: Vec::new(),
|
||||
name: account.name().to_owned(),
|
||||
email: account.email().map(str::to_owned),
|
||||
upstream_user_id: account.upstream_user_id().map(str::to_owned),
|
||||
@@ -808,6 +809,8 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
|
||||
access_token_expires_at: account.access_token_expires_at().map(DateTime::<Utc>::from),
|
||||
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
|
||||
enabled: account.enabled(),
|
||||
concurrency_limit: account.concurrency_limit(),
|
||||
weight: account.weight(),
|
||||
credential_state: account.credential_state(),
|
||||
credential_observed_at: now,
|
||||
quota: account.quota(),
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::collections::BTreeSet;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
@@ -8,10 +8,11 @@ use gateway_core::engine::continuation::{
|
||||
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
|
||||
};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountAttemptFeedback, AccountErrorReason, AccountFeedbackStats, AccountSelectionPolicy,
|
||||
AccountStateChange, CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId,
|
||||
ProviderAccountStore as _, QuotaAccessChange, QuotaAccessState, QuotaEvidence,
|
||||
QuotaObservation, QuotaState, QuotaWriteOutcome, RotationStrategy,
|
||||
AccountAttemptFeedback, AccountConcurrencyLimit, AccountErrorReason, AccountFeedbackStats,
|
||||
AccountSelectionPolicy, AccountStateChange, AccountWeight, CredentialState, OpaqueProviderData,
|
||||
ProviderAccount, ProviderAccountId, ProviderAccountStore as _, QuotaAccessChange,
|
||||
QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState, QuotaWriteOutcome,
|
||||
RotationStrategy,
|
||||
};
|
||||
use gateway_core::engine::{
|
||||
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, RequestAttemptContext,
|
||||
@@ -20,7 +21,9 @@ use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCooldownPort, ProviderSessionAffinityKey, ProviderSessionAffinityPort,
|
||||
};
|
||||
use gateway_core::routing::ProviderKind;
|
||||
use gateway_core::routing::{
|
||||
ClientRoutingScope, FrozenAccountScope, ProviderKind, RuntimeAccount, RuntimeAccountDirectory,
|
||||
};
|
||||
use provider_openai::OFFICIAL_CODEX_BASE_URL;
|
||||
use provider_openai::credential::{
|
||||
CodexAccountFailure, CodexCookiePolicy, CodexCredentialCatalogService, CodexCredentialCodec,
|
||||
@@ -48,6 +51,34 @@ fn create_account(store: &Arc<MemoryAccountStore>, id: &str, token: &str) {
|
||||
}));
|
||||
}
|
||||
|
||||
fn contract_account_scope() -> Arc<FrozenAccountScope> {
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let accounts = [
|
||||
"acct_available",
|
||||
"acct_fallback",
|
||||
"acct_fallback_signal",
|
||||
"acct_first",
|
||||
"acct_missing",
|
||||
"acct_original",
|
||||
"acct_original_signal",
|
||||
"acct_other",
|
||||
"acct_primary",
|
||||
"acct_second",
|
||||
]
|
||||
.into_iter()
|
||||
.map(|id| {
|
||||
(
|
||||
ProviderAccountId::new(id).expect("account"),
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
|
||||
)
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
Arc::new(FrozenAccountScope::new(
|
||||
Arc::new(RuntimeAccountDirectory::new(accounts)),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
))
|
||||
}
|
||||
|
||||
fn attempt(excluded_accounts: BTreeSet<ProviderAccountId>) -> AttemptContext {
|
||||
attempt_with_required(excluded_accounts, None)
|
||||
}
|
||||
@@ -64,7 +95,8 @@ fn attempt_with_required(
|
||||
NonZeroU32::new(1).expect("attempt"),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(excluded_accounts, required_account, None),
|
||||
AccountAttemptContext::new(excluded_accounts, required_account, None)
|
||||
.with_account_scope(contract_account_scope()),
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
@@ -83,7 +115,8 @@ fn round_robin_attempt() -> AttemptContext {
|
||||
NonZeroU32::new(2).expect("concurrency"),
|
||||
Duration::ZERO,
|
||||
),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, None),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, None)
|
||||
.with_account_scope(contract_account_scope()),
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
@@ -358,6 +391,32 @@ fn selector_uses_frozen_global_account_policy_for_lease() {
|
||||
assert_eq!(requests[0].deadline(), attempt.deadline());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_uses_the_account_concurrency_override_for_the_redis_lease() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_primary", "at-primary");
|
||||
store.set_scheduling(
|
||||
"acct_primary",
|
||||
Some(AccountConcurrencyLimit::new(7).expect("concurrency override")),
|
||||
AccountWeight::DEFAULT,
|
||||
);
|
||||
let leases = Arc::new(TestLeaseCoordinator::default());
|
||||
let selector = selector(&store, Arc::clone(&leases));
|
||||
let attempt = attempt(BTreeSet::new());
|
||||
|
||||
block_on(selector.select(&SelectCodexCredential {
|
||||
upstream_model: "gpt-5.4",
|
||||
request_url:
|
||||
&Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL"),
|
||||
attempt: &attempt,
|
||||
session_affinity_key: None,
|
||||
}))
|
||||
.expect("select account");
|
||||
|
||||
let requests = leases.requests.lock().expect("lease requests lock");
|
||||
assert_eq!(requests[0].max_concurrent().get(), 7);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selector_round_robin_cursor_advances_across_requests() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -1043,6 +1102,7 @@ fn native_continuation_surfaces_the_original_accounts_quota_status_to_the_coordi
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-response"),
|
||||
PreviousResponseId::new("upstream-response"),
|
||||
ClientApiKeyId::new("key_codex_contract").expect("client key id"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
original.id().clone(),
|
||||
);
|
||||
@@ -1054,7 +1114,8 @@ fn native_continuation_surfaces_the_original_accounts_quota_status_to_the_coordi
|
||||
NonZeroU32::new(1).expect("attempt"),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, None),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, None)
|
||||
.with_account_scope(contract_account_scope()),
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
);
|
||||
@@ -1112,6 +1173,7 @@ fn native_continuation_surfaces_the_original_accounts_quota_signal_to_the_coordi
|
||||
let continuation = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("previous-response"),
|
||||
PreviousResponseId::new("upstream-response"),
|
||||
ClientApiKeyId::new("key_codex_contract").expect("client key id"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
original.id().clone(),
|
||||
);
|
||||
@@ -1123,7 +1185,8 @@ fn native_continuation_surfaces_the_original_accounts_quota_signal_to_the_coordi
|
||||
NonZeroU32::new(1).expect("attempt"),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, None),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, None)
|
||||
.with_account_scope(contract_account_scope()),
|
||||
Some(ContinuationBinding::Pinned(continuation)),
|
||||
CancellationToken::new(),
|
||||
);
|
||||
|
||||
@@ -7,6 +7,7 @@ use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::engine::credential::{ProviderAccountId, ProviderAccountStore};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCredentialState, ProviderCredentialStatePort, ProviderLeaseAcquisition,
|
||||
ProviderLeasePort, ProviderLeaseRequest, ProviderRefreshPolicy, ProviderRuntimePolicyPort,
|
||||
@@ -86,6 +87,7 @@ struct RefreshLeases;
|
||||
impl ProviderLeasePort for RefreshLeases {
|
||||
fn load_state<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientApiKeyId,
|
||||
_: &'a ProviderKind,
|
||||
_: &'a [ProviderAccountId],
|
||||
) -> BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::collections::BTreeSet;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
@@ -21,7 +21,8 @@ use gateway_core::event::GatewayEvent;
|
||||
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload, ProviderSessionState};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::{
|
||||
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
|
||||
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
|
||||
ProviderModel, PublicModelId, RoutingContext, RuntimeAccount, RuntimeAccountDirectory,
|
||||
RuntimeSnapshot, UpstreamModelId,
|
||||
};
|
||||
use provider_openai::credential::{
|
||||
@@ -252,6 +253,13 @@ fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest
|
||||
let provider = ProviderKind::new(provider_name).expect("provider");
|
||||
let upstream_model = UpstreamModelId::new("gpt-5.4").expect("upstream model");
|
||||
let public_model = PublicModelId::new(upstream_model.as_str()).expect("public model");
|
||||
let account_scope = Arc::new(FrozenAccountScope::new(
|
||||
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
|
||||
ProviderAccountId::new("acct_provider_contract").expect("account"),
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
|
||||
)]))),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
));
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("revision"),
|
||||
account_policy(),
|
||||
@@ -265,12 +273,67 @@ fn planned_request(provider_name: &str, operation: Operation) -> ProviderRequest
|
||||
)
|
||||
.expect("snapshot");
|
||||
let plan = snapshot
|
||||
.plan(&public_model, &operation, &RoutingContext::default())
|
||||
.plan(
|
||||
&public_model,
|
||||
&operation,
|
||||
account_scope,
|
||||
&RoutingContext::default(),
|
||||
)
|
||||
.expect("routing plan");
|
||||
|
||||
ProviderRequest::new(operation, plan.candidates()[0].clone())
|
||||
}
|
||||
|
||||
fn contract_account_scope() -> Arc<FrozenAccountScope> {
|
||||
let provider = ProviderKind::new("openai").expect("provider");
|
||||
let accounts = [
|
||||
"acct_affinity",
|
||||
"acct_atomic_failure",
|
||||
"acct_bare_atomic_failure",
|
||||
"acct_bounded_replay_grace",
|
||||
"acct_bounded_session_state",
|
||||
"acct_capacity_busy",
|
||||
"acct_client_history",
|
||||
"acct_completed_affinity",
|
||||
"acct_continuation_prefetch",
|
||||
"acct_disabled_scheduling",
|
||||
"acct_first_event_latency",
|
||||
"acct_header_new",
|
||||
"acct_header_old",
|
||||
"acct_header_same",
|
||||
"acct_http_sse_exhausted",
|
||||
"acct_local_affinity",
|
||||
"acct_metadata_new",
|
||||
"acct_metadata_old",
|
||||
"acct_prefetch_limit",
|
||||
"acct_presentation",
|
||||
"acct_provider_contract",
|
||||
"acct_scope_new",
|
||||
"acct_scope_old",
|
||||
"acct_scope_same",
|
||||
"acct_semantic_failure",
|
||||
"acct_session_affinity",
|
||||
"acct_subagent",
|
||||
"acct_success_exhausted",
|
||||
"acct_thread_spawn_affinity",
|
||||
"acct_truncated_stream",
|
||||
"acct_usage_limit_request_path",
|
||||
"acct_websocket_close",
|
||||
]
|
||||
.into_iter()
|
||||
.map(|id| {
|
||||
(
|
||||
ProviderAccountId::new(id).expect("account"),
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
|
||||
)
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
Arc::new(FrozenAccountScope::new(
|
||||
Arc::new(RuntimeAccountDirectory::new(accounts)),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
))
|
||||
}
|
||||
|
||||
fn context(request_id: &str, cancellation: CancellationToken) -> AttemptContext {
|
||||
AttemptContext::new(
|
||||
RequestAttemptContext::new(
|
||||
@@ -280,7 +343,8 @@ fn context(request_id: &str, cancellation: CancellationToken) -> AttemptContext
|
||||
NonZeroU32::new(1).expect("attempt"),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(BTreeSet::<ProviderAccountId>::new(), None, None),
|
||||
AccountAttemptContext::new(BTreeSet::<ProviderAccountId>::new(), None, None)
|
||||
.with_account_scope(contract_account_scope()),
|
||||
None,
|
||||
cancellation,
|
||||
)
|
||||
@@ -318,7 +382,8 @@ fn context_with_state_owner(request_id: &str, owner_account_id: &str) -> Attempt
|
||||
NonZeroU32::new(1).expect("attempt"),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
account_policy(),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, Some(owner)),
|
||||
AccountAttemptContext::new(BTreeSet::new(), None, Some(owner))
|
||||
.with_account_scope(contract_account_scope()),
|
||||
None,
|
||||
CancellationToken::new(),
|
||||
)
|
||||
|
||||
@@ -9,12 +9,14 @@ use std::time::{Duration, SystemTime};
|
||||
use async_trait::async_trait;
|
||||
use futures::future::BoxFuture;
|
||||
use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountRuntimeSignals, AccountStateChange, CredentialCasOutcome,
|
||||
CredentialCasUpdate, CredentialRevision, CredentialState, LoadedCredential, NewProviderAccount,
|
||||
OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore,
|
||||
ProviderAccountUpdate, QuotaAccessChange, QuotaObservation, QuotaState, QuotaWriteOutcome,
|
||||
AccountConcurrencyLimit, AccountErrorReason, AccountRuntimeSignals, AccountStateChange,
|
||||
AccountWeight, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState,
|
||||
LoadedCredential, NewProviderAccount, OpaqueProviderData, ProviderAccount, ProviderAccountId,
|
||||
ProviderAccountStore, ProviderAccountUpdate, QuotaAccessChange, QuotaObservation, QuotaState,
|
||||
QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::error::{StoreError, StoreErrorKind};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCooldown, ProviderCooldownPort,
|
||||
ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort, ProviderLeaseRequest,
|
||||
@@ -70,6 +72,21 @@ impl MemoryAccountStore {
|
||||
.map(|stored| stored.account.clone())
|
||||
}
|
||||
|
||||
pub(crate) fn set_scheduling(
|
||||
&self,
|
||||
id: &str,
|
||||
concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
weight: AccountWeight,
|
||||
) {
|
||||
let id = ProviderAccountId::new(id).expect("valid account ID");
|
||||
let mut accounts = self.accounts.lock().expect("account store lock");
|
||||
let stored = accounts.get_mut(&id).expect("seeded account");
|
||||
stored.account = stored
|
||||
.account
|
||||
.clone()
|
||||
.with_scheduling(concurrency_limit, weight);
|
||||
}
|
||||
|
||||
pub(crate) fn quota_reads(&self) -> usize {
|
||||
self.quota_reads.load(Ordering::SeqCst)
|
||||
}
|
||||
@@ -450,6 +467,7 @@ fn rebuild_account(current: &ProviderAccount, rebuild: AccountRebuild) -> Provid
|
||||
rebuild.last_error_reason,
|
||||
rebuild.last_error_message,
|
||||
)
|
||||
.with_scheduling(current.concurrency_limit(), current.weight())
|
||||
.with_refresh_schedule(rebuild.has_refresh_token, rebuild.next_refresh_at)
|
||||
}
|
||||
|
||||
@@ -487,6 +505,7 @@ pub(crate) struct TestLeaseCoordinator {
|
||||
impl ProviderLeasePort for TestLeaseCoordinator {
|
||||
fn load_state<'a>(
|
||||
&'a self,
|
||||
_client_api_key_id: &'a ClientApiKeyId,
|
||||
_provider_kind: &'a ProviderKind,
|
||||
accounts: &'a [ProviderAccountId],
|
||||
) -> BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
|
||||
|
||||
@@ -89,9 +89,14 @@ impl GrokAccountSessionSelector {
|
||||
.list_accounts_for_provider()
|
||||
.await
|
||||
.map_err(|_| GrokSessionSelectorError::Unavailable)?;
|
||||
if !diagnostic {
|
||||
self.quota.prepare_scheduling(&accounts).await;
|
||||
}
|
||||
let accounts = if diagnostic {
|
||||
accounts
|
||||
} else {
|
||||
accounts
|
||||
.into_iter()
|
||||
.filter(|account| request.account_scope().allows(account.id()))
|
||||
.collect()
|
||||
};
|
||||
let catalog_eligible = if diagnostic {
|
||||
accounts
|
||||
} else {
|
||||
@@ -128,6 +133,9 @@ impl GrokAccountSessionSelector {
|
||||
if catalog_eligible.is_empty() {
|
||||
return Err(GrokSessionSelectorError::NoEligibleSession);
|
||||
}
|
||||
if !diagnostic {
|
||||
self.quota.prepare_scheduling(&catalog_eligible).await;
|
||||
}
|
||||
|
||||
let account_ids = catalog_eligible
|
||||
.iter()
|
||||
@@ -135,7 +143,11 @@ impl GrokAccountSessionSelector {
|
||||
.collect::<Vec<_>>();
|
||||
let scheduling = self
|
||||
.scheduling
|
||||
.load_state(&self.provider_kind, &account_ids)
|
||||
.load_state(
|
||||
request.client_api_key_id(),
|
||||
&self.provider_kind,
|
||||
&account_ids,
|
||||
)
|
||||
.await
|
||||
.map_err(|_| GrokSessionSelectorError::Unavailable)?;
|
||||
let runtime_cooldowns = if diagnostic {
|
||||
@@ -179,6 +191,7 @@ impl GrokAccountSessionSelector {
|
||||
preferred_account: request.required_account().cloned().or(affinity_account),
|
||||
round_robin_cursor: scheduling.round_robin_cursor(),
|
||||
eligibility: request.eligibility(),
|
||||
account_scope: (!diagnostic).then(|| Arc::clone(request.account_scope())),
|
||||
};
|
||||
let mut capacity_denied = false;
|
||||
let mut retry_after = None;
|
||||
@@ -194,9 +207,11 @@ impl GrokAccountSessionSelector {
|
||||
self.provider_kind.clone(),
|
||||
selected_id.clone(),
|
||||
selected_revision,
|
||||
request
|
||||
.account_selection_policy()
|
||||
.max_concurrent_per_account(),
|
||||
selected.account.effective_concurrency(
|
||||
request
|
||||
.account_selection_policy()
|
||||
.max_concurrent_per_account(),
|
||||
),
|
||||
request.account_selection_policy().request_interval(),
|
||||
request.deadline(),
|
||||
),
|
||||
|
||||
@@ -391,6 +391,8 @@ async fn select_grok_session(
|
||||
required_account.clone(),
|
||||
context.account_selection_policy(),
|
||||
context.deadline(),
|
||||
Arc::clone(candidate.account_scope()),
|
||||
context.client_api_key_ref().clone(),
|
||||
)
|
||||
.with_eligibility_policy(if context.is_diagnostic_required_account() {
|
||||
AccountEligibilityPolicy::BypassForDiagnostic
|
||||
|
||||
@@ -2,13 +2,15 @@ use std::collections::BTreeSet;
|
||||
use std::fmt;
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use gateway_core::engine::credential::{
|
||||
AccountEligibilityPolicy, AccountSelectionPolicy, CredentialRevision, ProviderAccountId,
|
||||
};
|
||||
use gateway_core::engine::provider::ProviderResource;
|
||||
use gateway_core::routing::UpstreamModelId;
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::routing::{FrozenAccountScope, UpstreamModelId};
|
||||
use sha2::{Digest as _, Sha256};
|
||||
|
||||
use crate::SecretValue;
|
||||
@@ -213,6 +215,8 @@ pub struct GrokSessionSelection {
|
||||
eligibility: AccountEligibilityPolicy,
|
||||
affinity: Option<GrokSessionAffinityKey>,
|
||||
deadline: SystemTime,
|
||||
account_scope: Arc<FrozenAccountScope>,
|
||||
client_api_key_id: ClientApiKeyId,
|
||||
}
|
||||
|
||||
impl GrokSessionSelection {
|
||||
@@ -224,6 +228,8 @@ impl GrokSessionSelection {
|
||||
required_account: Option<ProviderAccountId>,
|
||||
account_selection_policy: AccountSelectionPolicy,
|
||||
deadline: SystemTime,
|
||||
account_scope: Arc<FrozenAccountScope>,
|
||||
client_api_key_id: ClientApiKeyId,
|
||||
) -> Self {
|
||||
Self {
|
||||
upstream_model,
|
||||
@@ -233,6 +239,8 @@ impl GrokSessionSelection {
|
||||
eligibility: AccountEligibilityPolicy::Enforce,
|
||||
affinity: None,
|
||||
deadline,
|
||||
account_scope,
|
||||
client_api_key_id,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -293,6 +301,16 @@ impl GrokSessionSelection {
|
||||
pub const fn deadline(&self) -> SystemTime {
|
||||
self.deadline
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
|
||||
&self.account_scope
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn client_api_key_id(&self) -> &ClientApiKeyId {
|
||||
&self.client_api_key_id
|
||||
}
|
||||
}
|
||||
|
||||
/// Grok session selector 返回的 future。
|
||||
|
||||
@@ -460,6 +460,7 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
|
||||
AccountRecord {
|
||||
id: account.id().to_string(),
|
||||
provider_kind: account.provider().clone(),
|
||||
groups: Vec::new(),
|
||||
name: account.name().to_owned(),
|
||||
email: account.email().map(str::to_owned),
|
||||
upstream_user_id: account.upstream_user_id().map(str::to_owned),
|
||||
@@ -471,6 +472,8 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
|
||||
access_token_expires_at: account.access_token_expires_at().map(DateTime::<Utc>::from),
|
||||
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
|
||||
enabled: account.enabled(),
|
||||
concurrency_limit: account.concurrency_limit(),
|
||||
weight: account.weight(),
|
||||
credential_state: account.credential_state(),
|
||||
credential_observed_at: now,
|
||||
quota: account.quota(),
|
||||
@@ -486,6 +489,7 @@ struct TestLeases;
|
||||
impl ProviderLeasePort for TestLeases {
|
||||
fn load_state<'a>(
|
||||
&'a self,
|
||||
_: &'a gateway_core::policy::ClientApiKeyId,
|
||||
_: &'a gateway_core::routing::ProviderKind,
|
||||
accounts: &'a [ProviderAccountId],
|
||||
) -> BoxFuture<
|
||||
|
||||
@@ -110,6 +110,7 @@ impl ProviderRuntimePolicyPort for MutableRuntimePolicy {
|
||||
impl ProviderLeasePort for TestRefreshLeases {
|
||||
fn load_state<'a>(
|
||||
&'a self,
|
||||
_: &'a gateway_core::policy::ClientApiKeyId,
|
||||
_: &'a gateway_core::routing::ProviderKind,
|
||||
_: &'a [gateway_core::engine::credential::ProviderAccountId],
|
||||
) -> futures::future::BoxFuture<
|
||||
|
||||
@@ -4,16 +4,22 @@ use std::time::{Duration, SystemTime};
|
||||
|
||||
use chrono::Utc;
|
||||
use gateway_core::engine::credential::{
|
||||
AccountAttemptFeedback, AccountFeedbackStats, AccountRuntimeSignals, AccountSelectionPolicy,
|
||||
CredentialCasOutcome, CredentialRevision, CredentialState, OpaqueProviderData,
|
||||
ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, QuotaAccessState,
|
||||
QuotaObservation, QuotaState, QuotaWriteOutcome, RotationStrategy,
|
||||
AccountAttemptFeedback, AccountConcurrencyLimit, AccountFeedbackStats, AccountRuntimeSignals,
|
||||
AccountSelectionPolicy, AccountWeight, CredentialCasOutcome, CredentialRevision,
|
||||
CredentialState, OpaqueProviderData, ProviderAccountId, ProviderAccountStore,
|
||||
ProviderAccountUpdate, QuotaAccessState, QuotaObservation, QuotaState, QuotaWriteOutcome,
|
||||
RotationStrategy,
|
||||
};
|
||||
use gateway_core::policy::ClientApiKeyId;
|
||||
use gateway_core::provider_ports::{
|
||||
ProviderCooldownPort, ProviderCooldownScope, ProviderLeaseAcquisition, ProviderLeasePort,
|
||||
ProviderLeaseRequest, ProviderSchedulingState, ProviderStoreError,
|
||||
ProviderLeaseRequest, ProviderSchedulingLeaseRequest, ProviderSchedulingState,
|
||||
ProviderStoreError,
|
||||
};
|
||||
use gateway_core::routing::{
|
||||
ClientRoutingScope, FrozenAccountScope, RuntimeAccount, RuntimeAccountDirectory,
|
||||
UpstreamModelId,
|
||||
};
|
||||
use gateway_core::routing::UpstreamModelId;
|
||||
use provider_xai::{
|
||||
GrokAccountSessionSelector, GrokBillingRequest, GrokBillingTransport,
|
||||
GrokBillingTransportError, GrokBillingTransportErrorKind, GrokBillingTransportFuture,
|
||||
@@ -31,6 +37,7 @@ use crate::support::{
|
||||
struct SchedulingCoordinator {
|
||||
signals: Mutex<BTreeMap<ProviderAccountId, AccountRuntimeSignals>>,
|
||||
denied: Mutex<BTreeSet<ProviderAccountId>>,
|
||||
requests: Mutex<Vec<ProviderSchedulingLeaseRequest>>,
|
||||
}
|
||||
|
||||
struct UnavailableBillingTransport;
|
||||
@@ -48,6 +55,7 @@ impl GrokBillingTransport for UnavailableBillingTransport {
|
||||
impl ProviderLeasePort for SchedulingCoordinator {
|
||||
fn load_state<'a>(
|
||||
&'a self,
|
||||
_: &'a ClientApiKeyId,
|
||||
_: &'a gateway_core::routing::ProviderKind,
|
||||
_: &'a [ProviderAccountId],
|
||||
) -> futures::future::BoxFuture<'a, Result<ProviderSchedulingState, ProviderStoreError>> {
|
||||
@@ -67,20 +75,19 @@ impl ProviderLeasePort for SchedulingCoordinator {
|
||||
let ProviderLeaseRequest::Scheduling(request) = request else {
|
||||
panic!("expected scheduling lease request");
|
||||
};
|
||||
Ok(
|
||||
if self
|
||||
.denied
|
||||
.lock()
|
||||
.expect("denied")
|
||||
.contains(request.account_id())
|
||||
{
|
||||
ProviderLeaseAcquisition::Busy {
|
||||
retry_after: Some(Duration::from_millis(25)),
|
||||
}
|
||||
} else {
|
||||
ProviderLeaseAcquisition::Acquired(Box::new(()))
|
||||
},
|
||||
)
|
||||
let denied = self
|
||||
.denied
|
||||
.lock()
|
||||
.expect("denied")
|
||||
.contains(request.account_id());
|
||||
self.requests.lock().expect("requests").push(request);
|
||||
Ok(if denied {
|
||||
ProviderLeaseAcquisition::Busy {
|
||||
retry_after: Some(Duration::from_millis(25)),
|
||||
}
|
||||
} else {
|
||||
ProviderLeaseAcquisition::Acquired(Box::new(()))
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -141,6 +148,7 @@ impl SelectorFixture {
|
||||
let coordinator = Arc::new(SchedulingCoordinator {
|
||||
signals: Mutex::new(signals),
|
||||
denied: Mutex::new(BTreeSet::new()),
|
||||
requests: Mutex::new(Vec::new()),
|
||||
});
|
||||
let cooldowns = Arc::new(MemoryCooldownPort::default());
|
||||
let catalog_cache: Arc<dyn GrokCredentialCatalogCache> = cache.clone();
|
||||
@@ -210,6 +218,21 @@ impl SelectorFixture {
|
||||
required_account: Option<ProviderAccountId>,
|
||||
strategy: RotationStrategy,
|
||||
) -> GrokSessionSelection {
|
||||
let provider = gateway_core::routing::ProviderKind::new("xai").expect("provider");
|
||||
let accounts = self
|
||||
.coordinator
|
||||
.signals
|
||||
.lock()
|
||||
.expect("signals")
|
||||
.keys()
|
||||
.cloned()
|
||||
.map(|account_id| {
|
||||
(
|
||||
account_id,
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
|
||||
)
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
GrokSessionSelection::new(
|
||||
UpstreamModelId::new(upstream_model).expect("model"),
|
||||
excluded,
|
||||
@@ -220,6 +243,11 @@ impl SelectorFixture {
|
||||
Duration::ZERO,
|
||||
),
|
||||
SystemTime::now() + Duration::from_secs(30),
|
||||
Arc::new(FrozenAccountScope::new(
|
||||
Arc::new(RuntimeAccountDirectory::new(accounts)),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
)),
|
||||
ClientApiKeyId::new("key_xai_selector").expect("client key"),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -286,6 +314,26 @@ async fn required_account_overrides_smart_selection_without_fallback() {
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn selector_uses_the_account_concurrency_override_for_the_redis_lease() {
|
||||
let fixture = SelectorFixture::new(&["concurrency-override"]).await;
|
||||
let account = account_id("concurrency-override");
|
||||
fixture.store.set_scheduling(
|
||||
&account,
|
||||
Some(AccountConcurrencyLimit::new(7).expect("concurrency override")),
|
||||
AccountWeight::DEFAULT,
|
||||
);
|
||||
|
||||
fixture
|
||||
.selector
|
||||
.select(fixture.request(BTreeSet::new()))
|
||||
.await
|
||||
.expect("select account");
|
||||
|
||||
let requests = fixture.coordinator.requests.lock().expect("requests");
|
||||
assert_eq!(requests[0].max_concurrent().get(), 7);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unauthorized_feedback_records_runtime_cooldown_without_persisting_account_state() {
|
||||
let fixture = SelectorFixture::new(&["feedback-a", "feedback-b"]).await;
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::collections::{BTreeSet, VecDeque};
|
||||
use std::collections::{BTreeMap, BTreeSet, VecDeque};
|
||||
use std::num::NonZeroU32;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
@@ -26,7 +26,8 @@ use gateway_core::operation::{
|
||||
Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload, ProviderSessionState,
|
||||
};
|
||||
use gateway_core::routing::{
|
||||
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
|
||||
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
|
||||
ProviderModel, PublicModelId, RoutingContext, RuntimeAccount, RuntimeAccountDirectory,
|
||||
RuntimeSnapshot, SupportLevel, UpstreamModelId,
|
||||
};
|
||||
use provider_xai::{
|
||||
@@ -810,6 +811,13 @@ fn provider_request_with_operation(provider_kind: &str, operation: Operation) ->
|
||||
.with_feature(Feature::Tools, SupportLevel::Native)
|
||||
.with_feature(Feature::NativeContinuation, SupportLevel::Native),
|
||||
);
|
||||
let account_scope = Arc::new(FrozenAccountScope::new(
|
||||
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
|
||||
account_id("provider"),
|
||||
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
|
||||
)]))),
|
||||
ClientRoutingScope::all_accounts(),
|
||||
));
|
||||
let snapshot = RuntimeSnapshot::new(
|
||||
ConfigRevision::new(1).expect("revision"),
|
||||
selection_policy(),
|
||||
@@ -822,6 +830,7 @@ fn provider_request_with_operation(provider_kind: &str, operation: Operation) ->
|
||||
.plan(
|
||||
&PublicModelId::new(MODEL).expect("model"),
|
||||
&operation,
|
||||
account_scope,
|
||||
&RoutingContext::default(),
|
||||
)
|
||||
.expect("routing plan");
|
||||
@@ -1858,6 +1867,7 @@ async fn native_previous_response_is_rejected_before_selection() {
|
||||
let pin = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("resp_previous"),
|
||||
PreviousResponseId::new("resp_upstream_previous"),
|
||||
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
|
||||
ProviderKind::new("openai").expect("provider"),
|
||||
account_id("provider"),
|
||||
);
|
||||
@@ -1891,6 +1901,7 @@ async fn native_previous_response_pins_account_and_sends_upstream_handle() {
|
||||
let pin = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("resp_previous"),
|
||||
PreviousResponseId::new("resp_upstream_previous"),
|
||||
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
|
||||
ProviderKind::new("xai").expect("provider"),
|
||||
account_id("provider"),
|
||||
);
|
||||
@@ -1936,6 +1947,7 @@ async fn native_previous_response_does_not_allow_quota_or_rate_limit_account_rot
|
||||
let pin = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("resp_previous"),
|
||||
PreviousResponseId::new("resp_upstream_previous"),
|
||||
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
|
||||
ProviderKind::new("xai").expect("provider"),
|
||||
account_id("provider"),
|
||||
);
|
||||
@@ -2257,6 +2269,7 @@ async fn connection_state_inherits_session_and_recovers_reasoning_on_pinned_acco
|
||||
let pin = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("resp_gateway_first"),
|
||||
PreviousResponseId::new("resp_state"),
|
||||
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
|
||||
ProviderKind::new("xai").expect("provider"),
|
||||
account_id("provider"),
|
||||
);
|
||||
@@ -2381,6 +2394,7 @@ async fn replay_owner_should_reencode_custom_apply_patch_call_for_grok() {
|
||||
let pin = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("resp_gateway_patch"),
|
||||
PreviousResponseId::new("resp_custom_patch"),
|
||||
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
|
||||
ProviderKind::new("xai").expect("provider"),
|
||||
account_id("provider"),
|
||||
);
|
||||
@@ -2456,6 +2470,7 @@ async fn missing_native_response_should_be_replay_safe_for_the_same_account() {
|
||||
let pin = NativeContinuationPin::new(
|
||||
PreviousResponseId::new("resp_gateway_first"),
|
||||
PreviousResponseId::new("resp_upstream_first"),
|
||||
gateway_core::policy::ClientApiKeyId::new("key_xai_contract").expect("client key id"),
|
||||
ProviderKind::new("xai").expect("provider"),
|
||||
account_id("provider"),
|
||||
);
|
||||
|
||||
@@ -8,10 +8,11 @@ use std::time::{Duration, SystemTime};
|
||||
use async_trait::async_trait;
|
||||
use chrono::Utc;
|
||||
use gateway_core::engine::credential::{
|
||||
AccountErrorReason, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
|
||||
CredentialRevision, CredentialState, LoadedCredential, NewProviderAccount, PlaintextCredential,
|
||||
ProviderAccount, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
|
||||
QuotaAccessChange, QuotaObservation, QuotaState, QuotaWriteOutcome,
|
||||
AccountConcurrencyLimit, AccountErrorReason, AccountStateChange, AccountWeight,
|
||||
CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState,
|
||||
LoadedCredential, NewProviderAccount, PlaintextCredential, ProviderAccount, ProviderAccountId,
|
||||
ProviderAccountStore, ProviderAccountUpdate, QuotaAccessChange, QuotaObservation, QuotaState,
|
||||
QuotaWriteOutcome,
|
||||
};
|
||||
use gateway_core::error::{StoreError, StoreErrorKind};
|
||||
use gateway_core::provider_ports::{
|
||||
@@ -213,6 +214,20 @@ impl MemoryProviderAccountStore {
|
||||
.map(|stored| stored.account.clone())
|
||||
}
|
||||
|
||||
pub fn set_scheduling(
|
||||
&self,
|
||||
id: &ProviderAccountId,
|
||||
concurrency_limit: Option<AccountConcurrencyLimit>,
|
||||
weight: AccountWeight,
|
||||
) {
|
||||
let mut accounts = lock(&self.accounts);
|
||||
let stored = accounts.get_mut(id).expect("seeded account");
|
||||
stored.account = stored
|
||||
.account
|
||||
.clone()
|
||||
.with_scheduling(concurrency_limit, weight);
|
||||
}
|
||||
|
||||
pub fn last_error_message(&self, id: &ProviderAccountId) -> Option<String> {
|
||||
lock(&self.accounts)
|
||||
.get(id)
|
||||
@@ -543,6 +558,7 @@ fn rebuild_account(previous: &ProviderAccount, replacement: AccountReplacement)
|
||||
replacement.last_error_reason,
|
||||
replacement.last_error_message,
|
||||
)
|
||||
.with_scheduling(previous.concurrency_limit(), previous.weight())
|
||||
.with_refresh_schedule(replacement.has_refresh_token, replacement.next_refresh_at)
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
-- Account groups are provider-neutral account sets. Existing provider-scoped
|
||||
-- keys are placed in ordinary migration groups before their provider column is
|
||||
-- removed, preserving their pre-upgrade account scope.
|
||||
|
||||
create table account_groups (
|
||||
id text primary key,
|
||||
name text not null,
|
||||
description text,
|
||||
enabled boolean not null default true,
|
||||
created_at timestamptz not null,
|
||||
updated_at timestamptz not null,
|
||||
constraint account_groups_id_ck check (
|
||||
id ~ '^grp_[0-9a-f]{32}$'
|
||||
),
|
||||
constraint account_groups_name_ck check (
|
||||
char_length(btrim(name)) between 1 and 100
|
||||
and name = btrim(name)
|
||||
and name !~ '[[:cntrl:]]'
|
||||
),
|
||||
constraint account_groups_description_ck check (
|
||||
description is null or (
|
||||
octet_length(description) <= 4096
|
||||
and description !~ '[[:cntrl:]]'
|
||||
)
|
||||
),
|
||||
constraint account_groups_time_ck check (
|
||||
created_at <= updated_at
|
||||
)
|
||||
);
|
||||
|
||||
create unique index account_groups_name_uq
|
||||
on account_groups (lower(name));
|
||||
|
||||
create index account_groups_list_idx
|
||||
on account_groups (enabled, created_at desc, id desc);
|
||||
|
||||
create table account_group_accounts (
|
||||
account_group_id text not null,
|
||||
provider_account_id text not null,
|
||||
created_at timestamptz not null,
|
||||
primary key (account_group_id, provider_account_id),
|
||||
constraint account_group_accounts_group_fk foreign key (account_group_id)
|
||||
references account_groups (id)
|
||||
on update restrict
|
||||
on delete cascade,
|
||||
constraint account_group_accounts_account_fk foreign key (provider_account_id)
|
||||
references provider_accounts (id)
|
||||
on update restrict
|
||||
on delete cascade
|
||||
);
|
||||
|
||||
create index account_group_accounts_account_idx
|
||||
on account_group_accounts (provider_account_id, account_group_id);
|
||||
|
||||
create table client_api_key_groups (
|
||||
client_api_key_id text not null,
|
||||
account_group_id text not null,
|
||||
created_at timestamptz not null,
|
||||
primary key (client_api_key_id, account_group_id),
|
||||
constraint client_api_key_groups_key_fk foreign key (client_api_key_id)
|
||||
references client_api_keys (id)
|
||||
on update restrict
|
||||
on delete cascade,
|
||||
constraint client_api_key_groups_group_fk foreign key (account_group_id)
|
||||
references account_groups (id)
|
||||
on update restrict
|
||||
on delete restrict
|
||||
);
|
||||
|
||||
create index client_api_key_groups_group_idx
|
||||
on client_api_key_groups (account_group_id, client_api_key_id);
|
||||
|
||||
with provider_kinds as (
|
||||
select provider_kind from provider_accounts
|
||||
union
|
||||
select provider_kind from client_api_keys
|
||||
)
|
||||
insert into account_groups (id, name, description, enabled, created_at, updated_at)
|
||||
select
|
||||
'grp_' || md5('codex-proxy-rs:migration-group:' || provider_kind),
|
||||
'迁移池-' || provider_kind,
|
||||
'升级时按原 Provider 权限自动生成',
|
||||
true,
|
||||
now(),
|
||||
now()
|
||||
from provider_kinds;
|
||||
|
||||
insert into account_group_accounts (account_group_id, provider_account_id, created_at)
|
||||
select
|
||||
'grp_' || md5('codex-proxy-rs:migration-group:' || provider_kind),
|
||||
id,
|
||||
clock_timestamp()
|
||||
from provider_accounts;
|
||||
|
||||
insert into client_api_key_groups (client_api_key_id, account_group_id, created_at)
|
||||
select
|
||||
id,
|
||||
'grp_' || md5('codex-proxy-rs:migration-group:' || provider_kind),
|
||||
clock_timestamp()
|
||||
from client_api_keys;
|
||||
|
||||
alter table model_requests
|
||||
add column routing_scope text not null default 'legacy_provider',
|
||||
add column routing_group_refs text[] not null default '{}',
|
||||
add column routing_group_names_snapshot jsonb not null default '[]'::jsonb,
|
||||
add constraint model_requests_routing_scope_ck check (
|
||||
routing_scope in ('legacy_provider', 'all', 'groups')
|
||||
),
|
||||
add constraint model_requests_routing_group_names_ck check (
|
||||
jsonb_typeof(routing_group_names_snapshot) = 'array'
|
||||
and (
|
||||
(
|
||||
routing_scope in ('legacy_provider', 'all')
|
||||
and cardinality(routing_group_refs) = 0
|
||||
and jsonb_array_length(routing_group_names_snapshot) = 0
|
||||
)
|
||||
or (
|
||||
routing_scope = 'groups'
|
||||
and cardinality(routing_group_refs) > 0
|
||||
and array_position(routing_group_refs, null) is null
|
||||
and jsonb_array_length(routing_group_names_snapshot)
|
||||
= cardinality(routing_group_refs)
|
||||
)
|
||||
)
|
||||
);
|
||||
|
||||
create index model_requests_routing_group_refs_idx
|
||||
on model_requests using gin (routing_group_refs);
|
||||
|
||||
alter table model_requests
|
||||
alter column routing_scope drop default;
|
||||
|
||||
alter table client_api_keys
|
||||
drop constraint client_api_keys_provider_kind_ck,
|
||||
drop column provider_kind;
|
||||
@@ -0,0 +1,9 @@
|
||||
alter table account_groups
|
||||
add column color text;
|
||||
|
||||
update account_groups
|
||||
set color = '#2563EB';
|
||||
|
||||
alter table account_groups
|
||||
alter column color set not null,
|
||||
add constraint account_groups_color_ck check (color ~ '^#[0-9A-F]{6}$');
|
||||
@@ -0,0 +1,8 @@
|
||||
alter table account_groups
|
||||
drop constraint account_groups_color_ck;
|
||||
|
||||
update account_groups
|
||||
set color = color || 'FF';
|
||||
|
||||
alter table account_groups
|
||||
add constraint account_groups_color_ck check (color ~ '^#[0-9A-F]{8}$');
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user