Compare commits

..
22 Commits
Author SHA1 Message Date
ZyyandCodex 6074cff68c chore: bump version to 3.2.2
Co-authored-by: Codex <noreply@openai.com>
2026-08-14 00:43:13 +08:00
ZyyandCodex 6ab8e7d0a6 fix(frontend): grade model success rate colors
Co-authored-by: Codex <noreply@openai.com>
2026-08-14 00:42:20 +08:00
ZyyandCodex 17f9ef8915 fix(observability): omit inactive accounts from daily usage
Co-authored-by: Codex <noreply@openai.com>
2026-08-14 00:39:59 +08:00
ZyyandCodex 15e0a28bd0 fix(frontend): refine refresh loading and group alignment
Co-authored-by: Codex <noreply@openai.com>
2026-08-14 00:39:50 +08:00
ZyyandCodex 58bc438992 chore: bump version to 3.2.1
Co-authored-by: Codex <noreply@openai.com>
2026-08-14 00:13:58 +08:00
ZyyandCodex b04616f59f refactor(migrations): squash schema into initial baseline
Co-authored-by: Codex <noreply@openai.com>
2026-08-14 00:13:09 +08:00
ZyyandCodex 82ba4fcbba fix(observability): align dashboard and quota window usage
Co-authored-by: Codex <noreply@openai.com>
2026-08-14 00:11:56 +08:00
ZyyandCodex 819f428ac3 fix(frontend): align loading and daily usage display
Co-authored-by: Codex <noreply@openai.com>
2026-08-14 00:10:13 +08:00
ZyyandCodex 35f28e8aab chore: bump version to 3.2.0
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 22:43:48 +08:00
ZyyandCodex 65645c7403 fix(accounts): gate quota recovery and add manual reset
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 22:43:03 +08:00
ZyyandCodex 71159d5633 refactor(account-groups): remove unused members endpoint
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 17:28:43 +08:00
ZyyandCodex f2ebb09e34 docs: describe group routing and account scheduling
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 16:58:36 +08:00
ZyyandCodex 07acc67ce6 feat(frontend): add account group management
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 16:58:36 +08:00
ZyyandCodex 01ad8bba8a feat(routing): add account groups and weighted scheduling
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 16:58:36 +08:00
ZyyandCodex 986a128d15 fix(observability): separate live status from history range
Co-authored-by: Codex <noreply@openai.com>
2026-08-13 00:26:10 +08:00
ZyyandCodex 6cad2bb6a1 fix(openai): align account import with official token claims
Co-authored-by: Codex <noreply@openai.com>
2026-08-12 23:51:26 +08:00
ZyyandCodex 2b8b3de0e4 refactor(accounts): unify five-state account status
Co-authored-by: Codex <noreply@openai.com>
2026-08-12 23:50:31 +08:00
ZyyandCodex 4127886584 chore: bump version to 3.1.21
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:42:54 +08:00
ZyyandCodex ee3e5db6c2 fix(accounts): isolate model-specific quota usage
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:42:02 +08:00
ZyyandCodex 153cb67b9c refactor(login): refine panel details
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:21:31 +08:00
ZyyandCodex f370d355b6 fix(accounts): align quota ordering and summaries
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:21:23 +08:00
ZyyandCodex a31dad992c fix(frontend): stabilize button loading indicator
Co-authored-by: Codex <noreply@openai.com>
2026-08-11 17:21:16 +08:00
213 changed files with 14935 additions and 6349 deletions
+17 -8
View File
@@ -21,9 +21,9 @@ pub mod ports;
mod use_case;
pub use use_case::{
accounts::AccountsService, auth::AuthService, backup::BackupService,
client_keys::ClientKeyService, observability::ObservabilityService, openai::OpenAiService,
settings::SettingsService, system::SystemService, xai::XaiService,
account_groups::AccountGroupService, accounts::AccountsService, auth::AuthService,
backup::BackupService, client_keys::ClientKeyService, observability::ObservabilityService,
openai::OpenAiService, settings::SettingsService, system::SystemService, xai::XaiService,
};
use model::{AdminError, AdminErrorKind};
@@ -33,10 +33,10 @@ use ports::{
system::SystemOperations,
};
use use_case::{
accounts::DefaultAccountsService, auth::DefaultAuthService, backup::DefaultBackupService,
client_keys::DefaultClientKeyService, observability::DefaultObservabilityService,
openai::DefaultOpenAiService, settings::DefaultSettingsService, system::DefaultSystemService,
xai::DefaultXaiService,
account_groups::DefaultAccountGroupService, accounts::DefaultAccountsService,
auth::DefaultAuthService, backup::DefaultBackupService, client_keys::DefaultClientKeyService,
observability::DefaultObservabilityService, openai::DefaultOpenAiService,
settings::DefaultSettingsService, system::DefaultSystemService, xai::DefaultXaiService,
};
const OPENAI_PROVIDER_KIND: &str = "openai";
@@ -147,6 +147,7 @@ pub enum AdminConfigError {
pub struct AdminServices {
auth: Arc<dyn AuthService>,
accounts: Arc<dyn AccountsService>,
account_groups: Arc<dyn AccountGroupService>,
client_keys: Arc<dyn ClientKeyService>,
observability: Arc<dyn ObservabilityService>,
settings: Arc<dyn SettingsService>,
@@ -167,6 +168,11 @@ impl AdminServices {
self.accounts.as_ref()
}
#[must_use]
pub fn account_groups(&self) -> &dyn AccountGroupService {
self.account_groups.as_ref()
}
#[must_use]
pub fn client_keys(&self) -> &dyn ClientKeyService {
self.client_keys.as_ref()
@@ -255,7 +261,6 @@ pub async fn initialize(
let accounts = Arc::new(DefaultAccountsService::new(
store.accounts(),
store.settings(),
registry.clone(),
snapshot.clone(),
probe.clone(),
@@ -275,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,142 @@
//! Provider-neutral account group commands and projections.
use std::collections::BTreeMap;
use chrono::{DateTime, Utc};
use gateway_core::routing::AccountGroupId;
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,
}
/// 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>,
}
@@ -1,159 +1,19 @@
//! 多 Provider 账号目录与连接测试的公共事实。
use std::{pin::Pin, time::SystemTime};
use std::pin::Pin;
use chrono::{DateTime, Utc};
use futures::Stream;
use gateway_core::routing::ProviderKind;
use super::{PageSize, Revision, observability::TimeRange};
use super::{PageSize, Revision, account_groups::AccountGroupRef, observability::TimeRange};
/// Provider 账号当前可用性;权威定义与字符串编解码在 gateway-core,此处仅复用。
pub use gateway_core::engine::credential::AccountAvailability;
/// 管理页使用的归一化账号状态。
///
/// 只回答两个问题:是否可调度、是否需要人工干预。具体错误原因不进入此枚举,
/// 由 [`AccountErrorReason`] 承载;新增原因只扩展 reason,不改变分类。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountStatus {
/// 正常,可调度。
Normal,
/// 配额耗尽(持久化终态),等待额度恢复,无需干预。
QuotaExhausted,
/// 限流中(quota 快照级信号),自动恢复,无需干预。
RateLimited,
/// 管理员手动停用(`enabled = false`)。
Disabled,
/// 错误(过期 / 失效 / 封禁等),需要人工干预。具体原因见 [`AccountErrorReason`]。
Error,
}
impl AccountStatus {
/// 返回稳定 wire 值。
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Normal => "normal",
Self::QuotaExhausted => "quota_exhausted",
Self::RateLimited => "rate_limited",
Self::Disabled => "disabled",
Self::Error => "error",
}
}
/// 解析稳定 wire 值。
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
match value {
"normal" => Some(Self::Normal),
"quota_exhausted" => Some(Self::QuotaExhausted),
"rate_limited" => Some(Self::RateLimited),
"disabled" => Some(Self::Disabled),
"error" => Some(Self::Error),
_ => None,
}
}
}
/// 状态派生所需的最小事实;由 store / use_case 各自构造,派生规则唯一。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountStatusSignals {
pub enabled: bool,
pub availability: AccountAvailability,
pub access_token_expired: bool,
pub quota_limit_reached: bool,
/// 429 临时限流(Redis 冷却)到期时间;`Some` 表示账号当前处于临时限流。
pub rate_limited_until: Option<SystemTime>,
}
impl AccountStatusSignals {
/// 派生产出运营分类。终态优先于瞬时信号:停用 > 配额耗尽 > 错误 > 限流。
#[must_use]
pub const fn derive(self) -> AccountStatus {
if !self.enabled {
AccountStatus::Disabled
} else if matches!(self.availability, AccountAvailability::QuotaExhausted) {
AccountStatus::QuotaExhausted
} else if self.access_token_expired
|| matches!(
self.availability,
AccountAvailability::Expired
| AccountAvailability::Invalid
| AccountAvailability::Banned
| AccountAvailability::Unknown
)
{
AccountStatus::Error
} else if self.rate_limited_until.is_some() {
AccountStatus::RateLimited
} else if self.quota_limit_reached {
AccountStatus::QuotaExhausted
} else {
AccountStatus::Normal
}
}
}
/// 错误分类下的具体原因;仅在 `AccountStatus::Error` 时存在。
///
/// 受控枚举用于展示与聚合,扩展不改变状态分类。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountErrorReason {
Expired,
TokenExpired,
Invalid,
Banned,
Unknown,
}
impl AccountErrorReason {
/// 返回稳定 wire 值。
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Expired => "expired",
Self::TokenExpired => "token_expired",
Self::Invalid => "invalid",
Self::Banned => "banned",
Self::Unknown => "unknown",
}
}
/// 解析稳定 wire 值。
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
match value {
"expired" => Some(Self::Expired),
"token_expired" => Some(Self::TokenExpired),
"invalid" => Some(Self::Invalid),
"banned" => Some(Self::Banned),
"unknown" => Some(Self::Unknown),
_ => None,
}
}
}
/// 由信号派生错误原因;非错误分类(正常 / 配额耗尽 / 限流 / 停用)返回 `None`。
///
/// 与 [`AccountStatusSignals::derive`] 保持同一判定:`Error` 当且仅当此处返回 `Some`。
#[must_use]
pub fn derive_error_reason(signals: &AccountStatusSignals) -> Option<AccountErrorReason> {
if !signals.enabled {
return None;
}
match signals.availability {
AccountAvailability::QuotaExhausted => None,
AccountAvailability::Ready => signals
.access_token_expired
.then_some(AccountErrorReason::TokenExpired),
AccountAvailability::Expired => Some(AccountErrorReason::Expired),
AccountAvailability::Invalid => Some(AccountErrorReason::Invalid),
AccountAvailability::Banned => Some(AccountErrorReason::Banned),
AccountAvailability::Unknown => Some(AccountErrorReason::Unknown),
}
}
pub use gateway_core::engine::credential::{
AccountConcurrencyLimit, AccountErrorReason, AccountStatus, AccountStatusFacts,
AccountStatusProjection, AccountWeight, CredentialState, QuotaAccessState, QuotaEvidence,
QuotaState, resolve_account_status,
};
/// 账号列表排序字段。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -186,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>,
@@ -207,11 +76,13 @@ pub struct AccountRecord {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
/// 最近一次非 ready 状态变更的原始原因(上游错误原文 / 原因码);恢复 ready 后清空。
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
pub quota: QuotaState,
pub last_error_reason: Option<AccountErrorReason>,
pub last_error_message: Option<String>,
pub quota_observed_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
@@ -296,11 +167,18 @@ pub struct AccountUsageWindowResult {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountPage {
pub config_revision: Revision,
pub items: Vec<AccountRecord>,
pub items: Vec<AccountPageItem>,
pub total: u64,
pub summary: AccountSummary,
}
/// 同一状态快照下的账号事实与唯一状态投影。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountPageItem {
pub account: AccountRecord,
pub projection: AccountStatusProjection,
}
/// 统一账号目录的全局状态计数,不受当前筛选和分页影响。
///
/// 计数与 [`AccountStatus`] 一一对应,由 store 按派生状态聚合。
@@ -314,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>,
}
/// 账号批量删除命令。
@@ -342,9 +247,7 @@ pub enum AccountConnectionTestEvent {
Content {
text: String,
},
Completed {
account_status: AccountStatus,
},
Completed,
Failed {
message: String,
provider_error_code: Option<String>,
@@ -352,7 +255,6 @@ pub enum AccountConnectionTestEvent {
upstream_status: Option<u16>,
upstream_content_type: Option<String>,
upstream_body: Option<String>,
account_status: AccountStatus,
},
}
@@ -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;
@@ -439,25 +439,15 @@ pub struct RequestMetricPoint {
pub costs: Vec<CurrencyCost>,
}
/// 账号池容量统计。
///
/// 首页概览的 `unavailable` 表示不可参与调度,包含配额耗尽及过期、禁用和封禁账号。
///
/// `quota_exhausted`(402 确认耗尽)与 `rate_limited`(429 限流中)单列,
/// 以便状态详情展示和诊断。
/// 账号池的统一五态统计。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AccountPoolMetrics {
pub total: u64,
pub enabled: u64,
pub unavailable: u64,
pub active: u64,
pub rate_limited: u64,
pub expired: u64,
pub invalid: u64,
pub normal: u64,
pub quota_exhausted: u64,
pub refreshing: Option<u64>,
pub rate_limited: u64,
pub disabled: u64,
pub banned: u64,
pub error: u64,
}
/// Dashboard 中一个账号的模型级用量事实。
@@ -488,7 +478,7 @@ pub struct DashboardAccountRequestBucket {
pub request_count: u64,
}
/// Dashboard 中一个账号的完整用量与公共状态事实。
/// Dashboard 中一个账号的完整用量事实。
#[derive(Debug, Clone, PartialEq)]
pub struct DashboardAccountUsage {
pub account_id: String,
@@ -497,8 +487,6 @@ pub struct DashboardAccountUsage {
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
pub enabled: bool,
pub availability: String,
pub request_count: u64,
pub success_count: u64,
pub input_tokens: Option<u64>,
@@ -525,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>,
}
@@ -560,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,
@@ -11,10 +11,7 @@ use gateway_core::{
use super::{
AdminError, MutationActor, MutationContext, PageSize, Revision,
accounts::{
AccountAvailability, AccountErrorReason, AccountRecord, AccountStatus, AccountSummary,
AccountUsage,
},
accounts::{AccountRecord, AccountSummary, AccountUsage, CredentialState},
};
/// Provider-owned JSON;公共层只搬运且 Debug 不输出值。
@@ -55,24 +52,24 @@ pub struct CredentialCursor {
/// Provider credential 列表查询。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CredentialListQuery {
pub availability: Option<CredentialAvailabilityFilter>,
pub credential_state: Option<CredentialStateFilter>,
pub enabled: Option<bool>,
pub window: CredentialListWindow,
}
/// Credential 可用性筛选;支持一个 wire 值对应多个持久状态。
/// Credential 状态筛选。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CredentialAvailabilityFilter {
Exact(AccountAvailability),
AnyOf(Vec<AccountAvailability>),
pub enum CredentialStateFilter {
Exact(CredentialState),
AnyOf(Vec<CredentialState>),
}
impl CredentialAvailabilityFilter {
impl CredentialStateFilter {
#[must_use]
pub fn matches(&self, availability: AccountAvailability) -> bool {
pub fn matches(&self, credential_state: CredentialState) -> bool {
match self {
Self::Exact(expected) => *expected == availability,
Self::AnyOf(expected) => expected.contains(&availability),
Self::Exact(expected) => *expected == credential_state,
Self::AnyOf(expected) => expected.contains(&credential_state),
}
}
}
@@ -155,8 +152,8 @@ pub struct PreparedCredentialCreate {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
}
/// Provider 对一份导入文档的完整验证结果。
@@ -564,6 +561,15 @@ impl fmt::Debug for RotateCredential {
}
}
/// 通用账号用量能否可靠归属到一个 Provider quota 窗口。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QuotaLocalUsageAttribution {
/// 窗口覆盖账号的全部请求,可按账号与时间范围聚合。
AccountWide,
/// 窗口需要 Provider / 模型级归属,通用账号聚合不可用。
Unavailable,
}
/// 一个 Provider quota 窗口的公共投影。
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderQuotaWindow {
@@ -571,6 +577,7 @@ pub struct ProviderQuotaWindow {
pub group: String,
pub label: String,
pub source: Option<String>,
pub local_usage_attribution: QuotaLocalUsageAttribution,
pub window_seconds: Option<u64>,
pub used_percent: Option<f64>,
pub reset_at: Option<DateTime<Utc>>,
@@ -593,15 +600,31 @@ pub struct ProviderQuota {
pub observed_at: Option<DateTime<Utc>>,
pub refresh_token_expires_at: Option<DateTime<Utc>>,
pub windows: Vec<ProviderQuotaWindow>,
/// 快照级限流事实(顶层或任一窗口触顶);Admin 状态派生用,前端不再遍历窗口猜测。
/// 展示用快照级触顶事实(顶层或任一窗口触顶);不参与账号五态派生。
pub limit_reached: bool,
/// 429 临时限流(Redis 冷却)到期时间;`Some` 时账号处于临时限流,到期自动恢复。
/// 与 `limit_reached`(配额耗尽)是独立维度,不污染额度窗口。
pub rate_limited_until: Option<DateTime<Utc>>,
pub provider_data: Option<ProviderDocument>,
}
impl ProviderQuota {
/// 返回账号展开面板使用的当前额度窗口用量。
///
/// 只选择具备完整时间边界且可按账号归属的窗口;优先级与 Dashboard
/// 代表性额度窗口一致,同一优先级保持 Provider 投影顺序。
#[must_use]
pub fn representative_window_usage(&self) -> Option<&AccountUsage> {
self.windows
.iter()
.enumerate()
.filter(|(_, window)| {
window.local_usage_attribution == QuotaLocalUsageAttribution::AccountWide
&& window.window_seconds.is_some()
&& window.reset_at.is_some()
&& window.local_usage.is_some()
})
.min_by_key(|(index, window)| (quota_usage_priority(window), *index))
.and_then(|(_, window)| window.local_usage.as_ref())
}
/// 返回 Dashboard 使用的代表性额度比例。
///
/// 优先使用短周期窗口,再依次使用周、月和其它窗口;同一优先级取较高的已用比例。
@@ -693,9 +716,7 @@ impl fmt::Debug for ProviderExport {
#[derive(Debug, Clone, PartialEq)]
pub struct AccountDirectoryItem {
pub account: AccountRecord,
pub status: AccountStatus,
/// `status == Error` 时的具体原因;其余状态为 `None`。
pub error_reason: Option<AccountErrorReason>,
pub projection: gateway_core::engine::credential::AccountStatusProjection,
pub usage: Option<super::accounts::AccountUsage>,
pub quota: ProviderQuota,
}
@@ -10,9 +10,14 @@ use chrono::{DateTime, Utc};
use super::backup::BackupStorePorts;
use crate::model::{
MutationContext, Revision,
account_groups::{
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, DeleteAccountGroup,
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountListQuery, AccountPage, AccountPageItem, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -83,7 +88,7 @@ pub type AdminStoreResult<T> = Result<T, AdminStoreError>;
pub trait AccountStore: Send + Sync {
async fn list_accounts(&self, query: AccountListQuery) -> AdminStoreResult<AccountPage>;
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountRecord>>;
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountPageItem>>;
async fn load_account_usage(
&self,
@@ -138,11 +143,23 @@ 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 recover_account(
&self,
account_id: &gateway_core::engine::credential::ProviderAccountId,
context: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult>;
async fn batch_update_accounts(
&self,
command: BatchUpdateAccounts,
context: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult>;
async fn delete_accounts(
&self,
@@ -215,10 +232,48 @@ pub trait ClientKeyStore: Send + Sync {
) -> AdminStoreResult<Revision>;
}
/// Provider-neutral account group management transactions.
#[async_trait]
pub trait AccountGroupStore: Send + Sync {
async fn list_account_groups(
&self,
query: AccountGroupListQuery,
) -> AdminStoreResult<AccountGroupPage>;
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 {
async fn dashboard_summary(&self, range: TimeRange) -> AdminStoreResult<DashboardObservation>;
/// 返回历史统计区间和指定观测时刻下的实时账号状态。
async fn dashboard_summary(
&self,
range: TimeRange,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation>;
/// 返回 Dashboard 可选的实时槽位事实。
///
@@ -297,6 +352,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>,
@@ -308,6 +364,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>,
@@ -316,6 +373,7 @@ impl AdminStorePorts {
) -> Self {
Self {
accounts,
account_groups,
auth,
client_keys,
observability,
@@ -329,6 +387,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,130 @@
//! 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, 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 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 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
}
}
@@ -1,6 +1,6 @@
//! 统一账号目录与跨 Provider 动态分派。
use std::{collections::BTreeMap, sync::Arc, time::SystemTime};
use std::{collections::BTreeMap, sync::Arc};
use async_trait::async_trait;
use chrono::{Duration, Utc};
@@ -18,20 +18,17 @@ use crate::{
AdminError, MutationContext,
accounts::{
AccountConnectionTestEvent, AccountConnectionTestEventStream, AccountListQuery,
AccountRecord, AccountStatus, AccountStatusSignals, AccountUsageWindowQuery,
derive_error_reason,
AccountPageItem, AccountUpdateResult, AccountUsageWindowQuery, AccountsUpdateResult,
BatchUpdateAccounts, UpdateAccount,
},
observability::TimeRange,
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
PrepareCredentialRefresh, ProviderModels, ProviderQuota, ProviderQuotaRequest,
ProviderQuotaWindow,
ProviderQuotaWindow, QuotaLocalUsageAttribution,
},
},
ports::{
provider::ProviderAdminRegistry,
store::{AccountStore, SettingsStore},
},
ports::{provider::ProviderAdminRegistry, store::AccountStore},
};
use super::{
@@ -58,6 +55,24 @@ pub trait AccountsService: Send + Sync {
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError>;
async fn recover(
&self,
context: &MutationContext,
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,
@@ -79,7 +94,6 @@ pub trait AccountsService: Send + Sync {
pub(crate) struct DefaultAccountsService {
accounts: Arc<dyn AccountStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
probe: Arc<dyn AccountProbe>,
@@ -89,14 +103,12 @@ impl DefaultAccountsService {
#[must_use]
pub(crate) const fn new(
accounts: Arc<dyn AccountStore>,
settings: Arc<dyn SettingsStore>,
providers: ProviderAdminRegistry,
snapshot: Arc<dyn SnapshotControl>,
probe: Arc<dyn AccountProbe>,
) -> Self {
Self {
accounts,
settings,
providers,
snapshot,
probe,
@@ -106,7 +118,7 @@ impl DefaultAccountsService {
async fn load_account(
&self,
account_id: &ProviderAccountId,
) -> Result<AccountRecord, AdminError> {
) -> Result<AccountPageItem, AdminError> {
self.accounts
.load_account(account_id.as_str())
.await
@@ -119,32 +131,32 @@ impl DefaultAccountsService {
account_id: &ProviderAccountId,
) -> Result<
(
AccountRecord,
AccountPageItem,
Arc<dyn crate::ports::provider::ProviderAdmin>,
),
AdminError,
> {
let account = self.load_account(account_id).await?;
let item = self.load_account(account_id).await?;
let provider = self
.providers
.require(&account.provider_kind)
.require(&item.account.provider_kind)
.map_err(|error| map_provider_error(error, "provider account"))?;
Ok((account, provider))
Ok((item, provider))
}
async fn attach_quota_local_usage(
&self,
accounts: &[AccountRecord],
accounts: &[AccountPageItem],
quotas: &mut [ProviderQuota],
) -> Result<(), AdminError> {
let windows = accounts
.iter()
.zip(quotas.iter())
.flat_map(|(account, quota)| {
.flat_map(|(item, quota)| {
quota
.windows
.iter()
.filter_map(|window| quota_usage_window(&account.id, window))
.filter_map(|window| quota_usage_window(&item.account.id, window))
})
.collect::<Vec<_>>();
if windows.is_empty() {
@@ -158,10 +170,12 @@ impl DefaultAccountsService {
.into_iter()
.map(|result| ((result.account_id, result.key), result.usage))
.collect::<BTreeMap<_, _>>();
for (account, quota) in accounts.iter().zip(quotas) {
for (item, quota) in accounts.iter().zip(quotas) {
for window in &mut quota.windows {
if window.local_usage.is_none() {
let key = (account.id.clone(), window.key.clone());
if window.local_usage.is_none()
&& window.local_usage_attribution == QuotaLocalUsageAttribution::AccountWide
{
let key = (item.account.id.clone(), window.key.clone());
if let Some(usage) = usage_by_window.get(&key) {
window.local_usage = Some(usage.clone());
}
@@ -176,33 +190,19 @@ impl DefaultAccountsService {
account_id: &ProviderAccountId,
refresh_quota: bool,
) -> Result<AccountDirectoryItem, AdminError> {
let (account, provider) = self.provider_for_account(account_id).await?;
let settings = self
.settings
.load_runtime_settings()
.await
.map_err(|error| map_store_error(error, "runtime settings"))?;
let (stored, provider) = self.provider_for_account(account_id).await?;
let account = &stored.account;
let now = Utc::now();
let retained_range = retained_usage_range(now, settings.usage_retention_days);
let rolling_range = TimeRange {
start: now - Duration::hours(24),
end: now,
};
let ids = vec![account.id.clone()];
let (usage, rolling_usage) = futures::try_join!(
async {
self.accounts
.load_account_usage(retained_range, &ids)
.await
.map_err(|error| map_store_error(error, "account usage"))
},
async {
self.accounts
.load_account_usage(rolling_range, &ids)
.await
.map_err(|error| map_store_error(error, "rolling account usage"))
},
)?;
let rolling_usage = self
.accounts
.load_account_usage(rolling_range, &ids)
.await
.map_err(|error| map_store_error(error, "rolling account usage"))?;
let rolling_usage = rolling_usage.into_iter().next();
let mut quota = provider
.quota(ProviderQuotaRequest {
@@ -213,17 +213,15 @@ impl DefaultAccountsService {
.await
.map_err(|error| map_provider_error(error, "provider quota"))?;
self.attach_quota_local_usage(
std::slice::from_ref(&account),
std::slice::from_ref(&stored),
std::slice::from_mut(&mut quota),
)
.await?;
let signals = account_signals(&account, &quota, now);
let error_reason = derive_error_reason(&signals);
let usage = quota.representative_window_usage().cloned();
Ok(AccountDirectoryItem {
status: signals.derive(),
error_reason,
usage: usage.into_iter().next(),
account,
projection: stored.projection,
usage,
account: stored.account,
quota,
})
}
@@ -232,13 +230,12 @@ impl DefaultAccountsService {
#[async_trait]
impl AccountsService for DefaultAccountsService {
async fn list(&self, query: AccountListQuery) -> Result<AccountDirectoryPage, AdminError> {
let (page, settings) = futures::try_join!(
self.accounts.list_accounts(query),
self.settings.load_runtime_settings(),
)
.map_err(|error| map_store_error(error, "account directory"))?;
let page = self
.accounts
.list_accounts(query)
.await
.map_err(|error| map_store_error(error, "account directory"))?;
let now = Utc::now();
let retained_range = retained_usage_range(now, settings.usage_retention_days);
let rolling_range = TimeRange {
start: now - Duration::hours(24),
end: now,
@@ -246,22 +243,19 @@ impl AccountsService for DefaultAccountsService {
let ids = page
.items
.iter()
.map(|account| account.id.clone())
.map(|item| item.account.id.clone())
.collect::<Vec<_>>();
let (usage, rolling_usage) = futures::try_join!(
self.accounts.load_account_usage(retained_range, &ids),
self.accounts.load_account_usage(rolling_range, &ids),
)
.map_err(|error| map_store_error(error, "account usage"))?;
let usage = usage
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect::<BTreeMap<_, _>>();
let rolling_usage = self
.accounts
.load_account_usage(rolling_range, &ids)
.await
.map_err(|error| map_store_error(error, "rolling account usage"))?;
let rolling_usage = rolling_usage
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect::<BTreeMap<_, _>>();
let mut quotas = futures::future::join_all(page.items.iter().map(|account| async {
let mut quotas = futures::future::join_all(page.items.iter().map(|item| async {
let account = &item.account;
let account_id = ProviderAccountId::new(account.id.clone())
.map_err(|_| AdminError::invalid("Invalid provider account ID"))?;
// 单个账号的 quota 投影失败(Provider 未注册或 quota 读取失败)不拖垮整页:
@@ -305,14 +299,12 @@ impl AccountsService for DefaultAccountsService {
.items
.into_iter()
.zip(quotas)
.map(|(account, quota)| {
let signals = account_signals(&account, &quota, now);
let error_reason = derive_error_reason(&signals);
.map(|(item, quota)| {
let usage = quota.representative_window_usage().cloned();
AccountDirectoryItem {
status: signals.derive(),
error_reason,
usage: usage.get(&account.id).cloned(),
account,
usage,
account: item.account,
projection: item.projection,
quota,
}
})
@@ -340,7 +332,7 @@ impl AccountsService for DefaultAccountsService {
for account_id in account_ids {
let account = self.load_account(&account_id).await?;
grouped
.entry(account.provider_kind)
.entry(account.account.provider_kind)
.or_default()
.push(account_id);
}
@@ -383,7 +375,8 @@ impl AccountsService for DefaultAccountsService {
context: &MutationContext,
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError> {
let (account, provider) = self.provider_for_account(&account_id).await?;
let (stored, provider) = self.provider_for_account(&account_id).await?;
let account = stored.account;
let prepared = provider
.prepare_refresh(PrepareCredentialRefresh {
account: account.clone(),
@@ -409,6 +402,98 @@ impl AccountsService for DefaultAccountsService {
})
}
async fn recover(
&self,
context: &MutationContext,
account_id: ProviderAccountId,
) -> Result<AccountRefreshResult, AdminError> {
let (_, provider) = self.provider_for_account(&account_id).await?;
let result = self
.accounts
.recover_account(&account_id, context)
.await
.map_err(|error| map_store_error(error, "provider account recovery"))?;
provider
.account_facts_changed(std::slice::from_ref(&account_id))
.await;
publish_committed(self.snapshot.as_ref(), result.config_revision).await?;
let account = self.load_directory_item(&account_id, false).await?;
Ok(AccountRefreshResult {
config_revision: result.config_revision,
account,
})
}
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,
@@ -434,8 +519,8 @@ impl AccountsService for DefaultAccountsService {
account_id: ProviderAccountId,
upstream_model: UpstreamModelId,
) -> Result<AccountConnectionTestEventStream, AdminError> {
let (account, provider) = self.provider_for_account(&account_id).await?;
let initial_status = account_status(&account, false, Utc::now());
let (stored, provider) = self.provider_for_account(&account_id).await?;
let account = stored.account;
let model = upstream_model.as_str().to_owned();
let operation = provider
.connection_test_operation(&upstream_model, CONNECTION_TEST_INPUT)
@@ -452,8 +537,6 @@ impl AccountsService for DefaultAccountsService {
},
];
let probe = Arc::clone(&self.probe);
let accounts = Arc::clone(&self.accounts);
let stored_account_id = account.id.clone();
let terminal = futures::stream::once(async move {
let result = probe
.probe(AccountProbeRequest {
@@ -463,22 +546,12 @@ impl AccountsService for DefaultAccountsService {
operation,
})
.await;
let status = accounts
.load_account(&stored_account_id)
.await
.ok()
.flatten()
.map_or(initial_status, |account| {
account_status(&account, false, Utc::now())
});
match result {
Ok(result) => result
.text
.into_iter()
.map(|text| AccountConnectionTestEvent::Content { text })
.chain(std::iter::once(AccountConnectionTestEvent::Completed {
account_status: status,
}))
.chain(std::iter::once(AccountConnectionTestEvent::Completed))
.collect(),
Err(error) => {
let upstream_status = error
@@ -504,7 +577,6 @@ impl AccountsService for DefaultAccountsService {
upstream_status,
upstream_content_type,
upstream_body,
account_status: status,
}]
}
}
@@ -521,54 +593,17 @@ fn empty_quota() -> ProviderQuota {
refresh_token_expires_at: None,
windows: Vec::new(),
limit_reached: false,
rate_limited_until: None,
provider_data: None,
}
}
/// 构造状态派生所需的信号快照。
fn account_signals(
account: &AccountRecord,
quota: &ProviderQuota,
now: chrono::DateTime<Utc>,
) -> AccountStatusSignals {
AccountStatusSignals {
enabled: account.enabled,
availability: account.availability,
access_token_expired: account
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now),
quota_limit_reached: quota.limit_reached,
rate_limited_until: quota.rate_limited_until.map(SystemTime::from),
}
}
/// 账号运营分类(互斥展示状态)。
fn account_status(
account: &AccountRecord,
quota_limit_reached: bool,
now: chrono::DateTime<Utc>,
) -> AccountStatus {
let quota = ProviderQuota {
limit_reached: quota_limit_reached,
..empty_quota()
};
account_signals(account, &quota, now).derive()
}
fn retained_usage_range(now: chrono::DateTime<Utc>, retention_days: u32) -> TimeRange {
// 账号留存投影是内部查询,不受外部观测接口 366 天上限约束。
TimeRange {
start: now - Duration::days(i64::from(retention_days)),
end: now,
}
}
fn quota_usage_window(
account_id: &str,
window: &ProviderQuotaWindow,
) -> Option<AccountUsageWindowQuery> {
if window.local_usage.is_some() {
if window.local_usage.is_some()
|| window.local_usage_attribution != QuotaLocalUsageAttribution::AccountWide
{
return None;
}
let reset_at = window.reset_at?;
@@ -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,
@@ -98,10 +98,11 @@ impl ObservabilityService for DefaultObservabilityService {
range: TimeRange,
kind: TrendKind,
) -> Result<DashboardResult, AdminError> {
let observed_at = Utc::now();
let (mut observation, settings, runtime_slots) = futures::try_join!(
self.store.dashboard_summary(range),
self.store.dashboard_summary(range, observed_at),
self.settings.load_runtime_settings(),
self.store.dashboard_runtime_slots(range.end),
self.store.dashboard_runtime_slots(observed_at),
)
.map_err(|error| map_store_error(error, "dashboard"))?;
self.enrich_billing(&mut observation.recent_requests);
@@ -125,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.active, |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 {
@@ -484,8 +493,7 @@ impl DefaultObservabilityService {
});
let quota_used = futures::future::join_all(quota_reads).await;
for (account, used_percent) in accounts.iter_mut().zip(quota_used) {
account.quota_used_percent = used_percent
.or_else(|| (account.availability == "quota_exhausted").then_some(100.0));
account.quota_used_percent = used_percent;
}
}
@@ -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,
+2 -1
View File
@@ -11,6 +11,7 @@ use gateway_admin::{
observability::{DecimalAmount, RequestOutcome, TimeRange},
provider_credentials::{
CredentialCommitGuard, ProviderDocument, ProviderQuota, ProviderQuotaWindow,
QuotaLocalUsageAttribution,
},
settings::AdminApiKey,
},
@@ -128,7 +129,6 @@ fn representative_quota_should_prefer_short_window_and_highest_usage() {
quota_window("shortTerm", Some(14_400), Some(45.0)),
],
limit_reached: false,
rate_limited_until: None,
provider_data: None,
};
@@ -145,6 +145,7 @@ fn quota_window(
group: group.to_owned(),
label: group.to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds,
used_percent,
reset_at: None,
@@ -5,7 +5,10 @@ use chrono::{TimeDelta, Utc};
use futures::{StreamExt as _, future::BoxFuture};
use gateway_core::{
engine::{
credential::{OpaqueProviderData, ProviderAccountId},
credential::{
AccountStatusFacts, CredentialState, OpaqueProviderData, ProviderAccountId,
QuotaEvidence, QuotaState, resolve_account_status,
},
probe::{AccountProbe, AccountProbeError, AccountProbeRequest, AccountProbeResult},
},
error::{ClientVisibleUpstreamError, GatewayError, GatewayErrorKind},
@@ -18,9 +21,10 @@ use gateway_admin::{
model::{
AdminError, MutationContext, Revision,
accounts::{
AccountAvailability, AccountConnectionTestEvent, AccountListQuery, AccountPage,
AccountRecord, AccountSummary, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountConnectionTestEvent, AccountListQuery, AccountPage, AccountPageItem,
AccountRecord, AccountSummary, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
observability::TimeRange,
provider_credentials::{
@@ -33,7 +37,7 @@ use gateway_admin::{
PreparedAuthorizationCredential, PreparedCredentialCreate, PreparedCredentialImport,
PreparedCredentialRotation, PreparedCredentialRotationFacts, ProviderDocument,
ProviderExport, ProviderExportCredentialInput, ProviderModels, ProviderQuota,
ProviderQuotaRequest, ProviderQuotaWindow,
ProviderQuotaRequest, ProviderQuotaWindow, QuotaLocalUsageAttribution,
},
settings::{
AdminApiKey, AdminApiKeyMutation, ReplaceRuntimeSettings, RotationStrategy,
@@ -378,6 +382,7 @@ pub(super) struct FakeAccountStore {
fail_commit: Mutex<bool>,
audit_requests: Mutex<Vec<String>>,
quota_window_usage: Mutex<Vec<AccountUsageWindowResult>>,
quota_window_queries: Mutex<Vec<AccountUsageWindowQuery>>,
}
impl FakeAccountStore {
@@ -393,6 +398,7 @@ impl FakeAccountStore {
fail_commit: Mutex::new(false),
audit_requests: Mutex::new(Vec::new()),
quota_window_usage: Mutex::new(Vec::new()),
quota_window_queries: Mutex::new(Vec::new()),
})
}
@@ -408,6 +414,13 @@ impl FakeAccountStore {
*self.quota_window_usage.lock().expect("quota window usage") = usage;
}
fn quota_window_queries(&self) -> Vec<AccountUsageWindowQuery> {
self.quota_window_queries
.lock()
.expect("quota window queries")
.clone()
}
pub(super) fn set_accounts(&self, accounts: Vec<AccountRecord>) {
*self.accounts.lock().expect("accounts") = accounts;
}
@@ -438,6 +451,22 @@ impl FakeAccountStore {
.expect("audit requests")
.push(context.request_id.clone());
}
fn page_item(account: AccountRecord) -> AccountPageItem {
let facts = AccountStatusFacts {
enabled: account.enabled,
credential_state: account.credential_state,
access_token_expires_at: account.access_token_expires_at.map(Into::into),
quota: account.quota,
rate_limited_until: None,
last_error_reason: account.last_error_reason,
last_error_message: account.last_error_message.clone(),
};
AccountPageItem {
account,
projection: resolve_account_status(&facts, std::time::SystemTime::now()),
}
}
}
#[async_trait]
@@ -445,12 +474,13 @@ impl AccountStore for FakeAccountStore {
async fn list_accounts(&self, _: AccountListQuery) -> AdminStoreResult<AccountPage> {
self.record("store.list_accounts");
let accounts = self.accounts.lock().expect("accounts").clone();
let total = accounts.len() as u64;
Ok(AccountPage {
config_revision: revision(1),
items: accounts.clone(),
total: accounts.len() as u64,
items: accounts.into_iter().map(Self::page_item).collect(),
total,
summary: AccountSummary {
total: accounts.len() as u64,
total,
normal: 1,
quota_exhausted: 0,
rate_limited: 0,
@@ -460,7 +490,7 @@ impl AccountStore for FakeAccountStore {
})
}
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountRecord>> {
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountPageItem>> {
self.record("store.load_account");
// probe 后的账号状态覆盖只对同一 id 生效;其余按账号列表查询。
let account = self
@@ -477,7 +507,7 @@ impl AccountStore for FakeAccountStore {
.find(|account| account.id == account_id)
.cloned()
});
Ok(account)
Ok(account.map(Self::page_item))
}
async fn load_account_usage(
@@ -490,8 +520,12 @@ impl AccountStore for FakeAccountStore {
async fn load_account_usage_by_windows(
&self,
_: &[AccountUsageWindowQuery],
windows: &[AccountUsageWindowQuery],
) -> AdminStoreResult<Vec<AccountUsageWindowResult>> {
*self
.quota_window_queries
.lock()
.expect("quota window queries") = windows.to_vec();
Ok(self
.quota_window_usage
.lock()
@@ -628,15 +662,61 @@ 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 recover_account(
&self,
account_id: &ProviderAccountId,
context: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult> {
self.record("store.recover_account");
self.record_context(context);
self.require_commit()?;
let mut accounts = self.accounts.lock().expect("accounts");
if let Some(account) = accounts
.iter_mut()
.find(|account| account.id == account_id.as_str())
{
account.enabled = true;
account.credential_state = CredentialState::Ready;
account.quota = QuotaState::allowed(std::time::SystemTime::now());
account.last_error_reason = None;
account.last_error_message = None;
}
Ok(AccountUpdateResult {
config_revision: revision(2),
account_id: account_id.clone(),
})
}
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(
@@ -815,7 +895,11 @@ fn provider_registry_should_reject_duplicate_kind() {
async fn connection_test_should_probe_unavailable_account() {
let provider = FakeProviderAdmin::new("xai", events());
let mut account = account_record("xai");
account.availability = AccountAvailability::QuotaExhausted;
account.quota = QuotaState::exhausted(
QuotaEvidence::UsageLimitReached,
std::time::SystemTime::now(),
None,
);
let store = FakeAccountStore::with_account(account, events());
let services =
accounts_service_with_probe(provider, store, Arc::new(FailingAccountProbe)).await;
@@ -837,7 +921,6 @@ async fn connection_test_should_probe_unavailable_account() {
message,
provider_error_code: Some(code),
provider_error_type: Some(error_type),
account_status: gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
..
}) if message == "included usage exhausted"
&& code == "usage_exhausted"
@@ -846,7 +929,7 @@ async fn connection_test_should_probe_unavailable_account() {
}
#[tokio::test]
async fn connection_test_rate_limited_probe_preserves_normal_account_status() {
async fn connection_test_rate_limited_probe_returns_provider_failure() {
let events = events();
let provider = FakeProviderAdmin::new("xai", events.clone());
let store = FakeAccountStore::new("xai", events);
@@ -869,9 +952,10 @@ async fn connection_test_rate_limited_probe_preserves_normal_account_status() {
assert!(matches!(
events.last(),
Some(AccountConnectionTestEvent::Failed {
account_status: gateway_admin::model::accounts::AccountStatus::Normal,
message,
..
})
if message == "xAI free model quota is exhausted"
));
}
@@ -897,10 +981,7 @@ async fn connection_test_should_preserve_disabled_account_status() {
assert!(matches!(
events.last(),
Some(AccountConnectionTestEvent::Failed {
account_status: gateway_admin::model::accounts::AccountStatus::Disabled,
..
})
Some(AccountConnectionTestEvent::Failed { .. })
));
}
@@ -949,7 +1030,7 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
assert_eq!(result.config_revision, revision(2));
assert_eq!(result.account.account.provider_kind.as_str(), "openai");
assert_eq!(
result.account.status,
result.account.projection.status,
gateway_admin::model::accounts::AccountStatus::Normal
);
@@ -967,6 +1048,175 @@ async fn accounts_refresh_should_keep_guard_through_store_commit() {
assert_eq!(store.audit_requests(), ["refresh-request"]);
}
#[tokio::test]
async fn accounts_recover_should_commit_facts_then_return_normal_account() {
let events = events();
let provider = FakeProviderAdmin::new("openai", events.clone());
let mut account = account_record("openai");
account.enabled = false;
account.credential_state = CredentialState::Invalid;
account.quota = QuotaState::exhausted(
QuotaEvidence::UsageLimitReached,
std::time::SystemTime::now(),
None,
);
account.last_error_reason =
Some(gateway_core::engine::credential::AccountErrorReason::CredentialInvalid);
account.last_error_message = Some("invalid credential".to_owned());
let store = FakeAccountStore::with_account(account, events.clone());
let services = accounts_service(provider, store.clone()).await;
let result = services
.accounts()
.recover(
&context("recover-request"),
ProviderAccountId::new("acct_test").expect("account ID"),
)
.await
.expect("recover account");
assert_eq!(result.config_revision, revision(2));
assert!(result.account.account.enabled);
assert_eq!(
result.account.account.credential_state,
CredentialState::Ready
);
assert_eq!(
result.account.account.quota.access(),
gateway_core::engine::credential::QuotaAccessState::Allowed
);
assert_eq!(
result.account.projection.status,
gateway_admin::model::accounts::AccountStatus::Normal
);
assert_eq!(
recorded(&events),
[
"store.load_account",
"store.recover_account",
"provider.account_facts_changed",
"store.load_account",
]
);
assert_eq!(store.audit_requests(), ["recover-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());
@@ -978,6 +1228,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,
@@ -990,7 +1241,7 @@ async fn accounts_list_should_return_complete_directory_semantics() {
let account = page.items.first().expect("account item");
assert_eq!(account.account.provider_kind.as_str(), "openai");
assert_eq!(
account.status,
account.projection.status,
gateway_admin::model::accounts::AccountStatus::Normal
);
}
@@ -1007,6 +1258,7 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(97.0),
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
@@ -1015,7 +1267,6 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
provider_data: None,
}],
limit_reached: false,
rate_limited_until: None,
provider_data: None,
});
let failing = FakeProviderAdmin::new("xai", events.clone());
@@ -1039,6 +1290,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,
@@ -1063,10 +1315,14 @@ async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_drop
}
#[tokio::test]
async fn accounts_list_should_prefer_quota_exhaustion_over_stale_token_expiry() {
async fn accounts_list_should_prefer_credential_error_over_quota_exhaustion() {
let provider = FakeProviderAdmin::new("openai", events());
let mut account = account_record("openai");
account.availability = AccountAvailability::QuotaExhausted;
account.quota = QuotaState::exhausted(
QuotaEvidence::ProviderDenied,
std::time::SystemTime::now(),
None,
);
account.access_token_expires_at = Some(Utc::now() - TimeDelta::minutes(5));
let store = FakeAccountStore::with_account(account, events());
@@ -1077,6 +1333,7 @@ async fn accounts_list_should_prefer_quota_exhaustion_over_stale_token_expiry()
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,
@@ -1085,17 +1342,17 @@ async fn accounts_list_should_prefer_quota_exhaustion_over_stale_token_expiry()
.expect("usage-limited account directory");
assert_eq!(
page.items.first().expect("account item").status,
gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
page.items.first().expect("account item").projection.status,
gateway_admin::model::accounts::AccountStatus::Error,
);
}
#[tokio::test]
async fn accounts_list_should_map_unknown_availability_to_error_not_normal() {
async fn accounts_list_should_map_unknown_credential_to_error_not_normal() {
// Unknown 不可调度,Admin 不得显示为 normal。
let provider = FakeProviderAdmin::new("openai", events());
let mut account = account_record("openai");
account.availability = AccountAvailability::Unknown;
account.credential_state = CredentialState::Unknown;
let store = FakeAccountStore::with_account(account, events());
let page = accounts_service(provider, store)
@@ -1105,6 +1362,7 @@ async fn accounts_list_should_map_unknown_availability_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,
@@ -1113,15 +1371,13 @@ async fn accounts_list_should_map_unknown_availability_to_error_not_normal() {
.expect("unknown account directory");
assert_eq!(
page.items.first().expect("account item").status,
page.items.first().expect("account item").projection.status,
gateway_admin::model::accounts::AccountStatus::Error,
);
}
#[tokio::test]
async fn accounts_list_should_map_rate_limited_from_quota_limit_reached() {
// ready + 快照级 quota.limitReached → quota_exhausted(配额耗尽);
// rate_limited 只来自 ProviderQuota.rate_limited_until(429 Redis 冷却)。
async fn accounts_list_should_not_derive_rate_limited_from_provider_quota_view() {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
observed_at: Some(Utc::now()),
@@ -1131,6 +1387,7 @@ async fn accounts_list_should_map_rate_limited_from_quota_limit_reached() {
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(100.0),
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
@@ -1139,7 +1396,6 @@ async fn accounts_list_should_map_rate_limited_from_quota_limit_reached() {
provider_data: None,
}],
limit_reached: true,
rate_limited_until: Some(Utc::now() + TimeDelta::minutes(5)),
provider_data: None,
});
let store = FakeAccountStore::new("openai", events());
@@ -1151,22 +1407,22 @@ async fn accounts_list_should_map_rate_limited_from_quota_limit_reached() {
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,
})
.await
.expect("rate limited account directory");
.expect("account directory with display quota cooldown");
assert_eq!(
page.items.first().expect("account item").status,
gateway_admin::model::accounts::AccountStatus::RateLimited,
page.items.first().expect("account item").projection.status,
gateway_admin::model::accounts::AccountStatus::Normal,
);
}
#[tokio::test]
async fn accounts_list_should_map_quota_exhausted_when_only_limit_reached() {
// ready + 快照级 quota.limitReached、无 429 冷却 → quota_exhausted。
async fn accounts_list_should_not_derive_exhaustion_from_provider_quota_view() {
let provider = FakeProviderAdmin::new("openai", events());
provider.set_quota(ProviderQuota {
observed_at: Some(Utc::now()),
@@ -1176,6 +1432,7 @@ async fn accounts_list_should_map_quota_exhausted_when_only_limit_reached() {
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(100.0),
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
@@ -1184,7 +1441,6 @@ async fn accounts_list_should_map_quota_exhausted_when_only_limit_reached() {
provider_data: None,
}],
limit_reached: true,
rate_limited_until: None,
provider_data: None,
});
let store = FakeAccountStore::new("openai", events());
@@ -1196,16 +1452,17 @@ async fn accounts_list_should_map_quota_exhausted_when_only_limit_reached() {
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,
})
.await
.expect("quota exhausted account directory");
.expect("account directory with display quota limit");
assert_eq!(
page.items.first().expect("account item").status,
gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
page.items.first().expect("account item").projection.status,
gateway_admin::model::accounts::AccountStatus::Normal,
);
}
@@ -1221,6 +1478,7 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(97.0),
reset_at: Some(reset_at),
@@ -1229,7 +1487,6 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
provider_data: None,
}],
limit_reached: false,
rate_limited_until: None,
provider_data: None,
});
let store = FakeAccountStore::new("openai", events());
@@ -1239,13 +1496,14 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
usage: quota_local_usage("acct_test", 4_330_000),
}]);
let page = accounts_service(provider, store)
let page = accounts_service(provider, store.clone())
.await
.accounts()
.list(AccountListQuery {
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,
@@ -1253,11 +1511,120 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
.await
.expect("quota window usage");
let usage = page.items[0].quota.windows[0]
let item = &page.items[0];
let usage = item.quota.windows[0]
.local_usage
.as_ref()
.expect("quota window local usage");
assert_eq!(usage.total_tokens, Some(4_330_000));
assert_eq!(
item.usage.as_ref().and_then(|usage| usage.total_tokens),
Some(4_330_000),
);
let queries = store.quota_window_queries();
assert_eq!(queries.len(), 1);
assert_eq!(queries[0].account_id, "acct_test");
assert_eq!(queries[0].key, "primary");
assert_eq!(queries[0].range.start, reset_at - TimeDelta::hours(5));
assert_eq!(queries[0].range.end, reset_at);
}
#[tokio::test]
async fn accounts_list_should_not_attach_account_usage_to_model_specific_quota_windows() {
let provider = FakeProviderAdmin::new("openai", events());
let reset_at = Utc::now() + TimeDelta::days(7);
provider.set_quota(ProviderQuota {
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![
ProviderQuotaWindow {
key: "core-primary".to_owned(),
group: "weekly".to_owned(),
label: "周额度".to_owned(),
source: Some("codex".to_owned()),
local_usage_attribution: QuotaLocalUsageAttribution::AccountWide,
window_seconds: Some(7 * 24 * 60 * 60),
used_percent: Some(1.0),
reset_at: Some(reset_at),
limit_reached: false,
local_usage: None,
provider_data: None,
},
ProviderQuotaWindow {
key: "codex-bengalfox-primary".to_owned(),
group: "weekly".to_owned(),
label: "GPT-5.3-Codex-Spark · 周额度".to_owned(),
source: Some("codex_bengalfox".to_owned()),
local_usage_attribution: QuotaLocalUsageAttribution::Unavailable,
window_seconds: Some(7 * 24 * 60 * 60),
used_percent: Some(0.0),
reset_at: Some(reset_at),
limit_reached: false,
local_usage: None,
provider_data: None,
},
],
limit_reached: false,
provider_data: None,
});
let store = FakeAccountStore::new("openai", events());
store.set_quota_window_usage(vec![
AccountUsageWindowResult {
account_id: "acct_test".to_owned(),
key: "core-primary".to_owned(),
usage: quota_local_usage("acct_test", 12_818_806),
},
AccountUsageWindowResult {
account_id: "acct_test".to_owned(),
key: "codex-bengalfox-primary".to_owned(),
usage: quota_local_usage("acct_test", 127_926),
},
]);
let page = accounts_service(provider, store)
.await
.accounts()
.list(AccountListQuery {
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,
})
.await
.expect("quota window usage");
let local_tokens = page.items[0]
.quota
.windows
.iter()
.map(|window| {
(
window.key.as_str(),
window
.local_usage
.as_ref()
.and_then(|usage| usage.total_tokens),
)
})
.collect::<Vec<_>>();
assert_eq!(
local_tokens,
[
("core-primary", Some(12_818_806)),
("codex-bengalfox-primary", None),
],
);
assert_eq!(
page.items[0]
.usage
.as_ref()
.and_then(|usage| usage.total_tokens),
Some(12_818_806),
);
}
#[tokio::test]
@@ -1338,7 +1705,6 @@ fn empty_quota() -> ProviderQuota {
refresh_token_expires_at: None,
windows: Vec::new(),
limit_reached: false,
rate_limited_until: None,
provider_data: None,
}
}
@@ -1371,6 +1737,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()),
@@ -1382,10 +1749,13 @@ 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,
availability: AccountAvailability::Ready,
availability_observed_at: now,
concurrency_limit: None,
weight: gateway_core::engine::credential::AccountWeight::DEFAULT,
credential_state: CredentialState::Ready,
credential_observed_at: now,
quota: QuotaState::allowed(now.into()),
last_error_reason: None,
last_error_message: None,
quota_observed_at: None,
created_at: now,
updated_at: now,
}
@@ -1419,8 +1789,8 @@ fn prepared_create_with_id(
access_token_expires_at: Some(now + TimeDelta::hours(1)),
next_refresh_at: Some(now + TimeDelta::minutes(30)),
enabled: true,
availability: AccountAvailability::Ready,
availability_observed_at: now,
credential_state: CredentialState::Ready,
credential_observed_at: now,
}
}
@@ -1503,9 +1873,7 @@ impl AccountProbe for FreeModelQuotaProbe {
_: AccountProbeRequest,
) -> BoxFuture<'_, Result<AccountProbeResult, AccountProbeError>> {
Box::pin(async move {
let mut account = account_record("xai");
account.availability = AccountAvailability::Ready;
self.store.set_account_after_probe(account);
self.store.set_account_after_probe(account_record("xai"));
Err(GatewayError::new(
GatewayErrorKind::RateLimited,
"xAI free model quota is exhausted",
@@ -14,14 +14,20 @@ use std::{
};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use futures::future::BoxFuture;
use gateway_admin::{
AdminConfig, AdminServices, InitialAdminPassword,
model::{
MutationContext, Revision,
account_groups::{
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, DeleteAccountGroup,
NewAccountGroup, SetAccountGroupEnabled, UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountRecord, AccountUsage, AccountUsageWindowQuery,
AccountUsageWindowResult, DeleteAccounts, SetAccountEnabled,
AccountListQuery, AccountPage, AccountUpdateResult, AccountUsage,
AccountUsageWindowQuery, AccountUsageWindowResult, AccountsUpdateResult,
BatchUpdateAccounts, DeleteAccounts, UpdateAccount,
},
auth::{AdminAuditEvent, AdminSession},
client_keys::{
@@ -50,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,
@@ -170,6 +176,7 @@ impl AdminHarness {
},
AdminStorePorts::new(
self.accounts,
Arc::new(UnavailableAccountGroupStore),
self.auth,
self.client_keys,
self.observability,
@@ -234,13 +241,60 @@ 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 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> {
Err(unavailable("accounts"))
}
async fn load_account(&self, _: &str) -> AdminStoreResult<Option<AccountRecord>> {
async fn load_account(
&self,
_: &str,
) -> AdminStoreResult<Option<gateway_admin::model::accounts::AccountPageItem>> {
Err(unavailable("account"))
}
@@ -315,14 +369,30 @@ 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 recover_account(
&self,
_: &ProviderAccountId,
_: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult> {
Err(unavailable("account recovery"))
}
async fn batch_update_accounts(
&self,
_: BatchUpdateAccounts,
_: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
Err(unavailable("account batch update"))
}
async fn delete_accounts(
&self,
_: DeleteAccounts,
@@ -388,7 +458,11 @@ impl ClientKeyStore for UnavailableStore {
#[async_trait]
impl ObservabilityStore for UnavailableStore {
async fn dashboard_summary(&self, _: TimeRange) -> AdminStoreResult<DashboardObservation> {
async fn dashboard_summary(
&self,
_: TimeRange,
_: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
Err(unavailable("dashboard"))
}
@@ -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,28 @@ 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]
async fn dashboard_summary_should_share_one_current_observation_time_across_runtime_facts() {
let historical_end = Utc::now() - Duration::hours(1);
let store = Arc::new(FixtureObservabilityStore::new(observation_range(
historical_end,
)));
let services = observability_services(store.clone()).await;
services
.observability()
.dashboard_summary(observation_range(historical_end), TrendKind::Usage)
.await
.expect("dashboard summary");
let (summary_observed_at, slots_observed_at) = store.observed_times();
assert_eq!(summary_observed_at, slots_observed_at);
assert!(summary_observed_at.is_some_and(|value| value > historical_end));
}
#[tokio::test]
@@ -394,6 +414,8 @@ struct FixtureObservabilityStore {
calculated_billing_facts: Mutex<Vec<UsageCalculatedBillingFact>>,
diagnostics: Mutex<Vec<DiagnosticObservation>>,
runtime_slots: Mutex<Option<DashboardRuntimeSlots>>,
summary_observed_at: Mutex<Option<DateTime<Utc>>>,
slots_observed_at: Mutex<Option<DateTime<Utc>>>,
usage_records: Mutex<Vec<UsageRecord>>,
}
@@ -410,6 +432,8 @@ impl FixtureObservabilityStore {
calculated_billing_facts: Mutex::new(Vec::new()),
diagnostics: Mutex::new(Vec::new()),
runtime_slots: Mutex::new(None),
summary_observed_at: Mutex::new(None),
slots_observed_at: Mutex::new(None),
usage_records: Mutex::new(Vec::new()),
}
}
@@ -437,6 +461,16 @@ impl FixtureObservabilityStore {
*self.runtime_slots.lock().expect("runtime slots") = runtime_slots;
}
fn observed_times(&self) -> (Option<DateTime<Utc>>, Option<DateTime<Utc>>) {
(
*self
.summary_observed_at
.lock()
.expect("summary observed at"),
*self.slots_observed_at.lock().expect("slots observed at"),
)
}
fn replace_usage_records(&self, records: Vec<UsageRecord>) {
*self.usage_records.lock().expect("usage records") = records;
}
@@ -444,7 +478,15 @@ impl FixtureObservabilityStore {
#[async_trait]
impl ObservabilityStore for FixtureObservabilityStore {
async fn dashboard_summary(&self, range: TimeRange) -> AdminStoreResult<DashboardObservation> {
async fn dashboard_summary(
&self,
range: TimeRange,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
*self
.summary_observed_at
.lock()
.expect("summary observed at") = Some(observed_at);
Ok(DashboardObservation {
range,
requests: RequestMetrics::default(),
@@ -459,8 +501,9 @@ impl ObservabilityStore for FixtureObservabilityStore {
async fn dashboard_runtime_slots(
&self,
_: DateTime<Utc>,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<Option<DashboardRuntimeSlots>> {
*self.slots_observed_at.lock().expect("slots observed at") = Some(observed_at);
Ok(*self.runtime_slots.lock().expect("runtime slots"))
}
@@ -612,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,417 @@
//! 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,
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, 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 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>))
}
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 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")
}
+273 -99
View File
@@ -18,9 +18,10 @@ use futures::{Stream, StreamExt as _};
use gateway_admin::model::{
AdminError as AdminServiceError, PageSize,
accounts::{
AccountConnectionTestEvent as DomainConnectionTestEvent, AccountCost, AccountListQuery,
AccountModelUsage, AccountSort, AccountSortField, AccountStatus as DomainAccountStatus,
AccountUsage, SortDirection,
AccountConcurrencyLimit, AccountConnectionTestEvent as DomainConnectionTestEvent,
AccountCost, AccountGroupFilter, AccountListQuery, AccountModelUsage, AccountSort,
AccountSortField, AccountStatus as DomainAccountStatus, AccountUpdateResult, AccountUsage,
AccountWeight, AccountsUpdateResult, BatchUpdateAccounts, SortDirection, UpdateAccount,
},
provider_credentials::{
AccountDirectoryItem, AccountDirectoryPage, AccountExportBundle, AccountRefreshResult,
@@ -32,12 +33,12 @@ 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_number};
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, PageMeta,
WireValidationError,
@@ -51,8 +52,10 @@ const MAX_NAME_BYTES: usize = 512;
const MAX_IMPORT_DATA_BYTES: usize = 64 * 1024 * 1024;
const MAX_ACCESS_TOKEN_BYTES: usize = 16 * 1024;
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)]
@@ -61,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> {
@@ -79,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())
@@ -111,6 +184,7 @@ impl ListQuery {
)
.map_err(|_| WireValidationError::new("pageSize"))?,
provider_kind,
group_filter,
search,
status,
sort,
@@ -176,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>,
@@ -185,13 +260,13 @@ pub struct AccountView {
pub authentication_kind: String,
pub has_refresh_token: bool,
pub status: String,
pub display_status: String,
/// `status == "error"` 时的具体原因;其余状态为 `null`。
pub error_reason: Option<String>,
/// 最近一次失败的上游错误描述;仅错误状态存在。
pub error_message: Option<String>,
pub availability: 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>,
@@ -204,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")]
@@ -522,10 +606,9 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
DomainConnectionTestEvent::Content { text } => {
serde_json::json!({ "type": "content", "text": text })
}
DomainConnectionTestEvent::Completed { account_status } => serde_json::json!({
DomainConnectionTestEvent::Completed => serde_json::json!({
"type": "test_complete",
"success": true,
"accountStatus": account_status.as_str()
"success": true
}),
DomainConnectionTestEvent::Failed {
message,
@@ -534,7 +617,6 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
upstream_status,
upstream_content_type,
upstream_body,
account_status,
} => serde_json::json!({
"type": "error",
"error": message,
@@ -542,8 +624,7 @@ impl From<DomainConnectionTestEvent> for AccountConnectionTestEvent {
"providerErrorType": provider_error_type,
"upstreamStatus": upstream_status,
"upstreamContentType": upstream_content_type,
"upstreamBody": upstream_body,
"accountStatus": account_status.as_str()
"upstreamBody": upstream_body
}),
};
Self { data }
@@ -686,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(),
}
}
}
@@ -765,6 +864,7 @@ pub struct RotateAccountRequest {
pub account_id: String,
pub access_token: String,
pub refresh_token: Option<String>,
pub id_token: Option<String>,
}
impl RotateAccountRequest {
@@ -773,7 +873,11 @@ impl RotateAccountRequest {
return Err(WireValidationError::new("provider"));
}
require_account_id(&self.account_id, "accountId")?;
validate_oauth_material(&self.access_token, self.refresh_token.as_deref())
validate_oauth_material(
&self.access_token,
self.refresh_token.as_deref(),
self.id_token.as_deref(),
)
}
fn into_command(
@@ -787,6 +891,10 @@ impl RotateAccountRequest {
"refresh_token".to_owned(),
self.refresh_token.map_or(Value::Null, Value::String),
);
material.insert(
"id_token".to_owned(),
self.id_token.map_or(Value::Null, Value::String),
);
Ok(RotateCredential {
mutation: CredentialMutation {
context,
@@ -805,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()
@@ -883,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
@@ -894,10 +1042,14 @@ where
.route("/api/admin/accounts/export", get(export_accounts::<S>))
.route("/api/admin/accounts/import", post(import_accounts::<S>))
.route("/api/admin/accounts/refresh", post(refresh_account::<S>))
.route("/api/admin/accounts/recover", post(recover_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",
@@ -922,6 +1074,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>,
@@ -1011,7 +1184,7 @@ where
.map_err(map_service_error)?;
Ok(AdminResponse::new(
StatusCode::CREATED,
AdminEnvelope::ok(AccountImportData::from(result)),
AdminEnvelope::ok(AccountImportData::from_result(result)),
))
}
@@ -1106,64 +1279,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)),
))
}
@@ -1208,6 +1341,25 @@ where
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
}
async fn recover_account<S>(
auth: AdminAuth,
State(state): State<S>,
Json(request): Json<AccountActionRequest>,
) -> Result<impl IntoResponse, AdminError>
where
S: AdminSessionState + Send + Sync,
{
let account_id = request.into_id().map_err(map_wire_error)?;
let result = state
.admin_services()
.accounts()
.recover(&auth.context().mutation_context(), account_id)
.await
.map_err(map_service_error)?;
let data = account_refresh_data(result, Utc::now());
Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data)))
}
async fn account_quota<S>(
_auth: AdminAuth,
State(state): State<S>,
@@ -1365,20 +1517,30 @@ fn account_models_data(result: ProviderModels) -> AccountModelsData {
fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
let AccountDirectoryItem {
account,
status,
error_reason,
projection,
usage,
quota,
} = item;
let status = status.as_str().to_owned();
let status = projection.status.as_str().to_owned();
let rate_limited_until = projection.rate_limited_until.map(DateTime::<Utc>::from);
let expires_at = account.access_token_expires_at.as_ref().map(china_rfc3339);
let added_at = china_rfc3339(&account.created_at);
let updated_at = china_rfc3339(&account.updated_at);
let (quota, refresh_token_expires_at) = account_quota_view(quota, now);
let (quota, refresh_token_expires_at) = account_quota_view(quota, rate_limited_until, now);
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,
@@ -1387,12 +1549,14 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
plan_type: account.plan_type,
authentication_kind: account.authentication_kind,
has_refresh_token: account.has_refresh_token,
status: status.clone(),
display_status: status,
error_reason: error_reason.map(|reason| reason.as_str().to_owned()),
error_message: account.last_error_message,
availability: account.availability.as_str().to_owned(),
status,
error_reason: projection
.error_reason
.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
@@ -1411,6 +1575,7 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
fn account_quota_view(
quota: ProviderQuota,
rate_limited_until: Option<DateTime<Utc>>,
now: DateTime<Utc>,
) -> (AccountQuotaView, Option<String>) {
let refresh_token_expires_at = quota
@@ -1420,7 +1585,7 @@ fn account_quota_view(
.observed_at
.map_or_else(|| "—".to_owned(), |value| relative_time(value, now));
let windows = quota.windows.into_iter().map(quota_window_view).collect();
let rate_limited_until = quota.rate_limited_until.map(|until| china_datetime(&until));
let rate_limited_until = rate_limited_until.map(|until| china_datetime(&until));
(
AccountQuotaView {
refreshed_at_display,
@@ -1438,6 +1603,7 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
group,
label,
source: _,
local_usage_attribution: _,
window_seconds,
used_percent,
reset_at,
@@ -1589,7 +1755,7 @@ fn account_model_usage_view(usage: AccountModelUsage, now: DateTime<Utc>) -> Mod
billing_amount_usd: usd.map(|cost| cost.amount.as_str().to_owned()),
billing_amount_usd_display: usd.map_or_else(
|| "—".to_owned(),
|cost| format!("${}", cost.amount.as_str()),
|cost| format_decimal_currency(cost.amount.as_str(), "USD"),
),
cost_estimate_status: cost_estimate_status.to_owned(),
known_cost_count: known_count,
@@ -1684,6 +1850,7 @@ fn require_text(
fn validate_oauth_material(
access_token: &str,
refresh_token: Option<&str>,
id_token: Option<&str>,
) -> Result<(), WireValidationError> {
if access_token.len() > MAX_ACCESS_TOKEN_BYTES
|| !valid_visible_ascii(access_token)
@@ -1698,6 +1865,13 @@ fn validate_oauth_material(
}) {
return Err(WireValidationError::new("refreshToken"));
}
if id_token.is_some_and(|token| {
token.len() > MAX_ID_TOKEN_BYTES
|| !valid_visible_ascii(token)
|| !valid_compact_jwt_shape(token)
}) {
return Err(WireValidationError::new("idToken"));
}
Ok(())
}
@@ -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>())
@@ -15,7 +15,7 @@ use gateway_admin::model::{PageSize as DomainPageSize, observability as domain};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::presenter::{format_compact_number, format_number};
use super::presenter::{format_compact_number, format_decimal_currency, format_number};
use super::{
AdminAuth, AdminEnvelope, AdminError, AdminResponse, AdminSessionState, PageMeta,
WireValidationError, wire::map_admin_service_error,
@@ -127,20 +127,6 @@ pub(crate) fn token_details(record: &domain::UsageRecord) -> TokenDetailsView {
}
}
pub(crate) fn format_decimal_currency(amount: &str, currency: &str) -> String {
if currency == "USD" {
let value = amount.parse::<f64>().map_or(0.0, |value| value);
let precision = if value != 0.0 && value.abs() < 0.01 {
4
} else {
2
};
format!("${value:.precision$}")
} else {
format!("{currency} {amount}")
}
}
pub(crate) fn format_money(cost: &domain::CurrencyCost) -> String {
format_decimal_currency(cost.amount.as_str(), &cost.currency)
}
@@ -246,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,
@@ -747,10 +736,16 @@ pub(crate) fn dashboard_view(
credentials: DashboardCredentialsCardView {
total: format_compact_number(provider_accounts.total),
total_value: provider_accounts.total,
available: format_compact_number(provider_accounts.active),
available_value: provider_accounts.active,
unavailable: format_compact_number(provider_accounts.unavailable),
unavailable_value: provider_accounts.unavailable,
available: format_compact_number(provider_accounts.normal),
available_value: provider_accounts.normal,
unavailable: format_compact_number(
provider_accounts
.total
.saturating_sub(provider_accounts.normal),
),
unavailable_value: provider_accounts
.total
.saturating_sub(provider_accounts.normal),
},
traffic: DashboardTrafficCardView {
today_requests: format_compact_number(today.request_count),
@@ -783,14 +778,11 @@ pub(crate) fn dashboard_view(
usage_records: recent_requests.into_iter().map(usage_record_view).collect(),
pool_summary: DashboardPoolSummaryView {
total: provider_accounts.total,
active: provider_accounts.active,
rate_limited: provider_accounts.rate_limited,
expired: provider_accounts.expired,
invalid: provider_accounts.invalid,
normal: provider_accounts.normal,
quota_exhausted: provider_accounts.quota_exhausted,
refreshing: provider_accounts.refreshing,
rate_limited: provider_accounts.rate_limited,
disabled: provider_accounts.disabled,
banned: provider_accounts.banned,
error: provider_accounts.error,
},
capacity_info: DashboardCapacityInfoView {
max_concurrent_per_account: capacity.max_concurrent_per_account,
@@ -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>,
@@ -347,14 +350,11 @@ pub struct DashboardAccountRequestBucketView {
#[serde(rename_all = "camelCase")]
pub struct DashboardPoolSummaryView {
pub total: u64,
pub active: u64,
pub rate_limited: u64,
pub expired: u64,
pub invalid: u64,
pub normal: u64,
pub quota_exhausted: u64,
pub refreshing: Option<u64>,
pub rate_limited: u64,
pub disabled: u64,
pub banned: u64,
pub error: u64,
}
/// 同 target 账号调度容量;Redis 未提供聚合事实时运行中槽位保持空值。
@@ -3,6 +3,8 @@
//! 按 docs/project-redundancy-boundary-audit.md BE-02 的目标,display 字段最终
//! 迁移到前端;过渡期统一由本模块提供,避免 accounts/observability 各自维护。
use gateway_core::accounting::Decimal;
/// 千分位格式化。
#[must_use]
pub fn format_number(value: u64) -> String {
@@ -37,3 +39,33 @@ pub fn format_compact_number(value: u64) -> String {
}
format_number(value)
}
/// 货币展示格式化。USD 常规金额保留两位,小于 1 美元时最多保留四位。
#[must_use]
pub fn format_decimal_currency(amount: &str, currency: &str) -> String {
if currency != "USD" {
return format!("{currency} {amount}");
}
const DECIMAL_SCALE: u128 = 10_000_000_000;
let scaled = amount.parse::<Decimal>().unwrap_or_default().scaled();
let precision = if scaled != 0 && scaled < DECIMAL_SCALE {
4_u32
} else {
2_u32
};
let rounding_unit = 10_u128.pow(10 - precision);
let rounded = (scaled + rounding_unit / 2) / rounding_unit;
let display_scale = 10_u128.pow(precision);
let whole = rounded / display_scale;
let fraction = rounded % display_scale;
let mut display = format!("{whole}.{fraction:0width$}", width = precision as usize);
while display.ends_with('0')
&& display
.split_once('.')
.is_some_and(|(_, fraction)| fraction.len() > 2)
{
display.pop();
}
format!("${display}")
}
@@ -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,271 @@
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 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,8 +70,138 @@ 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;
#[test]
fn account_usage_view_should_keep_unobserved_numbers_null() {
@@ -112,12 +242,22 @@ mod response {
assert!(value["totalTokens"].is_null());
assert_eq!(value["createdTokensDisplay"], "-");
}
#[test]
fn account_usage_currency_should_limit_usd_to_four_fraction_digits() {
assert_eq!(format_decimal_currency("0.1204956", "USD"), "$0.1205");
assert_eq!(format_decimal_currency("0.99996", "USD"), "$1.00");
assert_eq!(format_decimal_currency("0.0300", "USD"), "$0.03");
assert_eq!(format_decimal_currency("12.3456", "USD"), "$12.35");
assert_eq!(format_decimal_currency("0.1204956", "CNY"), "CNY 0.1204956");
}
}
mod actions {
use base64::Engine as _;
use gateway_admin::model::{
Revision,
accounts::{AccountConnectionTestEvent as DomainConnectionTestEvent, AccountStatus},
accounts::AccountConnectionTestEvent as DomainConnectionTestEvent,
provider_credentials::{
CredentialDeletionResult, CredentialImportResult, CredentialMutationResult,
},
@@ -126,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;
@@ -207,7 +348,8 @@ mod actions {
"provider": "openai",
"accountId": "acct_1",
"accessToken": "header.payload.signature",
"refreshToken": "refresh-token"
"refreshToken": "refresh-token",
"idToken": "id-header.id-payload.id-signature"
}))
.expect("decode rotation");
assert!(rotation.validate().is_ok());
@@ -269,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]
@@ -305,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")],
});
@@ -343,9 +514,7 @@ mod actions {
DomainConnectionTestEvent::Content {
text: "OK".to_owned(),
},
DomainConnectionTestEvent::Completed {
account_status: AccountStatus::Normal,
},
DomainConnectionTestEvent::Completed {},
DomainConnectionTestEvent::Failed {
message: "upstream unavailable".to_owned(),
provider_error_code: Some("usage_exhausted".to_owned()),
@@ -353,7 +522,6 @@ mod actions {
upstream_status: Some(429),
upstream_content_type: Some("application/json".to_owned()),
upstream_body: Some(r#"{"error":{"type":"usage_limit_reached"}}"#.to_owned()),
account_status: AccountStatus::QuotaExhausted,
},
]
.map(|event| AccountConnectionTestEvent::from(event).data);
@@ -378,7 +546,7 @@ mod actions {
}
}),
json!({ "type": "content", "text": "OK" }),
json!({ "type": "test_complete", "success": true, "accountStatus": "normal" }),
json!({ "type": "test_complete", "success": true }),
json!({
"type": "error",
"error": "upstream unavailable",
@@ -386,8 +554,7 @@ mod actions {
"providerErrorType": "invalid_request_error",
"upstreamStatus": 429,
"upstreamContentType": "application/json",
"upstreamBody": r#"{"error":{"type":"usage_limit_reached"}}"#,
"accountStatus": "quota_exhausted"
"upstreamBody": r#"{"error":{"type":"usage_limit_reached"}}"#
}),
]
);
@@ -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());
+273 -14
View File
@@ -7,15 +7,22 @@ use std::{
};
use async_trait::async_trait;
use chrono::{Duration, Utc};
use chrono::{DateTime, Duration, Utc};
use futures::future::BoxFuture;
use gateway_admin::{
AdminConfig, AdminServices, InitialAdminPassword,
model::{
MutationContext, Revision,
account_groups::{
AccountGroupAccountSummary, AccountGroupCapacity, AccountGroupColor,
AccountGroupListQuery, AccountGroupMutation, AccountGroupPage, AccountGroupRecord,
AccountGroupUsage, DeleteAccountGroup, NewAccountGroup, SetAccountGroupEnabled,
UpdateAccountGroup,
},
accounts::{
AccountListQuery, AccountPage, AccountRecord, 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,204 @@ 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>,
}
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 groups = BTreeMap::from([
(
primary_id.clone(),
AccountGroupRecord {
id: primary_id,
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,
}),
}
}
}
#[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 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);
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"));
}
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 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 +574,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(),
@@ -426,7 +635,7 @@ impl AccountStore for UnusedStore {
Err(unavailable("account list"))
}
async fn load_account(&self, _: &str) -> AdminStoreResult<Option<AccountRecord>> {
async fn load_account(&self, _: &str) -> AdminStoreResult<Option<AccountPageItem>> {
Err(unavailable("account"))
}
@@ -501,14 +710,30 @@ 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 recover_account(
&self,
_: &ProviderAccountId,
_: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult> {
Err(unavailable("account recovery"))
}
async fn batch_update_accounts(
&self,
_: BatchUpdateAccounts,
_: &MutationContext,
) -> AdminStoreResult<AccountsUpdateResult> {
Err(unavailable("account batch update"))
}
async fn delete_accounts(
&self,
_: DeleteAccounts,
@@ -528,7 +753,11 @@ impl AccountStore for UnusedStore {
#[async_trait]
impl ObservabilityStore for UnusedStore {
async fn dashboard_summary(&self, range: TimeRange) -> AdminStoreResult<DashboardObservation> {
async fn dashboard_summary(
&self,
range: TimeRange,
_: DateTime<Utc>,
) -> AdminStoreResult<DashboardObservation> {
*self
.dashboard_summary_range
.lock()
@@ -788,6 +1017,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,
@@ -796,6 +1051,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)
}
@@ -222,16 +222,11 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
totals: Default::default(),
provider_accounts: AccountPoolMetrics {
total: 807,
enabled: 804,
unavailable: 792,
active: 15,
rate_limited: 0,
expired: 2,
invalid: 0,
normal: 15,
quota_exhausted: 785,
refreshing: None,
rate_limited: 0,
disabled: 3,
banned: 2,
error: 4,
},
trend: Vec::new(),
account_usage: Vec::new(),
@@ -679,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(),
@@ -805,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(),
@@ -924,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"],
@@ -954,6 +958,9 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
}),
serde_json::json!({
"route": "/v1/responses",
"routingScope": "groups",
"routingGroupRefs": ["grp_usage"],
"routingGroupNamesSnapshot": ["Usage group"],
"serviceTier": "priority",
"accountId": "acct_snapshot",
"accountName": "Snapshot Alpha",
@@ -33,6 +33,7 @@ fn source_tree_should_match_frozen_machine_manifest() {
let mut actual = rust_files(&root.join("src"));
actual.sort();
let mut expected = vec![
"src/admin/account_groups.rs",
"src/admin/accounts.rs",
"src/admin/auth.rs",
"src/admin/backups.rs",
@@ -73,6 +74,7 @@ fn test_tree_should_match_frozen_rust_mirror() {
let mut actual = rust_files(&root.join("tests"));
actual.sort();
let mut expected = vec![
"tests/admin/account_groups.rs",
"tests/admin/accounts.rs",
"tests/admin/auth.rs",
"tests/admin/client_keys.rs",
+16 -5
View File
@@ -4,7 +4,7 @@ mod models;
mod responses;
mod router;
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
@@ -18,7 +18,9 @@ use gateway_core::engine::admission::{
use gateway_core::engine::continuation::{
NativeContinuationPin, NativeContinuationPort, NativeContinuationStoreError, PreviousResponseId,
};
use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy};
use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
};
use gateway_core::engine::execution::{
AuthenticatedClient, ClientApiKeyUsageSink, DefaultExecutionService, ExecutionService,
ProviderCircuitDecision, ProviderCircuitError, ProviderCircuitPort,
@@ -34,8 +36,8 @@ use gateway_core::lifecycle::{ConnectionDraining, ConnectionGuard, ConnectionLif
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::RuntimeSnapshotHandle;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, RuntimeSnapshot,
UpstreamModelId,
ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities, ProviderKind,
ProviderModel, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
};
pub(super) async fn api_router(execution: Arc<dyn ExecutionService>) -> axum::Router {
@@ -94,6 +96,10 @@ impl ClientApiKeyUsageSink for IgnoredClientApiKeyUsage {
fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
let provider = ProviderKind::new(provider_name).expect("provider");
let account_directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
ProviderAccountId::new("acct_api_test").expect("account ID"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)])));
let capabilities = ModelCapabilities::new(
BTreeSet::from([gateway_core::operation::OperationKind::Generate]),
Some(16_000),
@@ -119,12 +125,16 @@ fn snapshot(plaintext: &str, provider_name: &str) -> RuntimeSnapshot {
vec![ClientPolicy::new(
ClientApiKeyId::new("key_api_test").expect("key ID"),
PlaintextClientApiKey::new(plaintext).expect("plaintext key"),
provider,
Arc::new(FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::all_accounts(),
)),
true,
RateLimits::unlimited(),
)],
)
.expect("runtime snapshot")
.with_account_directory(account_directory)
}
#[derive(Default)]
@@ -260,6 +270,7 @@ struct UnusedContinuation;
impl NativeContinuationPort for UnusedContinuation {
fn resolve<'a>(
&'a self,
_: &'a ClientApiKeyId,
_: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>> {
Box::pin(async { unreachable!("authentication fixture does not execute") })
@@ -451,10 +451,12 @@ fn decoder_should_not_reject_large_bodies_using_catalog_context_limits() {
.to_string();
let decoded = decode_request(body.as_bytes()).expect("large request should decode");
let plan = super::snapshot("sk_context_test", "openai")
let snapshot = super::snapshot("sk_context_test", "openai");
let plan = snapshot
.plan(
&PublicModelId::new(decoded.metadata().public_model()).expect("public model"),
decoded.operation(),
snapshot.all_account_scope(),
&RoutingContext::default(),
)
.expect("large body must not be locally context-gated");
@@ -9,6 +9,7 @@ use futures::future::BoxFuture;
use crate::engine::credential::ProviderAccountId;
use crate::operation::ProviderSessionState;
use crate::policy::ClientApiKeyId;
use crate::routing::ProviderKind;
/// 客户端或 Provider 传递的 opaque response handle。
@@ -48,14 +49,15 @@ pub enum NativeContinuationScope {
/// 从可丢失会话亲和存储恢复的 native previous-response pin。
///
/// 该值不依赖调用方 API Key;它只阻止 native handle 与不透明 Provider state
/// 被发送到错误的 Provider 或账号。
/// 该值同时绑定调用方 API Key、Provider 与账号,阻止 native handle 和不透明
/// Provider state 被跨租户或跨上游目标复用。
#[derive(Clone, PartialEq)]
pub struct NativeContinuationPin {
/// 客户端提交、用于查找可丢失会话亲和记录的 response ID。
previous_response_id: PreviousResponseId,
/// Provider 原生 response handle。
upstream_response_id: PreviousResponseId,
client_api_key_id: ClientApiKeyId,
provider: ProviderKind,
account: ProviderAccountId,
scope: NativeContinuationScope,
@@ -67,12 +69,14 @@ impl NativeContinuationPin {
pub const fn new(
previous_response_id: PreviousResponseId,
upstream_response_id: PreviousResponseId,
client_api_key_id: ClientApiKeyId,
provider: ProviderKind,
account: ProviderAccountId,
) -> Self {
Self {
previous_response_id,
upstream_response_id,
client_api_key_id,
provider,
account,
scope: NativeContinuationScope::ConnectionLocal,
@@ -105,6 +109,16 @@ impl NativeContinuationPin {
&self.upstream_response_id
}
#[must_use]
pub const fn client_api_key_id(&self) -> &ClientApiKeyId {
&self.client_api_key_id
}
#[must_use]
pub fn matches_client(&self, client_api_key_id: &ClientApiKeyId) -> bool {
self.client_api_key_id == *client_api_key_id
}
#[must_use]
pub const fn provider(&self) -> &ProviderKind {
&self.provider
@@ -139,6 +153,7 @@ impl fmt::Debug for NativeContinuationPin {
.debug_struct("NativeContinuationPin")
.field("previous_response_id", &"<redacted>")
.field("upstream_response_id", &"<redacted>")
.field("client_api_key_id", &self.client_api_key_id)
.field("provider", &self.provider)
.field("account", &self.account)
.field("scope", &self.scope)
@@ -190,21 +205,57 @@ impl fmt::Debug for ContinuationBinding {
}
}
/// Native continuation lookup 的稳定失败分类。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum NativeContinuationStoreErrorKind {
/// Redis 等可丢失协调存储暂不可用;调用方可以按外部 continuation fail-open。
Unavailable,
/// 存储中的记录无法还原为可信 pin;不得将原始 handle 继续透传上游。
InvalidData,
/// 已存在的记录属于另一个客户端 API Key;必须 fail closed。
OwnershipMismatch,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("native continuation store is unavailable: {detail}")]
#[error("native continuation store failed: {kind:?}: {detail}")]
pub struct NativeContinuationStoreError {
kind: NativeContinuationStoreErrorKind,
detail: String,
}
impl NativeContinuationStoreError {
/// 构造带可排查细节的亲和存储错误;`detail` 不应包含 response handle。
/// 构造可 fail-open 的协调存储不可用错误;`detail` 不应包含 response handle。
#[must_use]
pub fn new(detail: impl Into<String>) -> Self {
pub fn unavailable(detail: impl Into<String>) -> Self {
Self {
kind: NativeContinuationStoreErrorKind::Unavailable,
detail: detail.into(),
}
}
/// 构造必须 fail closed 的无效记录错误;`detail` 不应包含 response handle。
#[must_use]
pub fn invalid_data(detail: impl Into<String>) -> Self {
Self {
kind: NativeContinuationStoreErrorKind::InvalidData,
detail: detail.into(),
}
}
/// 构造不泄露记录归属细节的跨调用方错误。
#[must_use]
pub fn ownership_mismatch() -> Self {
Self {
kind: NativeContinuationStoreErrorKind::OwnershipMismatch,
detail: "continuation belongs to a different client API key".to_owned(),
}
}
#[must_use]
pub const fn kind(&self) -> NativeContinuationStoreErrorKind {
self.kind
}
#[must_use]
pub fn detail(&self) -> &str {
&self.detail
@@ -213,11 +264,12 @@ impl NativeContinuationStoreError {
/// 可丢失的 previous-response 亲和存储端口。
///
/// 亲和记录不按调用方 API Key 隔离:与旧版一样,response ID 是会话级 opaque
/// handle;Redis 不可用或超时必须由调用方 fail-open。
/// 亲和记录按调用方 API Key 隔离;Redis 不可用或超时必须由调用方 fail-open,
/// 但命中其他 Key 的记录必须 fail closed。
pub trait NativeContinuationPort: Send + Sync {
fn resolve<'a>(
&'a self,
client_api_key_id: &'a ClientApiKeyId,
previous_response_id: &'a PreviousResponseId,
) -> BoxFuture<'a, Result<Option<NativeContinuationPin>, NativeContinuationStoreError>>;
@@ -46,6 +46,10 @@ impl AccountSelection {
Self::Diagnostic(account) => Some(account),
}
}
const fn is_diagnostic(&self) -> bool {
matches!(self, Self::Diagnostic(_))
}
}
impl<S: ?Sized> AttemptCoordinator<S>
@@ -126,6 +130,15 @@ where
.map(ProviderAccountStateOwner::from_continuation);
let continuation_attempt =
initial_continuation_attempt(&operation, &plan, continuation.as_ref());
let candidate_index = continuation
.as_ref()
.and_then(ContinuationBinding::pinned)
.map_or(Ok(0), |pin| {
plan.candidates()
.iter()
.position(|candidate| candidate.provider() == pin.provider())
.ok_or(EngineError::ContinuationPinMismatch)
})?;
let image_generation_requested = operation.image_generation_requested();
let mut session = ResponseExecutionSession {
engine: Arc::clone(&self.engine),
@@ -149,6 +162,7 @@ where
account_state_owner,
cancellation,
attempts: 0,
candidate_index,
excluded_accounts: BTreeSet::new(),
credential_recovery_attempted_accounts: BTreeSet::new(),
recovery_account: None,
@@ -225,6 +239,7 @@ pub struct ResponseExecutionSession<S: ?Sized> {
account_state_owner: Option<ProviderAccountStateOwner>,
cancellation: CancellationToken,
attempts: u32,
candidate_index: usize,
excluded_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
credential_recovery_attempted_accounts: BTreeSet<crate::engine::credential::ProviderAccountId>,
/// 凭据恢复后的一次性同账号钉选;只约束紧随其后的 replay attempt,
@@ -461,6 +476,7 @@ where
NativeContinuationPin::new(
previous_response_id,
upstream_response_id,
self.client_api_key_ref.clone(),
provider,
account,
)
@@ -565,7 +581,7 @@ where
if self.attempts >= self.plan.max_attempts().get() {
return Err(EngineError::EmptyRoutingPlan);
}
let Some(candidate) = self.plan.candidates().first().cloned() else {
let Some(candidate) = self.plan.candidates().get(self.candidate_index).cloned() else {
let error = GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream Provider is available",
@@ -606,7 +622,8 @@ where
self.excluded_accounts.clone(),
pinned_account.clone(),
self.account_state_owner.clone(),
),
)
.with_account_scope(Arc::clone(self.plan.account_scope())),
}
.with_credential_recovery_attempted(pinned_account.as_ref().is_some_and(|account| {
self.credential_recovery_attempted_accounts
@@ -663,6 +680,20 @@ where
if self.prepare_unavailable_native_continuation_replay(&error) {
return Ok(Some(PullOutcome::AttemptDiscarded));
}
if matches!(
error.kind(),
ProviderErrorKind::AccountCapacityUnavailable
| ProviderErrorKind::NoEligibleAccount
| ProviderErrorKind::ProviderInfrastructureUnavailable
) && error.send_state() == UpstreamSendState::NotSent
&& matches!(
self.continuation_attempt,
ContinuationAttempt::None | ContinuationAttempt::ReplayAny
)
&& self.advance_provider_candidate()
{
return Ok(Some(PullOutcome::AttemptDiscarded));
}
if matches!(
error.kind(),
ProviderErrorKind::AccountCapacityUnavailable
@@ -716,6 +747,27 @@ where
}
let metadata = stream.metadata().clone();
if !self.account_selection.is_diagnostic()
&& metadata
.provider_account_id()
.is_none_or(|account| !candidate.account_scope().allows(account))
{
let error = GatewayError::new(
GatewayErrorKind::Internal,
"provider selected an account outside the frozen client scope",
);
self.finish_failure(FailureFinalization {
outcome: ExecutionOutcome::Failed,
send_state: self.current_send_state(),
error,
upstream_status_code: None,
provider_error_code: None,
retry_after_ms: None,
provider_response_id: None,
})
.await?;
return Err(EngineError::AccountOutsideClientScope);
}
if pinned_account
.as_ref()
.is_some_and(|required| metadata.provider_account_id() != Some(required))
@@ -819,6 +871,17 @@ where
Ok(None)
}
fn advance_provider_candidate(&mut self) -> bool {
let Some(next) = self.candidate_index.checked_add(1) else {
return false;
};
if next >= self.plan.candidates().len() {
return false;
}
self.candidate_index = next;
true
}
async fn observe_event(&mut self, event: &GatewayEvent) {
self.observe_event_facts(event);
self.mark_send_observed().await;
@@ -3,7 +3,7 @@
use std::cmp::Reverse;
use std::collections::{BTreeSet, HashMap};
use std::fmt;
use std::num::{NonZeroU32, NonZeroU64};
use std::num::{NonZeroU16, NonZeroU32, NonZeroU64};
use std::sync::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, SystemTime};
@@ -38,6 +38,58 @@ impl ProviderAccountId {
}
}
/// Per-account scheduling concurrency override.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountConcurrencyLimit(NonZeroU32);
impl AccountConcurrencyLimit {
#[must_use]
pub const fn new(value: u32) -> Option<Self> {
match NonZeroU32::new(value) {
Some(value) => Some(Self(value)),
None => None,
}
}
#[must_use]
pub const fn get(self) -> u32 {
self.0.get()
}
#[must_use]
pub const fn into_non_zero(self) -> NonZeroU32 {
self.0
}
}
/// Relative scheduling priority for an account.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountWeight(NonZeroU16);
impl AccountWeight {
pub const DEFAULT: Self = Self(NonZeroU16::MIN);
pub const MAX: u16 = 100;
#[must_use]
pub const fn new(value: u16) -> Option<Self> {
match NonZeroU16::new(value) {
Some(value) if value.get() <= Self::MAX => Some(Self(value)),
_ => None,
}
}
#[must_use]
pub const fn get(self) -> u16 {
self.0.get()
}
}
impl Default for AccountWeight {
fn default() -> Self {
Self::DEFAULT
}
}
impl fmt::Display for ProviderAccountId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
@@ -187,45 +239,41 @@ impl fmt::Debug for OpaqueProviderData {
}
}
/// 数据库中固定的账号运行状态。
/// 已持久化的凭据/账号身份状态。
///
/// 只承载终态或需要外部干预的状态;临时限流(429)不进入此状态,由
/// Provider quota 数据驱动(`AccountQuotaSignals::limit_reached`)。
/// 这里只保存凭据与账号身份事实。额度耗尽属于 [`QuotaState`],临时 429 属于
/// 运行时冷却;二者都不得写入此枚举。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AccountAvailability {
pub enum CredentialState {
Unknown,
Ready,
QuotaExhausted,
Expired,
Banned,
Invalid,
}
impl AccountAvailability {
impl CredentialState {
/// 返回数据库稳定值。
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Unknown => "unknown",
Self::Ready => "ready",
Self::QuotaExhausted => "quota_exhausted",
Self::Expired => "expired",
Self::Banned => "banned",
Self::Invalid => "invalid",
}
}
/// 返回账号进入不可恢复错误态、但 Provider 未给出更具体上游原因时的稳定兜底。
///
/// Provider 应优先持久化已脱敏的上游 message 或原因码;该值只避免错误态沿用
/// 上一次无关的错误信息。
/// 返回该凭据状态对应的稳定错误原因。
#[must_use]
pub const fn fallback_error_reason(self) -> Option<&'static str> {
pub const fn error_reason(self) -> Option<AccountErrorReason> {
match self {
Self::Expired => Some("credential_expired"),
Self::Banned => Some("account_banned"),
Self::Invalid => Some("credential_invalid"),
Self::Unknown | Self::Ready | Self::QuotaExhausted => None,
Self::Unknown => Some(AccountErrorReason::AccountUnverified),
Self::Expired => Some(AccountErrorReason::CredentialExpired),
Self::Banned => Some(AccountErrorReason::AccountBanned),
Self::Invalid => Some(AccountErrorReason::CredentialInvalid),
Self::Ready => None,
}
}
@@ -235,7 +283,6 @@ impl AccountAvailability {
match value {
"unknown" => Some(Self::Unknown),
"ready" => Some(Self::Ready),
"quota_exhausted" => Some(Self::QuotaExhausted),
"expired" => Some(Self::Expired),
"banned" => Some(Self::Banned),
"invalid" => Some(Self::Invalid),
@@ -244,6 +291,373 @@ impl AccountAvailability {
}
}
/// 额度访问结论;百分比和余额等展示值不进入此枚举。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum QuotaAccessState {
Unknown,
Allowed,
Exhausted,
}
impl QuotaAccessState {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Unknown => "unknown",
Self::Allowed => "allowed",
Self::Exhausted => "exhausted",
}
}
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
match value {
"unknown" => Some(Self::Unknown),
"allowed" => Some(Self::Allowed),
"exhausted" => Some(Self::Exhausted),
_ => None,
}
}
}
/// 权威额度耗尽证据。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum QuotaEvidence {
ProviderDenied,
AccountLimitReached,
UsageLimitReached,
PaymentRequired,
}
impl QuotaEvidence {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::ProviderDenied => "provider_denied",
Self::AccountLimitReached => "account_limit_reached",
Self::UsageLimitReached => "usage_limit_reached",
Self::PaymentRequired => "payment_required",
}
}
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
match value {
"provider_denied" => Some(Self::ProviderDenied),
"account_limit_reached" => Some(Self::AccountLimitReached),
"usage_limit_reached" => Some(Self::UsageLimitReached),
"payment_required" => Some(Self::PaymentRequired),
_ => None,
}
}
}
/// 可审计的额度访问事实。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct QuotaState {
access: QuotaAccessState,
evidence: Option<QuotaEvidence>,
observed_at: Option<SystemTime>,
reset_at: Option<SystemTime>,
}
impl Default for QuotaState {
fn default() -> Self {
Self::unknown()
}
}
impl QuotaState {
#[must_use]
pub const fn unknown() -> Self {
Self {
access: QuotaAccessState::Unknown,
evidence: None,
observed_at: None,
reset_at: None,
}
}
#[must_use]
pub const fn observed_unknown(observed_at: SystemTime) -> Self {
Self {
access: QuotaAccessState::Unknown,
evidence: None,
observed_at: Some(observed_at),
reset_at: None,
}
}
#[must_use]
pub const fn allowed(observed_at: SystemTime) -> Self {
Self {
access: QuotaAccessState::Allowed,
evidence: None,
observed_at: Some(observed_at),
reset_at: None,
}
}
#[must_use]
pub const fn exhausted(
evidence: QuotaEvidence,
observed_at: SystemTime,
reset_at: Option<SystemTime>,
) -> Self {
Self {
access: QuotaAccessState::Exhausted,
evidence: Some(evidence),
observed_at: Some(observed_at),
reset_at,
}
}
/// 从持久化列恢复额度状态;非法组合返回 `None`。
#[must_use]
pub const fn from_persisted(
access: QuotaAccessState,
evidence: Option<QuotaEvidence>,
observed_at: Option<SystemTime>,
reset_at: Option<SystemTime>,
) -> Option<Self> {
match access {
QuotaAccessState::Unknown if evidence.is_none() && reset_at.is_none() => Some(Self {
access,
evidence,
observed_at,
reset_at,
}),
QuotaAccessState::Allowed
if evidence.is_none() && observed_at.is_some() && reset_at.is_none() =>
{
Some(Self {
access,
evidence,
observed_at,
reset_at,
})
}
QuotaAccessState::Exhausted if evidence.is_some() && observed_at.is_some() => {
Some(Self {
access,
evidence,
observed_at,
reset_at,
})
}
_ => None,
}
}
#[must_use]
pub const fn access(self) -> QuotaAccessState {
self.access
}
#[must_use]
pub const fn evidence(self) -> Option<QuotaEvidence> {
self.evidence
}
#[must_use]
pub const fn observed_at(self) -> Option<SystemTime> {
self.observed_at
}
#[must_use]
pub const fn reset_at(self) -> Option<SystemTime> {
self.reset_at
}
/// 返回当前已确认的额度耗尽结论。
///
/// `reset_at` 只表示何时应重新向 Provider 求证;时间到期本身不是额度恢复证据。
#[must_use]
pub fn is_exhausted(self) -> bool {
self.access == QuotaAccessState::Exhausted
}
/// 判断已耗尽额度是否到达 Provider 的下一次复核时间。
///
/// 有明确 `reset_at` 时严格等待该时刻;没有时由 Provider 传入自己的保守复核周期。
#[must_use]
pub fn exhaustion_refresh_due(self, now: SystemTime, fallback_interval: Duration) -> bool {
if !self.is_exhausted() {
return false;
}
let due_at = self.reset_at.or_else(|| {
self.observed_at
.and_then(|observed_at| observed_at.checked_add(fallback_interval))
});
due_at.is_none_or(|due_at| due_at <= now)
}
/// 合并一次额度访问观察;`Unknown` 不能擦除已经确认的访问结论。
#[must_use]
pub fn merge_observation(self, observation: Self) -> Self {
if observation.access == QuotaAccessState::Unknown
&& self.access != QuotaAccessState::Unknown
{
self
} else {
observation
}
}
}
/// 对外唯一的五态账号状态。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AccountStatus {
Normal,
QuotaExhausted,
RateLimited,
Disabled,
Error,
}
impl AccountStatus {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Normal => "normal",
Self::QuotaExhausted => "quota_exhausted",
Self::RateLimited => "rate_limited",
Self::Disabled => "disabled",
Self::Error => "error",
}
}
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
match value {
"normal" => Some(Self::Normal),
"quota_exhausted" => Some(Self::QuotaExhausted),
"rate_limited" => Some(Self::RateLimited),
"disabled" => Some(Self::Disabled),
"error" => Some(Self::Error),
_ => None,
}
}
/// 管理页稳定业务排序:正常、限流、耗尽、错误、停用。
#[must_use]
pub const fn sort_rank(self) -> u8 {
match self {
Self::Normal => 0,
Self::RateLimited => 1,
Self::QuotaExhausted => 2,
Self::Error => 3,
Self::Disabled => 4,
}
}
}
/// `error` 状态下的稳定原因码。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AccountErrorReason {
AccountUnverified,
AccessTokenExpired,
CredentialExpired,
CredentialInvalid,
AccountBanned,
}
impl AccountErrorReason {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::AccountUnverified => "account_unverified",
Self::AccessTokenExpired => "access_token_expired",
Self::CredentialExpired => "credential_expired",
Self::CredentialInvalid => "credential_invalid",
Self::AccountBanned => "account_banned",
}
}
#[must_use]
pub fn parse(value: &str) -> Option<Self> {
match value {
"account_unverified" => Some(Self::AccountUnverified),
"access_token_expired" => Some(Self::AccessTokenExpired),
"credential_expired" => Some(Self::CredentialExpired),
"credential_invalid" => Some(Self::CredentialInvalid),
"account_banned" => Some(Self::AccountBanned),
_ => None,
}
}
}
/// 唯一状态解析器的完整输入事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountStatusFacts {
pub enabled: bool,
pub credential_state: CredentialState,
pub access_token_expires_at: Option<SystemTime>,
pub quota: QuotaState,
pub rate_limited_until: Option<SystemTime>,
pub last_error_reason: Option<AccountErrorReason>,
pub last_error_message: Option<String>,
}
/// 唯一状态解析器的输出。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountStatusProjection {
pub status: AccountStatus,
pub error_reason: Option<AccountErrorReason>,
pub error_message: Option<String>,
/// 仅 `rate_limited` 状态携带仍有效的运行时冷却截止时间。
pub rate_limited_until: Option<SystemTime>,
}
/// 从独立事实派生唯一、互斥的对外状态。
#[must_use]
pub fn resolve_account_status(
facts: &AccountStatusFacts,
now: SystemTime,
) -> AccountStatusProjection {
if !facts.enabled {
return status_projection(AccountStatus::Disabled);
}
let credential_error = facts.credential_state.error_reason().or_else(|| {
facts
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now)
.then_some(AccountErrorReason::AccessTokenExpired)
});
if let Some(default_reason) = credential_error {
return AccountStatusProjection {
status: AccountStatus::Error,
error_reason: Some(facts.last_error_reason.unwrap_or(default_reason)),
error_message: facts.last_error_message.clone(),
rate_limited_until: None,
};
}
if facts.quota.is_exhausted() {
return status_projection(AccountStatus::QuotaExhausted);
}
if facts
.rate_limited_until
.is_some_and(|rate_limited_until| rate_limited_until > now)
{
return AccountStatusProjection {
status: AccountStatus::RateLimited,
error_reason: None,
error_message: None,
rate_limited_until: facts.rate_limited_until,
};
}
status_projection(AccountStatus::Normal)
}
const fn status_projection(status: AccountStatus) -> AccountStatusProjection {
AccountStatusProjection {
status,
error_reason: None,
error_message: None,
rate_limited_until: None,
}
}
/// 不含 secret 的账号持久事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccount {
@@ -257,7 +671,12 @@ pub struct ProviderAccount {
authentication_kind: String,
revision: CredentialRevision,
enabled: bool,
availability: AccountAvailability,
concurrency_limit: Option<AccountConcurrencyLimit>,
weight: AccountWeight,
credential_state: CredentialState,
quota: QuotaState,
last_error_reason: Option<AccountErrorReason>,
last_error_message: Option<String>,
access_token_expires_at: Option<SystemTime>,
next_refresh_at: Option<SystemTime>,
has_refresh_token: bool,
@@ -286,7 +705,12 @@ impl ProviderAccount {
authentication_kind,
revision,
enabled: true,
availability: AccountAvailability::Unknown,
concurrency_limit: None,
weight: AccountWeight::DEFAULT,
credential_state: CredentialState::Unknown,
quota: QuotaState::unknown(),
last_error_reason: None,
last_error_message: None,
access_token_expires_at,
next_refresh_at: None,
has_refresh_token: false,
@@ -307,17 +731,34 @@ impl ProviderAccount {
}
#[must_use]
pub const fn with_runtime_state(
pub fn with_account_facts(
mut self,
enabled: bool,
availability: AccountAvailability,
credential_state: CredentialState,
quota: QuotaState,
last_error_reason: Option<AccountErrorReason>,
last_error_message: Option<String>,
) -> Self {
self.enabled = enabled;
self.availability = if self.upstream_user_id.is_some() {
availability
self.credential_state = if self.upstream_user_id.is_some() {
credential_state
} else {
AccountAvailability::Unknown
CredentialState::Unknown
};
self.quota = quota;
self.last_error_reason = last_error_reason;
self.last_error_message = last_error_message;
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
}
@@ -382,8 +823,23 @@ impl ProviderAccount {
}
#[must_use]
pub const fn availability(&self) -> AccountAvailability {
self.availability
pub const fn credential_state(&self) -> CredentialState {
self.credential_state
}
#[must_use]
pub const fn quota(&self) -> QuotaState {
self.quota
}
#[must_use]
pub const fn last_error_reason(&self) -> Option<AccountErrorReason> {
self.last_error_reason
}
#[must_use]
pub fn last_error_message(&self) -> Option<&str> {
self.last_error_message.as_deref()
}
#[must_use]
@@ -391,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
@@ -407,23 +881,25 @@ impl ProviderAccount {
self.has_refresh_token
}
/// 判断账号当前能否进入同 target 的候选池。
/// 组合持久事实与请求级冷却,交给唯一解析器派生状态。
#[must_use]
pub fn is_schedulable(&self, now: SystemTime) -> bool {
let available = match self.availability {
AccountAvailability::Ready => true,
AccountAvailability::Unknown
| AccountAvailability::QuotaExhausted
| AccountAvailability::Expired
| AccountAvailability::Banned
| AccountAvailability::Invalid => false,
};
self.upstream_user_id.is_some()
&& self.enabled
&& available
&& self
.access_token_expires_at
.is_none_or(|expires_at| expires_at > now)
pub fn status_projection(
&self,
now: SystemTime,
rate_limited_until: Option<SystemTime>,
) -> AccountStatusProjection {
resolve_account_status(
&AccountStatusFacts {
enabled: self.enabled,
credential_state: self.credential_state,
access_token_expires_at: self.access_token_expires_at,
quota: self.quota,
rate_limited_until,
last_error_reason: self.last_error_reason,
last_error_message: self.last_error_message.clone(),
},
now,
)
}
}
@@ -600,10 +1076,11 @@ pub enum CredentialCasOutcome {
pub struct QuotaObservation {
pub account_id: ProviderAccountId,
pub expected_revision: CredentialRevision,
pub quota: Option<OpaqueProviderData>,
pub observed_at: Option<SystemTime>,
/// 快照级限流事实(Provider 解析后物化);None 表示不更新物化列。
pub limit_reached: Option<bool>,
pub quota: OpaqueProviderData,
/// Provider 原始 quota document 的观察时间;不代表访问结论发生变化。
pub observed_at: SystemTime,
/// Provider 已从私有 JSON 归一化出的额度访问事实。
pub state: QuotaState,
}
impl fmt::Debug for QuotaObservation {
@@ -614,6 +1091,7 @@ impl fmt::Debug for QuotaObservation {
.field("expected_revision", &self.expected_revision)
.field("quota", &self.quota)
.field("observed_at", &self.observed_at)
.field("state", &self.state)
.finish()
}
}
@@ -624,14 +1102,24 @@ pub enum QuotaWriteOutcome {
Conflict,
}
/// 不改 Provider 原始 quota JSON 的额度访问事实写入。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QuotaAccessChange {
pub account_id: ProviderAccountId,
pub expected_revision: CredentialRevision,
pub state: QuotaState,
}
/// 账号状态的 revision-fenced 写入。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountStateChange {
pub account_id: ProviderAccountId,
pub expected_revision: CredentialRevision,
pub availability: AccountAvailability,
pub credential_state: CredentialState,
pub observed_at: SystemTime,
/// 上游原因(错误原文 / 原因码);`Some` 写入持久化错误信息,`None` 不更新,恢复 `Ready` 时清空。
/// 受控错误原因;恢复 Ready 时必须清空。
pub error_reason: Option<AccountErrorReason>,
/// 已脱敏的上游诊断信息;恢复 Ready 时必须清空。
pub message: Option<String>,
}
@@ -682,6 +1170,11 @@ pub trait ProviderAccountStore: Send + Sync {
observation: QuotaObservation,
) -> Result<QuotaWriteOutcome, StoreError>;
async fn apply_quota_access(
&self,
change: QuotaAccessChange,
) -> Result<QuotaWriteOutcome, StoreError>;
async fn apply_state_change(&self, change: AccountStateChange) -> Result<(), StoreError>;
async fn update_account(&self, update: ProviderAccountUpdate) -> Result<(), StoreError>;
@@ -772,7 +1265,7 @@ pub struct AccountRuntimeSignals {
pub last_started_at: Option<SystemTime>,
pub quota_reset_at: Option<SystemTime>,
pub quota_remaining_rank: Option<u64>,
pub quota_limit_reached: bool,
pub rate_limited_until: Option<SystemTime>,
pub failure_rate_basis_points: Option<u16>,
pub first_output_latency_ms: Option<u64>,
}
@@ -920,11 +1413,16 @@ impl AccountRuntimeSignals {
if let Some(quota) = quota {
self.quota_reset_at = quota.reset_at;
self.quota_remaining_rank = quota.remaining_rank;
self.quota_limit_reached = quota.limit_reached;
}
self
}
#[must_use]
pub const fn with_rate_limit(mut self, rate_limited_until: Option<SystemTime>) -> Self {
self.rate_limited_until = rate_limited_until;
self
}
#[must_use]
pub const fn with_runtime_health(
mut self,
@@ -942,20 +1440,14 @@ impl AccountRuntimeSignals {
pub struct AccountQuotaSignals {
reset_at: Option<SystemTime>,
remaining_rank: Option<u64>,
limit_reached: bool,
}
impl AccountQuotaSignals {
#[must_use]
pub const fn new(
reset_at: Option<SystemTime>,
remaining_rank: Option<u64>,
limit_reached: bool,
) -> Self {
pub const fn new(reset_at: Option<SystemTime>, remaining_rank: Option<u64>) -> Self {
Self {
reset_at,
remaining_rank,
limit_reached,
}
}
@@ -968,12 +1460,6 @@ impl AccountQuotaSignals {
pub const fn remaining_rank(self) -> Option<u64> {
self.remaining_rank
}
/// 任一计量窗口已触顶(`limit_reached`/`used_percent >= 100`);调度据此排除候选。
#[must_use]
pub const fn limit_reached(self) -> bool {
self.limit_reached
}
}
/// 账号持久事实与可重建运行信号的请求级组合。
@@ -991,24 +1477,25 @@ pub struct AccountSelectionContext {
pub excluded_accounts: BTreeSet<ProviderAccountId>,
pub preferred_account: Option<ProviderAccountId>,
pub round_robin_cursor: u64,
pub availability: AccountAvailabilityPolicy,
pub eligibility: AccountEligibilityPolicy,
pub account_scope: Option<std::sync::Arc<crate::routing::FrozenAccountScope>>,
}
/// 选择账号时是否信任本地可用性投影。
/// 选择账号时是否执行本地调度资格投影。
///
/// 管理端对指定账号执行诊断时,会直接向上游确认实际状态;该模式跳过
/// `enabled`、可用性、冷却和 token 到期等本地投影,仍保留租约、并发和
/// 请求间隔约束,且只能与固定账号约束组合使用。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum AccountAvailabilityPolicy {
pub enum AccountEligibilityPolicy {
#[default]
Enforce,
BypassForDiagnostic,
}
impl AccountAvailabilityPolicy {
impl AccountEligibilityPolicy {
#[must_use]
pub const fn bypasses_local_availability(self) -> bool {
pub const fn bypasses_local_eligibility(self) -> bool {
matches!(self, Self::BypassForDiagnostic)
}
}
@@ -1016,10 +1503,12 @@ impl AccountAvailabilityPolicy {
/// 候选账号未进入调度池的约束。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountSchedulingBlocker {
OutsideClientScope,
LocalAvailability,
Excluded,
ConcurrencyLimit,
RequestInterval,
LowerWeight,
}
/// 优先账号在本次选择中的处理结果。
@@ -1062,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())
@@ -1089,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 => {
@@ -1132,15 +1628,31 @@ impl AccountSelector {
candidate: &AccountCandidate,
context: &AccountSelectionContext,
) -> Option<AccountSchedulingBlocker> {
if !context.availability.bypasses_local_availability()
&& !candidate.account.is_schedulable(context.now)
if context
.account_scope
.as_ref()
.is_some_and(|scope| !scope.allows(candidate.account.id()))
{
return Some(AccountSchedulingBlocker::LocalAvailability);
return Some(AccountSchedulingBlocker::OutsideClientScope);
}
if !context.eligibility.bypasses_local_eligibility() {
let status = candidate
.account
.status_projection(context.now, candidate.signals.rate_limited_until)
.status;
if status != AccountStatus::Normal {
return Some(AccountSchedulingBlocker::LocalAvailability);
}
}
if context.excluded_accounts.contains(candidate.account.id()) {
return Some(AccountSchedulingBlocker::Excluded);
}
if candidate.signals.in_flight >= context.policy.max_concurrent_per_account().get() {
if candidate.signals.in_flight
>= candidate
.account
.effective_concurrency(context.policy.max_concurrent_per_account())
.get()
{
return Some(AccountSchedulingBlocker::ConcurrencyLimit);
}
if candidate
@@ -15,7 +15,8 @@ use crate::engine::admission::{
ClientAdmissionDecision, ClientAdmissionPort, ClientAdmissionRejection, ClientAdmissionRequest,
};
use crate::engine::continuation::{
ContinuationBinding, NativeContinuationPin, NativeContinuationPort, PreviousResponseId,
ContinuationBinding, NativeContinuationPin, NativeContinuationPort,
NativeContinuationStoreErrorKind, PreviousResponseId,
};
use crate::engine::coordinator::ResponseExecutionSession;
use crate::engine::probe::{
@@ -261,22 +262,32 @@ impl DefaultExecutionService {
})?;
let request_id = new_request_id()?;
let routing_context = self
.route_context(request.client.policy.provider_kind())
.route_context(request.client.policy.account_scope().provider_kinds())
.await?;
let plan = request
.client
.snapshot
.plan(&request.public_model, &request.operation, &routing_context)
.plan(
&request.public_model,
&request.operation,
Arc::clone(request.client.policy.account_scope()),
&routing_context,
)
.map_err(map_routing_error)?;
let continuation = match request.metadata.previous_response_id.as_ref() {
Some(previous) => {
let resolve = self.continuation.resolve(previous).fuse();
let resolve = self
.continuation
.resolve(request.client.policy.key_id(), previous)
.fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(resolve, timeout);
let pin = select_biased! {
result = resolve => match result {
Ok(pin) => pin,
Err(error) => {
Err(error)
if error.kind() == NativeContinuationStoreErrorKind::Unavailable =>
{
tracing::debug!(
request_id = request_id.as_str(),
%error,
@@ -284,10 +295,53 @@ impl DefaultExecutionService {
);
None
}
Err(error)
if error.kind()
== NativeContinuationStoreErrorKind::OwnershipMismatch =>
{
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"continuation does not belong to this client API key",
));
}
Err(error) => {
tracing::warn!(
request_id = request_id.as_str(),
%error,
"Continuation affinity 记录无效,已拒绝续接"
);
return Err(GatewayError::new(
GatewayErrorKind::Internal,
"continuation state is invalid",
));
}
},
_ = timeout => None,
};
match pin {
Some(pin) if !pin.matches_client(request.client.policy.key_id()) => {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"continuation does not belong to this client API key",
));
}
Some(pin) if !request.client.policy.account_scope().allows(pin.account()) => {
return Err(GatewayError::new(
GatewayErrorKind::PolicyDenied,
"continuation account is outside the client account scope",
));
}
Some(pin)
if !plan
.candidates()
.iter()
.any(|candidate| candidate.provider() == pin.provider()) =>
{
return Err(GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"continuation provider is not available",
));
}
Some(pin) => {
attach_continuation_session_state(&mut request.operation, &pin);
ContinuationBinding::Pinned(pin)
@@ -360,14 +414,19 @@ impl DefaultExecutionService {
client_api_key_id: request.client.policy.key_id().clone(),
model_request_id: request_id.clone(),
};
let observation = self
.providers
.request_observation(request.client.policy.provider_kind(), &request.operation);
let observation = plan
.candidates()
.first()
.map_or_else(Default::default, |candidate| {
self.providers
.request_observation(candidate.provider(), &request.operation)
});
let new_request = NewModelRequest {
id: request_id.clone(),
client_api_key_id: Some(request.client.policy.key_id().clone()),
client_api_key_ref: request.client.policy.key_id().clone(),
config_revision: plan.config_revision(),
routing: request.client.policy.account_scope().routing_snapshot(),
protocol: request.metadata.protocol,
operation: request.operation.kind(),
endpoint: request.metadata.endpoint,
@@ -417,38 +476,42 @@ impl DefaultExecutionService {
async fn route_context(
&self,
provider_kind: &ProviderKind,
provider_kinds: &BTreeSet<ProviderKind>,
) -> Result<RoutingContext, GatewayError> {
let decision = self.circuits.decision(provider_kind).fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(decision, timeout);
let decision = select_biased! {
result = decision => Some(result),
_ = timeout => None,
};
let blocked_providers = match decision {
Some(Ok(ProviderCircuitDecision::Allow)) => BTreeSet::new(),
Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => {
BTreeSet::from([provider_kind.clone()])
let decisions = futures::future::join_all(provider_kinds.iter().map(|provider_kind| {
let circuits = Arc::clone(&self.circuits);
async move {
let decision = circuits.decision(provider_kind).fuse();
let timeout = Delay::new(COORDINATION_TIMEOUT).fuse();
pin_mut!(decision, timeout);
let decision = select_biased! {
result = decision => Some(result),
_ = timeout => None,
};
(provider_kind, decision)
}
Some(Err(error)) => {
tracing::warn!(
}))
.await;
let mut blocked_providers = BTreeSet::new();
for (provider_kind, decision) in decisions {
match decision {
Some(Ok(ProviderCircuitDecision::BlockedUntil(_))) => {
blocked_providers.insert(provider_kind.clone());
}
Some(Err(error)) => tracing::warn!(
provider = provider_kind.as_str(),
%error,
"Provider circuit 读取失败,按可重建协调状态 fail-open"
);
BTreeSet::new()
}
None => {
tracing::warn!(
),
None => tracing::warn!(
provider = provider_kind.as_str(),
"Provider circuit 读取超时,按可重建协调状态 fail-open"
);
BTreeSet::new()
),
Some(Ok(ProviderCircuitDecision::Allow)) => {}
}
};
}
Ok(RoutingContext {
provider_kind: Some(provider_kind.clone()),
required_provider: None,
blocked_providers,
})
}
@@ -479,11 +542,16 @@ impl DefaultExecutionService {
GatewayError::new(GatewayErrorKind::Unsupported, "requested model is invalid")
})?;
let routing_context = RoutingContext {
provider_kind: Some(provider_kind),
required_provider: Some(provider_kind),
..RoutingContext::default()
};
let plan = snapshot
.plan(&public_model, &operation, &routing_context)
.plan(
&public_model,
&operation,
snapshot.all_account_scope(),
&routing_context,
)
.map_err(map_routing_error)?;
let started_at = SystemTime::now();
let deadline_at = started_at
@@ -499,6 +567,7 @@ impl DefaultExecutionService {
client_api_key_id: None,
client_api_key_ref: actor,
config_revision: plan.config_revision(),
routing: crate::routing::AccountRoutingSnapshot::all(),
protocol: "admin_connection_test".to_owned(),
operation: operation.kind(),
endpoint: "/api/admin/accounts/connection-test".to_owned(),
@@ -693,19 +762,19 @@ impl ExecutionService for DefaultExecutionService {
fn public_models(&self, client: &AuthenticatedClient) -> Vec<PublicModelId> {
client
.snapshot
.public_models_for_provider(client.policy.provider_kind())
.public_models_for_scope(client.policy.account_scope())
}
fn public_model_profiles(&self, client: &AuthenticatedClient) -> Vec<PublicModelProfile> {
client
.snapshot
.public_model_profiles_for_provider(client.policy.provider_kind())
.public_model_profiles_for_scope(client.policy.account_scope())
}
fn contains_public_model(&self, client: &AuthenticatedClient, model: &PublicModelId) -> bool {
client
.snapshot
.contains_public_model_for_provider(model, client.policy.provider_kind())
.contains_public_model_for_scope(model, client.policy.account_scope())
}
fn start(
@@ -1006,7 +1075,8 @@ fn new_request_id() -> Result<ModelRequestId, GatewayError> {
fn map_routing_error(error: crate::error::RoutingError) -> GatewayError {
match error {
crate::error::RoutingError::NoCapableProvider { .. } => GatewayError::new(
crate::error::RoutingError::NoCapableProvider { .. }
| crate::error::RoutingError::EmptyAccountScope => GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no provider can execute this request",
),
@@ -233,6 +233,7 @@ pub struct AccountAttemptContext {
state_owner: Option<ProviderAccountStateOwner>,
credential_recovery_attempted: bool,
diagnostic_required_account: bool,
account_scope: Option<Arc<crate::routing::FrozenAccountScope>>,
}
impl AccountAttemptContext {
@@ -248,6 +249,7 @@ impl AccountAttemptContext {
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: false,
account_scope: None,
}
}
@@ -266,9 +268,17 @@ impl AccountAttemptContext {
state_owner,
credential_recovery_attempted: false,
diagnostic_required_account: true,
account_scope: None,
}
}
/// 附着普通请求认证时冻结的账号范围。
#[must_use]
pub fn with_account_scope(mut self, scope: Arc<crate::routing::FrozenAccountScope>) -> Self {
self.account_scope = Some(scope);
self
}
/// 标记本请求已对即将选择的固定账号执行过一次凭据恢复。
#[must_use]
pub const fn with_credential_recovery_attempted(mut self, attempted: bool) -> Self {
@@ -301,6 +311,11 @@ impl AccountAttemptContext {
pub const fn is_diagnostic_required_account(&self) -> bool {
self.diagnostic_required_account
}
#[must_use]
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
self.account_scope.as_ref()
}
}
/// 请求中 Provider 账号绑定状态的唯一归属。
@@ -506,6 +521,12 @@ impl AttemptContext {
self.account.is_diagnostic_required_account()
}
/// 普通请求认证时冻结的账号范围;管理端诊断为 `None`。
#[must_use]
pub const fn account_scope(&self) -> Option<&Arc<crate::routing::FrozenAccountScope>> {
self.account.account_scope()
}
#[must_use]
pub const fn continuation(&self) -> Option<&ContinuationBinding> {
self.continuation.as_ref()
@@ -529,6 +550,7 @@ pub struct NewModelRequest {
pub client_api_key_id: Option<ClientApiKeyId>,
pub client_api_key_ref: ClientApiKeyId,
pub config_revision: ConfigRevision,
pub routing: crate::routing::AccountRoutingSnapshot,
pub protocol: String,
pub operation: OperationKind,
pub endpoint: String,
@@ -749,6 +771,8 @@ pub enum EngineError {
ContinuationPinMismatch,
#[error("provider did not use the account required by this execution")]
RequiredAccountMismatch,
#[error("provider selected an account outside the frozen client scope")]
AccountOutsideClientScope,
#[error("provider execution failed")]
Provider(ProviderError),
#[error("request was cancelled")]
+9
View File
@@ -33,6 +33,9 @@ pub enum IdentifierError {
/// 规定前缀。
expected: &'static str,
},
/// 标识不满足该领域值对象的固定格式。
#[error("identifier has an invalid format")]
InvalidFormat,
}
/// Operation 构造或校验失败。
@@ -67,6 +70,12 @@ pub enum RoutingError {
/// 配置 revision 必须为正数。
#[error("config revision must be greater than zero")]
InvalidRevision,
/// Restricted scope 必须保留至少一个持久化分组 binding。
#[error("account routing scope is invalid")]
InvalidAccountScope,
/// Key 的账号范围中当前没有任何可参与路由的账号。
#[error("account routing scope is empty")]
EmptyAccountScope,
/// 快照中存在重复实体。
#[error("duplicate {entity} `{id}`")]
DuplicateEntity {
@@ -1,11 +1,12 @@
//! 下游 Client API Key 的准入策略。
//!
//! Client API Key 绑定一个 Provider 平台;模型名称不参与改写或 allowlist 判断。
//! Client API Key 冻结账号分组权限;模型名称不参与权限判断。
use std::fmt;
use std::sync::Arc;
use crate::error::{IdentifierError, PolicyError, validate_text};
use crate::routing::ProviderKind;
use crate::routing::FrozenAccountScope;
/// `client_api_keys.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -88,7 +89,7 @@ impl RateLimits {
pub struct ClientPolicy {
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: ProviderKind,
account_scope: Arc<FrozenAccountScope>,
enabled: bool,
limits: RateLimits,
}
@@ -98,14 +99,14 @@ impl ClientPolicy {
pub const fn new(
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: ProviderKind,
account_scope: Arc<FrozenAccountScope>,
enabled: bool,
limits: RateLimits,
) -> Self {
Self {
key_id,
plaintext_key,
provider_kind,
account_scope,
enabled,
limits,
}
@@ -122,8 +123,8 @@ impl ClientPolicy {
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
#[must_use]
@@ -9,10 +9,11 @@ use std::time::{Duration, SystemTime};
use futures::future::BoxFuture;
use crate::engine::credential::{
AccountAvailability, AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision,
AccountFeedbackStats, AccountRuntimeSignals, CredentialRevision, CredentialState,
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>>;
@@ -411,7 +413,7 @@ pub struct ProviderCredentialState {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
enabled: bool,
availability: AccountAvailability,
credential_state: CredentialState,
observed_at: SystemTime,
}
@@ -421,14 +423,14 @@ impl ProviderCredentialState {
account_id: ProviderAccountId,
credential_revision: CredentialRevision,
enabled: bool,
availability: AccountAvailability,
credential_state: CredentialState,
observed_at: SystemTime,
) -> Self {
Self {
account_id,
credential_revision,
enabled,
availability,
credential_state,
observed_at,
}
}
@@ -449,8 +451,8 @@ impl ProviderCredentialState {
}
#[must_use]
pub const fn availability(&self) -> AccountAvailability {
self.availability
pub const fn credential_state(&self) -> CredentialState {
self.credential_state
}
#[must_use]
+296 -3
View File
@@ -9,12 +9,293 @@ use std::fmt;
use std::num::{NonZeroU32, NonZeroU64};
use std::sync::Arc;
use crate::engine::credential::AccountSelectionPolicy;
use crate::engine::credential::{AccountSelectionPolicy, ProviderAccountId};
use crate::error::{IdentifierError, RoutingError, validate_text};
use crate::operation::{CapabilityRequirements, Feature, OperationKind};
const MAX_REQUEST_ATTEMPTS: u32 = 32;
/// `account_groups.id` 的核心值对象。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AccountGroupId(String);
impl AccountGroupId {
/// 校验并创建账号分组 ID。
pub fn new(value: impl Into<String>) -> Result<Self, IdentifierError> {
let value = value.into();
let Some(suffix) = value.strip_prefix("grp_") else {
return Err(IdentifierError::MissingPrefix { expected: "grp_" });
};
if suffix.len() != 32
|| !suffix
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
{
return Err(IdentifierError::InvalidFormat);
}
Ok(Self(value))
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for AccountGroupId {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(&self.0)
}
}
/// 快照中一个账号的 Provider 与分组归属。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RuntimeAccount {
provider_kind: ProviderKind,
group_ids: Arc<BTreeSet<AccountGroupId>>,
}
impl RuntimeAccount {
#[must_use]
pub fn new(provider_kind: ProviderKind, group_ids: BTreeSet<AccountGroupId>) -> Self {
Self {
provider_kind,
group_ids: Arc::new(group_ids),
}
}
#[must_use]
pub const fn provider_kind(&self) -> &ProviderKind {
&self.provider_kind
}
#[must_use]
pub fn group_ids(&self) -> &BTreeSet<AccountGroupId> {
&self.group_ids
}
}
/// 全快照共享的账号、Provider 与分组反向索引。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct RuntimeAccountDirectory {
accounts: BTreeMap<ProviderAccountId, RuntimeAccount>,
providers_with_accounts: BTreeSet<ProviderKind>,
providers_by_group: BTreeMap<AccountGroupId, BTreeSet<ProviderKind>>,
}
impl RuntimeAccountDirectory {
#[must_use]
pub fn new(accounts: BTreeMap<ProviderAccountId, RuntimeAccount>) -> Self {
let mut providers_with_accounts = BTreeSet::new();
let mut providers_by_group = BTreeMap::<AccountGroupId, BTreeSet<ProviderKind>>::new();
for account in accounts.values() {
providers_with_accounts.insert(account.provider_kind.clone());
for group_id in account.group_ids.iter() {
providers_by_group
.entry(group_id.clone())
.or_default()
.insert(account.provider_kind.clone());
}
}
Self {
accounts,
providers_with_accounts,
providers_by_group,
}
}
#[must_use]
pub fn account(&self, account_id: &ProviderAccountId) -> Option<&RuntimeAccount> {
self.accounts.get(account_id)
}
#[must_use]
pub fn providers_with_accounts(&self) -> &BTreeSet<ProviderKind> {
&self.providers_with_accounts
}
#[must_use]
pub fn providers_for_groups<'a>(
&self,
group_ids: impl IntoIterator<Item = &'a AccountGroupId>,
) -> BTreeSet<ProviderKind> {
group_ids
.into_iter()
.filter_map(|group_id| self.providers_by_group.get(group_id))
.flatten()
.cloned()
.collect()
}
}
/// 历史请求保存的账号范围种类。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountRoutingScopeKind {
All,
Groups,
}
impl AccountRoutingScopeKind {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::All => "all",
Self::Groups => "groups",
}
}
}
/// 请求开始时冻结的分组 ID 与名称。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RoutingGroupSnapshot {
id: AccountGroupId,
name: String,
}
impl RoutingGroupSnapshot {
#[must_use]
pub fn new(id: AccountGroupId, name: String) -> Self {
Self { id, name }
}
#[must_use]
pub const fn id(&self) -> &AccountGroupId {
&self.id
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
}
/// 请求历史所需的完整、稳定账号范围快照。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountRoutingSnapshot {
kind: AccountRoutingScopeKind,
groups: Arc<[RoutingGroupSnapshot]>,
}
impl AccountRoutingSnapshot {
#[must_use]
pub fn all() -> Self {
Self {
kind: AccountRoutingScopeKind::All,
groups: Arc::from([]),
}
}
#[must_use]
pub fn groups(groups: Vec<RoutingGroupSnapshot>) -> Self {
Self {
kind: AccountRoutingScopeKind::Groups,
groups: Arc::from(groups),
}
}
#[must_use]
pub const fn kind(&self) -> AccountRoutingScopeKind {
self.kind
}
#[must_use]
pub fn groups_snapshot(&self) -> &[RoutingGroupSnapshot] {
&self.groups
}
}
/// Key 持久 binding 编译出的账号权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClientRoutingScope {
AllAccounts,
Restricted {
bound_groups: Arc<[RoutingGroupSnapshot]>,
enabled_group_ids: Arc<BTreeSet<AccountGroupId>>,
provider_kinds: Arc<BTreeSet<ProviderKind>>,
},
}
impl ClientRoutingScope {
#[must_use]
pub fn all_accounts() -> Self {
Self::AllAccounts
}
pub fn restricted(
bound_groups: Vec<RoutingGroupSnapshot>,
enabled_group_ids: BTreeSet<AccountGroupId>,
provider_kinds: BTreeSet<ProviderKind>,
) -> Result<Self, RoutingError> {
if bound_groups.is_empty() {
return Err(RoutingError::InvalidAccountScope);
}
Ok(Self::Restricted {
bound_groups: Arc::from(bound_groups),
enabled_group_ids: Arc::new(enabled_group_ids),
provider_kinds: Arc::new(provider_kinds),
})
}
}
/// 一次认证随 RuntimeSnapshot 冻结的账号目录与 Key 权限。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FrozenAccountScope {
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
}
impl FrozenAccountScope {
#[must_use]
pub const fn new(
directory: Arc<RuntimeAccountDirectory>,
client_scope: ClientRoutingScope,
) -> Self {
Self {
directory,
client_scope,
}
}
#[must_use]
pub fn allows(&self, account_id: &ProviderAccountId) -> bool {
let Some(account) = self.directory.account(account_id) else {
return false;
};
match &self.client_scope {
ClientRoutingScope::AllAccounts => true,
ClientRoutingScope::Restricted {
enabled_group_ids, ..
} => account
.group_ids()
.iter()
.any(|group_id| enabled_group_ids.contains(group_id)),
}
}
#[must_use]
pub fn provider_kinds(&self) -> &BTreeSet<ProviderKind> {
match &self.client_scope {
ClientRoutingScope::AllAccounts => self.directory.providers_with_accounts(),
ClientRoutingScope::Restricted { provider_kinds, .. } => provider_kinds,
}
}
#[must_use]
pub fn routing_snapshot(&self) -> AccountRoutingSnapshot {
match &self.client_scope {
ClientRoutingScope::AllAccounts => AccountRoutingSnapshot::all(),
ClientRoutingScope::Restricted { bound_groups, .. } => {
AccountRoutingSnapshot::groups(bound_groups.to_vec())
}
}
}
#[must_use]
pub const fn directory(&self) -> &Arc<RuntimeAccountDirectory> {
&self.directory
}
}
/// 编译进二进制的 Provider adapter slug。
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ProviderKind(String);
@@ -494,8 +775,8 @@ impl ProviderModel {
/// 本次请求选择 Provider 时使用的动态过滤事实。
#[derive(Debug, Clone, Default)]
pub struct RoutingContext {
/// 已认证 Client API Key 绑定的平台;模型名称不参与平台猜测。
pub provider_kind: Option<ProviderKind>,
/// 管理端 connection test 显式限制的 Provider;普通请求留空。
pub required_provider: Option<ProviderKind>,
pub blocked_providers: BTreeSet<ProviderKind>,
}
@@ -505,6 +786,7 @@ pub struct ProviderCandidate {
provider: ProviderKind,
upstream_model: UpstreamModelId,
emulated_features: BTreeSet<Feature>,
account_scope: Arc<FrozenAccountScope>,
}
impl ProviderCandidate {
@@ -522,6 +804,11 @@ impl ProviderCandidate {
pub const fn emulated_features(&self) -> &BTreeSet<Feature> {
&self.emulated_features
}
#[must_use]
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
}
/// 一次请求冻结的 Provider 尝试顺序。
@@ -532,6 +819,7 @@ pub struct RoutingPlan {
public_model: PublicModelId,
operation: OperationKind,
max_attempts: NonZeroU32,
account_scope: Arc<FrozenAccountScope>,
candidates: Arc<[ProviderCandidate]>,
}
@@ -561,6 +849,11 @@ impl RoutingPlan {
self.max_attempts
}
#[must_use]
pub const fn account_scope(&self) -> &Arc<FrozenAccountScope> {
&self.account_scope
}
#[must_use]
pub fn candidates(&self) -> &[ProviderCandidate] {
&self.candidates
@@ -11,7 +11,7 @@ use futures::{FutureExt as _, Stream, StreamExt as _, pin_mut, select_biased};
use futures_timer::Delay;
use crate::engine::CancellationToken;
use crate::engine::credential::{AccountSelectionPolicy, RotationStrategy};
use crate::engine::credential::{AccountSelectionPolicy, ProviderAccountId, RotationStrategy};
use crate::engine::provider::{ProviderCatalogGeneration, ProviderRegistry};
use crate::error::RoutingError;
use crate::health::{HealthProbe, HealthState};
@@ -24,8 +24,9 @@ use crate::task::{
};
use super::{
ConfigRevision, ModelCapabilities, ProviderCandidate, ProviderKind, ProviderModel,
PublicModelId, RoutingContext, RoutingPlan, UpstreamModelId,
AccountGroupId, ClientRoutingScope, ConfigRevision, FrozenAccountScope, ModelCapabilities,
ProviderCandidate, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RoutingGroupSnapshot, RoutingPlan, RuntimeAccount, RuntimeAccountDirectory, UpstreamModelId,
};
const RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5);
@@ -66,7 +67,7 @@ impl SnapshotSettingsFacts {
pub struct SnapshotClientPolicyFacts {
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: String,
group_ids: Vec<AccountGroupId>,
limits: RateLimits,
}
@@ -75,18 +76,67 @@ impl SnapshotClientPolicyFacts {
pub fn new(
key_id: ClientApiKeyId,
plaintext_key: PlaintextClientApiKey,
provider_kind: impl Into<String>,
group_ids: Vec<AccountGroupId>,
limits: RateLimits,
) -> Self {
Self {
key_id,
plaintext_key,
provider_kind: provider_kind.into(),
group_ids,
limits,
}
}
}
/// Store 读取到的账号分组事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotAccountGroupFacts {
id: AccountGroupId,
name: String,
enabled: bool,
}
impl SnapshotAccountGroupFacts {
#[must_use]
pub fn new(id: AccountGroupId, name: String, enabled: bool) -> Self {
Self { id, name, enabled }
}
}
/// Store 读取到的账号及其固有 Provider 事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotProviderAccountFacts {
account_id: ProviderAccountId,
provider_kind: String,
}
impl SnapshotProviderAccountFacts {
#[must_use]
pub fn new(account_id: ProviderAccountId, provider_kind: impl Into<String>) -> Self {
Self {
account_id,
provider_kind: provider_kind.into(),
}
}
}
/// Store 读取到的一条分组成员关系。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotAccountGroupMemberFacts {
group_id: AccountGroupId,
account_id: ProviderAccountId,
}
impl SnapshotAccountGroupMemberFacts {
#[must_use]
pub const fn new(group_id: AccountGroupId, account_id: ProviderAccountId) -> Self {
Self {
group_id,
account_id,
}
}
}
/// 一次一致性读取产生的全部 RuntimeSnapshot 持久事实。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SnapshotFacts {
@@ -94,6 +144,9 @@ pub struct SnapshotFacts {
observed_current_revision: ConfigRevision,
settings: SnapshotSettingsFacts,
client_policies: Vec<SnapshotClientPolicyFacts>,
account_groups: Vec<SnapshotAccountGroupFacts>,
provider_accounts: Vec<SnapshotProviderAccountFacts>,
group_memberships: Vec<SnapshotAccountGroupMemberFacts>,
}
impl SnapshotFacts {
@@ -103,12 +156,18 @@ impl SnapshotFacts {
observed_current_revision: ConfigRevision,
settings: SnapshotSettingsFacts,
client_policies: Vec<SnapshotClientPolicyFacts>,
account_groups: Vec<SnapshotAccountGroupFacts>,
provider_accounts: Vec<SnapshotProviderAccountFacts>,
group_memberships: Vec<SnapshotAccountGroupMemberFacts>,
) -> Self {
Self {
config_revision,
observed_current_revision,
settings,
client_policies,
account_groups,
provider_accounts,
group_memberships,
}
}
@@ -223,13 +282,16 @@ async fn compile_runtime_snapshot(
providers: &ProviderRegistry,
) -> Result<RuntimeSnapshot, RuntimeSnapshotCompileError> {
let provider_kinds = providers.provider_kinds().cloned().collect::<Vec<_>>();
let registered_providers = provider_kinds.iter().cloned().collect::<BTreeSet<_>>();
// 实时目录只提供公开模型与能力提示;查询失败时保留 Provider 透传语义。
// 目录查询失败表示未知;查询成功后,即使为空,也必须与“已知缺少模型”区分。
let mut provider_models = Vec::new();
let mut known_provider_catalogs = BTreeSet::new();
for provider in &provider_kinds {
let Ok(models) = providers.query_model_capabilities(provider).await else {
continue;
};
known_provider_catalogs.insert(provider.clone());
provider_models.extend(models.into_iter().map(|model| {
let compiled = ProviderModel::new(
provider.clone(),
@@ -243,6 +305,56 @@ async fn compile_runtime_snapshot(
}));
}
let mut groups = BTreeMap::new();
for group in facts.account_groups {
if group.name.trim() != group.name
|| group.name.is_empty()
|| group.name.chars().count() > 100
|| group.name.chars().any(char::is_control)
|| groups.insert(group.id.clone(), group).is_some()
{
return Err(RuntimeSnapshotCompileError::InvalidData);
}
}
let mut account_groups = BTreeMap::<ProviderAccountId, BTreeSet<AccountGroupId>>::new();
let mut accounts = BTreeMap::new();
for account in facts.provider_accounts {
let provider_kind = ProviderKind::new(account.provider_kind)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
if !registered_providers.contains(&provider_kind)
|| accounts
.insert(
account.account_id.clone(),
RuntimeAccount::new(provider_kind, BTreeSet::new()),
)
.is_some()
{
return Err(RuntimeSnapshotCompileError::InvalidData);
}
account_groups.insert(account.account_id, BTreeSet::new());
}
let mut memberships = BTreeSet::new();
for membership in facts.group_memberships {
if !groups.contains_key(&membership.group_id)
|| !accounts.contains_key(&membership.account_id)
|| !memberships.insert((membership.group_id.clone(), membership.account_id.clone()))
{
return Err(RuntimeSnapshotCompileError::InvalidData);
}
account_groups
.get_mut(&membership.account_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?
.insert(membership.group_id);
}
for (account_id, group_ids) in account_groups {
let account = accounts
.get_mut(&account_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
*account = RuntimeAccount::new(account.provider_kind().clone(), group_ids);
}
let account_directory = Arc::new(RuntimeAccountDirectory::new(accounts));
let model_mappings = facts.settings.model_mappings;
let rotation_strategy = RotationStrategy::parse(facts.settings.rotation_strategy.as_str())
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
@@ -252,21 +364,48 @@ async fn compile_runtime_snapshot(
.ok_or(RuntimeSnapshotCompileError::InvalidData)?,
Duration::from_millis(facts.settings.request_interval_ms),
);
let client_policies = facts
.client_policies
.into_iter()
.map(|policy| {
let provider = ProviderKind::new(policy.provider_kind)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?;
Ok(ClientPolicy::new(
policy.key_id,
policy.plaintext_key,
provider,
true,
policy.limits,
))
})
.collect::<Result<Vec<_>, RuntimeSnapshotCompileError>>()?;
let mut client_policies = Vec::with_capacity(facts.client_policies.len());
for policy in facts.client_policies {
let account_scope = if policy.group_ids.is_empty() {
FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::all_accounts(),
)
} else {
let mut seen = BTreeSet::new();
let mut bound_groups = Vec::with_capacity(policy.group_ids.len());
let mut enabled_group_ids = BTreeSet::new();
for group_id in policy.group_ids {
if !seen.insert(group_id.clone()) {
return Err(RuntimeSnapshotCompileError::InvalidData);
}
let group = groups
.get(&group_id)
.ok_or(RuntimeSnapshotCompileError::InvalidData)?;
bound_groups.push(RoutingGroupSnapshot::new(
group.id.clone(),
group.name.clone(),
));
if group.enabled {
enabled_group_ids.insert(group_id);
}
}
bound_groups.sort_by(|left, right| left.id().cmp(right.id()));
let provider_kinds = account_directory.providers_for_groups(&enabled_group_ids);
FrozenAccountScope::new(
Arc::clone(&account_directory),
ClientRoutingScope::restricted(bound_groups, enabled_group_ids, provider_kinds)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)?,
)
};
client_policies.push(ClientPolicy::new(
policy.key_id,
policy.plaintext_key,
Arc::new(account_scope),
true,
policy.limits,
));
}
RuntimeSnapshot::new(
facts.config_revision,
@@ -276,7 +415,12 @@ async fn compile_runtime_snapshot(
client_policies,
)
.map_err(|_| RuntimeSnapshotCompileError::InvalidData)
.map(|snapshot| snapshot.with_model_mappings(model_mappings))
.map(|snapshot| {
snapshot
.with_model_mappings(model_mappings)
.with_account_directory(account_directory)
.with_known_provider_catalogs(known_provider_catalogs)
})
}
/// 数据面使用的不可变配置快照。
@@ -290,6 +434,8 @@ pub struct RuntimeSnapshot {
Arc<BTreeMap<ProviderKind, BTreeMap<UpstreamModelId, super::ModelPresentation>>>,
model_mappings: Arc<BTreeMap<String, String>>,
provider_catalog_generations: Arc<BTreeMap<ProviderKind, ProviderCatalogGeneration>>,
known_provider_catalogs: Arc<BTreeSet<ProviderKind>>,
account_directory: Arc<RuntimeAccountDirectory>,
client_policies: Arc<BTreeMap<ClientApiKeyId, ClientPolicy>>,
}
@@ -312,6 +458,7 @@ impl RuntimeSnapshot {
}
}
let mut known_provider_catalogs = BTreeSet::new();
let mut model_map =
BTreeMap::<ProviderKind, BTreeMap<UpstreamModelId, ModelCapabilities>>::new();
let mut presentation_map =
@@ -323,6 +470,7 @@ impl RuntimeSnapshot {
capabilities,
presentation,
} = model;
known_provider_catalogs.insert(provider.clone());
if !provider_set.contains(&provider) {
return Err(RoutingError::NotFound {
entity: "provider",
@@ -367,6 +515,8 @@ impl RuntimeSnapshot {
provider_model_presentations: Arc::new(presentation_map),
model_mappings: Arc::new(BTreeMap::new()),
provider_catalog_generations: Arc::new(BTreeMap::new()),
known_provider_catalogs: Arc::new(known_provider_catalogs),
account_directory: Arc::new(RuntimeAccountDirectory::default()),
client_policies: Arc::new(client_policy_map),
})
}
@@ -377,6 +527,26 @@ impl RuntimeSnapshot {
self
}
#[must_use]
pub fn with_account_directory(mut self, directory: Arc<RuntimeAccountDirectory>) -> Self {
self.account_directory = directory;
self
}
#[must_use]
fn with_known_provider_catalogs(mut self, providers: BTreeSet<ProviderKind>) -> Self {
self.known_provider_catalogs = Arc::new(providers);
self
}
#[must_use]
pub fn all_account_scope(&self) -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
Arc::clone(&self.account_directory),
ClientRoutingScope::all_accounts(),
))
}
#[must_use]
fn with_provider_catalog_generations(
mut self,
@@ -459,8 +629,39 @@ impl RuntimeSnapshot {
.collect()
}
/// 已取得目录时按目录与全局别名判断模型可用性;目录暂不可用时沿用旧版
/// fail-open 语义,把具体 wire 可接受性留给上游。
/// 合并冻结账号范围内实际存在 Provider 的公开模型。
#[must_use]
pub fn public_models_for_scope(&self, scope: &FrozenAccountScope) -> Vec<PublicModelId> {
scope
.provider_kinds()
.iter()
.flat_map(|provider| self.public_models_for_provider(provider))
.collect::<BTreeSet<_>>()
.into_iter()
.collect()
}
/// 合并冻结账号范围内实际存在 Provider 的公开模型画像。
#[must_use]
pub fn public_model_profiles_for_scope(
&self,
scope: &FrozenAccountScope,
) -> Vec<super::PublicModelProfile> {
let mut profiles = BTreeMap::new();
for provider in scope.provider_kinds() {
for profile in self.public_model_profiles_for_provider(provider) {
profiles
.entry(profile.model().clone())
.or_insert_with(|| profile.presentation().clone());
}
}
profiles
.into_iter()
.map(|(model, presentation)| super::PublicModelProfile::new(model, presentation))
.collect()
}
/// 已取得目录时以映射后的上游模型为准;目录不可用时由 Provider 在发送前判定。
#[must_use]
pub fn contains_public_model_for_provider(
&self,
@@ -470,17 +671,25 @@ impl RuntimeSnapshot {
if !self.providers.contains(provider) {
return false;
}
if self.model_mappings.contains_key(public_model.as_str()) {
return true;
}
let upstream_model = self.mapped_model(public_model.as_str());
match self.provider_models.get(provider) {
Some(models) if !models.is_empty() => models
.keys()
.any(|model| model.as_str() == public_model.as_str()),
Some(_) | None => true,
Some(models) => models.keys().any(|model| model.as_str() == upstream_model),
None => !self.known_provider_catalogs.contains(provider),
}
}
#[must_use]
pub fn contains_public_model_for_scope(
&self,
public_model: &PublicModelId,
scope: &FrozenAccountScope,
) -> bool {
scope
.provider_kinds()
.iter()
.any(|provider| self.contains_public_model_for_provider(public_model, provider))
}
#[must_use]
pub fn mapped_model(&self, requested: &str) -> String {
let original = requested;
@@ -506,14 +715,26 @@ impl RuntimeSnapshot {
&self,
public_model: &PublicModelId,
operation: &Operation,
account_scope: Arc<FrozenAccountScope>,
context: &RoutingContext,
) -> Result<RoutingPlan, RoutingError> {
let requirements = operation.capability_requirements();
let mut candidates = Vec::new();
for provider in self.providers.iter() {
if context.required_provider.is_none() && account_scope.provider_kinds().is_empty() {
return Err(RoutingError::EmptyAccountScope);
}
let providers = context.required_provider.as_ref().map_or_else(
|| account_scope.provider_kinds().clone(),
|provider| BTreeSet::from([provider.clone()]),
);
for provider in &providers {
if !self.providers.contains(provider) {
continue;
}
if context
.provider_kind
.required_provider
.as_ref()
.is_some_and(|expected| expected != provider)
|| context.blocked_providers.contains(provider)
@@ -539,12 +760,14 @@ impl RuntimeSnapshot {
};
emulated
}
None if self.known_provider_catalogs.contains(provider) => continue,
None => BTreeSet::new(),
};
candidates.push(ProviderCandidate {
provider: provider.clone(),
upstream_model,
emulated_features,
account_scope: Arc::clone(&account_scope),
});
}
@@ -561,6 +784,7 @@ impl RuntimeSnapshot {
operation: operation.kind(),
max_attempts: NonZeroU32::new(super::MAX_REQUEST_ATTEMPTS)
.expect("constant request attempt limit is non-zero"),
account_scope,
candidates: Arc::from(candidates),
})
}
@@ -2,12 +2,14 @@ use gateway_core::engine::continuation::{
ContinuationBinding, NativeContinuationPin, PreviousResponseId,
};
use gateway_core::engine::credential::ProviderAccountId;
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::ProviderKind;
fn pin() -> NativeContinuationPin {
NativeContinuationPin::new(
PreviousResponseId::new("response-private"),
PreviousResponseId::new("upstream-private"),
ClientApiKeyId::new("key_client").expect("valid client key"),
ProviderKind::new("openai").expect("valid provider"),
ProviderAccountId::new("acct_codex").expect("valid account"),
)
@@ -43,3 +45,8 @@ fn native_pin_should_reject_different_account() {
&ProviderAccountId::new("acct_other").expect("valid account"),
));
}
#[test]
fn native_pin_should_reject_different_client_api_key() {
assert!(!pin().matches_client(&ClientApiKeyId::new("key_other").expect("valid client key")));
}
@@ -43,8 +43,9 @@ use gateway_core::operation::{
};
use gateway_core::policy::ClientApiKeyId;
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RoutingPlan, RuntimeSnapshot, UpstreamModelId,
AccountRoutingSnapshot, ClientRoutingScope, ConfigRevision, FrozenAccountScope,
ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext, RoutingPlan,
RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId,
};
use serde_json::{Map, Value, json};
@@ -531,7 +532,32 @@ fn plan_with_policy(
let public_model = PublicModelId::new("gpt-5").expect("public model");
let capabilities = ModelCapabilities::new(BTreeSet::from([operation.kind()]), Some(32_000))
.with_upstream_feature_validation();
RuntimeSnapshot::new(
let directory = Arc::new(RuntimeAccountDirectory::new(
[
"acct_failed",
"acct_first",
"acct_image",
"acct_observation_mismatch",
"acct_observed",
"acct_one",
"acct_only",
"acct_other",
"acct_required",
"acct_second",
"acct_tool",
"acct_two",
"acct_wrong",
]
.into_iter()
.map(|id| {
(
ProviderAccountId::new(id).expect("account"),
RuntimeAccount::new(provider.clone(), BTreeSet::new()),
)
})
.collect(),
));
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("config revision"),
account_selection_policy,
vec![provider.clone()],
@@ -543,15 +569,21 @@ fn plan_with_policy(
Vec::new(),
)
.expect("snapshot")
.plan(
&public_model,
operation,
&RoutingContext {
provider_kind: Some(provider),
..RoutingContext::default()
},
)
.expect("routing plan")
.with_account_directory(Arc::clone(&directory));
snapshot
.plan(
&public_model,
operation,
Arc::new(FrozenAccountScope::new(
directory,
ClientRoutingScope::all_accounts(),
)),
&RoutingContext {
required_provider: Some(provider),
..RoutingContext::default()
},
)
.expect("routing plan")
}
fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest {
@@ -561,6 +593,7 @@ fn model_request(operation: &Operation, deadline: SystemTime) -> NewModelRequest
client_api_key_id: Some(client_key.clone()),
client_api_key_ref: client_key,
config_revision: ConfigRevision::new(1).expect("config revision"),
routing: AccountRoutingSnapshot::all(),
protocol: "openai".to_owned(),
operation: operation.kind(),
endpoint: "responses".to_owned(),
@@ -1354,6 +1387,7 @@ fn authenticated_native_continuation_reaches_every_attempt_context() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-secret-id"),
PreviousResponseId::new("provider-native-id"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -1411,6 +1445,7 @@ fn streaming_native_continuation_delivers_started_event_before_later_output() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -1487,6 +1522,7 @@ fn native_continuation_replays_owner_before_safely_switching_account() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-secret-id"),
PreviousResponseId::new("provider-native-id"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_one").expect("account"),
);
@@ -1552,6 +1588,7 @@ fn unavailable_native_continuation_replays_with_the_configured_account_policy()
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-response"),
PreviousResponseId::new("upstream-response"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
original.clone(),
);
@@ -2498,6 +2535,7 @@ fn native_continuation_explicit_429_is_not_retried() {
let continuation = NativeContinuationPin::new(
PreviousResponseId::new("previous-secret-id"),
PreviousResponseId::new("provider-native-id"),
ClientApiKeyId::new("key_client_1").expect("client key"),
ProviderKind::new("openai").expect("provider"),
ProviderAccountId::new("acct_first").expect("account"),
);
@@ -1,18 +1,21 @@
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use serde_json::{Map, Value};
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountAvailability, AccountAvailabilityPolicy, AccountCandidate,
AccountAttemptFeedback, AccountCandidate, AccountConcurrencyLimit, AccountEligibilityPolicy,
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker,
AccountSelectionContext, AccountSelectionPolicy, AccountSelector, CredentialCasUpdate,
CredentialRevision, OpaqueProviderData, PlaintextCredential, PreferredAccountSelection,
ProviderAccount, ProviderAccountId, ProviderAccountIdentity, ProviderAccountUpdate,
RotationStrategy,
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(
@@ -24,7 +27,13 @@ fn account(id: &str) -> ProviderAccount {
CredentialRevision::new(1).expect("valid revision"),
Some(SystemTime::now() + Duration::from_secs(3600)),
)
.with_runtime_state(true, AccountAvailability::Ready)
.with_account_facts(
true,
CredentialState::Ready,
QuotaState::unknown(),
None,
None,
)
}
fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandidate {
@@ -35,13 +44,41 @@ fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandida
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: remaining,
quota_limit_reached: false,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
}
}
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(
@@ -53,7 +90,8 @@ fn context(strategy: RotationStrategy) -> AccountSelectionContext {
excluded_accounts: BTreeSet::new(),
preferred_account: None,
round_robin_cursor: 0,
availability: AccountAvailabilityPolicy::Enforce,
eligibility: AccountEligibilityPolicy::Enforce,
account_scope: None,
}
}
@@ -101,43 +139,99 @@ fn provider_account_identity_debug_should_redact_upstream_ids() {
}
#[test]
fn availability_should_round_trip_all_database_values() {
fn account_fact_values_should_round_trip() {
assert_eq!(
AccountAvailability::parse(AccountAvailability::QuotaExhausted.as_str()),
Some(AccountAvailability::QuotaExhausted)
CredentialState::parse(CredentialState::Banned.as_str()),
Some(CredentialState::Banned)
);
assert_eq!(
QuotaEvidence::parse("provider_denied"),
Some(QuotaEvidence::ProviderDenied)
);
}
#[test]
fn diagnostic_selection_bypasses_all_local_account_availability() {
fn exhausted_quota_refresh_is_due_at_reset_or_provider_fallback() {
let now = SystemTime::now();
let observed_at = now - Duration::from_secs(60);
let before_reset = QuotaState::exhausted(
QuotaEvidence::ProviderDenied,
observed_at,
Some(now + Duration::from_secs(1)),
);
let after_reset = QuotaState::exhausted(
QuotaEvidence::ProviderDenied,
observed_at,
Some(now - Duration::from_secs(1)),
);
let without_reset = QuotaState::exhausted(QuotaEvidence::ProviderDenied, observed_at, None);
assert!(!before_reset.exhaustion_refresh_due(now, Duration::from_secs(30)));
assert!(after_reset.exhaustion_refresh_due(now, Duration::from_secs(30)));
assert!(without_reset.exhaustion_refresh_due(now, Duration::from_secs(30)));
assert!(!without_reset.exhaustion_refresh_due(now, Duration::from_secs(120)));
}
#[test]
fn inconclusive_quota_observation_cannot_erase_confirmed_access() {
let now = SystemTime::now();
let exhausted = QuotaState::exhausted(QuotaEvidence::ProviderDenied, now, None);
assert_eq!(
exhausted
.merge_observation(QuotaState::observed_unknown(now + Duration::from_secs(1)))
.access(),
gateway_core::engine::credential::QuotaAccessState::Exhausted
);
assert_eq!(
exhausted
.merge_observation(QuotaState::allowed(now + Duration::from_secs(1)))
.access(),
gateway_core::engine::credential::QuotaAccessState::Allowed
);
}
#[test]
fn diagnostic_selection_bypasses_all_local_account_eligibility() {
let observed_at = SystemTime::now();
let exhausted = AccountCandidate {
account: account("acct_exhausted")
.with_runtime_state(true, AccountAvailability::QuotaExhausted),
account: account("acct_exhausted").with_account_facts(
true,
CredentialState::Ready,
QuotaState::exhausted(QuotaEvidence::ProviderDenied, observed_at, None),
None,
None,
),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
};
let disabled = AccountCandidate {
account: account("acct_disabled")
.with_runtime_state(false, AccountAvailability::QuotaExhausted),
account: account("acct_disabled").with_account_facts(
false,
CredentialState::Ready,
QuotaState::exhausted(QuotaEvidence::ProviderDenied, observed_at, None),
None,
None,
),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
};
let mut context = context(RotationStrategy::Sticky);
context.availability = AccountAvailabilityPolicy::BypassForDiagnostic;
context.eligibility = AccountEligibilityPolicy::BypassForDiagnostic;
assert_eq!(
AccountSelector
@@ -154,10 +248,63 @@ fn diagnostic_selection_bypasses_all_local_account_availability() {
}
#[test]
fn disabled_account_should_not_be_schedulable() {
let account = account("acct_disabled").with_runtime_state(false, AccountAvailability::Ready);
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(),
)));
assert!(!account.is_schedulable(SystemTime::now()));
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(
false,
CredentialState::Ready,
QuotaState::unknown(),
None,
None,
);
assert_eq!(
account.status_projection(SystemTime::now(), None).status,
AccountStatus::Disabled
);
}
#[test]
fn rate_limited_projection_carries_only_the_active_cooldown_deadline() {
let now = SystemTime::now();
let active_until = now + Duration::from_secs(60);
let account = account("acct_rate_limited");
let active = account.status_projection(now, Some(active_until));
assert_eq!(active.status, AccountStatus::RateLimited);
assert_eq!(active.rate_limited_until, Some(active_until));
let elapsed = account.status_projection(now, Some(now - Duration::from_secs(1)));
assert_eq!(elapsed.status, AccountStatus::Normal);
assert_eq!(elapsed.rate_limited_until, None);
}
#[test]
@@ -171,14 +318,20 @@ fn account_without_upstream_identity_should_stay_unknown_and_not_schedulable() {
CredentialRevision::new(1).expect("valid revision"),
Some(SystemTime::now() + Duration::from_secs(3600)),
)
.with_runtime_state(true, AccountAvailability::Ready);
.with_account_facts(
true,
CredentialState::Ready,
QuotaState::unknown(),
None,
None,
);
assert_eq!(
(
account.availability(),
account.is_schedulable(SystemTime::now())
account.credential_state(),
account.status_projection(SystemTime::now(), None).status
),
(AccountAvailability::Unknown, false)
(CredentialState::Unknown, AccountStatus::Error)
);
}
@@ -195,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));
@@ -308,15 +525,11 @@ fn provider_quota_overlay_should_preserve_store_concurrency_facts() {
last_started_at: Some(last_started_at),
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
}
.with_provider_quota(Some(AccountQuotaSignals::new(
Some(reset_at),
Some(75),
false,
)));
.with_provider_quota(Some(AccountQuotaSignals::new(Some(reset_at), Some(75))));
assert_eq!(
(
@@ -633,21 +846,31 @@ fn credential_cas_should_reject_refresh_schedule_without_refresh_token() {
}
#[test]
fn enforced_availability_excludes_terminal_states_despite_quota_signals() {
// availability 终态(Banned/Expired/Invalid)始终排除,
// 不受 quota 信号或 skip_exhausted 放宽影响。
fn enforced_eligibility_excludes_credential_errors_despite_quota_signals() {
let mut banned = candidate("acct_banned", 0, Some(100));
banned.account = banned
.account
.with_runtime_state(true, AccountAvailability::Banned);
banned.account = banned.account.with_account_facts(
true,
CredentialState::Banned,
QuotaState::unknown(),
None,
None,
);
let mut expired = candidate("acct_expired", 0, Some(100));
expired.account = expired
.account
.with_runtime_state(true, AccountAvailability::Expired);
expired.account = expired.account.with_account_facts(
true,
CredentialState::Expired,
QuotaState::unknown(),
None,
None,
);
let mut invalid = candidate("acct_invalid", 0, Some(100));
invalid.account = invalid
.account
.with_runtime_state(true, AccountAvailability::Invalid);
invalid.account = invalid.account.with_account_facts(
true,
CredentialState::Invalid,
QuotaState::unknown(),
None,
None,
);
let healthy = candidate("acct_healthy", 0, Some(100));
let candidates = vec![banned, expired, invalid, healthy];
@@ -659,14 +882,15 @@ fn enforced_availability_excludes_terminal_states_despite_quota_signals() {
}
#[test]
fn enforced_availability_allows_quota_exhausted_with_skip_exhausted_projection() {
// QuotaExhausted 是可恢复额度耗尽,skip_exhausted=false 的
// 放宽投影(provider 层)允许其参与调度;core 的 Enforce 策略本身
// 只排除终态。这里验证 Enforce 下 QuotaExhausted 仍被排除(严格模式)。
fn enforced_eligibility_excludes_authoritative_quota_exhaustion() {
let mut exhausted = candidate("acct_exhausted", 0, Some(100));
exhausted.account = exhausted
.account
.with_runtime_state(true, AccountAvailability::QuotaExhausted);
exhausted.account = exhausted.account.with_account_facts(
true,
CredentialState::Ready,
QuotaState::exhausted(QuotaEvidence::ProviderDenied, SystemTime::now(), None),
None,
None,
);
let healthy = candidate("acct_healthy", 0, Some(100));
let candidates = vec![exhausted, healthy];
@@ -676,3 +900,31 @@ fn enforced_availability_allows_quota_exhausted_with_skip_exhausted_projection()
assert_eq!(selected.candidate().account.id().as_str(), "acct_healthy");
}
#[test]
fn elapsed_quota_reset_does_not_fabricate_recovery() {
let now = SystemTime::now();
let mut exhausted = candidate("acct_exhausted", 0, Some(100));
exhausted.account = exhausted.account.with_account_facts(
true,
CredentialState::Ready,
QuotaState::exhausted(
QuotaEvidence::ProviderDenied,
now - Duration::from_secs(2),
Some(now - Duration::from_secs(1)),
),
None,
None,
);
let healthy = candidate("acct_healthy", 0, Some(100));
assert_eq!(
exhausted.account.status_projection(now, None).status,
AccountStatus::QuotaExhausted
);
let candidates = [exhausted, healthy];
let selected = AccountSelector
.select(&candidates, &context(RotationStrategy::Smart))
.expect("healthy candidate available");
assert_eq!(selected.candidate().account.id().as_str(), "acct_healthy");
}
@@ -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(),
);
@@ -92,7 +92,7 @@ fn scheduling_state_preserves_provider_neutral_signals() {
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: Some(7),
quota_limit_reached: false,
rate_limited_until: None,
failure_rate_basis_points: Some(125),
first_output_latency_ms: Some(250),
},
+168 -16
View File
@@ -1,15 +1,20 @@
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::Duration;
use gateway_core::engine::credential::{AccountSelectionPolicy, RotationStrategy};
use gateway_core::engine::credential::{
AccountSelectionPolicy, ProviderAccountId, RotationStrategy,
};
use gateway_core::operation::{
CapabilityRequirements, Feature, GenerateRequest, Operation, OperationKind, ProtocolPayload,
};
use gateway_core::policy::{ClientApiKeyId, ClientPolicy, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::{
ConfigRevision, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId, RoutingContext,
RuntimeSnapshot, UpstreamModelId,
AccountGroupId, AccountRoutingScopeKind, ClientRoutingScope, ConfigRevision,
FrozenAccountScope, ModelCapabilities, ProviderKind, ProviderModel, PublicModelId,
RoutingContext, RoutingGroupSnapshot, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot,
UpstreamModelId,
};
mod snapshot;
@@ -45,11 +50,34 @@ fn operation() -> Operation {
))
}
fn account_directory() -> Arc<RuntimeAccountDirectory> {
Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
(
ProviderAccountId::new("acct_openai").expect("account"),
RuntimeAccount::new(
ProviderKind::new("openai").expect("provider"),
BTreeSet::new(),
),
),
(
ProviderAccountId::new("acct_xai").expect("account"),
RuntimeAccount::new(ProviderKind::new("xai").expect("provider"), BTreeSet::new()),
),
])))
}
fn account_scope() -> Arc<FrozenAccountScope> {
Arc::new(FrozenAccountScope::new(
account_directory(),
ClientRoutingScope::all_accounts(),
))
}
fn client_policy(id: &str, plaintext: &str, enabled: bool) -> ClientPolicy {
ClientPolicy::new(
ClientApiKeyId::new(id).expect("client key ID"),
PlaintextClientApiKey::new(plaintext).expect("plaintext client key"),
ProviderKind::new("openai").expect("valid provider"),
account_scope(),
enabled,
RateLimits::unlimited(),
)
@@ -70,6 +98,7 @@ fn snapshot() -> RuntimeSnapshot {
Vec::new(),
)
.expect("snapshot")
.with_account_directory(account_directory())
.with_model_mappings(BTreeMap::from([
("gpt-5.4".to_owned(), "gpt-5.5".to_owned()),
("grok-latest".to_owned(), "grok-4.5".to_owned()),
@@ -81,6 +110,92 @@ fn config_revision_should_reject_zero() {
assert!(ConfigRevision::new(0).is_err());
}
#[test]
fn account_group_id_should_match_the_database_contract_exactly() {
let valid = "grp_0123456789abcdef0123456789abcdef";
assert_eq!(
AccountGroupId::new(valid).expect("valid group ID").as_str(),
valid
);
for invalid in [
"0123456789abcdef0123456789abcdef",
"grp_0123456789abcdef0123456789abcde",
"grp_0123456789abcdef0123456789abcdef0",
"grp_0123456789abcdef0123456789abcdeF",
"grp_0123456789abcdef0123456789abcdeg",
] {
assert!(
AccountGroupId::new(invalid).is_err(),
"invalid group ID was accepted: {invalid}"
);
}
}
#[test]
fn restricted_scope_should_use_enabled_group_union_and_preserve_all_bindings_for_history() {
let enabled =
AccountGroupId::new("grp_11111111111111111111111111111111").expect("enabled group");
let disabled =
AccountGroupId::new("grp_22222222222222222222222222222222").expect("disabled group");
let provider = ProviderKind::new("openai").expect("provider");
let allowed_account = ProviderAccountId::new("acct_allowed").expect("account");
let disabled_only_account = ProviderAccountId::new("acct_disabled").expect("account");
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([
(
allowed_account.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([enabled.clone()])),
),
(
disabled_only_account.clone(),
RuntimeAccount::new(provider.clone(), BTreeSet::from([disabled.clone()])),
),
])));
let scope = FrozenAccountScope::new(
directory,
ClientRoutingScope::restricted(
vec![
RoutingGroupSnapshot::new(enabled.clone(), "Enabled".to_owned()),
RoutingGroupSnapshot::new(disabled, "Disabled".to_owned()),
],
BTreeSet::from([enabled]),
BTreeSet::from([provider]),
)
.expect("restricted scope"),
);
assert!(scope.allows(&allowed_account));
assert!(!scope.allows(&disabled_only_account));
let history = scope.routing_snapshot();
assert_eq!(history.kind(), AccountRoutingScopeKind::Groups);
assert_eq!(history.groups_snapshot().len(), 2);
assert_eq!(history.groups_snapshot()[1].name(), "Disabled");
}
#[test]
fn restricted_scope_with_no_enabled_group_should_fail_closed() {
let group = AccountGroupId::new("grp_33333333333333333333333333333333").expect("group");
let account = ProviderAccountId::new("acct_group_member").expect("account");
let directory = Arc::new(RuntimeAccountDirectory::new(BTreeMap::from([(
account.clone(),
RuntimeAccount::new(
ProviderKind::new("openai").expect("provider"),
BTreeSet::from([group.clone()]),
),
)])));
let scope = FrozenAccountScope::new(
directory,
ClientRoutingScope::restricted(
vec![RoutingGroupSnapshot::new(group, "Disabled".to_owned())],
BTreeSet::new(),
BTreeSet::new(),
)
.expect("restricted scope"),
);
assert!(scope.provider_kinds().is_empty());
assert!(!scope.allows(&account));
}
#[test]
fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
let snapshot = RuntimeSnapshot::new(
@@ -98,7 +213,12 @@ fn snapshot_should_publish_only_enabled_plaintext_client_policies() {
let policies = snapshot.client_policies().collect::<Vec<_>>();
assert_eq!(policies.len(), 1);
assert_eq!(policies[0].key_id().as_str(), "key_enabled");
assert_eq!(policies[0].provider_kind().as_str(), "openai");
assert!(
policies[0]
.account_scope()
.provider_kinds()
.contains(&ProviderKind::new("openai").expect("provider"))
);
}
#[test]
@@ -137,8 +257,9 @@ fn selected_provider_should_use_global_model_mapping() {
.plan(
&PublicModelId::new("gpt-5.4").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
@@ -170,26 +291,37 @@ fn cyclic_model_mapping_should_fall_back_to_the_original_name() {
}
#[test]
fn global_mapping_should_apply_to_every_provider() {
fn known_provider_catalog_missing_mapped_model_should_be_filtered() {
let snapshot = snapshot();
let plan = snapshot
let error = snapshot
.plan(
&PublicModelId::new("gpt-5.4").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("xai").expect("provider")),
required_provider: Some(ProviderKind::new("xai").expect("provider")),
..RoutingContext::default()
},
)
.expect("transparent plan");
.expect_err("known xAI catalog does not contain the mapped OpenAI model");
assert_eq!(plan.candidates()[0].provider().as_str(), "xai");
assert_eq!(plan.candidates()[0].upstream_model().as_str(), "gpt-5.5");
assert!(matches!(
error,
gateway_core::error::RoutingError::NoCapableProvider { .. }
));
}
#[test]
fn unmapped_model_should_pass_through_unchanged() {
let snapshot = snapshot();
let snapshot = RuntimeSnapshot::new(
ConfigRevision::new(1).expect("revision"),
scheduling(),
vec![ProviderKind::new("openai").expect("provider")],
Vec::new(),
Vec::new(),
)
.expect("snapshot")
.with_account_directory(account_directory());
for requested in [
"future-openai-model".to_owned(),
format!("future-{}", "x".repeat(512)),
@@ -200,8 +332,9 @@ fn unmapped_model_should_pass_through_unchanged() {
.plan(
&PublicModelId::from_client_wire(requested.clone()).expect("client model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
@@ -229,8 +362,9 @@ fn blocked_provider_should_be_filtered() {
.plan(
&PublicModelId::new("gpt-5.5").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
blocked_providers: BTreeSet::from([ProviderKind::new("openai").expect("provider")]),
},
)
@@ -262,8 +396,9 @@ fn known_unsupported_operation_should_not_be_bypassed() {
.plan(
&PublicModelId::new("gpt-known-unsupported").expect("model"),
&operation(),
snapshot.all_account_scope(),
&RoutingContext {
provider_kind: Some(ProviderKind::new("openai").expect("provider")),
required_provider: Some(ProviderKind::new("openai").expect("provider")),
..RoutingContext::default()
},
)
@@ -316,3 +451,20 @@ fn unknown_model_should_be_hidden_when_provider_catalog_is_available() {
&ProviderKind::new("openai").expect("provider"),
));
}
#[test]
fn alias_should_only_be_available_from_a_provider_with_its_mapped_model() {
let snapshot = snapshot();
let alias = PublicModelId::new("grok-latest").expect("model");
assert!(!snapshot.contains_public_model_for_provider(
&alias,
&ProviderKind::new("openai").expect("provider"),
));
assert!(
snapshot.contains_public_model_for_provider(
&alias,
&ProviderKind::new("xai").expect("provider"),
)
);
}
@@ -19,7 +19,8 @@ use gateway_core::operation::OperationKind;
use gateway_core::policy::{ClientApiKeyId, PlaintextClientApiKey, RateLimits};
use gateway_core::routing::snapshot::{
RuntimeSnapshotCompileError, RuntimeSnapshotCompiler, RuntimeSnapshotHandle,
RuntimeSnapshotPublisher, SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts,
RuntimeSnapshotPublisher, SnapshotAccountGroupFacts, SnapshotAccountGroupMemberFacts,
SnapshotClientPolicyFacts, SnapshotControl, SnapshotFacts, SnapshotProviderAccountFacts,
SnapshotRevisionStream, SnapshotSettingsFacts, SnapshotStoreError, SnapshotStorePort,
SnapshotSubscriptionError, SnapshotSubscriptionPort, runtime_revision_needs_refresh,
};
@@ -348,9 +349,12 @@ fn facts(config_revision: u64, observed_current_revision: u64) -> SnapshotFacts
vec![SnapshotClientPolicyFacts::new(
ClientApiKeyId::new("key_one").expect("key ID"),
PlaintextClientApiKey::new("sk_test").expect("plaintext key"),
"alpha",
Vec::new(),
RateLimits::unlimited(),
)],
Vec::<SnapshotAccountGroupFacts>::new(),
Vec::<SnapshotProviderAccountFacts>::new(),
Vec::<SnapshotAccountGroupMemberFacts>::new(),
)
}
@@ -90,6 +90,11 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
pool.clone(),
Some(Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>),
)),
Arc::new(postgres::PgAccountGroupRepository::with_runtime_state(
pool.clone(),
credential_leases.clone(),
Arc::clone(&cooldowns) as Arc<dyn ProviderCooldownPort>,
)),
Arc::new(AdminAuthStoreAdapter {
security: postgres::PgAdminSecurityAuditRepository::new(pool.clone()),
settings: postgres::PgRuntimeSettingsRepository::new(pool.clone()),
@@ -0,0 +1,730 @@
//! 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, 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,
};
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 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 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 } => {
@@ -1,12 +1,15 @@
//! Pg 观测 adapter:实现 `ObservabilityRepository` 与 `AdminObservabilityStore`。
use std::sync::Arc;
use std::time::SystemTime;
use gateway_core::provider_ports::ProviderCooldownPort;
use std::sync::Arc;
use super::*;
use crate::postgres::{
PgProviderAccountRepository, ProviderAccountRepository, ProviderAccountSummary,
account_status_projection, load_rate_limited_until,
};
use crate::redis::{CredentialLeaseRepository as _, RedisCredentialLeaseRepository};
#[derive(Clone)]
@@ -21,73 +24,48 @@ impl PgObservabilityRepository {
Self { pool, cooldowns }
}
/// 账号池指标:SQL 聚合后,仅用 Redis 429 冷却重分类仍可调度的 ready 账号。
/// 冷却中从 `active` 或 `quota_exhausted` 移到 `rate_limited`;冷却不可用时保留 SQL 分类。
async fn provider_account_metrics_with_cooldowns(
/// 从账号事实和当前冷却一次性派生 Dashboard 五态;不在 SQL 中复制状态机。
async fn account_status_snapshot(
&self,
observed_at: DateTime<Utc>,
) -> StoreResult<ProviderAccountMetrics> {
let base = provider_account_metrics(&self.pool, observed_at).await?;
let Some(cooldowns) = &self.cooldowns else {
return Ok(base);
) -> StoreResult<(ProviderAccountMetrics, Vec<ProviderAccountSummary>)> {
let accounts = PgProviderAccountRepository::new(self.pool.clone())
.list_provider_accounts(None, true)
.await?;
let now = observed_at.into();
let rate_limited_until =
load_rate_limited_until(self.cooldowns.as_deref(), &accounts, now).await;
let mut metrics = ProviderAccountMetrics {
total: u64::try_from(accounts.len()).unwrap_or(u64::MAX),
..ProviderAccountMetrics::default()
};
let candidates = schedulable_metric_candidates(&self.pool, observed_at).await?;
if candidates.is_empty() {
return Ok(base);
}
// Redis 只能把 SQL 已归入 active/quota_exhausted 的候选移到冷却桶;
// 其余状态不重建,避免 Redis 失败或脏 ID 时丢失持久化计数。
let mut active = base.active;
let mut rate_limited = base.rate_limited;
let mut quota_exhausted = base.quota_exhausted;
let now = SystemTime::now();
for (account_id, revision, quota_reached) in candidates {
let Ok(account_id) =
gateway_core::engine::credential::ProviderAccountId::new(account_id)
else {
continue;
};
let cooling = cooldowns
.read(&account_id)
.await
.ok()
.flatten()
.filter(|cooldown| cooldown.credential_revision().get() as i64 == revision)
.filter(|cooldown| cooldown.until() > now)
.is_some();
if !cooling {
continue;
}
if quota_reached {
if quota_exhausted == 0 {
continue;
let mut normal_accounts = Vec::new();
for account in &accounts {
let projection = account_status_projection(
account,
now,
rate_limited_until.get(&account.id).copied(),
);
match projection.status {
gateway_core::engine::credential::AccountStatus::Normal => {
metrics.normal = metrics.normal.saturating_add(1);
normal_accounts.push(account.clone());
}
quota_exhausted -= 1;
} else {
if active == 0 {
continue;
gateway_core::engine::credential::AccountStatus::QuotaExhausted => {
metrics.quota_exhausted = metrics.quota_exhausted.saturating_add(1);
}
gateway_core::engine::credential::AccountStatus::RateLimited => {
metrics.rate_limited = metrics.rate_limited.saturating_add(1);
}
gateway_core::engine::credential::AccountStatus::Disabled => {
metrics.disabled = metrics.disabled.saturating_add(1);
}
gateway_core::engine::credential::AccountStatus::Error => {
metrics.error = metrics.error.saturating_add(1);
}
active -= 1;
}
rate_limited += 1;
}
Ok(ProviderAccountMetrics {
total: base.total,
enabled: base.enabled,
unavailable: rate_limited
+ quota_exhausted
+ base.expired
+ base.invalid
+ base.disabled
+ base.banned,
active,
rate_limited,
expired: base.expired,
invalid: base.invalid,
quota_exhausted,
disabled: base.disabled,
banned: base.banned,
})
Ok((metrics, normal_accounts))
}
}
@@ -120,14 +98,11 @@ impl ObservabilityRepository for PgObservabilityRepository {
async fn dashboard_summary(
&self,
range: ObservabilityRange,
observed_at: DateTime<Utc>,
) -> StoreResult<DashboardObservation> {
let filter = UsageRecordFilter::default();
let account_usage_range = ObservabilityRange::new(
range.end - TimeDelta::hours(ACCOUNT_USAGE_TIMELINE_HOURS),
range.end,
)?;
let account_usage_query =
ProviderAccountUsageQuery::recent(account_usage_range, DASHBOARD_ACCOUNT_LIMIT)?
ProviderAccountUsageQuery::recent(range, DASHBOARD_ACCOUNT_LIMIT)?
.with_hourly_request_buckets()?;
let recent_query = UsageRecordQuery {
range,
@@ -145,9 +120,7 @@ impl ObservabilityRepository for PgObservabilityRepository {
attempt_metrics(&self.pool, range, &filter),
dashboard_totals(&self.pool),
)?;
let provider_accounts = self
.provider_account_metrics_with_cooldowns(range.end)
.await?;
let (provider_accounts, _) = self.account_status_snapshot(observed_at).await?;
let (trend, account_usage, recent_requests) = futures::try_join!(
request_metric_series(&self.pool, range, &filter),
provider_account_usage(&self.pool, account_usage_query),
@@ -239,10 +212,11 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
async fn dashboard_summary(
&self,
range: admin_observability::TimeRange,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<admin_observability::DashboardObservation> {
let observation = self
.repository
.dashboard_summary(store_range(range)?)
.dashboard_summary(store_range(range)?, observed_at)
.await
.map_err(observability_error)?;
admin_dashboard_observation(observation)
@@ -252,28 +226,52 @@ impl AdminObservabilityStore for PgAdminObservabilityStore {
&self,
observed_at: DateTime<Utc>,
) -> AdminStoreResult<Option<admin_observability::DashboardRuntimeSlots>> {
let active_account_ids = active_provider_account_ids(&self.repository.pool, observed_at)
let (_, normal_accounts) = self
.repository
.account_status_snapshot(observed_at)
.await
.map_err(observability_error)?;
let active_accounts = u64::try_from(active_account_ids.len())
.map_err(|_| observability_error(invalid("active account count overflows u64")))?;
if active_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(&active_account_ids)
.credential_runtime_signals(&normal_account_ids)
.await
{
Ok(signals) => signals,
Err(_) => {
return Ok(Some(admin_observability::DashboardRuntimeSlots {
active_accounts,
inherited_accounts,
overridden_slots,
used_slots: None,
}));
}
@@ -282,7 +280,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),
}))
}
@@ -274,16 +274,11 @@ pub(crate) const fn admin_account_pool_metrics(
) -> admin_observability::AccountPoolMetrics {
admin_observability::AccountPoolMetrics {
total: metrics.total,
enabled: metrics.enabled,
unavailable: metrics.unavailable,
active: metrics.active,
rate_limited: metrics.rate_limited,
expired: metrics.expired,
invalid: metrics.invalid,
normal: metrics.normal,
quota_exhausted: metrics.quota_exhausted,
refreshing: None,
rate_limited: metrics.rate_limited,
disabled: metrics.disabled,
banned: metrics.banned,
error: metrics.error,
}
}
@@ -297,8 +292,6 @@ pub(crate) fn admin_dashboard_account_usage(
name: usage.name,
email: usage.email,
plan_type: usage.plan_type,
enabled: usage.enabled,
availability: usage.availability,
request_count: usage.request_count,
success_count: usage.success_count,
input_tokens: usage.input_tokens,
@@ -443,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,
@@ -662,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")?,
@@ -363,15 +363,11 @@ pub struct CalculatedUsageBillingFact {
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ProviderAccountMetrics {
pub total: u64,
pub enabled: u64,
pub unavailable: u64,
pub active: u64,
pub rate_limited: u64,
pub expired: u64,
pub invalid: u64,
pub normal: u64,
pub quota_exhausted: u64,
pub rate_limited: u64,
pub disabled: u64,
pub banned: u64,
pub error: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -388,8 +384,6 @@ pub struct ProviderAccountUsageObservation {
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
pub enabled: bool,
pub availability: String,
pub request_count: u64,
pub success_count: u64,
pub input_tokens: Option<u64>,
@@ -503,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,
@@ -689,6 +686,7 @@ pub trait ObservabilityRepository: Send + Sync {
async fn dashboard_summary(
&self,
range: ObservabilityRange,
observed_at: DateTime<Utc>,
) -> StoreResult<DashboardObservation>;
async fn dashboard_trend(
&self,
@@ -15,7 +15,7 @@ pub(crate) async fn provider_account_usage(
let mut statement = QueryBuilder::<Postgres>::new(
"select pa.id, pa.provider_kind, pa.authentication_kind, pa.name, pa.email,
pa.plan_type, pa.enabled, pa.availability,
pa.plan_type,
count(mr.id)::bigint as request_count,
count(mr.id) filter (where mr.outcome = 'succeeded')::bigint as success_count,
sum(mr.input_tokens)::bigint as input_tokens,
@@ -52,9 +52,12 @@ pub(crate) async fn provider_account_usage(
}
statement.push(
" group by pa.id, pa.provider_kind, pa.authentication_kind, pa.name, pa.email,
pa.plan_type, pa.enabled, pa.availability
order by max(mr.started_at) desc nulls last, pa.name, pa.id limit ",
pa.plan_type",
);
if query.account_ids.is_none() {
statement.push(" having count(mr.id) > 0");
}
statement.push(" order by max(mr.started_at) desc nulls last, pa.name, pa.id limit ");
statement.push_bind(i64::from(query.limit));
let rows = statement
.build()
@@ -313,8 +316,6 @@ pub(crate) fn provider_account_from_row(
name: get(row, "name")?,
email: get(row, "email")?,
plan_type: get(row, "plan_type")?,
enabled: get(row, "enabled")?,
availability: get(row, "availability")?,
request_count: unsigned(row, "request_count")?,
success_count: unsigned(row, "success_count")?,
input_tokens: optional_unsigned(row, "input_tokens")?,
@@ -369,118 +369,6 @@ pub(crate) async fn request_costs(
.collect()
}
pub(crate) async fn provider_account_metrics(
pool: &PgPool,
observed_at: DateTime<Utc>,
) -> StoreResult<ProviderAccountMetrics> {
let row = sqlx::query(
"with normalized as (
select enabled, availability, quota_limit_reached,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
when availability = 'quota_exhausted' then 'quota_exhausted'
when availability = 'invalid' then 'invalid'
when availability = 'expired'
or (access_token_expires_at is not null and access_token_expires_at <= $1)
then 'expired'
when availability = 'unknown' then 'invalid'
when quota_limit_reached then 'quota_exhausted'
else 'active'
end as status
from provider_accounts
)
select count(*)::bigint as total,
count(*) filter (where enabled)::bigint as enabled,
count(*) filter (where status in ('rate_limited','quota_exhausted','expired','invalid','disabled','banned'))::bigint as unavailable,
count(*) filter (where status = 'active')::bigint as active,
count(*) filter (where status = 'rate_limited')::bigint as rate_limited,
count(*) filter (where status = 'expired')::bigint as expired,
count(*) filter (where status = 'invalid')::bigint as invalid,
count(*) filter (where status = 'quota_exhausted')::bigint as quota_exhausted,
count(*) filter (where status = 'disabled')::bigint as disabled,
count(*) filter (where status = 'banned')::bigint as banned
from normalized",
)
.bind(observed_at)
.fetch_one(pool)
.await
.map_err(|_| postgres_unavailable("load provider account metrics"))?;
Ok(ProviderAccountMetrics {
total: unsigned(&row, "total")?,
enabled: unsigned(&row, "enabled")?,
unavailable: unsigned(&row, "unavailable")?,
active: unsigned(&row, "active")?,
rate_limited: unsigned(&row, "rate_limited")?,
expired: unsigned(&row, "expired")?,
invalid: unsigned(&row, "invalid")?,
quota_exhausted: unsigned(&row, "quota_exhausted")?,
disabled: unsigned(&row, "disabled")?,
banned: unsigned(&row, "banned")?,
})
}
pub(crate) async fn active_provider_account_ids(
pool: &PgPool,
observed_at: DateTime<Utc>,
) -> StoreResult<Vec<String>> {
sqlx::query(
"with normalized as (
select id,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
when availability = 'quota_exhausted' then 'quota_exhausted'
when availability in ('expired', 'invalid')
or (access_token_expires_at is not null and access_token_expires_at <= $1)
then 'expired'
else 'active'
end as status
from provider_accounts
)
select id
from normalized
where status = 'active'
order by id",
)
.bind(observed_at)
.fetch_all(pool)
.await
.map_err(|_| postgres_unavailable("load active provider account IDs"))?
.iter()
.map(|row| get(row, "id"))
.collect()
}
/// 查询可被 Redis 冷却重新归类的 ready 账号:id、credential revision 与 quota 标记。
/// 已过期或非 ready 的账号由 [`provider_account_metrics`] 直接计数。
pub(crate) async fn schedulable_metric_candidates(
pool: &PgPool,
observed_at: DateTime<Utc>,
) -> StoreResult<Vec<(String, i64, bool)>> {
sqlx::query(
"select id, credential_revision, quota_limit_reached
from provider_accounts
where enabled
and availability = 'ready'
and (access_token_expires_at is null or access_token_expires_at > $1)
order by id",
)
.bind(observed_at)
.fetch_all(pool)
.await
.map_err(|_| postgres_unavailable("load schedulable metric candidates"))?
.into_iter()
.map(|row| {
Ok((
get(&row, "id")?,
get(&row, "credential_revision")?,
get(&row, "quota_limit_reached")?,
))
})
.collect::<StoreResult<Vec<_>>>()
}
pub(crate) async fn provider_observations(
pool: &PgPool,
range: ObservabilityRange,
@@ -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,
@@ -2,7 +2,6 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::SystemTime;
use gateway_core::provider_ports::ProviderCooldownPort;
@@ -32,34 +31,6 @@ impl PgAdminAccountStore {
}
}
/// 批量读取 Redis 429 冷却,返回 `account_id → rate_limited_until` 映射。
/// 冷却不可用(Redis 挂/无 cooldown 端口)时返回空映射,账号归入配额耗尽(近似)。
async fn load_rate_limited_until(
&self,
accounts: &[ProviderAccountSummary],
) -> AdminStoreResult<BTreeMap<String, SystemTime>> {
let Some(cooldowns) = &self.cooldowns else {
return Ok(BTreeMap::new());
};
let mut map = BTreeMap::new();
for account in accounts {
let Ok(account_id) = CoreProviderAccountId::new(account.id.clone()) else {
continue;
};
let Some(cooldown) = cooldowns.read(&account_id).await.ok().flatten() else {
continue;
};
if cooldown.credential_revision().get() != account.credential_revision.get() {
continue;
}
let until = cooldown.until();
if until > SystemTime::now() {
map.insert(account.id.clone(), until);
}
}
Ok(map)
}
async fn usage_observations(
&self,
range: ObservabilityRange,
@@ -126,22 +97,59 @@ impl PgAdminAccountStore {
.iter()
.map(|window| window.account_id.clone())
.collect::<Vec<_>>();
sqlx::query(ACCOUNT_USAGE_BY_WINDOWS_SQL)
.bind(account_ids)
.bind(keys)
.bind(starts)
.bind(ends)
.fetch_all(&self.pool)
.await
.map_err(|_| {
admin_store_error(
ENTITY,
postgres_unavailable("load provider account quota window usage"),
)
})?
let (usage_rows, model_rows, model_cost_rows) = futures::try_join!(
sqlx::query(ACCOUNT_USAGE_BY_WINDOWS_SQL)
.bind(account_ids.clone())
.bind(keys.clone())
.bind(starts.clone())
.bind(ends.clone())
.fetch_all(&self.pool),
sqlx::query(ACCOUNT_USAGE_MODELS_BY_WINDOWS_SQL)
.bind(account_ids.clone())
.bind(keys.clone())
.bind(starts.clone())
.bind(ends.clone())
.fetch_all(&self.pool),
sqlx::query(ACCOUNT_USAGE_MODEL_COSTS_BY_WINDOWS_SQL)
.bind(account_ids)
.bind(keys)
.bind(starts)
.bind(ends)
.fetch_all(&self.pool),
)
.map_err(|_| {
admin_store_error(
ENTITY,
postgres_unavailable("load provider account quota window usage"),
)
})?;
let mut model_costs = BTreeMap::<(String, String, String), Vec<AccountCost>>::new();
for row in &model_cost_rows {
let (key, cost) = admin_account_usage_window_model_cost(row)?;
model_costs.entry(key).or_default().push(cost);
}
let mut models_by_window = BTreeMap::<(String, String), Vec<AccountModelUsage>>::new();
for row in &model_rows {
let ((account_id, window_key, model), mut usage) =
admin_account_usage_window_model(row)?;
usage.costs = model_costs
.remove(&(account_id.clone(), window_key.clone(), model))
.unwrap_or_default();
models_by_window
.entry((account_id, window_key))
.or_default()
.push(usage);
}
let mut results = usage_rows
.iter()
.map(admin_account_usage_window)
.collect()
.collect::<AdminStoreResult<Vec<_>>>()?;
for result in &mut results {
result.usage.models = models_by_window
.remove(&(result.account_id.clone(), result.key.clone()))
.unwrap_or_default();
}
Ok(results)
}
async fn required_scope(
@@ -260,6 +268,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]
@@ -278,24 +391,26 @@ impl AccountStore for PgAdminAccountStore {
)
.map_err(|error| admin_store_error(ENTITY, error))?;
let now = Utc::now();
// 批量读取 Redis 429 冷却,供列表状态派生与筛选使用。
let rate_limited_until = self
.load_rate_limited_until(&accounts)
.await
.unwrap_or_default();
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 status = admin_account_status(&account, now, until);
account_matches_admin_query(&account, status, &query).then_some(
AdminAccountListItem {
account,
status,
usage: None,
},
)
let projection = account_status_projection(&account, now.into(), until);
(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<_>>();
@@ -330,7 +445,23 @@ impl AccountStore for PgAdminAccountStore {
.into_iter()
.skip(offset)
.take(page_size)
.map(|item| admin_account_record(item.account))
.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,
projection: item.projection,
})
})
.collect::<AdminStoreResult<Vec<_>>>()?;
Ok(AccountPage {
config_revision: admin_revision(control_plane.settings.config_revision)?,
@@ -340,13 +471,35 @@ impl AccountStore for PgAdminAccountStore {
})
}
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountRecord>> {
self.accounts
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountPageItem>> {
let record = self
.accounts
.load_provider_account(account_id)
.await
.map_err(|error| admin_store_error(ENTITY, error))?
.map(|record| admin_account_record(record.summary))
.transpose()
.map_err(|error| admin_store_error(ENTITY, error))?;
let Some(record) = record else {
return Ok(None);
};
let now = Utc::now();
let rate_limited_until = load_rate_limited_until(
self.cooldowns.as_deref(),
std::slice::from_ref(&record.summary),
now.into(),
)
.await
.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,
projection,
}))
}
async fn load_account_usage(
@@ -386,9 +539,9 @@ impl AccountStore for PgAdminAccountStore {
.filter(|account| account.provider_kind == provider_kind.as_str())
.filter(|account| {
query
.availability
.credential_state
.as_ref()
.is_none_or(|expected| expected.matches(account.availability))
.is_none_or(|expected| expected.matches(account.credential_state))
})
.filter(|account| {
query
@@ -583,29 +736,134 @@ 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 recover_account(
&self,
account_id: &CoreProviderAccountId,
context: &MutationContext,
) -> AdminStoreResult<AccountUpdateResult> {
if let Some(cooldowns) = self.cooldowns.as_deref() {
cooldowns.clear_all(account_id).await.map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
ENTITY,
"provider account cooldown cleanup failed",
)
})?;
}
let config_revision = self
.accounts
.recover_provider_account_admin(RecoverProviderAccount {
account_id: account_id.as_str().to_owned(),
audit: mutation_audit(
context,
"recover",
"provider_account",
account_id.as_str(),
vec!["status".to_owned(), "quota".to_owned()],
),
})
.await
.map_err(|error| admin_store_error(ENTITY, error))
.and_then(admin_revision)?;
Ok(AccountUpdateResult {
config_revision,
account_id: account_id.clone(),
})
}
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,8 +31,10 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.map(DateTime::<Utc>::from),
next_refresh_at: account.account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.account.enabled(),
availability: account.account.availability(),
availability_observed_at: Utc::now(),
concurrency_limit: account.account.concurrency_limit(),
weight: account.account.weight(),
credential_state: account.account.credential_state(),
credential_observed_at: Utc::now(),
})
.await
.map_err(core_store_error)
@@ -164,9 +166,9 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.collect::<Vec<_>>();
let rows = sqlx::query(
"select id, credential_revision, provider_quota_json, quota_observed_at, \
quota_limit_reached \
quota_access_state, quota_evidence, quota_access_observed_at, quota_reset_at \
from provider_accounts \
where id = any($1) and provider_quota_json is not null",
where id = any($1) and quota_observed_at is not null",
)
.bind(account_ids)
.fetch_all(&self.pool)
@@ -184,24 +186,49 @@ impl ProviderAccountStore for PgProviderAccountRepository {
.and_then(|value| CoreCredentialRevision::new(value).ok())
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let quota = row
.try_get::<serde_json::Value, _>("provider_quota_json")
.ok()
.try_get::<Option<serde_json::Value>, _>("provider_quota_json")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?
.and_then(|value| value.as_object().cloned())
.map(OpaqueProviderData::new)
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
.unwrap_or_else(|| OpaqueProviderData::new(serde_json::Map::new()));
let observed_at = row
.try_get::<DateTime<Utc>, _>("quota_observed_at")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let limit_reached = row
.try_get::<bool, _>("quota_limit_reached")
let access = row
.try_get::<String, _>("quota_access_state")
.ok()
.and_then(|value| QuotaAccessState::parse(&value))
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let evidence = row
.try_get::<Option<String>, _>("quota_evidence")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let evidence = match evidence {
Some(value) => Some(
QuotaEvidence::parse(&value)
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?,
),
None => None,
};
let reset_at = row
.try_get::<Option<DateTime<Utc>>, _>("quota_reset_at")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let access_observed_at = row
.try_get::<Option<DateTime<Utc>>, _>("quota_access_observed_at")
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
let state = QuotaState::from_persisted(
access,
evidence,
access_observed_at.map(Into::into),
reset_at.map(Into::into),
)
.ok_or_else(|| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?;
Ok(QuotaObservation {
account_id: CoreProviderAccountId::new(account_id)
.map_err(|_| CoreStoreError::new(CoreStoreErrorKind::InvalidData))?,
expected_revision: revision,
quota: Some(quota),
observed_at: Some(observed_at.into()),
limit_reached: Some(limit_reached),
quota,
observed_at: observed_at.into(),
state,
})
})
.collect()
@@ -211,24 +238,38 @@ impl ProviderAccountStore for PgProviderAccountRepository {
&self,
observation: QuotaObservation,
) -> Result<QuotaWriteOutcome, CoreStoreError> {
let quota = observation
.quota
.map(|value| {
JsonObject::try_from_value(
"provider_quota_json",
serde_json::Value::Object(value.into_inner()),
QUOTA_MAX_BYTES,
)
})
.transpose()
.map_err(core_store_error)?;
let quota = JsonObject::try_from_value(
"provider_quota_json",
serde_json::Value::Object(observation.quota.into_inner()),
QUOTA_MAX_BYTES,
)
.map_err(core_store_error)?;
let updated = self
.compare_and_swap_provider_quota(
observation.account_id.as_str(),
Revision::new(observation.expected_revision.get()).map_err(core_store_error)?,
quota,
observation.observed_at.map(DateTime::<Utc>::from),
observation.limit_reached,
DateTime::<Utc>::from(observation.observed_at),
observation.state,
)
.await
.map_err(core_store_error)?;
Ok(if updated {
QuotaWriteOutcome::Updated
} else {
QuotaWriteOutcome::Conflict
})
}
async fn apply_quota_access(
&self,
change: QuotaAccessChange,
) -> Result<QuotaWriteOutcome, CoreStoreError> {
let updated = self
.apply_provider_quota_access(
change.account_id.as_str(),
Revision::new(change.expected_revision.get()).map_err(core_store_error)?,
change.state,
)
.await
.map_err(core_store_error)?;
@@ -245,8 +286,9 @@ impl ProviderAccountStore for PgProviderAccountRepository {
account_id: change.account_id.as_str().to_owned(),
expected_revision: Revision::new(change.expected_revision.get())
.map_err(core_store_error)?,
availability: change.availability,
availability_observed_at: DateTime::<Utc>::from(change.observed_at),
credential_state: change.credential_state,
credential_observed_at: DateTime::<Utc>::from(change.observed_at),
error_reason: change.error_reason,
message: change.message,
})
.await
@@ -1,12 +1,12 @@
//! Admin 账号列表、用量与成本的领域映射。
use std::time::SystemTime;
use super::*;
use gateway_admin::model::accounts::AccountStatusSignals;
pub(crate) struct AdminAccountListItem {
pub(crate) account: ProviderAccountSummary,
pub(crate) status: AdminAccountStatus,
pub(crate) projection: gateway_core::engine::credential::AccountStatusProjection,
pub(crate) usage: Option<ProviderAccountUsageObservation>,
}
@@ -57,24 +57,6 @@ pub(crate) fn account_matches_admin_query(
provider_matches && search_matches && status_matches
}
pub(crate) fn admin_account_status(
account: &ProviderAccountSummary,
now: DateTime<Utc>,
rate_limited_until: Option<SystemTime>,
) -> AdminAccountStatus {
// 派生规则唯一归属 gateway-admin 域层,此处只做类型适配。
AccountStatusSignals {
enabled: account.enabled,
availability: account.availability,
access_token_expired: account
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now),
quota_limit_reached: account.quota_limit_reached,
rate_limited_until,
}
.derive()
}
pub(crate) fn admin_account_summary(
accounts: &[ProviderAccountSummary],
now: DateTime<Utc>,
@@ -90,7 +72,7 @@ pub(crate) fn admin_account_summary(
};
for account in accounts {
let until = rate_limited_until.get(&account.id).copied();
match admin_account_status(account, now, until) {
match account_status_projection(account, now.into(), until).status {
AdminAccountStatus::Normal => summary.normal = summary.normal.saturating_add(1),
AdminAccountStatus::QuotaExhausted => {
summary.quota_exhausted = summary.quota_exhausted.saturating_add(1);
@@ -118,7 +100,11 @@ pub(crate) fn sort_admin_account_items(
items.sort_by(|left, right| {
let ordering = match sort.field {
AdminAccountSortField::Email => left.account.email.cmp(&right.account.email),
AdminAccountSortField::Status => left.status.as_str().cmp(right.status.as_str()),
AdminAccountSortField::Status => left
.projection
.status
.sort_rank()
.cmp(&right.projection.status.sort_rank()),
AdminAccountSortField::PlanType => left.account.plan_type.cmp(&right.account.plan_type),
AdminAccountSortField::Usage => left
.usage
@@ -160,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,
@@ -171,10 +158,13 @@ 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,
availability: summary.availability,
availability_observed_at: summary.availability_observed_at,
concurrency_limit: summary.concurrency_limit,
weight: summary.weight,
credential_state: summary.credential_state,
credential_observed_at: summary.credential_observed_at,
quota: summary.quota,
last_error_reason: summary.last_error_reason,
last_error_message: summary.last_error_message,
quota_observed_at: summary.quota_observed_at,
created_at: summary.created_at,
updated_at: summary.updated_at,
})
@@ -197,8 +187,10 @@ pub(crate) fn prepared_account(
access_token_expires_at: credential.access_token_expires_at,
next_refresh_at: credential.next_refresh_at,
enabled: credential.enabled,
availability: credential.availability,
availability_observed_at: credential.availability_observed_at,
concurrency_limit: None,
weight: AccountWeight::DEFAULT,
credential_state: credential.credential_state,
credential_observed_at: credential.credential_observed_at,
})
}
@@ -284,6 +276,69 @@ pub(crate) fn admin_account_usage_window(
})
}
pub(crate) fn admin_account_usage_window_model(
row: &sqlx::postgres::PgRow,
) -> AdminStoreResult<((String, String, String), AccountModelUsage)> {
let account_id: String = window_usage_value(row, "account_id")?;
let window_key: String = window_usage_value(row, "window_key")?;
let model: String = window_usage_value(row, "model")?;
let provider_reported_count = window_usage_count(row, "provider_reported_count")?;
let calculated_count = window_usage_count(row, "calculated_count")?;
let unavailable_count = window_usage_count(row, "unavailable_count")?;
Ok((
(account_id, window_key, model.clone()),
AccountModelUsage {
model,
request_count: window_usage_count(row, "request_count")?,
success_count: window_usage_count(row, "success_count")?,
input_tokens: optional_window_usage_count(row, "input_tokens")?,
output_tokens: optional_window_usage_count(row, "output_tokens")?,
cached_tokens: optional_window_usage_count(row, "cached_tokens")?,
cache_write_tokens: optional_window_usage_count(row, "cache_write_tokens")?,
reasoning_tokens: optional_window_usage_count(row, "reasoning_tokens")?,
image_input_tokens: optional_window_usage_count(row, "image_input_tokens")?,
image_output_tokens: optional_window_usage_count(row, "image_output_tokens")?,
image_request_count: window_usage_count(row, "image_request_count")?,
image_request_failed_count: window_usage_count(row, "image_request_failed_count")?,
total_tokens: optional_window_usage_count(row, "total_tokens")?,
cost_coverage: AdminCostCoverage {
provider_reported_count,
calculated_count,
partial_count: 0,
unavailable_count,
not_billable_count: 0,
},
costs: Vec::new(),
last_used_at: window_usage_value(row, "last_used_at")?,
},
))
}
pub(crate) fn admin_account_usage_window_model_cost(
row: &sqlx::postgres::PgRow,
) -> AdminStoreResult<((String, String, String), AccountCost)> {
let key = (
window_usage_value(row, "account_id")?,
window_usage_value(row, "window_key")?,
window_usage_value(row, "model")?,
);
let amount = window_usage_value::<String>(row, "amount")?;
let amount = AdminDecimalAmount::from_str(&amount).map_err(|_| {
AdminStoreError::new(
AdminStoreErrorKind::Invalid,
ENTITY,
"persisted account model cost is invalid",
)
})?;
Ok((
key,
AccountCost {
currency: window_usage_value(row, "cost_currency")?,
amount,
},
))
}
pub(crate) fn window_usage_count(
row: &sqlx::postgres::PgRow,
column: &'static str,
@@ -4,7 +4,6 @@ use std::{
collections::{BTreeMap, BTreeSet},
fmt,
str::FromStr,
time::SystemTime,
};
use async_trait::async_trait;
@@ -12,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,
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,
@@ -35,15 +37,16 @@ use gateway_admin::{
use sqlx::{PgPool, Postgres, Row, Transaction};
use gateway_core::engine::credential::{
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision as CoreCredentialRevision, LoadedCredential,
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
ProviderAccountIdentity, ProviderAccountStore,
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome,
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,
@@ -63,8 +66,10 @@ mod core_adapter;
mod mapping;
mod repository;
mod rows;
mod runtime;
pub use admin_adapter::*;
pub(crate) use mapping::*;
pub use repository::*;
pub use rows::*;
pub(crate) use runtime::*;
@@ -25,9 +25,15 @@ pub trait ProviderAccountRepository: Send + Sync {
&self,
account_id: &str,
expected_revision: Revision,
quota: Option<JsonObject>,
observed_at: Option<DateTime<Utc>>,
limit_reached: Option<bool>,
quota: JsonObject,
observed_at: DateTime<Utc>,
state: QuotaState,
) -> StoreResult<bool>;
async fn apply_provider_quota_access(
&self,
account_id: &str,
expected_revision: Revision,
state: QuotaState,
) -> StoreResult<bool>;
async fn delete_provider_account(&self, id: &str) -> StoreResult<bool>;
}
@@ -50,9 +56,14 @@ 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 recover_provider_account_admin(
&self,
command: RecoverProviderAccount,
) -> StoreResult<Revision>;
async fn delete_provider_accounts_admin(
@@ -93,9 +104,10 @@ 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, availability,
availability_observed_at,
quota_observed_at, quota_limit_reached, last_error_message, created_at, updated_at
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
from provider_accounts
where ($1::text is null or provider_kind = $1) and ($2 or enabled)
order by provider_kind, name, id",
@@ -110,21 +122,21 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
async fn insert_provider_account(&self, account: NewProviderAccount) -> StoreResult<()> {
account.validate()?;
let availability = if account.upstream_user_id.is_some() {
account.availability
let credential_state = if account.upstream_user_id.is_some() {
account.credential_state
} else {
AccountAvailability::Unknown
CredentialState::Unknown
};
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,
availability, provider_quota_json,
availability_observed_at, quota_observed_at, created_at, updated_at
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, now(), greatest(now(), $15)
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
)",
)
.bind(account.id)
@@ -140,8 +152,10 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
.bind(account.access_token_expires_at)
.bind(account.next_refresh_at)
.bind(account.enabled)
.bind(availability.as_str())
.bind(account.availability_observed_at)
.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)
.await
.map_err(|_| postgres_unavailable("insert provider account"))?;
@@ -214,27 +228,33 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
update.validate()?;
let result = sqlx::query(
"update provider_accounts
set availability = case
set credential_state = case
when enabled and upstream_user_id is not null then $3
else availability
else credential_state
end,
availability_observed_at = case
credential_observed_at = case
when enabled and upstream_user_id is not null then $4
else availability_observed_at
else credential_observed_at
end,
last_error_reason = case
when enabled and upstream_user_id is not null then
case when $3 = 'ready' then null else $5 end
else last_error_reason
end,
last_error_message = case
when enabled and upstream_user_id is not null then
case when $3 = 'ready' then null else coalesce($5, last_error_message) end
case when $3 = 'ready' then null else $6 end
else last_error_message
end,
updated_at = case when enabled then greatest(now(), $4) else updated_at end
where id = $1 and credential_revision = $2
and (availability_observed_at is null or availability_observed_at <= $4)",
and (credential_observed_at is null or credential_observed_at <= $4)",
)
.bind(update.account_id)
.bind(to_i64(update.expected_revision.get())?)
.bind(update.availability.as_str())
.bind(update.availability_observed_at)
.bind(update.credential_state.as_str())
.bind(update.credential_observed_at)
.bind(update.error_reason.map(AccountErrorReason::as_str))
.bind(update.message.as_deref())
.execute(&self.pool)
.await
@@ -259,36 +279,81 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
&self,
account_id: &str,
expected_revision: Revision,
quota: Option<JsonObject>,
observed_at: Option<DateTime<Utc>>,
limit_reached: Option<bool>,
quota: JsonObject,
observed_at: DateTime<Utc>,
state: QuotaState,
) -> StoreResult<bool> {
require_nonempty(ENTITY, "account_id", account_id)?;
if quota.is_some() != observed_at.is_some() {
return Err(invalid("quota JSON and observed_at must agree"));
}
if let Some(quota) = quota.as_ref() {
validate_object_size("provider_quota_json", quota, QUOTA_MAX_BYTES)?;
}
validate_object_size("provider_quota_json", &quota, QUOTA_MAX_BYTES)?;
let access_observed_at = state.observed_at().map(DateTime::<Utc>::from);
let result = sqlx::query(
"update provider_accounts
set provider_quota_json = $3, quota_observed_at = $4,
quota_limit_reached = coalesce($5, quota_limit_reached),
updated_at = greatest(now(), coalesce($4, now()))
quota_access_state = case
when $7::timestamptz is not null
and (quota_access_observed_at is null or quota_access_observed_at <= $7)
then $5 else quota_access_state end,
quota_evidence = case
when $7::timestamptz is not null
and (quota_access_observed_at is null or quota_access_observed_at <= $7)
then $6 else quota_evidence end,
quota_access_observed_at = case
when $7::timestamptz is not null
and (quota_access_observed_at is null or quota_access_observed_at <= $7)
then $7 else quota_access_observed_at end,
quota_reset_at = case
when $7::timestamptz is not null
and (quota_access_observed_at is null or quota_access_observed_at <= $7)
then $8 else quota_reset_at end,
updated_at = greatest(now(), $4)
where id = $1 and credential_revision = $2
and ($4 is null or quota_observed_at is null or quota_observed_at <= $4)",
and (quota_observed_at is null or quota_observed_at <= $4)",
)
.bind(account_id)
.bind(to_i64(expected_revision.get())?)
.bind(quota.map(|value| value.as_value()))
.bind(quota.as_value())
.bind(observed_at)
.bind(limit_reached)
.bind(state.access().as_str())
.bind(state.evidence().map(QuotaEvidence::as_str))
.bind(access_observed_at)
.bind(state.reset_at().map(DateTime::<Utc>::from))
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("compare and swap provider quota"))?;
Ok(result.rows_affected() == 1)
}
async fn apply_provider_quota_access(
&self,
account_id: &str,
expected_revision: Revision,
state: QuotaState,
) -> StoreResult<bool> {
require_nonempty(ENTITY, "account_id", account_id)?;
let observed_at = state
.observed_at()
.map(DateTime::<Utc>::from)
.ok_or_else(|| invalid("quota access change requires observed_at"))?;
let result = sqlx::query(
"update provider_accounts
set quota_access_observed_at = $3, quota_access_state = $4,
quota_evidence = $5, quota_reset_at = $6,
updated_at = greatest(now(), $3)
where id = $1 and credential_revision = $2
and (quota_access_observed_at is null or quota_access_observed_at <= $3)",
)
.bind(account_id)
.bind(to_i64(expected_revision.get())?)
.bind(observed_at)
.bind(state.access().as_str())
.bind(state.evidence().map(QuotaEvidence::as_str))
.bind(state.reset_at().map(DateTime::<Utc>::from))
.execute(&self.pool)
.await
.map_err(|_| postgres_unavailable("apply provider quota access"))?;
Ok(result.rows_affected() == 1)
}
async fn delete_provider_account(&self, id: &str) -> StoreResult<bool> {
require_nonempty(ENTITY, "id", id)?;
let result = sqlx::query("delete from provider_accounts where id = $1 and not enabled")
@@ -413,12 +478,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()
@@ -426,11 +491,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)
@@ -441,6 +513,57 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
finish_admin_transaction(transaction, result, "provider account admin state change").await
}
async fn recover_provider_account_admin(
&self,
command: RecoverProviderAccount,
) -> StoreResult<Revision> {
validate_admin_account_ids(std::slice::from_ref(&command.account_id))?;
let mut transaction = self
.pool
.begin()
.await
.map_err(|_| postgres_unavailable("begin provider account admin recovery"))?;
let result = async {
let revision = bump_config_revision_in_transaction(&mut transaction).await?;
let recovered = sqlx::query_scalar::<_, String>(
"update provider_accounts
set enabled = true,
credential_state = 'ready',
credential_observed_at = now(),
access_token_expires_at = case
when access_token_expires_at <= now() then null
else access_token_expires_at
end,
provider_quota_json = null,
quota_observed_at = null,
quota_access_state = 'allowed',
quota_evidence = null,
quota_access_observed_at = now(),
quota_reset_at = null,
last_error_reason = null,
last_error_message = null,
updated_at = now()
where id = $1
returning id",
)
.bind(&command.account_id)
.fetch_optional(&mut *transaction)
.await
.map_err(|_| postgres_unavailable("recover provider account state"))?
.ok_or_else(|| StoreError::NotFound {
entity: ENTITY,
id: command.account_id.clone(),
})?;
append_admin_audit_event_in_transaction(&mut transaction, command.audit, revision)
.await?;
Ok((revision, recovered))
}
.await;
finish_admin_transaction(transaction, result, "provider account admin recovery")
.await
.map(|(revision, _)| revision)
}
async fn delete_provider_accounts_admin(
&self,
command: DeleteProviderAccounts,
@@ -469,26 +592,73 @@ 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,
) -> StoreResult<String> {
account.validate()?;
let availability = if account.upstream_user_id.is_some() {
account.availability
let credential_state = if account.upstream_user_id.is_some() {
account.credential_state
} else {
AccountAvailability::Unknown
CredentialState::Unknown
};
let imported_id = sqlx::query_scalar::<_, String>(
"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,
availability, provider_quota_json,
availability_observed_at, quota_observed_at, created_at, updated_at
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, now(), greatest(now(), $15)
$14, $15, $16, null, $17, null, null, now(), greatest(now(), $17)
)
on conflict (
provider_kind,
@@ -505,11 +675,17 @@ pub(crate) async fn upsert_provider_account_in_transaction(
access_token_expires_at = excluded.access_token_expires_at,
next_refresh_at = excluded.next_refresh_at,
enabled = excluded.enabled,
availability = excluded.availability,
credential_state = excluded.credential_state,
provider_quota_json = null,
availability_observed_at = excluded.availability_observed_at,
quota_access_state = 'unknown',
quota_evidence = null,
quota_access_observed_at = null,
quota_reset_at = null,
credential_observed_at = excluded.credential_observed_at,
quota_observed_at = null,
updated_at = greatest(now(), excluded.availability_observed_at)
last_error_reason = null,
last_error_message = null,
updated_at = greatest(now(), excluded.credential_observed_at)
returning id",
)
.bind(&account.id)
@@ -525,8 +701,10 @@ pub(crate) async fn upsert_provider_account_in_transaction(
.bind(account.access_token_expires_at)
.bind(account.next_refresh_at)
.bind(account.enabled)
.bind(availability.as_str())
.bind(account.availability_observed_at)
.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)
.await
.map_err(|error| {
@@ -574,17 +752,17 @@ pub(crate) async fn rotate_provider_account_in_transaction(
next_refresh_at = $10,
upstream_user_id = case when $11::boolean then $12::text else upstream_user_id end,
upstream_account_id = case when $11::boolean then $13::text else upstream_account_id end,
availability = case
when not enabled then availability
when availability = 'quota_exhausted' then availability
credential_state = case
when not enabled then credential_state
when coalesce($12::text, upstream_user_id) is not null then 'ready'
else 'unknown'
end,
availability_observed_at = case
when not enabled then availability_observed_at
when availability <> 'quota_exhausted' then now()
else availability_observed_at
credential_observed_at = case
when not enabled then credential_observed_at
else now()
end,
last_error_reason = case when enabled then null else last_error_reason end,
last_error_message = case when enabled then null else last_error_message end,
updated_at = now()
where id = $1 and provider_kind = $2
and credential_revision = $3
@@ -627,23 +805,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(
@@ -685,20 +877,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(
@@ -715,6 +893,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;
@@ -50,6 +51,74 @@ pub(crate) const ACCOUNT_USAGE_BY_WINDOWS_SQL: &str = "with requested_windows as
group by requested.account_id, requested.window_key
order by requested.account_id, requested.window_key";
pub(crate) const ACCOUNT_USAGE_MODELS_BY_WINDOWS_SQL: &str = "with requested_windows as (
select *
from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[])
as requested(account_id, window_key, window_start, window_end)
)
select requested.account_id,
requested.window_key,
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
count(mr.id)::bigint as request_count,
count(mr.id) filter (where mr.outcome = 'succeeded')::bigint as success_count,
sum(mr.input_tokens)::bigint as input_tokens,
sum(mr.output_tokens)::bigint as output_tokens,
sum(mr.cached_tokens)::bigint as cached_tokens,
sum(mr.cache_write_tokens)::bigint as cache_write_tokens,
sum(mr.reasoning_tokens)::bigint as reasoning_tokens,
sum(mr.image_input_tokens)::bigint as image_input_tokens,
sum(mr.image_output_tokens)::bigint as image_output_tokens,
count(mr.id) filter (where mr.image_generation_succeeded is true)::bigint
as image_request_count,
count(mr.id) filter (where mr.image_generation_succeeded is false)::bigint
as image_request_failed_count,
coalesce(sum(coalesce(
mr.total_tokens,
coalesce(mr.input_tokens, 0) + coalesce(mr.output_tokens, 0)
)), 0)::bigint as total_tokens,
count(mr.id) filter (where mr.cost_source = 'provider_reported')::bigint
as provider_reported_count,
count(mr.id) filter (where mr.cost_source = 'calculated')::bigint
as calculated_count,
count(mr.id) filter (where mr.cost_source = 'unavailable')::bigint
as unavailable_count,
max(mr.started_at) as last_used_at
from requested_windows requested
join model_requests mr
on mr.provider_account_ref = requested.account_id
and mr.started_at >= requested.window_start
and mr.started_at < requested.window_end
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
group by requested.account_id, requested.window_key,
coalesce(mr.upstream_model_id, mr.requested_model_id)
order by requested.account_id, requested.window_key, request_count desc, model";
pub(crate) const ACCOUNT_USAGE_MODEL_COSTS_BY_WINDOWS_SQL: &str = "with requested_windows as (
select *
from unnest($1::text[], $2::text[], $3::timestamptz[], $4::timestamptz[])
as requested(account_id, window_key, window_start, window_end)
)
select requested.account_id,
requested.window_key,
coalesce(mr.upstream_model_id, mr.requested_model_id) as model,
mr.cost_currency,
sum(mr.cost_amount)::text as amount
from requested_windows requested
join model_requests mr
on mr.provider_account_ref = requested.account_id
and mr.started_at >= requested.window_start
and mr.started_at < requested.window_end
and mr.outcome = 'succeeded'
and mr.downstream_committed_at is not null
and mr.client_status_code between 200 and 399
and mr.cost_amount is not null
and mr.cost_currency is not null
group by requested.account_id, requested.window_key,
coalesce(mr.upstream_model_id, mr.requested_model_id), mr.cost_currency
order by requested.account_id, requested.window_key, model, mr.cost_currency";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountAdminScope {
pub provider_kind: String,
@@ -65,9 +134,29 @@ impl ProviderAccountAdminScope {
}
}
/// 解析数据库 availability 列;字符串映射本体唯一归属 gateway-core。
pub(crate) fn parse_availability(value: &str) -> StoreResult<AccountAvailability> {
AccountAvailability::parse(value).ok_or_else(|| invalid("unknown availability value"))
pub(crate) fn parse_credential_state(value: &str) -> StoreResult<CredentialState> {
CredentialState::parse(value).ok_or_else(|| invalid("unknown credential_state value"))
}
pub(crate) fn parse_quota_access_state(value: &str) -> StoreResult<QuotaAccessState> {
QuotaAccessState::parse(value).ok_or_else(|| invalid("unknown quota_access_state value"))
}
pub(crate) fn parse_quota_evidence(value: Option<String>) -> StoreResult<Option<QuotaEvidence>> {
value
.map(|value| {
QuotaEvidence::parse(&value).ok_or_else(|| invalid("unknown quota_evidence value"))
})
.transpose()
}
pub(crate) fn parse_error_reason(value: Option<String>) -> StoreResult<Option<AccountErrorReason>> {
value
.map(|value| {
AccountErrorReason::parse(&value)
.ok_or_else(|| invalid("unknown last_error_reason value"))
})
.transpose()
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -85,12 +174,12 @@ pub struct ProviderAccountSummary {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
/// 快照级限流事实(Provider 物化)。
pub quota_limit_reached: bool,
/// 最近一次非 ready 状态变更的原始原因(上游错误原文 / 原因码);恢复 ready 后清空。
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
pub quota: QuotaState,
pub last_error_reason: Option<AccountErrorReason>,
pub last_error_message: Option<String>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
@@ -132,8 +221,10 @@ pub struct NewProviderAccount {
pub access_token_expires_at: Option<DateTime<Utc>>,
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
pub concurrency_limit: Option<AccountConcurrencyLimit>,
pub weight: AccountWeight,
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -152,7 +243,7 @@ impl fmt::Debug for NewProviderAccount {
.field("provider_kind", &self.provider_kind)
.field("name", &self.name)
.field("enabled", &self.enabled)
.field("availability", &self.availability)
.field("credential_state", &self.credential_state)
.field("provider_credentials_json", &"[REDACTED]")
.finish_non_exhaustive()
}
@@ -256,10 +347,18 @@ 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,
}
#[derive(Debug, Clone)]
pub struct RecoverProviderAccount {
pub account_id: String,
pub audit: AdminAuditEvent,
}
@@ -297,24 +396,13 @@ impl fmt::Debug for ProviderCredentialUpdate {
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderAccountObservation {
pub account_id: String,
pub availability: AccountAvailability,
pub provider_quota_json: Option<JsonObject>,
pub availability_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
/// 快照级限流事实(Provider 物化)。
pub quota_limit_reached: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountStateUpdate {
pub account_id: String,
pub expected_revision: Revision,
pub availability: AccountAvailability,
pub availability_observed_at: DateTime<Utc>,
/// 上游原因(错误原文 / 原因码);`None` 不更新错误信息列,恢复 `Ready` 时清空。
pub credential_state: CredentialState,
pub credential_observed_at: DateTime<Utc>,
pub error_reason: Option<AccountErrorReason>,
pub message: Option<String>,
}
@@ -325,31 +413,20 @@ impl ProviderAccountStateUpdate {
}
}
impl ProviderAccountObservation {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if self.provider_quota_json.is_some() != self.quota_observed_at.is_some() {
return Err(invalid("quota JSON and quota_observed_at must agree"));
}
if let Some(quota) = &self.provider_quota_json {
validate_object_size("provider_quota_json", quota, QUOTA_MAX_BYTES)?;
}
Ok(())
}
}
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, availability,
provider_quota_json, quota_limit_reached, last_error_message,
availability_observed_at, quota_observed_at, created_at, updated_at
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
from provider_accounts where id = $1";
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, availability,
provider_quota_json, quota_limit_reached, last_error_message,
availability_observed_at, quota_observed_at, created_at, updated_at
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
from provider_accounts
where id = any($1::text[]) and provider_kind = $2
order by id";
@@ -398,7 +475,14 @@ pub(crate) fn core_account_from_summary(
summary.upstream_account_id,
summary.plan_type,
)
.with_runtime_state(summary.enabled, summary.availability)
.with_account_facts(
summary.enabled,
summary.credential_state,
summary.quota,
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),
@@ -430,9 +514,32 @@ pub(crate) fn account_summary_from_row(
let revision = row
.try_get::<i64, _>("credential_revision")
.map_err(|_| invalid("invalid credential revision"))?;
let availability = row
.try_get::<String, _>("availability")
.map_err(|_| invalid("invalid availability"))?;
let credential_state = row
.try_get::<String, _>("credential_state")
.map_err(|_| invalid("invalid credential_state"))?;
let quota_access_state = parse_quota_access_state(&get::<String>(&row, "quota_access_state")?)?;
let quota_evidence = parse_quota_evidence(get(&row, "quota_evidence")?)?;
let quota_access_observed_at = get::<Option<DateTime<Utc>>>(&row, "quota_access_observed_at")?;
let quota_reset_at = get::<Option<DateTime<Utc>>>(&row, "quota_reset_at")?;
let quota = QuotaState::from_persisted(
quota_access_state,
quota_evidence,
quota_access_observed_at.map(Into::into),
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")?,
@@ -447,10 +554,12 @@ 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")?,
availability: parse_availability(&availability)?,
availability_observed_at: get(&row, "availability_observed_at")?,
quota_observed_at: get(&row, "quota_observed_at")?,
quota_limit_reached: get(&row, "quota_limit_reached")?,
concurrency_limit,
weight,
credential_state: parse_credential_state(&credential_state)?,
credential_observed_at: get(&row, "credential_observed_at")?,
quota,
last_error_reason: parse_error_reason(get(&row, "last_error_reason")?)?,
last_error_message: get(&row, "last_error_message")?,
created_at: get(&row, "created_at")?,
updated_at: get(&row, "updated_at")?,
@@ -0,0 +1,53 @@
//! 持久账号事实与可丢失冷却事实的统一状态投影。
use std::{collections::BTreeMap, time::SystemTime};
use futures::future::join_all;
use gateway_core::{
engine::credential::{
AccountStatusFacts, AccountStatusProjection, ProviderAccountId, resolve_account_status,
},
provider_ports::ProviderCooldownPort,
};
use super::ProviderAccountSummary;
#[must_use]
pub(crate) fn account_status_projection(
account: &ProviderAccountSummary,
now: SystemTime,
rate_limited_until: Option<SystemTime>,
) -> AccountStatusProjection {
resolve_account_status(
&AccountStatusFacts {
enabled: account.enabled,
credential_state: account.credential_state,
access_token_expires_at: account.access_token_expires_at.map(Into::into),
quota: account.quota,
rate_limited_until,
last_error_reason: account.last_error_reason,
last_error_message: account.last_error_message.clone(),
},
now,
)
}
/// 尽力读取账号当前有效的 429 冷却。
///
/// 冷却属于可丢失的账号级运行时事实:端口不可用、单条读取失败或脏账号 ID
/// 都视为没有冷却,不污染持久账号状态。凭据轮换不会解除上游对同一账号的限流。
pub(crate) async fn load_rate_limited_until(
cooldowns: Option<&dyn ProviderCooldownPort>,
accounts: &[ProviderAccountSummary],
now: SystemTime,
) -> BTreeMap<String, SystemTime> {
let Some(cooldowns) = cooldowns else {
return BTreeMap::new();
};
let reads = accounts.iter().map(|account| async move {
let account_id = ProviderAccountId::new(account.id.clone()).ok()?;
let cooldown = cooldowns.read(&account_id).await.ok().flatten()?;
(cooldown.until() > now).then(|| (account.id.clone(), cooldown.until()))
});
join_all(reads).await.into_iter().flatten().collect()
}
@@ -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,
},
};
@@ -31,7 +33,7 @@ pub struct SnapshotRuntimeSettings {
pub struct CredentialRevisionVectorEntry {
pub provider_account_id: String,
pub credential_revision: Revision,
pub availability_observed_at: DateTime<Utc>,
pub credential_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
}
@@ -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,
})
}
@@ -108,7 +138,7 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
async fn credential_revision_vector(&self) -> StoreResult<Vec<CredentialRevisionVectorEntry>> {
let rows = sqlx::query_as::<_, (String, i64, DateTime<Utc>, Option<DateTime<Utc>>)>(
"select id, credential_revision, availability_observed_at, quota_observed_at
"select id, credential_revision, credential_observed_at, quota_observed_at
from provider_accounts order by id",
)
.fetch_all(&self.pool)
@@ -119,7 +149,7 @@ impl RuntimeSnapshotRepository for PgRuntimeSnapshotRepository {
Ok(CredentialRevisionVectorEntry {
provider_account_id: row.0,
credential_revision: revision_from_i64(row.1)?,
availability_observed_at: row.2,
credential_observed_at: row.2,
quota_observed_at: row.3,
})
})
@@ -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"))
}
@@ -512,7 +524,7 @@ impl RedisProviderLeaseCoordinator {
last_started_at: signal.last_started_at.map(Into::into),
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
rate_limited_until: None,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
@@ -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,4 +1,4 @@
//! Provider account availability 的可重建 Redis cache fence。
//! Provider credential state 的可重建 Redis cache fence。
use std::time::Duration;
@@ -7,7 +7,7 @@ use chrono::{DateTime, Utc};
use redis::{Script, aio::ConnectionManager};
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, OpaqueProviderData, ProviderAccountId,
CredentialRevision, CredentialState, OpaqueProviderData, ProviderAccountId,
};
use gateway_core::provider_ports::{
ProviderCatalogCacheKey, ProviderCatalogCachePort, ProviderCredentialState,
@@ -24,7 +24,7 @@ if current >= tonumber(ARGV[1]) then return 0 end
redis.call('HSET', KEYS[1],
'revision', ARGV[1],
'enabled', ARGV[2],
'availability', ARGV[3],
'credential_state', ARGV[3],
'observed_at_ms', ARGV[4])
redis.call('PEXPIRE', KEYS[1], ARGV[5])
return 1
@@ -44,7 +44,7 @@ pub struct CredentialStateCache {
pub provider_account_id: String,
pub revision: Revision,
pub enabled: bool,
pub availability: String,
pub credential_state: String,
pub observed_at: DateTime<Utc>,
}
@@ -55,8 +55,8 @@ impl CredentialStateCache {
"provider_account_id",
&self.provider_account_id,
)?;
if AccountAvailability::parse(&self.availability).is_none() {
return Err(invalid("availability is not registered"));
if CredentialState::parse(&self.credential_state).is_none() {
return Err(invalid("credential_state is not registered"));
}
Ok(())
}
@@ -172,7 +172,7 @@ impl CredentialStateRepository for RedisCredentialStateRepository {
.key(self.key(&state.provider_account_id)?)
.arg(state.revision.get())
.arg(u8::from(state.enabled))
.arg(&state.availability)
.arg(&state.credential_state)
.arg(state.observed_at.timestamp_millis())
.arg(ttl_ms)
.invoke_async::<i64>(&mut connection)
@@ -195,12 +195,12 @@ impl CredentialStateRepository for RedisCredentialStateRepository {
.arg(self.key(provider_account_id)?)
.arg("revision")
.arg("enabled")
.arg("availability")
.arg("credential_state")
.arg("observed_at_ms")
.query_async(&mut connection)
.await
.map_err(|_| redis_unavailable("read credential state"))?;
let (Some(revision), Some(enabled), Some(availability), Some(observed_at_ms)) = values
let (Some(revision), Some(enabled), Some(credential_state), Some(observed_at_ms)) = values
else {
return Ok(None);
};
@@ -208,7 +208,7 @@ impl CredentialStateRepository for RedisCredentialStateRepository {
provider_account_id: provider_account_id.to_owned(),
revision: Revision::new(revision)?,
enabled: enabled == 1,
availability,
credential_state,
observed_at: DateTime::from_timestamp_millis(observed_at_ms)
.ok_or_else(|| invalid("observed_at is invalid"))?,
};
@@ -299,7 +299,7 @@ impl ProviderCredentialStatePort for RedisCredentialStateRepository {
revision: Revision::new(state.credential_revision().get())
.map_err(|_| provider_invalid("encode credential state"))?,
enabled: state.enabled(),
availability: state.availability().as_str().to_owned(),
credential_state: state.credential_state().as_str().to_owned(),
observed_at,
};
CredentialStateRepository::cache_credential_state(
@@ -329,13 +329,13 @@ impl ProviderCredentialStatePort for RedisCredentialStateRepository {
.map_err(|_| provider_invalid("decode credential state"))?;
let revision = CredentialRevision::new(state.revision.get())
.map_err(|_| provider_invalid("decode credential state"))?;
let availability = AccountAvailability::parse(&state.availability)
let credential_state = CredentialState::parse(&state.credential_state)
.ok_or_else(|| provider_invalid("decode credential state"))?;
Ok(ProviderCredentialState::new(
account_id,
revision,
state.enabled,
availability,
credential_state,
state.observed_at.into(),
))
})
@@ -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,299 @@
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 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,
@@ -55,8 +58,8 @@ async fn merged_first_attempt_insert_should_match_sequential_semantics() {
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, enabled, availability,
availability_observed_at, created_at, updated_at
has_refresh_token, access_token_expires_at, enabled, credential_state,
credential_observed_at, created_at, updated_at
) values (
'acct_merged', 'openai', 'merged', null, 'user-merged', null, null, 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', $1, $1, $1
@@ -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]
@@ -134,26 +134,27 @@ async fn dashboard_account_metrics_should_partition_account_statuses() {
"insert into provider_accounts (
id, provider_kind, name, upstream_user_id, authentication_kind,
provider_credentials_json, credential_revision, has_refresh_token,
access_token_expires_at, enabled, availability,
availability_observed_at, created_at, updated_at
access_token_expires_at, enabled, credential_state,
quota_access_state, quota_evidence, quota_access_observed_at,
credential_observed_at, created_at, updated_at
) values
('acct_available', 'openai', 'available', 'user-available', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_expired', 'openai', 'expired', 'user-expired', 'oauth',
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready',
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_rate_cooldown', 'openai', 'rate-cooldown', 'user-rate-cooldown', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_usage_limit', 'xai', 'usage-limit', 'user-usage-limit', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'quota_exhausted',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'exhausted', 'usage_limit_reached', $1,
$1, $1, $1),
('acct_banned', 'xai', 'banned', 'user-banned', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned', 'allowed', null, $1,
$1, $1, $1),
('acct_disabled', 'xai', 'disabled', 'user-disabled', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready',
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', 'allowed', null, $1,
$1, $1, $1)",
)
.bind(now)
@@ -165,7 +166,7 @@ async fn dashboard_account_metrics_should_partition_account_statuses() {
.expect("dashboard range");
let metrics = repository
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("dashboard summary")
.provider_accounts;
@@ -173,22 +174,22 @@ async fn dashboard_account_metrics_should_partition_account_statuses() {
assert_eq!(
(
metrics.total,
metrics.enabled,
metrics.active,
metrics.unavailable,
metrics.expired,
metrics.normal,
metrics.quota_exhausted,
metrics.rate_limited,
metrics.disabled,
metrics.banned,
metrics.error,
),
(6, 5, 2, 4, 1, 1, 1, 1),
(6, 2, 1, 0, 1, 2),
);
// 首页按调度可用性聚合:配额耗尽也属于不可用。
assert_eq!(
metrics.unavailable,
metrics.quota_exhausted + metrics.expired + metrics.disabled + metrics.banned
metrics.total,
metrics.normal
+ metrics.quota_exhausted
+ metrics.rate_limited
+ metrics.disabled
+ metrics.error
);
assert_eq!(metrics.total, metrics.active + metrics.unavailable);
database.close().await;
}
@@ -202,35 +203,36 @@ async fn dashboard_account_metrics_with_cooldowns_should_only_reclassify_eligibl
"insert into provider_accounts (
id, provider_kind, name, upstream_user_id, authentication_kind,
provider_credentials_json, credential_revision, has_refresh_token,
access_token_expires_at, enabled, availability, quota_limit_reached,
availability_observed_at, created_at, updated_at
access_token_expires_at, enabled, credential_state,
quota_access_state, quota_evidence, quota_access_observed_at,
credential_observed_at, created_at, updated_at
) values
('acct_metrics_active', 'openai', 'active', 'user-active', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', false,
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_metrics_active_cooling', 'openai', 'active-cooling', 'user-active-cooling', 'oauth',
'{}'::jsonb, 2, false, $1 + interval '1 day', true, 'ready', false,
'{}'::jsonb, 2, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_metrics_ready_quota', 'openai', 'ready-quota', 'user-ready-quota', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', true,
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'exhausted', 'provider_denied', $1,
$1, $1, $1),
('acct_metrics_cooling', 'openai', 'cooling', 'user-cooling', 'oauth',
'{}'::jsonb, 3, false, $1 + interval '1 day', true, 'ready', true,
'{}'::jsonb, 3, false, $1 + interval '1 day', true, 'ready', 'exhausted', 'provider_denied', $1,
$1, $1, $1),
('acct_metrics_stale_cooldown', 'openai', 'stale-cooldown', 'user-stale-cooldown', 'oauth',
'{}'::jsonb, 4, false, $1 + interval '1 day', true, 'ready', true,
'{}'::jsonb, 4, false, $1 + interval '1 day', true, 'ready', 'allowed', null, $1,
$1, $1, $1),
('acct_metrics_persistent_quota', 'openai', 'persistent-quota', 'user-persistent-quota', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'quota_exhausted', true,
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', 'exhausted', 'provider_denied', $1,
$1, $1, $1),
('acct_metrics_expired', 'openai', 'expired', 'user-expired', 'oauth',
'{}'::jsonb, 5, false, $1 - interval '1 second', true, 'ready', true,
'{}'::jsonb, 5, false, $1 - interval '1 second', true, 'ready', 'exhausted', 'provider_denied', $1,
$1, $1, $1),
('acct_metrics_unknown', 'openai', 'unknown', 'user-unknown', 'oauth',
'{}'::jsonb, 6, false, $1 + interval '1 day', true, 'unknown', false,
'{}'::jsonb, 6, false, $1 + interval '1 day', true, 'unknown', 'unknown', null, null,
$1, $1, $1),
('acct_metrics_disabled', 'openai', 'disabled', 'user-disabled', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', false,
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', 'allowed', null, $1,
$1, $1, $1)",
)
.bind(now)
@@ -250,7 +252,7 @@ async fn dashboard_account_metrics_with_cooldowns_should_only_reclassify_eligibl
.expect("dashboard range");
let metrics = repository
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("dashboard summary")
.provider_accounts;
@@ -258,36 +260,21 @@ async fn dashboard_account_metrics_with_cooldowns_should_only_reclassify_eligibl
assert_eq!(
(
metrics.total,
metrics.enabled,
metrics.active,
metrics.normal,
metrics.rate_limited,
metrics.quota_exhausted,
metrics.expired,
metrics.invalid,
metrics.disabled,
metrics.banned,
metrics.unavailable,
metrics.error,
),
(9, 8, 1, 2, 3, 1, 1, 1, 0, 8),
(9, 1, 2, 3, 1, 2),
);
assert_eq!(
metrics.total,
metrics.active
metrics.normal
+ metrics.rate_limited
+ metrics.quota_exhausted
+ metrics.expired
+ metrics.invalid
+ metrics.disabled
+ metrics.banned
);
assert_eq!(
metrics.unavailable,
metrics.rate_limited
+ metrics.quota_exhausted
+ metrics.expired
+ metrics.invalid
+ metrics.disabled
+ metrics.banned
+ metrics.error
);
database.close().await;
}
@@ -429,7 +416,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
let store = PgAdminObservabilityStore::new(database.pool.clone(), None, None);
let dashboard = store
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("admin dashboard summary");
assert_eq!(dashboard.range, range);
@@ -446,7 +433,7 @@ async fn admin_observability_adapter_preserves_utc_queries_metrics_costs_and_det
assert_eq!(dashboard.attempts.costs[0].amount.as_str(), "1.25");
assert_eq!(dashboard.provider_accounts.total, 1);
assert_eq!(dashboard.account_usage[0].request_count, 1);
assert_eq!(dashboard.account_usage[0].request_buckets.len(), 24);
assert_eq!(dashboard.account_usage[0].request_buckets.len(), 2);
assert_eq!(
dashboard.account_usage[0]
.request_buckets
@@ -563,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(),
@@ -597,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,
@@ -715,7 +714,7 @@ async fn dashboard_summary_totals_include_history_outside_selected_range() {
.expect("seed observability facts");
sqlx::query(
"update model_requests
set started_at = $1 - interval '2 days'
set started_at = $1 - interval '2 hours'
where id = 'req_observe_success'",
)
.bind(now)
@@ -727,10 +726,22 @@ async fn dashboard_summary_totals_include_history_outside_selected_range() {
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
let dashboard = repository
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("dashboard summary");
assert_eq!(dashboard.requests.request_count, 2);
assert!(dashboard.account_usage.is_empty());
let explicit_account_usage = repository
.provider_account_usage(
ProviderAccountUsageQuery::for_accounts(range, vec!["acct_observe".to_owned()])
.expect("account usage query"),
)
.await
.expect("provider account usage");
assert_eq!(explicit_account_usage.len(), 1);
assert_eq!(explicit_account_usage[0].request_count, 0);
assert_eq!(
(
dashboard.totals.request_count,
@@ -766,7 +777,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
let repository = PgObservabilityRepository::new(database.pool.clone(), None);
let dashboard = repository
.dashboard_summary(range)
.dashboard_summary(range, now)
.await
.expect("dashboard summary");
assert_eq!(dashboard.requests.request_count, 3);
@@ -814,7 +825,7 @@ async fn observability_queries_preserve_request_account_cost_and_diagnostic_fact
);
assert_eq!(dashboard.provider_accounts.total, 1);
assert_eq!(dashboard.account_usage[0].request_count, 1);
assert_eq!(dashboard.account_usage[0].request_buckets.len(), 24);
assert_eq!(dashboard.account_usage[0].request_buckets.len(), 2);
assert_eq!(
dashboard.account_usage[0]
.request_buckets
@@ -1013,8 +1024,8 @@ async fn seed_observability_facts(
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, enabled, availability,
availability_observed_at, created_at, updated_at
has_refresh_token, access_token_expires_at, enabled, credential_state,
credential_observed_at, created_at, updated_at
) values (
'acct_observe', 'openai', 'primary', 'account@example.invalid',
'user-observe', null, 'pro', 'oauth', '{}'::jsonb, 1, false, $1 + interval '1 day',
@@ -1034,7 +1045,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',
@@ -1042,7 +1054,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)
@@ -1064,7 +1077,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',
@@ -1075,7 +1089,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)
@@ -1093,7 +1108,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',
@@ -1102,7 +1118,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)
@@ -1144,7 +1161,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',
@@ -1152,7 +1170,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)
@@ -1167,7 +1186,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',
@@ -1175,7 +1194,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)
File diff suppressed because it is too large Load Diff
@@ -108,7 +108,7 @@ async fn insert_refreshable_account(
"insert into provider_accounts (
id, provider_kind, name, upstream_user_id, authentication_kind,
provider_credentials_json, has_refresh_token, access_token_expires_at,
next_refresh_at, availability, availability_observed_at, created_at, updated_at
next_refresh_at, credential_state, credential_observed_at, created_at, updated_at
) values ($1, 'openai', $1, $1, 'oauth', '{}'::jsonb, true, $2, $3,
'ready', now(), now(), now())",
)
@@ -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,
@@ -722,8 +725,8 @@ async fn seed_account(
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, enabled, availability,
availability_observed_at, created_at, updated_at
has_refresh_token, access_token_expires_at, enabled, credential_state,
credential_observed_at, created_at, updated_at
) values (
$1, 'openai', $2, $3, 'user-' || $1, null, null, $4,
'{}'::jsonb, 1, false, $5 + interval '1 day', true, 'ready',
@@ -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)
@@ -19,7 +19,7 @@ fn credential_state_rejects_provider_specific_status() {
provider_account_id: "account-1".to_owned(),
revision: Revision::new(1).expect("positive revision"),
enabled: true,
availability: "codex_special".to_owned(),
credential_state: "codex_special".to_owned(),
observed_at: Utc::now(),
};
assert!(state.validate().is_err());
@@ -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}:*"))
+33 -71
View File
@@ -19,15 +19,14 @@ use gateway_admin::model::provider_credentials::{
PreparedCredentialCreate, PreparedCredentialImport, PreparedCredentialRotation,
PreparedCredentialRotationFacts, ProviderDocument, ProviderExport,
ProviderExportCredentialInput, ProviderModel, ProviderModels, ProviderQuota,
ProviderQuotaRequest, ProviderQuotaWindow,
ProviderQuotaRequest, ProviderQuotaWindow, QuotaLocalUsageAttribution,
};
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
use gateway_core::accounting::Money;
use gateway_core::engine::credential::{
AccountAvailability, CredentialRevision, LoadedCredential, NewProviderAccount,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountStore,
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
};
use gateway_core::error::StoreErrorKind;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
@@ -161,6 +160,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
if account_ids.is_empty() {
return;
}
self.quota.invalidate_scheduling(account_ids);
if let Err(error) = self.catalog.invalidate() {
tracing::warn!(
account_count = account_ids.len(),
@@ -369,7 +369,15 @@ impl ProviderAdmin for OpenAiAdminProvider {
if !account_matches_record(&current.account, &command.account) {
return Err(provider_admin_error(ProviderAdminErrorKind::Conflict));
}
let secret = rotation_secret(command.provider_material)?;
let mut secret = rotation_secret(command.provider_material)?;
if secret.id_token.is_none() {
let runtime = CodexCredentialCodec::decode(&current.credential)
.map_err(|_| provider_admin_error(ProviderAdminErrorKind::Invalid))?;
secret.id_token = runtime
.authentication
.oauth()
.and_then(|oauth| oauth.id_token.clone());
}
let prepared = CodexCredentialAdmin
.prepare_refreshed_oauth_rotation(current, secret, None, None)
.map_err(map_credential_admin_error)?;
@@ -424,14 +432,7 @@ impl ProviderAdmin for OpenAiAdminProvider {
.await
.map_err(map_quota_error)?
};
let mut quota = project_quota(snapshot, &account);
quota.rate_limited_until = self
.quota
.rate_limited_until(account.id(), account.revision())
.await
.map_err(map_quota_error)?
.map(DateTime::<Utc>::from);
Ok(quota)
Ok(project_quota(snapshot, &account))
}
async fn models(
@@ -519,7 +520,6 @@ fn prepared_create(
account,
credential,
} = prepared;
let availability = account.availability();
Ok(PreparedCredentialCreate {
account_id: account.id().clone(),
provider_kind: account.provider().clone(),
@@ -534,8 +534,8 @@ fn prepared_create(
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(),
availability,
availability_observed_at: observed_at,
credential_state: account.credential_state(),
credential_observed_at: observed_at,
})
}
@@ -583,6 +583,7 @@ fn prepared_rotation(
struct RotationDocument {
access_token: String,
refresh_token: Option<String>,
id_token: Option<String>,
}
fn rotation_secret(document: ProviderDocument) -> Result<CodexOAuthSecret, ProviderAdminError> {
@@ -592,7 +593,7 @@ fn rotation_secret(document: ProviderDocument) -> Result<CodexOAuthSecret, Provi
Ok(CodexOAuthSecret {
access_token: SecretString::from(document.access_token),
refresh_token: document.refresh_token.map(SecretString::from),
id_token: None,
id_token: document.id_token.map(SecretString::from),
})
}
@@ -627,7 +628,13 @@ fn account_from_record(account: &AccountRecord) -> Result<ProviderAccount, Provi
account.upstream_account_id.clone(),
account.plan_type.clone(),
)
.with_runtime_state(account.enabled, account.availability)
.with_account_facts(
account.enabled,
account.credential_state,
account.quota,
account.last_error_reason,
account.last_error_message.clone(),
)
.with_refresh_schedule(
account.has_refresh_token,
account.next_refresh_at.map(SystemTime::from),
@@ -687,7 +694,6 @@ fn empty_quota() -> ProviderQuota {
refresh_token_expires_at: None,
windows: Vec::new(),
limit_reached: false,
rate_limited_until: None,
provider_data: None,
}
}
@@ -696,13 +702,10 @@ fn project_quota(
snapshot: Option<CodexAccountQuotaSnapshot>,
account: &ProviderAccount,
) -> ProviderQuota {
// 展示前先滚动已过期窗口,避免 reset 后的旧窗口值悬挂。
let mut quota = snapshot
.map(|snapshot| project_quota_snapshot(snapshot.roll_expired_windows(SystemTime::now())))
.map(project_quota_snapshot)
.unwrap_or_else(empty_quota);
if account.availability() == AccountAvailability::QuotaExhausted {
force_confirmed_exhaustion_projection(&mut quota);
}
quota.limit_reached = account.quota().is_exhausted();
quota
}
@@ -717,7 +720,7 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
);
provider_data.insert(
"exhausted".to_owned(),
Value::Bool(snapshot.fact().exhausted()),
Value::Bool(snapshot.quota().is_exhausted()),
);
let windows: Vec<ProviderQuotaWindow> = snapshot
.windows()
@@ -744,6 +747,11 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
window.window_seconds(),
),
source: Some(window.source().to_owned()),
local_usage_attribution: if window.is_account_wide() {
QuotaLocalUsageAttribution::AccountWide
} else {
QuotaLocalUsageAttribution::Unavailable
},
window_seconds: window.window_seconds(),
used_percent: window.used_percent(),
reset_at: window.reset_at(),
@@ -761,7 +769,6 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
refresh_token_expires_at: None,
windows,
limit_reached,
rate_limited_until: None,
provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(
provider_data,
))),
@@ -783,51 +790,6 @@ fn quota_windows_limit_reached(windows: &[ProviderQuotaWindow]) -> bool {
windows.iter().any(|window| window.limit_reached)
}
fn force_confirmed_exhaustion_projection(quota: &mut ProviderQuota) {
// 402 已确认耗尽:只覆盖 Admin 展示投影,原始 Provider JSON 仍作为
// 后续判断窗口是否真正恢复的基线。
quota.limit_reached = true;
quota.provider_data = Some(ProviderDocument::new(OpaqueProviderData::new(
Map::from_iter([
(
"remaining_percent".to_owned(),
Value::Number(Number::from(0)),
),
("exhausted".to_owned(), Value::Bool(true)),
]),
)));
if quota.windows.is_empty() {
quota.windows.push(ProviderQuotaWindow {
key: "confirmed-quota-exhaustion".to_owned(),
group: "other".to_owned(),
label: "额度".to_owned(),
source: None,
window_seconds: None,
used_percent: Some(100.0),
reset_at: None,
limit_reached: true,
local_usage: None,
provider_data: None,
});
return;
}
// 选已用比例最高的窗口强制 100%,保留其自身 reset_at。
let index = quota
.windows
.iter()
.enumerate()
.max_by(|(_, left), (_, right)| {
left.used_percent
.unwrap_or(-1.0)
.total_cmp(&right.used_percent.unwrap_or(-1.0))
})
.map(|(index, _)| index)
.unwrap_or(0);
quota.windows[index].used_percent = Some(100.0);
quota.windows[index].limit_reached = true;
}
const fn quota_group(kind: CodexQuotaWindowKind) -> &'static str {
match kind {
CodexQuotaWindowKind::Monthly => "monthly",
@@ -68,11 +68,6 @@ impl OpenAiConfig {
self.quota.refresh_policy()
}
#[must_use]
pub const fn quota_skip_exhausted(&self) -> bool {
self.quota.skip_exhausted
}
#[must_use]
pub fn token_client_config(&self) -> TokenClientConfig {
TokenClientConfig {
@@ -209,14 +204,12 @@ impl CodexWebSocketPoolSettings {
#[serde(deny_unknown_fields)]
pub struct CodexQuotaSettings {
pub refresh_interval_minutes: u64,
pub skip_exhausted: bool,
}
impl Default for CodexQuotaSettings {
fn default() -> Self {
Self {
refresh_interval_minutes: 15,
skip_exhausted: true,
}
}
}
@@ -10,9 +10,9 @@ use std::time::SystemTime;
use chrono::{DateTime, FixedOffset, Utc};
use gateway_core::engine::credential::{
AccountAvailability, CredentialCasUpdate, CredentialRevision, LoadedCredential,
AccountErrorReason, CredentialCasUpdate, CredentialRevision, CredentialState, LoadedCredential,
NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountIdentity,
ProviderAccountUpdate,
ProviderAccountUpdate, QuotaEvidence, QuotaState,
};
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest,
@@ -31,14 +31,13 @@ use super::types::{
CODEX_AUTHENTICATION_KIND_AGENT_IDENTITY, CODEX_AUTHENTICATION_KIND_OAUTH, CodexAccountProfile,
CodexAgentIdentityAuthMode, CodexAgentIdentityCredentialData, CodexCredentialData,
CodexCredentialPrincipal, CodexOAuthMetadata, CodexOAuthSecret, parse_access_token_expiration,
parse_chatgpt_jwt_claims,
};
const PROVIDER_NAME: &str = "openai";
const CODEX_CREDENTIAL_SCHEMA_VERSION: u32 = 1;
const MAX_BATCH: usize = 200;
const MAX_IMPORT_DOCUMENT_BYTES: usize = 64 * 1024 * 1024;
const MAX_IMPORT_EMAIL_BYTES: usize = 512;
const MAX_IMPORT_PLAN_BYTES: usize = 128;
pub struct ImportCodexOAuthCredential {
pub account_id: String,
@@ -51,8 +50,8 @@ pub struct ImportCodexOAuthCredential {
/// OAuth credential 的最小创建输入。
///
/// metadata 来自 ID token 的本地 payload 解码,或导入文档的展示字段;不包含
/// 签名、issuer、audience、nonce 或 access/refresh token 验证。
/// metadata 来自 ID token/access token 的官方本地 payload 解析;不包含签名、
/// issuer、audience、nonce 或 token 验证。
pub(crate) struct UnresolvedCodexOAuthCredential {
pub(crate) account_id: String,
pub(crate) name: String,
@@ -270,8 +269,9 @@ struct CodexCprExportCommon {
struct CodexCprOAuthExportAccount {
#[serde(flatten)]
common: CodexCprExportCommon,
token: String,
access_token: String,
refresh_token: Option<String>,
id_token: Option<String>,
access_token_expires_at: Option<String>,
}
@@ -433,7 +433,13 @@ impl CodexCredentialAdmin {
Some(input.verified_account.chatgpt_account_id),
input.verified_account.plan_type,
)
.with_runtime_state(input.enabled, AccountAvailability::Ready)
.with_account_facts(
input.enabled,
CredentialState::Ready,
QuotaState::unknown(),
None,
None,
)
.with_refresh_schedule(
input.secret.refresh_token.is_some(),
optional_time(input.next_refresh_at),
@@ -479,13 +485,16 @@ impl CodexCredentialAdmin {
optional_time(input.access_token_expires_at),
)
.with_profile(email, upstream_account_id, plan_type)
.with_runtime_state(
.with_account_facts(
input.enabled,
if has_upstream_user_id {
AccountAvailability::Ready
CredentialState::Ready
} else {
AccountAvailability::Unknown
CredentialState::Unknown
},
QuotaState::unknown(),
(!has_upstream_user_id).then_some(AccountErrorReason::AccountUnverified),
None,
)
.with_refresh_schedule(
input.secret.refresh_token.is_some(),
@@ -537,8 +546,9 @@ impl CodexCredentialAdmin {
}
CodexCprExportAccount::OAuth(CodexCprOAuthExportAccount {
common,
token: data.access_token,
access_token: data.access_token,
refresh_token: data.refresh_token,
id_token: data.id_token,
access_token_expires_at: account
.access_token_expires_at()
.map(DateTime::<Utc>::from)
@@ -785,23 +795,37 @@ impl CodexCredentialAdminService {
.refresh(refresh_token.expose_secret())
.await
.map_err(map_refresh_failure)?;
let access_token_rotated = tokens.access_token.is_some();
let secret = CodexOAuthSecret {
access_token: SecretString::from(tokens.access_token),
access_token: tokens
.access_token
.map(SecretString::from)
.unwrap_or_else(|| oauth.access_token.clone()),
refresh_token: tokens
.refresh_token
.map(SecretString::from)
.or_else(|| oauth.refresh_token.clone()),
id_token: oauth.id_token.clone(),
id_token: tokens
.id_token
.map(SecretString::from)
.or_else(|| oauth.id_token.clone()),
};
Self::record_recovery_log(
CodexOAuthRecoveryOperation::ManualRefresh,
Some(account_id.as_str()),
&secret,
);
let access_token_expires_at = tokens
.expires_in
.and_then(|expires_in| SystemTime::now().checked_add(expires_in))
.map(DateTime::<Utc>::from);
let access_token_expires_at = if access_token_rotated {
tokens
.expires_in
.and_then(|expires_in| SystemTime::now().checked_add(expires_in))
.map(DateTime::<Utc>::from)
} else {
current
.account
.access_token_expires_at()
.map(DateTime::<Utc>::from)
};
// 正常预刷新由 worker 读取当时的 runtime policy 动态判断;这里仅清除
// 之前瞬态失败留下的 retry-not-before。
let next_refresh_at = None;
@@ -820,8 +844,8 @@ impl CodexCredentialAdminService {
/// 归一导入 OAuth/agent-identity 凭据到唯一 `NewProviderAccount` 写入路径。
///
/// OAuth 导入只要取得 access token(直接提供或 RT exchange)即创建 Unknown
/// 账号;资格资料由后续 worker 异步补齐。
/// OAuth 导入先取得 access token(直接提供或 RT exchange),再按官方
/// `parse_chatgpt_jwt_claims` 从 ID token/access token 本地投影账号资料。
pub async fn prepare_import_document(
&self,
payload: Value,
@@ -861,16 +885,28 @@ impl CodexCredentialAdminService {
id_token.clone(),
)
.await?;
// 导入材料按管理员输入原样保存;ID token 仅尽力补展示字段,
// 解析失败不阻塞 AT/RT 入库。
let mut metadata = secret
// 与官方 `parse_chatgpt_jwt_claims` 一致:先读 ID token,缺失字段
// 再由 access token 补齐。解析失败不阻塞 token 入库。
let id_metadata = secret
.id_token
.as_ref()
.and_then(|token| {
super::types::parse_id_token_metadata(token.expose_secret()).ok()
})
.and_then(|token| parse_chatgpt_jwt_claims(token.expose_secret()).ok())
.unwrap_or_default();
apply_import_metadata(&mut metadata, &candidate);
let access_metadata =
parse_chatgpt_jwt_claims(secret.access_token.expose_secret())
.unwrap_or_default();
let metadata = CodexOAuthMetadata {
email: id_metadata.email.or(access_metadata.email),
chatgpt_plan_type: id_metadata
.chatgpt_plan_type
.or(access_metadata.chatgpt_plan_type),
chatgpt_user_id: id_metadata
.chatgpt_user_id
.or(access_metadata.chatgpt_user_id),
chatgpt_account_id: id_metadata
.chatgpt_account_id
.or(access_metadata.chatgpt_account_id),
};
CodexCredentialAdmin.prepare_unresolved_oauth(
UnresolvedCodexOAuthCredential {
account_id,
@@ -898,10 +934,10 @@ impl CodexCredentialAdminService {
if !account_ids.insert(account_id.clone()) {
return Err(CodexCredentialAdminError::InvalidInput);
}
let (enabled, availability) =
import_runtime_state(candidate.status.as_deref())?;
let facts =
parse_cpr_account_facts(candidate.status.as_deref(), SystemTime::now())?;
let mut account =
self.prepare_agent_identity_import(&candidate, account_id, enabled)?;
self.prepare_agent_identity_import(&candidate, account_id, facts.enabled)?;
let upstream_user_id = account
.account
.upstream_user_id()
@@ -915,7 +951,7 @@ impl CodexCredentialAdminService {
if !upstream_identities.insert((upstream_user_id, upstream_account_id)) {
return Err(CodexCredentialAdminError::InvalidInput);
}
account.account = account.account.with_runtime_state(enabled, availability);
account.account = facts.apply(account.account);
account
}
};
@@ -952,14 +988,17 @@ impl CodexCredentialAdminService {
.refresh(&refresh_token)
.await
.map_err(map_refresh_failure)?;
let access_token = tokens
.access_token
.ok_or(CodexCredentialAdminError::InvalidCredential)?;
let secret = CodexOAuthSecret {
access_token: SecretString::from(tokens.access_token),
access_token: SecretString::from(access_token),
// RT 未轮换时仍保留导入提供的 RT,保证本次补全不会丢失后续刷新能力。
refresh_token: tokens
.refresh_token
.map(SecretString::from)
.or_else(|| Some(SecretString::from(refresh_token))),
id_token,
id_token: tokens.id_token.map(SecretString::from).or(id_token),
};
Self::record_recovery_log(
CodexOAuthRecoveryOperation::ImportRefreshToken,
@@ -1048,7 +1087,13 @@ impl CodexCredentialAdminService {
Some(upstream_account_id),
candidate.plan_type.clone(),
)
.with_runtime_state(enabled, AccountAvailability::Ready)
.with_account_facts(
enabled,
CredentialState::Ready,
QuotaState::unknown(),
None,
None,
)
.with_refresh_schedule(false, None);
Ok(NewProviderAccount {
account,
@@ -1065,11 +1110,13 @@ fn cpr_status(account: &ProviderAccount) -> &'static str {
if !account.enabled() {
return "disabled";
}
match account.availability() {
AccountAvailability::QuotaExhausted => "quota_exhausted",
AccountAvailability::Expired | AccountAvailability::Invalid => "expired",
AccountAvailability::Banned => "banned",
AccountAvailability::Unknown | AccountAvailability::Ready => "active",
if account.quota().is_exhausted() {
return "quota_exhausted";
}
match account.credential_state() {
CredentialState::Expired | CredentialState::Invalid => "expired",
CredentialState::Banned => "banned",
CredentialState::Unknown | CredentialState::Ready => "active",
}
}
@@ -1204,9 +1251,9 @@ fn parse_oauth_import_tokens(
Value::Object(object) => {
for (key, value) in object {
let token = match key.as_str() {
"access_token" | "accessToken" | "at" | "token" => &mut access_token,
"refresh_token" | "refreshToken" | "rt" => &mut refresh_token,
"id_token" | "idToken" => &mut id_token,
"accessToken" => &mut access_token,
"refreshToken" => &mut refresh_token,
"idToken" => &mut id_token,
_ => {
pending.push(value);
continue;
@@ -1271,41 +1318,65 @@ fn first_string(value: &Value, keys: &[&str]) -> Option<String> {
.map(str::to_owned)
}
fn apply_import_metadata(metadata: &mut CodexOAuthMetadata, candidate: &ParsedCodexImportAccount) {
if metadata.email.is_none() {
metadata.email = candidate
.email
.as_deref()
.and_then(|value| normalize_import_text(value, MAX_IMPORT_EMAIL_BYTES));
}
if metadata.chatgpt_plan_type.is_none() {
metadata.chatgpt_plan_type = candidate
.plan_type
.as_deref()
.and_then(|value| normalize_import_text(value, MAX_IMPORT_PLAN_BYTES));
struct ImportedCprAccountFacts {
enabled: bool,
credential_state: CredentialState,
quota: QuotaState,
error_reason: Option<AccountErrorReason>,
}
impl ImportedCprAccountFacts {
fn apply(self, account: ProviderAccount) -> ProviderAccount {
account.with_account_facts(
self.enabled,
self.credential_state,
self.quota,
self.error_reason,
None,
)
}
}
fn normalize_import_text(value: &str, max_bytes: usize) -> Option<String> {
let value = value.trim();
(!value.is_empty() && value.len() <= max_bytes && !value.chars().any(char::is_control))
.then(|| value.to_owned())
}
fn import_runtime_state(
fn parse_cpr_account_facts(
status: Option<&str>,
) -> Result<(bool, AccountAvailability), CodexCredentialAdminError> {
observed_at: SystemTime,
) -> Result<ImportedCprAccountFacts, CodexCredentialAdminError> {
match status
.map(str::trim)
.unwrap_or("active")
.to_ascii_lowercase()
.as_str()
{
"active" => Ok((true, AccountAvailability::Ready)),
"disabled" => Ok((false, AccountAvailability::Ready)),
"expired" => Ok((true, AccountAvailability::Expired)),
"quota_exhausted" => Ok((true, AccountAvailability::QuotaExhausted)),
"banned" => Ok((true, AccountAvailability::Banned)),
"active" => Ok(ImportedCprAccountFacts {
enabled: true,
credential_state: CredentialState::Ready,
quota: QuotaState::unknown(),
error_reason: None,
}),
"disabled" => Ok(ImportedCprAccountFacts {
enabled: false,
credential_state: CredentialState::Ready,
quota: QuotaState::unknown(),
error_reason: None,
}),
"expired" => Ok(ImportedCprAccountFacts {
enabled: true,
credential_state: CredentialState::Expired,
quota: QuotaState::unknown(),
error_reason: Some(AccountErrorReason::CredentialExpired),
}),
"quota_exhausted" => Ok(ImportedCprAccountFacts {
enabled: true,
credential_state: CredentialState::Ready,
quota: QuotaState::exhausted(QuotaEvidence::ProviderDenied, observed_at, None),
error_reason: None,
}),
"banned" => Ok(ImportedCprAccountFacts {
enabled: true,
credential_state: CredentialState::Banned,
quota: QuotaState::unknown(),
error_reason: Some(AccountErrorReason::AccountBanned),
}),
_ => Err(CodexCredentialAdminError::InvalidInput),
}
}
@@ -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> {
@@ -7,7 +7,7 @@ use std::time::{Duration, SystemTime};
use chrono::{DateTime, Utc};
use gateway_core::engine::credential::{
AccountAvailability, OpaqueProviderData, ProviderAccount, ProviderAccountId,
CredentialState, OpaqueProviderData, ProviderAccount, ProviderAccountId,
};
use gateway_core::engine::provider::ProviderCatalogGeneration;
use gateway_core::provider_ports::{
@@ -780,12 +780,8 @@ fn eligible_catalog_account(account: &ProviderAccount, now: SystemTime) -> bool
&& account
.access_token_expires_at()
.is_none_or(|expires_at| expires_at > now)
&& match account.availability() {
AccountAvailability::Unknown
| AccountAvailability::Ready
| AccountAvailability::QuotaExhausted => true,
AccountAvailability::Expired
| AccountAvailability::Banned
| AccountAvailability::Invalid => false,
}
&& matches!(
account.credential_state(),
CredentialState::Unknown | CredentialState::Ready
)
}
@@ -35,7 +35,7 @@ use super::token_client::{
AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant,
OFFICIAL_CODEX_OAUTH_CLIENT_ID, OFFICIAL_CODEX_REDIRECT_URI,
};
use super::types::parse_id_token_metadata;
use super::types::parse_chatgpt_jwt_claims;
const AUTHORIZATION_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize";
const AUTHORIZATION_SCOPE: &str = "openid profile email offline_access";
@@ -504,7 +504,7 @@ impl CodexOAuthAdminService {
let access_token_expires_at = expires_in
.and_then(|expires_in| SystemTime::now().checked_add(expires_in))
.map(DateTime::<Utc>::from);
let metadata = parse_id_token_metadata(id_token.expose_secret())
let metadata = parse_chatgpt_jwt_claims(id_token.expose_secret())
.map_err(|_| CodexOAuthAdminError::TokenRejected)?;
secret.id_token = Some(id_token);
let credential = if let Some(current) = current {
@@ -0,0 +1,46 @@
//! 额度接口失败到规范化账号事实的唯一分类边界。
use gateway_core::engine::credential::{AccountErrorReason, CredentialState, QuotaEvidence};
use reqwest::StatusCode;
use serde_json::Value;
use crate::transport::CodexClientError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum QuotaEndpointFailure {
Credential {
state: CredentialState,
reason: AccountErrorReason,
},
Exhausted(QuotaEvidence),
}
/// 额度 endpoint 只有明确 402 能改变账号事实。401/403、429、5xx 与传输失败
/// 均不足以判断凭据或账号状态,留给 OAuth 与真实推理请求确认。
pub(crate) fn classify_quota_endpoint_failure(
error: &CodexClientError,
) -> Option<QuotaEndpointFailure> {
let CodexClientError::Upstream { status, body, .. } = error else {
return None;
};
if *status != StatusCode::PAYMENT_REQUIRED {
return None;
}
if is_deactivated_workspace(body) {
return Some(QuotaEndpointFailure::Credential {
state: CredentialState::Banned,
reason: AccountErrorReason::AccountBanned,
});
}
Some(QuotaEndpointFailure::Exhausted(
QuotaEvidence::PaymentRequired,
))
}
fn is_deactivated_workspace(body: &str) -> bool {
serde_json::from_str::<Value>(body)
.ok()
.is_some_and(|value| {
value.pointer("/detail/code").and_then(Value::as_str) == Some("deactivated_workspace")
})
}
@@ -1,9 +1,9 @@
//! Codex quota 服务编排:请求级额度投影、被动/主动同步与 429 合成。
//! Codex quota 服务编排:权威额度事实、展示投影、主动/被动同步与 429 冷却。
//!
//! - [`snapshot`]:快照/窗口解析、聚合、O(1) 滚动与调度信号。
//! - [`recovery`]:402 恢复证据与 availability 转换。
//! - [`evidence`]:额度接口错误到凭据/额度事实的分类。
pub(crate) mod recovery;
pub(crate) mod evidence;
pub(crate) mod snapshot;
pub use snapshot::{
@@ -13,12 +13,14 @@ pub use snapshot::{
use std::collections::{BTreeMap, BTreeSet};
use std::sync::{Arc, RwLock};
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use std::time::{Duration, Instant, SystemTime};
use chrono::Utc;
use gateway_core::engine::credential::{
AccountAvailability, AccountQuotaSignals, CredentialRevision, OpaqueProviderData,
ProviderAccount, ProviderAccountId, ProviderAccountStore, QuotaObservation, QuotaWriteOutcome,
AccountErrorReason, AccountQuotaSignals, CredentialRevision, CredentialState,
OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore,
QuotaAccessChange, QuotaAccessState, QuotaEvidence, QuotaObservation, QuotaState,
QuotaWriteOutcome,
};
use gateway_core::provider_ports::{ProviderCooldown, ProviderCooldownPort};
use gateway_protocol::openai::events::{
@@ -36,10 +38,7 @@ use crate::transport::{CodexBackendClient, CodexClientError, CodexRequestContext
use super::agent_identity::{CodexAgentIdentityTaskService, PreparedCodexRuntimeCredential};
use super::repository::{CodexCredentialRepository, CredentialRepositoryError};
use recovery::{
authoritative_quota_snapshot_state, quota_state_reason, quota_state_transition,
quota_success_state,
};
use evidence::{QuotaEndpointFailure, classify_quota_endpoint_failure};
use snapshot::{
parse_account_quota_snapshot, quota_projection_ttl, quota_snapshot_from_observation,
scheduling_signals_from_snapshot,
@@ -48,7 +47,9 @@ use snapshot::{
const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
pub(crate) const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
const EXHAUSTED_QUOTA_FALLBACK_RECHECK_INTERVAL: Duration = Duration::from_secs(30 * 60);
const EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL: Duration = Duration::from_secs(30 * 60);
const RESET_RECOVERY_MAX_USED_PERCENT: f64 = 10.0;
/// 首次 OAuth 异步观察失败时,由既有 quota worker 兜底重试的单轮上限。
const INITIAL_QUOTA_SYNC_BATCH: usize = 100;
// 5xx 上游拒绝的短退避重试预算;指数退避 1s/2s,吞掉瞬时抖动。
@@ -130,6 +131,12 @@ pub struct CodexCredentialQuotaService {
scheduling: CodexQuotaSchedulingProjection,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum QuotaRefreshAuthority {
ObserveAccess,
PreserveAccess,
}
#[derive(Clone, Default)]
struct CodexQuotaSchedulingProjection {
state: Arc<RwLock<CodexQuotaProjectionState>>,
@@ -177,6 +184,17 @@ enum CodexQuotaFetchError {
}
impl CodexQuotaSchedulingProjection {
fn invalidate(&self, account_ids: &[ProviderAccountId]) {
let mut state = self
.state
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
for account_id in account_ids {
state.entries.remove(account_id);
state.last_periodic_refresh_at.remove(account_id);
}
}
fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec<CodexQuotaHydrationTarget> {
let state = self
.state
@@ -299,12 +317,7 @@ impl CodexQuotaSchedulingProjection {
) -> Vec<ProviderAccount> {
let candidates = accounts
.into_iter()
.filter_map(|account| {
let limit_reached = self
.signals(&account)
.is_some_and(|signals| signals.limit_reached());
quota_refresh_candidate(account, now, limit_reached)
})
.filter_map(|account| quota_refresh_candidate(account, now))
.collect::<Vec<_>>();
let candidate_ids = candidates
.iter()
@@ -319,8 +332,8 @@ impl CodexQuotaSchedulingProjection {
.last_periodic_refresh_at
.retain(|account_id, _| candidate_ids.contains(account_id));
// 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。定时器只兜底
// 已耗尽账号,且不依据 reset 时间跳过复核;每个账号仍至少间隔 30 分钟。
// 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。定时器只复核
// 已耗尽且已到 reset/fallback 时刻的账号;重复复核至少间隔 30 分钟。
let mut reserved = Vec::new();
for account in candidates {
if !periodic_quota_refresh_due(&state, account.id(), refreshed_at) {
@@ -335,16 +348,13 @@ impl CodexQuotaSchedulingProjection {
}
}
fn quota_refresh_candidate(
account: ProviderAccount,
now: SystemTime,
limit_reached: bool,
) -> Option<ProviderAccount> {
// 定时器只兜底已耗尽或任一窗口触顶的账号;
// 正常账号由真实请求的响应头被动同步。
fn quota_refresh_candidate(account: ProviderAccount, now: SystemTime) -> Option<ProviderAccount> {
// reset 到期是“应刷新”的信号,不是“已恢复”的证据。
(eligible_periodic_quota_refresh(&account, now)
&& (account.availability() == AccountAvailability::QuotaExhausted || limit_reached))
.then_some(account)
&& account
.quota()
.exhaustion_refresh_due(now, EXHAUSTED_QUOTA_FALLBACK_RECHECK_INTERVAL))
.then_some(account)
}
fn periodic_quota_refresh_due(
@@ -356,7 +366,7 @@ fn periodic_quota_refresh_due(
.last_periodic_refresh_at
.get(account_id)
.is_none_or(|last| {
now.saturating_duration_since(*last) >= QUOTA_PERIODIC_REFRESH_MIN_INTERVAL
now.saturating_duration_since(*last) >= EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL
})
}
@@ -400,70 +410,55 @@ impl CodexCredentialQuotaService {
}
}
/// 写入新快照前读取 core primary 已确认耗尽窗口的 reset,用于判断该窗口是否真的推进。
async fn previous_confirmed_reset(&self, account: &ProviderAccount) -> Option<SystemTime> {
let observations = self
.store
.get_quotas(std::slice::from_ref(account.id()))
.await
.ok()?;
let observation = observations.into_iter().find(|observation| {
observation.account_id == *account.id()
&& observation.expected_revision == account.revision()
})?;
let snapshot = quota_snapshot_from_observation(&observation)?;
snapshot
.core_primary_window()
.and_then(|window| window.reset_at())
.map(SystemTime::from)
}
/// 将已确认的额度拒绝投影为 100%,使展示、调度和恢复共用同一窗口事实。
///
/// `reset_at` 来自官方 `usage_limit_reached` 错误体;缺失时保留当前快照的 reset。
/// 真实推理错误只更新额度访问事实,不伪造 Provider JSON 或展示百分比。
pub(crate) async fn record_confirmed_exhaustion(
&self,
account: &ProviderAccount,
reset_at: Option<i64>,
evidence: QuotaEvidence,
reset_at: Option<SystemTime>,
observed_at: SystemTime,
) -> Result<(), CodexCredentialQuotaError> {
let existing = self
.store
.get_quotas(std::slice::from_ref(account.id()))
.await?
.into_iter()
.find(|observation| {
observation.account_id == *account.id()
&& observation.expected_revision == account.revision()
})
.and_then(|observation| observation.quota)
.map(OpaqueProviderData::into_inner)
.unwrap_or_default();
let quota = confirmed_exhaustion_projection(
normalize_quota_window_placeholders(existing),
reset_at,
);
let observed_at = SystemTime::now();
let outcome = self
.store
.compare_and_swap_quota(QuotaObservation {
.apply_quota_access(QuotaAccessChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(quota.clone())),
observed_at: Some(observed_at),
limit_reached: Some(true),
state: QuotaState::exhausted(evidence, observed_at, reset_at),
})
.await?;
if outcome == QuotaWriteOutcome::Conflict {
return Err(CodexCredentialQuotaError::RevisionConflict);
}
if let Ok(snapshot) = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
&Value::Object(quota),
) {
self.scheduling.observe(&snapshot);
Ok(())
}
/// 成功推理是额度可访问的权威证据,同时解除账号级 429 冷却。
pub(crate) async fn record_successful_inference(
&self,
account: &ProviderAccount,
observed_at: SystemTime,
) -> Result<(), CodexCredentialQuotaError> {
// 已经发出的并发请求可能晚于耗尽事实成功返回;真实成功仍是独立的
// Allowed 权威证据,但调度不会为了试探恢复而放行耗尽账号。
if account.quota().access() == QuotaAccessState::Exhausted {
let outcome = self
.store
.apply_quota_access(QuotaAccessChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
state: QuotaState::allowed(observed_at),
})
.await?;
if outcome == QuotaWriteOutcome::Conflict {
return Err(CodexCredentialQuotaError::RevisionConflict);
}
}
self.cooldowns
.clear(account.id(), account.revision())
.await
.map_err(|error| CodexCredentialQuotaError::Store {
detail: error.to_string(),
})?;
Ok(())
}
@@ -511,6 +506,10 @@ impl CodexCredentialQuotaService {
self.scheduling.signals(account)
}
pub(crate) fn invalidate_scheduling(&self, account_ids: &[ProviderAccountId]) {
self.scheduling.invalidate(account_ids);
}
pub async fn synchronize(&self) -> Result<CodexQuotaSyncSummary, CodexCredentialQuotaError> {
let accounts = self.repository.list_for_provider().await?;
let mut summary = CodexQuotaSyncSummary::default();
@@ -562,30 +561,27 @@ impl CodexCredentialQuotaService {
);
}
Err(CodexQuotaFetchError::Upstream { account, error }) => {
match quota_state_transition(&error) {
// 402:真额度耗尽,写入 QuotaExhausted(worker 按 reset_at 自动恢复)。
Some(AccountAvailability::QuotaExhausted) => {
match classify_quota_endpoint_failure(&error) {
Some(QuotaEndpointFailure::Exhausted(evidence)) => {
summary.exhausted += 1;
self.persist_upstream_state(
&account,
AccountAvailability::QuotaExhausted,
&error,
observed_at,
)
.await;
if let Err(write_error) = self
.record_confirmed_exhaustion(&account, evidence, None, observed_at)
.await
{
summary.transient += 1;
tracing::warn!(
account_id = %account.id(),
error = %write_error,
"OpenAI quota exhaustion fact write failed"
);
}
}
Some(AccountAvailability::Banned) => {
Some(QuotaEndpointFailure::Credential { state, reason }) => {
summary.banned += 1;
self.persist_upstream_state(
&account,
AccountAvailability::Banned,
&error,
observed_at,
)
.await;
self.persist_credential_failure(&account, state, reason, observed_at)
.await;
}
// usage 401/403 与 429/503 都不写 availability;前者等待 OAuth 或真实请求确认。
None | Some(_) => {
None => {
summary.transient += 1;
tracing::warn!(
account_id = %account.id(),
@@ -613,7 +609,7 @@ impl CodexCredentialQuotaService {
let observed = self.store.get_quotas(&account_ids).await?;
let observed_ids = observed
.into_iter()
.filter_map(|observation| observation.observed_at.map(|_| observation.account_id))
.map(|observation| observation.account_id)
.collect::<BTreeSet<_>>();
Ok(accounts
.iter()
@@ -633,79 +629,53 @@ impl CodexCredentialQuotaService {
observed_at: SystemTime,
summary: &mut CodexQuotaSyncSummary,
) -> Result<(), CodexCredentialQuotaError> {
let refreshed_snapshot = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
value,
)?;
let previous_reset_at = self.previous_confirmed_reset(account).await;
let object = normalize_quota_window_placeholders(
value
.as_object()
.cloned()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?,
);
let object = if account.availability() == AccountAvailability::QuotaExhausted
&& quota_success_state(
account.availability(),
&refreshed_snapshot,
previous_reset_at,
)
.is_none()
{
confirmed_exhaustion_projection(
object,
confirmed_exhaustion_reset_at(previous_reset_at, &refreshed_snapshot),
)
} else {
object
};
let snapshot = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
&Value::Object(object.clone()),
)?;
let fact = snapshot.fact();
let previous = if account.quota().is_exhausted() {
self.read_snapshot_for(account).await?
} else {
None
};
if !quota_refresh_can_update(account, &snapshot, previous.as_ref()) {
summary.exhausted += 1;
return Ok(());
}
let state = if account.quota().is_exhausted() {
QuotaState::allowed(observed_at)
} else {
account.quota().merge_observation(snapshot.quota())
};
let snapshot = snapshot.with_quota_state(state);
let outcome = self
.store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(object)),
observed_at: Some(observed_at),
limit_reached: Some(fact.exhausted()),
quota: OpaqueProviderData::new(object),
observed_at,
state,
})
.await?;
if outcome == QuotaWriteOutcome::Conflict {
summary.stale += 1;
return Ok(());
}
let Some(current) = self.store.get_account(account.id()).await? else {
summary.stale += 1;
return Ok(());
};
if current.revision() != account.revision() {
summary.stale += 1;
return Ok(());
}
self.scheduling.observe(&snapshot);
if fact.exhausted() {
if snapshot.quota().is_exhausted() {
summary.exhausted += 1;
} else {
summary.updated += 1;
}
if let Some(availability) = authoritative_quota_snapshot_state(
current.availability(),
&refreshed_snapshot,
previous_reset_at,
) {
let _ = self
.repository
.apply_state(&current, availability, observed_at)
.await;
}
Ok(())
}
@@ -722,11 +692,6 @@ impl CodexCredentialQuotaService {
.limits
.values()
.any(|details| passive_rate_limit_snapshot(details).is_some());
let previous_reset_at = if account.availability() == AccountAvailability::QuotaExhausted {
self.previous_confirmed_reset(account).await
} else {
None
};
let existing = self
.store
.get_quotas(std::slice::from_ref(account.id()))
@@ -736,17 +701,14 @@ impl CodexCredentialQuotaService {
observation.account_id == *account.id()
&& observation.expected_revision == account.revision()
});
let existing_observed_at = existing
.as_ref()
.and_then(|observation| observation.observed_at);
let existing_state = existing.as_ref().map(|observation| observation.state);
let existing = existing
.and_then(|observation| observation.quota)
.map(OpaqueProviderData::into_inner)
.map(|observation| observation.quota.into_inner())
.unwrap_or_default();
// 套餐、credits 等元数据可以更新,但没有额度窗口事实时必须保留旧观察时刻,
// 也不能借旧快照重新推导 quota state。
if !has_quota_facts {
let Some(observed_at) = existing_observed_at else {
let Some(state) = existing_state else {
return Ok(false);
};
let outcome = self
@@ -754,62 +716,42 @@ impl CodexCredentialQuotaService {
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(merge_passive_quota(
existing,
&rate_limits,
))),
observed_at: Some(observed_at),
limit_reached: None,
quota: OpaqueProviderData::new(merge_passive_quota(existing, &rate_limits)),
observed_at: SystemTime::now(),
state,
})
.await?;
return Ok(outcome != QuotaWriteOutcome::Conflict);
}
let quota = merge_passive_quota(existing, &rate_limits);
let observed_at = SystemTime::now();
let refreshed_snapshot = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
&Value::Object(quota.clone()),
)?;
let quota = if account.availability() == AccountAvailability::QuotaExhausted {
// 被动成功响应不能解除确认额度耗尽;继续展示 100%,直到权威 usage 快照
// 证明旧 reset 已跨过且新窗口已经推进。
confirmed_exhaustion_projection(
quota,
confirmed_exhaustion_reset_at(previous_reset_at, &refreshed_snapshot),
)
} else {
quota
};
let fact = parse_codex_quota_usage(&Value::Object(quota.clone()))?;
let snapshot = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
&Value::Object(quota.clone()),
);
)?;
// 这些 headers 来自一次成功推理,访问结论优先于可能滞后的百分比。
let state = QuotaState::allowed(observed_at);
let outcome = self
.store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(quota)),
observed_at: Some(observed_at),
limit_reached: Some(fact.exhausted()),
quota: OpaqueProviderData::new(quota),
observed_at,
state,
})
.await?;
if outcome == QuotaWriteOutcome::Conflict {
return Ok(false);
}
if let Ok(snapshot) = snapshot {
self.scheduling.observe(&snapshot);
}
self.scheduling.observe(&snapshot);
Ok(true)
}
/// 真实 429 的单一事实入口:写入 Redis 临时限流冷却(`until = now + retry_after`)。
/// availability 与 quota 主窗口(额度重置时间)都不改变——临时限流是独立维度,
/// 凭据与额度主窗口(额度重置时间)都不改变——临时限流是独立维度,
/// 到期由 Redis key 过期自动解除,不污染配额耗尽状态。
/// 已有更晚的冷却不会被缩短(put_if_later)。
pub async fn apply_rate_limit_429(
@@ -838,7 +780,6 @@ impl CodexCredentialQuotaService {
pub async fn rate_limited_until(
&self,
account_id: &ProviderAccountId,
credential_revision: CredentialRevision,
) -> Result<Option<SystemTime>, CodexCredentialQuotaError> {
let Some(cooldown) = self.cooldowns.read(account_id).await.map_err(|error| {
CodexCredentialQuotaError::Store {
@@ -848,9 +789,6 @@ impl CodexCredentialQuotaService {
else {
return Ok(None);
};
if cooldown.credential_revision() != credential_revision {
return Ok(None);
}
let until = cooldown.until();
if until <= SystemTime::now() {
return Ok(None);
@@ -894,18 +832,17 @@ impl CodexCredentialQuotaService {
{
return Err(CodexCredentialQuotaError::RevisionConflict);
}
let Some(data) = observation.quota else {
if observation.quota.expose_to_provider().is_empty() {
return Ok(None);
};
let observed_at = observation
.observed_at
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?;
}
let observed_at = observation.observed_at;
let snapshot = parse_account_quota_snapshot(
account_id.clone(),
account.revision(),
observed_at,
&Value::Object(data.expose_to_provider().clone()),
)?;
&Value::Object(observation.quota.expose_to_provider().clone()),
)?
.with_quota_state(observation.state);
self.scheduling.observe(&snapshot);
Ok(Some(snapshot))
}
@@ -915,7 +852,8 @@ impl CodexCredentialQuotaService {
&self,
account_id: &ProviderAccountId,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
self.refresh_account_snapshot(account_id).await
self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::ObserveAccess)
.await
}
/// 真实限额失败后的异步刷新只补齐展示快照,不允许 usage 快照撤销已确认的失败状态。
@@ -923,12 +861,14 @@ impl CodexCredentialQuotaService {
&self,
account_id: &ProviderAccountId,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
self.refresh_account_snapshot(account_id).await
self.refresh_account_snapshot(account_id, QuotaRefreshAuthority::PreserveAccess)
.await
}
async fn refresh_account_snapshot(
&self,
account_id: &ProviderAccountId,
authority: QuotaRefreshAuthority,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
let account = self
.store
@@ -957,104 +897,89 @@ impl CodexCredentialQuotaService {
});
}
Err(CodexQuotaFetchError::Upstream { account, error }) => {
// 402 与明确的停用错误写对应状态;usage 401/403 不写 availability。
if let Some(availability) = quota_state_transition(&error) {
self.persist_upstream_state(&account, availability, &error, observed_at)
.await;
match classify_quota_endpoint_failure(&error) {
Some(QuotaEndpointFailure::Exhausted(evidence)) => {
self.record_confirmed_exhaustion(&account, evidence, None, observed_at)
.await?;
}
Some(QuotaEndpointFailure::Credential { state, reason }) => {
self.persist_credential_failure(&account, state, reason, observed_at)
.await;
}
None => {}
}
return Err(CodexCredentialQuotaError::Upstream {
detail: error.to_string(),
});
}
};
let refreshed_snapshot = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
&value,
)?;
let previous_reset_at = self.previous_confirmed_reset(&account).await;
let object = normalize_quota_window_placeholders(
value
.as_object()
.cloned()
.ok_or(CodexCredentialQuotaError::InvalidCredentialData)?,
);
let object = if account.availability() == AccountAvailability::QuotaExhausted
&& quota_success_state(
account.availability(),
&refreshed_snapshot,
previous_reset_at,
)
.is_none()
{
confirmed_exhaustion_projection(
object,
confirmed_exhaustion_reset_at(previous_reset_at, &refreshed_snapshot),
)
} else {
object
};
let snapshot = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
&Value::Object(object.clone()),
)?;
let previous = if account.quota().is_exhausted() {
self.read_snapshot_for(&account).await?
} else {
None
};
if !quota_refresh_can_update(&account, &snapshot, previous.as_ref()) {
return Ok(previous.unwrap_or_else(|| snapshot.with_quota_state(account.quota())));
}
let state = if account.quota().is_exhausted() {
QuotaState::allowed(observed_at)
} else {
match authority {
QuotaRefreshAuthority::ObserveAccess => {
account.quota().merge_observation(snapshot.quota())
}
QuotaRefreshAuthority::PreserveAccess => account.quota(),
}
};
let snapshot = snapshot.with_quota_state(state);
if self
.store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(object)),
observed_at: Some(observed_at),
limit_reached: Some(snapshot.fact().exhausted()),
quota: OpaqueProviderData::new(object),
observed_at,
state,
})
.await?
== QuotaWriteOutcome::Conflict
{
return Err(CodexCredentialQuotaError::RevisionConflict);
}
let current = self
.store
.get_account(account.id())
.await?
.ok_or(CodexCredentialQuotaError::NotFound)?;
if current.revision() != account.revision() {
return Err(CodexCredentialQuotaError::RevisionConflict);
}
self.scheduling.observe(&snapshot);
if let Some(availability) = authoritative_quota_snapshot_state(
current.availability(),
&refreshed_snapshot,
previous_reset_at,
) {
self.repository
.apply_state(&current, availability, observed_at)
.await?;
}
Ok(snapshot)
}
async fn persist_upstream_state(
async fn persist_credential_failure(
&self,
account: &ProviderAccount,
availability: AccountAvailability,
upstream_error: &CodexClientError,
credential_state: CredentialState,
reason: AccountErrorReason,
observed_at: SystemTime,
) {
let reason = quota_state_reason(upstream_error).map(str::to_owned);
if let Err(error) = self
.repository
.apply_state_with_message(account, availability, observed_at, reason.clone())
.apply_state_with_reason(account, credential_state, observed_at, Some(reason), None)
.await
{
tracing::warn!(
account_id = %account.id(),
?availability,
reason = reason.as_deref().unwrap_or_default(),
?credential_state,
reason = reason.as_str(),
error = %error,
"OpenAI quota upstream state write failed"
"OpenAI quota credential fact write failed"
);
}
}
@@ -1102,6 +1027,44 @@ impl CodexCredentialQuotaService {
}
}
fn quota_refresh_can_update(
account: &ProviderAccount,
refreshed: &CodexAccountQuotaSnapshot,
previous: Option<&CodexAccountQuotaSnapshot>,
) -> bool {
if !account.quota().is_exhausted() {
return true;
}
if refreshed.quota().is_exhausted() {
return false;
}
let previous_reset_at = account.quota().reset_at().or_else(|| {
previous?
.windows()
.iter()
.find(|window| {
window.source() == "codex" && window.role() == CodexQuotaWindowRole::Primary
})?
.reset_at()
.map(SystemTime::from)
});
let Some(previous_reset_at) = previous_reset_at else {
return false;
};
let Some(primary) = refreshed.windows().iter().find(|window| {
window.source() == "codex" && window.role() == CodexQuotaWindowRole::Primary
}) else {
return false;
};
let Some(refreshed_reset_at) = primary.reset_at().map(SystemTime::from) else {
return false;
};
let Some(used_percent) = primary.used_percent() else {
return false;
};
refreshed_reset_at > previous_reset_at && used_percent < RESET_RECOVERY_MAX_USED_PERCENT
}
async fn fetch_usage_once(
client: &CodexBackendClient,
prepared: &PreparedCodexRuntimeCredential,
@@ -1163,10 +1126,8 @@ fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -
fn eligible_initial_quota_sync(account: &ProviderAccount, now: SystemTime) -> bool {
// 首次观察只兜底刚入库、尚无 quota 快照的账号。已耗尽等运行时状态必须由
// periodic 路径处理,才能保留同一账号的最小复核间隔。
matches!(
account.availability(),
AccountAvailability::Ready | AccountAvailability::Unknown
) && eligible_periodic_quota_refresh(account, now)
account.credential_state() == CredentialState::Ready
&& eligible_periodic_quota_refresh(account, now)
}
fn access_token_is_current(account: &ProviderAccount, now: SystemTime) -> bool {
@@ -1297,71 +1258,6 @@ fn secondary_window_is_placeholder(window: &Map<String, Value>) -> bool {
})
}
/// 用已确认的上游额度拒绝覆盖 core 限额投影。
///
/// 上游 usage 快照可能在结算期间短暂滞后为 98/99% 或 `allowed=true`;保留其他
/// 原始 metadata,但把 core 窗口固定为 100%,直到权威快照证明进入下一窗口。
fn confirmed_exhaustion_projection(
mut quota: Map<String, Value>,
reset_at: Option<i64>,
) -> Map<String, Value> {
let mut rate_limit = quota
.remove("rate_limit")
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
rate_limit.insert("allowed".to_owned(), Value::Bool(false));
rate_limit.insert("limit_reached".to_owned(), Value::Bool(true));
for field in ["primary_window", "secondary_window"] {
let existed = rate_limit.contains_key(field);
let mut window = rate_limit
.remove(field)
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
if field == "primary_window" || existed {
window.insert("used_percent".to_owned(), Value::from(100_u64));
if field == "primary_window"
&& let Some(reset_at) = reset_at
{
window.insert("reset_at".to_owned(), Value::from(reset_at));
}
rate_limit.insert(field.to_owned(), Value::Object(window));
}
}
quota.insert("rate_limit".to_owned(), Value::Object(rate_limit));
quota
}
/// 选择应继续锁定的 reset:新鲜快照仍确认耗尽时可把锁推进到更晚窗口;
/// 否则保留已确认失败的旧窗口,避免 98/99% 快照反向解锁。
fn confirmed_exhaustion_reset_at(
previous_reset_at: Option<SystemTime>,
refreshed_snapshot: &CodexAccountQuotaSnapshot,
) -> Option<i64> {
let refreshed_reset_at = refreshed_snapshot
.core_primary_window()
.and_then(|window| window.reset_at())
.map(SystemTime::from);
let reset_at = match (previous_reset_at, refreshed_reset_at) {
(Some(previous), Some(refreshed))
if refreshed_snapshot.fact().exhausted() && refreshed > previous =>
{
refreshed
}
(Some(previous), _) => previous,
(None, refreshed) => refreshed?,
};
system_time_to_unix_seconds(reset_at)
}
fn system_time_to_unix_seconds(value: SystemTime) -> Option<i64> {
value
.duration_since(UNIX_EPOCH)
.ok()?
.as_secs()
.try_into()
.ok()
}
fn passive_rate_limit_snapshot(details: &RateLimitDetails) -> Option<Map<String, Value>> {
// 响应头的窗口属于同一次上游观测;跨响应合并会制造不存在的额度窗口。
let mut snapshot = Map::new();
@@ -1,102 +0,0 @@
//! 权威额度快照的 availability 转换。
//!
//! 已触顶的 core primary 会立即把可调度账号标为 `QuotaExhausted`;
//! `QuotaExhausted -> Ready` 则只能由同一 credential revision 的权威 usage 快照证明:
//! 新快照未耗尽、core primary reset 已推进到下一窗口,且该窗口用量低于 10%。
use std::time::SystemTime;
use gateway_core::engine::credential::AccountAvailability;
use reqwest::StatusCode;
use serde_json::Value;
use super::snapshot::CodexAccountQuotaSnapshot;
use crate::transport::CodexClientError;
const RESET_RECOVERY_MAX_USED_PERCENT: f64 = 10.0;
/// 基于成功的权威 usage 快照收敛可调度账号的额度状态。
///
/// 只接受上游明确拒绝使用且 core primary 触顶的证据;secondary、附加桶、
/// credits 与 spend control 都不能单独把整个账号标成额度耗尽。
pub(crate) fn authoritative_quota_snapshot_state(
current: AccountAvailability,
snapshot: &CodexAccountQuotaSnapshot,
previous_reset_at: Option<SystemTime>,
) -> Option<AccountAvailability> {
match current {
AccountAvailability::Ready | AccountAvailability::Unknown
if snapshot.authoritative_core_primary_exhausted() =>
{
Some(AccountAvailability::QuotaExhausted)
}
AccountAvailability::QuotaExhausted => {
quota_success_state(current, snapshot, previous_reset_at)
}
AccountAvailability::Ready | AccountAvailability::Unknown => None,
AccountAvailability::Expired
| AccountAvailability::Banned
| AccountAvailability::Invalid => None,
}
}
/// 只处理 `QuotaExhausted` 的恢复。`allowed=true`、成功响应和百分比回落都不能单独覆盖
/// 一次已确认的额度拒绝;它们在 reset 未推进时可能只是上游滞后的观测。
///
/// `previous_reset_at` 是写入新快照前读取的 core primary 已确认耗尽窗口。上游可能只
/// 调整 reset 时间而尚未重置额度,因此还要求同一新鲜 primary 窗口的用量低于 10%。
pub(crate) fn quota_success_state(
current: AccountAvailability,
snapshot: &CodexAccountQuotaSnapshot,
previous_reset_at: Option<SystemTime>,
) -> Option<AccountAvailability> {
if current != AccountAvailability::QuotaExhausted || snapshot.fact().exhausted() {
return None;
}
let previous_reset_at = previous_reset_at?;
let primary_window = snapshot.core_primary_window()?;
let refreshed_reset_at = primary_window.reset_at().map(SystemTime::from)?;
let used_percent = primary_window.used_percent()?;
(refreshed_reset_at > previous_reset_at && used_percent < RESET_RECOVERY_MAX_USED_PERCENT)
.then_some(AccountAvailability::Ready)
}
pub(crate) fn quota_state_transition(error: &CodexClientError) -> Option<AccountAvailability> {
let CodexClientError::Upstream { status, body, .. } = error else {
return None;
};
if *status == StatusCode::PAYMENT_REQUIRED && is_deactivated_workspace(body) {
return Some(AccountAvailability::Banned);
}
match *status {
// 402:真额度耗尽 → QuotaExhausted(worker 按 reset_at 自动恢复)。
StatusCode::PAYMENT_REQUIRED => Some(AccountAvailability::QuotaExhausted),
// usage 401/403 不能单独判定账号不可用:保持现有状态,交给 OAuth 或真实请求确认。
// 429/503 同样不写 availability,仅记失败待重试。
StatusCode::UNAUTHORIZED
| StatusCode::FORBIDDEN
| StatusCode::TOO_MANY_REQUESTS
| StatusCode::SERVICE_UNAVAILABLE => None,
_ => None,
}
}
/// 配额接口将账号变为非 ready 状态时可安全持久化的上游原因码。
pub(crate) fn quota_state_reason(error: &CodexClientError) -> Option<&'static str> {
match quota_state_transition(error)? {
AccountAvailability::Banned => Some("deactivated_workspace"),
AccountAvailability::QuotaExhausted => Some("payment_required"),
AccountAvailability::Unknown
| AccountAvailability::Ready
| AccountAvailability::Expired
| AccountAvailability::Invalid => None,
}
}
fn is_deactivated_workspace(body: &str) -> bool {
serde_json::from_str::<Value>(body)
.ok()
.is_some_and(|value| {
value.pointer("/detail/code").and_then(Value::as_str) == Some("deactivated_workspace")
})
}

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