refactor(quota): separate rate limiting from account availability

Remove the persistent Cooldown state and cooldown_until/availability_reason
columns; rate limiting (429/503) now lives entirely in provider quota data
with projected window rolling, while availability carries only terminal or
intervention-required states (QuotaExhausted for 402, worker auto-recovers
on quota reset). xAI migrated to the same model (runtime cooldown via Redis,
no persistent state). Frontend derives 限流中 from quota window limitReached.

Align quota window parsing with the official codex schema: drop the invented
monthly_limit/spend_control windows, add the top-level code_review_rate_limit
window, prefer upstream limit_name for display labels, and keep duration
guessing only for grouping.

Also fix: account directory degrades one failing quota to an empty projection
instead of failing the whole page; dashboard rate_limited bucket SQL; fake
store Mutex re-entrancy deadlock; remove dead tokenRefreshing field.

Migration 0002 drops cooldown columns and rewrites the availability CHECK.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Zyy
2026-08-03 00:13:56 +08:00
co-authored by Claude Fable 5
parent 3d1e33558a
commit cf03c56cfa
61 changed files with 1007 additions and 2117 deletions
@@ -75,8 +75,6 @@ pub struct AccountRecord {
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
@@ -443,13 +443,15 @@ pub struct RequestMetricPoint {
///
/// 首页概览的 `unavailable` 表示不可参与调度,包含配额耗尽及过期、禁用和封禁账号。
///
/// `quota_exhausted` 仍单列,以便状态详情展示和诊断。
/// `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 quota_exhausted: u64,
pub refreshing: Option<u64>,
@@ -153,8 +153,6 @@ pub struct PreparedCredentialCreate {
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
}
@@ -573,6 +571,7 @@ pub struct ProviderQuotaWindow {
pub window_seconds: Option<u64>,
pub used_percent: Option<f64>,
pub reset_at: Option<DateTime<Utc>>,
pub limit_reached: bool,
pub local_usage: Option<AccountUsage>,
pub provider_data: Option<ProviderDocument>,
}
@@ -260,18 +260,37 @@ impl AccountsService for DefaultAccountsService {
let mut quotas = futures::future::join_all(page.items.iter().map(|account| async {
let account_id = ProviderAccountId::new(account.id.clone())
.map_err(|_| AdminError::invalid("Invalid provider account ID"))?;
let provider = self
.providers
.require(&account.provider_kind)
.map_err(|error| map_provider_error(error, "provider quota"))?;
provider
// 单个账号的 quota 投影失败(Provider 未注册或 quota 读取失败)不拖垮整页:
// 该账号降级为空额度投影,其余账号与页面状态照常返回。
let provider = match self.providers.require(&account.provider_kind) {
Ok(provider) => provider,
Err(error) => {
tracing::warn!(
account_id = %account.id,
error = %error,
"account directory provider is not registered; showing empty quota"
);
return Ok(empty_quota());
}
};
match provider
.quota(ProviderQuotaRequest {
account_id,
refresh: false,
rolling_usage: rolling_usage.get(&account.id).cloned(),
})
.await
.map_err(|error| map_provider_error(error, "provider quota"))
{
Ok(quota) => Ok(quota),
Err(error) => {
tracing::warn!(
account_id = %account.id,
error = %error,
"account directory quota projection failed; showing empty quota"
);
Ok(empty_quota())
}
}
}))
.await
.into_iter()
@@ -486,16 +505,22 @@ impl AccountsService for DefaultAccountsService {
}
}
/// 账号目录中单个账号 quota 读取失败时使用的空额度投影。
fn empty_quota() -> ProviderQuota {
ProviderQuota {
observed_at: None,
refresh_token_expires_at: None,
windows: Vec::new(),
provider_data: None,
}
}
fn account_status(account: &AccountRecord, now: chrono::DateTime<Utc>) -> AccountStatus {
if !account.enabled {
AccountStatus::Disabled
} else if account.availability == crate::model::accounts::AccountAvailability::Banned {
AccountStatus::Banned
} else if account.availability == crate::model::accounts::AccountAvailability::QuotaExhausted
|| (account.availability == crate::model::accounts::AccountAvailability::Cooldown
&& account.availability_reason.as_deref() == Some("usage_limit_exhausted")
&& account.cooldown_until.is_none_or(|until| until > now))
{
} else if account.availability == crate::model::accounts::AccountAvailability::QuotaExhausted {
AccountStatus::QuotaExhausted
} else if matches!(
account.availability,
@@ -146,6 +146,7 @@ fn quota_window(
window_seconds,
used_percent,
reset_at: None,
limit_reached: false,
local_usage: None,
provider_data: None,
}
@@ -373,7 +373,7 @@ impl ProviderAdmin for FakeProviderAdmin {
pub(super) struct FakeAccountStore {
events: EventLog,
account: AccountRecord,
accounts: Mutex<Vec<AccountRecord>>,
account_after_probe: Mutex<Option<AccountRecord>>,
fail_commit: Mutex<bool>,
audit_requests: Mutex<Vec<String>>,
@@ -388,7 +388,7 @@ impl FakeAccountStore {
fn with_account(account: AccountRecord, events: EventLog) -> Arc<Self> {
Arc::new(Self {
events,
account,
accounts: Mutex::new(vec![account]),
account_after_probe: Mutex::new(None),
fail_commit: Mutex::new(false),
audit_requests: Mutex::new(Vec::new()),
@@ -408,6 +408,10 @@ impl FakeAccountStore {
*self.quota_window_usage.lock().expect("quota window usage") = usage;
}
pub(super) fn set_accounts(&self, accounts: Vec<AccountRecord>) {
*self.accounts.lock().expect("accounts") = accounts;
}
fn set_account_after_probe(&self, account: AccountRecord) {
*self
.account_after_probe
@@ -439,12 +443,14 @@ impl FakeAccountStore {
#[async_trait]
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();
Ok(AccountPage {
config_revision: revision(1),
items: vec![self.account.clone()],
total: 1,
items: accounts.clone(),
total: accounts.len() as u64,
summary: AccountSummary {
total: 1,
total: accounts.len() as u64,
active: 1,
quota_exhausted: 0,
unavailable: 0,
@@ -454,13 +460,22 @@ impl AccountStore for FakeAccountStore {
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountRecord>> {
self.record("store.load_account");
// probe 后的账号状态覆盖只对同一 id 生效;其余按账号列表查询。
let account = self
.account_after_probe
.lock()
.expect("account after probe")
.clone()
.unwrap_or_else(|| self.account.clone());
Ok((account_id == account.id.as_str()).then_some(account))
.filter(|account| account.id == account_id)
.or_else(|| {
self.accounts
.lock()
.expect("accounts")
.iter()
.find(|account| account.id == account_id)
.cloned()
});
Ok(account)
}
async fn load_account_usage(
@@ -488,11 +503,12 @@ impl AccountStore for FakeAccountStore {
_: CredentialListQuery,
) -> AdminStoreResult<CredentialPage> {
self.record("store.list_credentials");
let accounts = self.accounts.lock().expect("accounts").clone();
Ok(CredentialPage {
config_revision: revision(1),
items: (provider_kind == &self.account.provider_kind)
.then(|| self.account.clone())
items: accounts
.into_iter()
.filter(|account| &account.provider_kind == provider_kind)
.collect(),
next_cursor: None,
})
@@ -504,14 +520,19 @@ impl AccountStore for FakeAccountStore {
account_id: &ProviderAccountId,
) -> AdminStoreResult<Option<CredentialDetails>> {
self.record("store.credential_details");
Ok(
(provider_kind == &self.account.provider_kind
&& account_id.as_str() == self.account.id)
.then(|| CredentialDetails {
config_revision: revision(1),
credential: self.account.clone(),
}),
)
let account = self
.accounts
.lock()
.expect("accounts")
.iter()
.find(|account| {
&account.provider_kind == provider_kind && account.id == account_id.as_str()
})
.cloned();
Ok(account.map(|credential| CredentialDetails {
config_revision: revision(1),
credential,
}))
}
async fn load_credentials_for_export(
@@ -520,21 +541,26 @@ impl AccountStore for FakeAccountStore {
account_ids: &[ProviderAccountId],
) -> AdminStoreResult<Vec<ProviderExportCredentialInput>> {
self.record("store.load_credentials_for_export");
if provider_kind != &self.account.provider_kind
|| account_ids
.iter()
.any(|account_id| account_id.as_str() != self.account.id)
{
let accounts = self.accounts.lock().expect("accounts").clone();
if accounts.iter().any(|account| {
&account.provider_kind != provider_kind
|| !account_ids
.iter()
.any(|account_id| account_id.as_str() == account.id)
}) {
return Err(AdminStoreError::new(
AdminStoreErrorKind::NotFound,
"test account",
"credential not found",
));
}
Ok(vec![ProviderExportCredentialInput {
account: self.account.clone(),
provider_material: document(),
}])
Ok(accounts
.into_iter()
.map(|account| ProviderExportCredentialInput {
account,
provider_material: document(),
})
.collect())
}
async fn commit_credential_import(
@@ -818,7 +844,7 @@ async fn connection_test_should_probe_unavailable_account() {
}
#[tokio::test]
async fn connection_test_reloads_resettable_usage_limit_status_after_probe() {
async fn connection_test_rate_limited_probe_preserves_active_account_status() {
let events = events();
let provider = FakeProviderAdmin::new("xai", events.clone());
let store = FakeAccountStore::new("xai", events);
@@ -841,7 +867,7 @@ async fn connection_test_reloads_resettable_usage_limit_status_after_probe() {
assert!(matches!(
events.last(),
Some(AccountConnectionTestEvent::Failed {
account_status: gateway_admin::model::accounts::AccountStatus::QuotaExhausted,
account_status: gateway_admin::model::accounts::AccountStatus::Active,
..
})
));
@@ -880,7 +906,7 @@ async fn connection_test_should_preserve_disabled_account_status() {
async fn accounts_export_should_pass_store_loaded_timestamps_and_material_to_provider() {
let provider = FakeProviderAdmin::new("openai", events());
let store = FakeAccountStore::new("openai", events());
let expected = store.account.clone();
let expected = store.accounts.lock().expect("accounts")[0].clone();
accounts_service(provider.clone(), store)
.await
.accounts()
@@ -968,12 +994,75 @@ async fn accounts_list_should_return_complete_directory_semantics() {
}
#[tokio::test]
async fn accounts_list_should_prefer_usage_limit_exhaustion_over_stale_token_expiry() {
async fn accounts_list_should_degrade_quota_failure_to_empty_window_without_dropping_page() {
let events = events();
let openai = FakeProviderAdmin::new("openai", events.clone());
openai.set_quota(ProviderQuota {
observed_at: Some(Utc::now()),
refresh_token_expires_at: None,
windows: vec![ProviderQuotaWindow {
key: "primary".to_owned(),
group: "shortTerm".to_owned(),
label: "5小时限额".to_owned(),
source: None,
window_seconds: Some(5 * 60 * 60),
used_percent: Some(97.0),
reset_at: Some(Utc::now() + TimeDelta::hours(1)),
limit_reached: false,
local_usage: None,
provider_data: None,
}],
provider_data: None,
});
let failing = FakeProviderAdmin::new("xai", events.clone());
failing.fail_next_quota(ProviderAdminErrorKind::Invalid);
let store = FakeAccountStore::new("openai", events.clone());
let mut xai_account = account_record("xai");
xai_account.id = "acct_bad_xai".to_owned();
store.set_accounts(vec![account_record("openai"), xai_account]);
let services = super::AdminHarness::new()
.accounts(store.clone())
.settings(Arc::new(StaticSettingsStore))
.provider(openai)
.provider(failing)
.probe(Arc::new(SuccessfulAccountProbe))
.build()
.await;
let page = services
.accounts()
.list(AccountListQuery {
page: 1,
page_size: gateway_admin::model::PageSize::new(20).expect("page size"),
provider_kind: None,
search: None,
status: None,
sort: None,
})
.await
.expect("one failing quota must not fail the directory page");
assert_eq!(page.items.len(), 2);
let by_id = page
.items
.iter()
.map(|item| (item.account.id.as_str(), item))
.collect::<std::collections::HashMap<_, _>>();
let healthy = by_id.get("acct_test").expect("healthy account item");
assert_eq!(healthy.quota.windows.len(), 1);
assert_eq!(healthy.quota.windows[0].key, "primary");
let degraded = by_id.get("acct_bad_xai").expect("degraded account item");
assert!(degraded.quota.windows.is_empty());
assert!(degraded.quota.observed_at.is_none());
assert!(degraded.quota.provider_data.is_none());
assert_eq!(page.summary.total, 2);
}
#[tokio::test]
async fn accounts_list_should_prefer_quota_exhaustion_over_stale_token_expiry() {
let provider = FakeProviderAdmin::new("openai", events());
let mut account = account_record("openai");
account.availability = AccountAvailability::Cooldown;
account.availability_reason = Some("usage_limit_exhausted".to_owned());
account.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
account.availability = AccountAvailability::QuotaExhausted;
account.access_token_expires_at = Some(Utc::now() - TimeDelta::minutes(5));
let store = FakeAccountStore::with_account(account, events());
@@ -1012,6 +1101,7 @@ async fn accounts_list_should_attach_local_usage_to_quota_windows() {
window_seconds: Some(5 * 60 * 60),
used_percent: Some(97.0),
reset_at: Some(reset_at),
limit_reached: false,
local_usage: None,
provider_data: None,
}],
@@ -1166,8 +1256,6 @@ pub(super) fn account_record(kind: &str) -> AccountRecord {
next_refresh_at: Some(now + TimeDelta::minutes(30)),
enabled: true,
availability: AccountAvailability::Ready,
availability_reason: None,
cooldown_until: None,
availability_observed_at: now,
quota_observed_at: None,
created_at: now,
@@ -1204,8 +1292,6 @@ fn prepared_create_with_id(
next_refresh_at: Some(now + TimeDelta::minutes(30)),
enabled: true,
availability: AccountAvailability::Ready,
availability_reason: None,
cooldown_until: None,
availability_observed_at: now,
}
}
@@ -1290,9 +1376,7 @@ impl AccountProbe for FreeModelQuotaProbe {
) -> BoxFuture<'_, Result<AccountProbeResult, AccountProbeError>> {
Box::pin(async move {
let mut account = account_record("xai");
account.availability = AccountAvailability::Cooldown;
account.availability_reason = Some("usage_limit_exhausted".to_owned());
account.cooldown_until = Some(Utc::now() + TimeDelta::hours(24));
account.availability = AccountAvailability::Ready;
self.store.set_account_after_probe(account);
Err(GatewayError::new(
GatewayErrorKind::RateLimited,
@@ -194,7 +194,6 @@ pub struct AccountView {
pub has_refresh_token: bool,
pub status: String,
pub display_status: String,
pub token_refreshing: bool,
pub availability: String,
pub enabled: bool,
pub access_token_expires_at: Option<String>,
@@ -227,6 +226,7 @@ pub struct AccountQuotaWindowView {
pub label_display: String,
pub used_percent: Option<f64>,
pub used_percent_display: String,
pub limit_reached: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub local_usage: Option<serde_json::Value>,
pub reset_at_display: String,
@@ -1387,7 +1387,6 @@ fn account_view(item: AccountDirectoryItem, now: DateTime<Utc>) -> AccountView {
has_refresh_token: account.has_refresh_token,
status: status.clone(),
display_status: status,
token_refreshing: false,
availability: account.availability.as_str().to_owned(),
enabled: account.enabled,
access_token_expires_at: expires_at,
@@ -1435,6 +1434,7 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
window_seconds,
used_percent,
reset_at,
limit_reached,
local_usage,
provider_data: _,
} = window;
@@ -1446,6 +1446,7 @@ fn quota_window_view(window: ProviderQuotaWindow) -> AccountQuotaWindowView {
used_percent,
used_percent_display: used_percent
.map_or_else(|| "—".to_owned(), |value| format!("{value:.1}%")),
limit_reached,
local_usage: local_usage.as_ref().map(quota_local_usage),
reset_at_display: reset_at.map_or_else(|| "—".to_owned(), |value| china_datetime(&value)),
}
@@ -784,6 +784,7 @@ pub(crate) fn dashboard_view(
pool_summary: DashboardPoolSummaryView {
total: provider_accounts.total,
active: provider_accounts.active,
rate_limited: provider_accounts.rate_limited,
expired: provider_accounts.expired,
quota_exhausted: provider_accounts.quota_exhausted,
refreshing: provider_accounts.refreshing,
@@ -348,6 +348,7 @@ pub struct DashboardAccountRequestBucketView {
pub struct DashboardPoolSummaryView {
pub total: u64,
pub active: u64,
pub rate_limited: u64,
pub expired: u64,
pub quota_exhausted: u64,
pub refreshing: Option<u64>,
@@ -224,6 +224,7 @@ async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headli
enabled: 804,
unavailable: 792,
active: 15,
rate_limited: 0,
expired: 2,
quota_exhausted: 785,
refreshing: None,
@@ -188,11 +188,13 @@ impl fmt::Debug for OpaqueProviderData {
}
/// 数据库中固定的账号运行状态。
///
/// 只承载终态或需要外部干预的状态;临时限流(429)不进入此状态,由
/// Provider quota 数据驱动(`AccountQuotaSignals::limit_reached`)。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AccountAvailability {
Unknown,
Ready,
Cooldown,
QuotaExhausted,
Expired,
Banned,
@@ -206,7 +208,6 @@ impl AccountAvailability {
match self {
Self::Unknown => "unknown",
Self::Ready => "ready",
Self::Cooldown => "cooldown",
Self::QuotaExhausted => "quota_exhausted",
Self::Expired => "expired",
Self::Banned => "banned",
@@ -220,7 +221,6 @@ impl AccountAvailability {
match value {
"unknown" => Some(Self::Unknown),
"ready" => Some(Self::Ready),
"cooldown" => Some(Self::Cooldown),
"quota_exhausted" => Some(Self::QuotaExhausted),
"expired" => Some(Self::Expired),
"banned" => Some(Self::Banned),
@@ -244,7 +244,6 @@ pub struct ProviderAccount {
revision: CredentialRevision,
enabled: bool,
availability: AccountAvailability,
cooldown_until: Option<SystemTime>,
access_token_expires_at: Option<SystemTime>,
next_refresh_at: Option<SystemTime>,
has_refresh_token: bool,
@@ -274,7 +273,6 @@ impl ProviderAccount {
revision,
enabled: true,
availability: AccountAvailability::Unknown,
cooldown_until: None,
access_token_expires_at,
next_refresh_at: None,
has_refresh_token: false,
@@ -299,11 +297,9 @@ impl ProviderAccount {
mut self,
enabled: bool,
availability: AccountAvailability,
cooldown_until: Option<SystemTime>,
) -> Self {
self.enabled = enabled;
self.availability = availability;
self.cooldown_until = cooldown_until;
self
}
@@ -373,11 +369,6 @@ impl ProviderAccount {
self.enabled
}
#[must_use]
pub const fn cooldown_until(&self) -> Option<SystemTime> {
self.cooldown_until
}
#[must_use]
pub const fn access_token_expires_at(&self) -> Option<SystemTime> {
self.access_token_expires_at
@@ -398,7 +389,6 @@ impl ProviderAccount {
pub fn is_schedulable(&self, now: SystemTime) -> bool {
let available = match self.availability {
AccountAvailability::Ready => true,
AccountAvailability::Cooldown => self.cooldown_until.is_some_and(|until| until <= now),
AccountAvailability::Unknown
| AccountAvailability::QuotaExhausted
| AccountAvailability::Expired
@@ -614,8 +604,6 @@ pub struct AccountStateChange {
pub account_id: ProviderAccountId,
pub expected_revision: CredentialRevision,
pub availability: AccountAvailability,
pub reason: Option<String>,
pub cooldown_until: Option<SystemTime>,
pub observed_at: SystemTime,
}
@@ -756,6 +744,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 failure_rate_basis_points: Option<u16>,
pub first_output_latency_ms: Option<u64>,
}
@@ -903,6 +892,7 @@ 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
}
@@ -924,14 +914,20 @@ 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>) -> Self {
pub const fn new(
reset_at: Option<SystemTime>,
remaining_rank: Option<u64>,
limit_reached: bool,
) -> Self {
Self {
reset_at,
remaining_rank,
limit_reached,
}
}
@@ -944,6 +940,12 @@ 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
}
}
/// 账号持久事实与可重建运行信号的请求级组合。
@@ -24,7 +24,7 @@ fn account(id: &str) -> ProviderAccount {
CredentialRevision::new(1).expect("valid revision"),
Some(SystemTime::now() + Duration::from_secs(3600)),
)
.with_runtime_state(true, AccountAvailability::Ready, None)
.with_runtime_state(true, AccountAvailability::Ready)
}
fn candidate(id: &str, in_flight: u32, remaining: Option<u64>) -> AccountCandidate {
@@ -35,6 +35,7 @@ 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,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
@@ -110,31 +111,27 @@ fn availability_should_round_trip_all_database_values() {
#[test]
fn diagnostic_selection_bypasses_all_local_account_availability() {
let exhausted = AccountCandidate {
account: account("acct_exhausted").with_runtime_state(
true,
AccountAvailability::QuotaExhausted,
None,
),
account: account("acct_exhausted")
.with_runtime_state(true, AccountAvailability::QuotaExhausted),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
};
let disabled = AccountCandidate {
account: account("acct_disabled").with_runtime_state(
false,
AccountAvailability::QuotaExhausted,
None,
),
account: account("acct_disabled")
.with_runtime_state(false, AccountAvailability::QuotaExhausted),
signals: AccountRuntimeSignals {
in_flight: 0,
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
@@ -158,33 +155,11 @@ 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, None);
let account = account("acct_disabled").with_runtime_state(false, AccountAvailability::Ready);
assert!(!account.is_schedulable(SystemTime::now()));
}
#[test]
fn cooldown_account_should_become_schedulable_at_its_deadline() {
let now = SystemTime::now();
let account =
account("acct_cooldown").with_runtime_state(true, AccountAvailability::Cooldown, Some(now));
assert!(account.is_schedulable(now));
}
#[test]
fn cooldown_account_should_remain_isolated_before_its_deadline() {
let now = SystemTime::now();
let account = account("acct_cooldown").with_runtime_state(
true,
AccountAvailability::Cooldown,
Some(now + Duration::from_secs(1)),
);
assert!(!account.is_schedulable(now));
}
#[test]
fn smart_selector_should_prefer_lower_inflight_count() {
let candidates = vec![
@@ -311,10 +286,15 @@ 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,
failure_rate_basis_points: None,
first_output_latency_ms: None,
}
.with_provider_quota(Some(AccountQuotaSignals::new(Some(reset_at), Some(75))));
.with_provider_quota(Some(AccountQuotaSignals::new(
Some(reset_at),
Some(75),
false,
)));
assert_eq!(
(
@@ -101,6 +101,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,
failure_rate_basis_points: Some(125),
first_output_latency_ms: Some(250),
},
+4 -7
View File
@@ -1,8 +1,10 @@
//! 完成连接、迁移与 hydration 的 Store 能力集合与启动屏障。
use gateway_core::engine::credential::ProviderAccountStore;
use super::*;
/// 已完成连接、迁移与 cooldown hydration 的 Store 能力集合。
/// 已完成连接、迁移与 hydration 的 Store 能力集合。
pub struct StoreBundle {
admin_ports: AdminStorePorts,
core_ports: CoreStorePorts,
@@ -61,12 +63,7 @@ pub async fn initialize(mut config: StoreConfig) -> StoreResult<StoreBundle> {
redis_connection.clone(),
REDIS_NAMESPACE,
)?);
let account_store = Arc::new(redis::CooldownCachingProviderAccountStore::new(
provider_accounts,
cooldowns.clone(),
));
let hydrated = account_store.hydrate(std::time::SystemTime::now()).await;
tracing::info!(hydrated, "账号 cooldown Redis 热缓存重建完成");
let account_store: Arc<dyn ProviderAccountStore> = provider_accounts;
let credential_leases =
redis::RedisCredentialLeaseRepository::new(redis_connection.clone(), REDIS_NAMESPACE)?;
@@ -263,6 +263,7 @@ pub(crate) const fn admin_account_pool_metrics(
enabled: metrics.enabled,
unavailable: metrics.unavailable,
active: metrics.active,
rate_limited: metrics.rate_limited,
expired: metrics.expired,
quota_exhausted: metrics.quota_exhausted,
refreshing: None,
@@ -366,6 +366,7 @@ pub struct ProviderAccountMetrics {
pub enabled: u64,
pub unavailable: u64,
pub active: u64,
pub rate_limited: u64,
pub expired: u64,
pub quota_exhausted: u64,
pub disabled: u64,
@@ -348,17 +348,11 @@ pub(crate) async fn provider_account_metrics(
) -> StoreResult<ProviderAccountMetrics> {
let row = sqlx::query(
"with normalized as (
select enabled,
select enabled, availability, provider_quota_json,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
when availability = 'quota_exhausted'
or (
availability = 'cooldown'
and availability_reason = 'usage_limit_exhausted'
and (cooldown_until is null or cooldown_until > $1)
)
then 'quota_exhausted'
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'
@@ -370,6 +364,17 @@ pub(crate) async fn provider_account_metrics(
count(*) filter (where enabled)::bigint as enabled,
count(*) filter (where status <> 'active')::bigint as unavailable,
count(*) filter (where status = 'active')::bigint as active,
count(*) filter (
where status = 'active'
and (
provider_quota_json is not null
and (
provider_quota_json->'rate_limit'->>'limit_reached' = 'true'
or (provider_quota_json->'rate_limit'->'primary_window'->>'used_percent')::numeric >= 100
or (provider_quota_json->'rate_limit'->'secondary_window'->>'used_percent')::numeric >= 100
)
)
)::bigint as rate_limited,
count(*) filter (where status = 'expired')::bigint as expired,
count(*) filter (where status = 'quota_exhausted')::bigint as quota_exhausted,
count(*) filter (where status = 'disabled')::bigint as disabled,
@@ -385,6 +390,7 @@ pub(crate) async fn provider_account_metrics(
enabled: unsigned(&row, "enabled")?,
unavailable: unsigned(&row, "unavailable")?,
active: unsigned(&row, "active")?,
rate_limited: unsigned(&row, "rate_limited")?,
expired: unsigned(&row, "expired")?,
quota_exhausted: unsigned(&row, "quota_exhausted")?,
disabled: unsigned(&row, "disabled")?,
@@ -402,13 +408,7 @@ pub(crate) async fn active_provider_account_ids(
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
when availability = 'quota_exhausted'
or (
availability = 'cooldown'
and availability_reason = 'usage_limit_exhausted'
and (cooldown_until is null or cooldown_until > $1)
)
then 'quota_exhausted'
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'
@@ -32,8 +32,6 @@ impl ProviderAccountStore for PgProviderAccountRepository {
next_refresh_at: account.account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.account.enabled(),
availability: account.account.availability(),
availability_reason: None,
cooldown_until: account.account.cooldown_until().map(DateTime::<Utc>::from),
availability_observed_at: Utc::now(),
})
.await
@@ -242,8 +240,6 @@ impl ProviderAccountStore for PgProviderAccountRepository {
expected_revision: Revision::new(change.expected_revision.get())
.map_err(core_store_error)?,
availability: change.availability,
availability_reason: change.reason,
cooldown_until: change.cooldown_until.map(DateTime::<Utc>::from),
availability_observed_at: DateTime::<Utc>::from(change.observed_at),
})
.await
@@ -63,11 +63,7 @@ pub(crate) fn admin_account_status(
AdminAccountStatus::Disabled
} else if account.availability == AccountAvailability::Banned {
AdminAccountStatus::Banned
} else if account.availability == AccountAvailability::QuotaExhausted
|| (account.availability == AccountAvailability::Cooldown
&& account.availability_reason.as_deref() == Some("usage_limit_exhausted")
&& account.cooldown_until.is_none_or(|until| until > now))
{
} else if account.availability == AccountAvailability::QuotaExhausted {
AdminAccountStatus::QuotaExhausted
} else if account.availability == AccountAvailability::Expired
|| account.availability == AccountAvailability::Invalid
@@ -184,8 +180,6 @@ pub(crate) fn admin_account_record(
next_refresh_at: summary.next_refresh_at,
enabled: summary.enabled,
availability: summary.availability,
availability_reason: summary.availability_reason,
cooldown_until: summary.cooldown_until,
availability_observed_at: summary.availability_observed_at,
quota_observed_at: summary.quota_observed_at,
created_at: summary.created_at,
@@ -211,8 +205,6 @@ pub(crate) fn prepared_account(
next_refresh_at: credential.next_refresh_at,
enabled: credential.enabled,
availability: credential.availability,
availability_reason: credential.availability_reason,
cooldown_until: credential.cooldown_until,
availability_observed_at: credential.availability_observed_at,
})
}
@@ -93,7 +93,7 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
"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_reason, cooldown_until, availability_observed_at,
availability_observed_at,
quota_observed_at, created_at, updated_at
from provider_accounts
where ($1::text is null or provider_kind = $1) and ($2 or enabled)
@@ -114,11 +114,11 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
availability, availability_reason, cooldown_until, provider_quota_json,
availability, provider_quota_json,
availability_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, $16, null, now(), greatest(now(), $16)
$14, null, $15, null, now(), greatest(now(), $15)
)",
)
.bind(account.id)
@@ -135,7 +135,6 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
.bind(account.next_refresh_at)
.bind(account.enabled)
.bind(account.availability.as_str())
.bind(account.cooldown_until)
.bind(account.availability_observed_at)
.execute(&self.pool)
.await
@@ -210,21 +209,17 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
let result = sqlx::query(
"update provider_accounts
set availability = case when enabled then $3 else availability end,
availability_reason = case when enabled then $4 else availability_reason end,
cooldown_until = case when enabled then $5 else cooldown_until end,
availability_observed_at = case
when enabled then $6
when enabled then $4
else availability_observed_at
end,
updated_at = case when enabled then greatest(now(), $6) else updated_at 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 <= $6)",
and (availability_observed_at is null or availability_observed_at <= $4)",
)
.bind(update.account_id)
.bind(to_i64(update.expected_revision.get())?)
.bind(update.availability.as_str())
.bind(update.availability_reason)
.bind(update.cooldown_until)
.bind(update.availability_observed_at)
.execute(&self.pool)
.await
@@ -466,11 +461,11 @@ pub(crate) async fn upsert_provider_account_in_transaction(
id, provider_kind, name, email, upstream_user_id,
upstream_account_id, plan_type, authentication_kind, provider_credentials_json, credential_revision,
has_refresh_token, access_token_expires_at, next_refresh_at, enabled,
availability, availability_reason, cooldown_until, provider_quota_json,
availability, provider_quota_json,
availability_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, $15, $16, null, $17, null, now(), greatest(now(), $17)
$14, null, $15, null, now(), greatest(now(), $15)
)
on conflict (
provider_kind,
@@ -488,8 +483,6 @@ pub(crate) async fn upsert_provider_account_in_transaction(
next_refresh_at = excluded.next_refresh_at,
enabled = excluded.enabled,
availability = excluded.availability,
availability_reason = excluded.availability_reason,
cooldown_until = excluded.cooldown_until,
provider_quota_json = null,
availability_observed_at = excluded.availability_observed_at,
quota_observed_at = null,
@@ -510,8 +503,6 @@ pub(crate) async fn upsert_provider_account_in_transaction(
.bind(account.next_refresh_at)
.bind(account.enabled)
.bind(account.availability.as_str())
.bind(&account.availability_reason)
.bind(account.cooldown_until)
.bind(account.availability_observed_at)
.fetch_optional(&mut **transaction)
.await
@@ -565,16 +556,6 @@ pub(crate) async fn rotate_provider_account_in_transaction(
when availability <> 'quota_exhausted' then 'ready'
else availability
end,
availability_reason = case
when not enabled then availability_reason
when availability <> 'quota_exhausted' then null
else availability_reason
end,
cooldown_until = case
when not enabled then cooldown_until
when availability <> 'quota_exhausted' then null
else cooldown_until
end,
availability_observed_at = case
when not enabled then availability_observed_at
when availability <> 'quota_exhausted' then now()
@@ -70,26 +70,6 @@ pub(crate) fn parse_availability(value: &str) -> StoreResult<AccountAvailability
AccountAvailability::parse(value).ok_or_else(|| invalid("unknown availability value"))
}
pub(crate) fn validate_availability_cooldown(
availability: AccountAvailability,
has_cooldown_until: bool,
) -> StoreResult<()> {
if availability == AccountAvailability::Cooldown && !has_cooldown_until {
return Err(invalid("cooldown availability requires cooldown_until"));
}
if has_cooldown_until
&& !matches!(
availability,
AccountAvailability::Cooldown | AccountAvailability::QuotaExhausted
)
{
return Err(invalid(
"cooldown_until is only allowed for cooldown or quota_exhausted availability",
));
}
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountSummary {
pub id: String,
@@ -106,8 +86,6 @@ pub struct ProviderAccountSummary {
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
pub created_at: DateTime<Utc>,
@@ -151,8 +129,6 @@ pub struct NewProviderAccount {
pub next_refresh_at: Option<DateTime<Utc>>,
pub enabled: bool,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
}
@@ -173,7 +149,6 @@ impl fmt::Debug for NewProviderAccount {
.field("name", &self.name)
.field("enabled", &self.enabled)
.field("availability", &self.availability)
.field("cooldown_until", &self.cooldown_until)
.field("provider_credentials_json", &"[REDACTED]")
.finish_non_exhaustive()
}
@@ -189,7 +164,6 @@ impl NewProviderAccount {
if !self.has_refresh_token && self.next_refresh_at.is_some() {
return Err(invalid("next_refresh_at requires a refresh token"));
}
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
validate_object_size(
"provider_credentials_json",
&self.provider_credentials_json,
@@ -321,8 +295,6 @@ impl fmt::Debug for ProviderCredentialUpdate {
pub struct ProviderAccountObservation {
pub account_id: String,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub provider_quota_json: Option<JsonObject>,
pub availability_observed_at: DateTime<Utc>,
pub quota_observed_at: Option<DateTime<Utc>>,
@@ -333,15 +305,12 @@ pub struct ProviderAccountStateUpdate {
pub account_id: String,
pub expected_revision: Revision,
pub availability: AccountAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub availability_observed_at: DateTime<Utc>,
}
impl ProviderAccountStateUpdate {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
Ok(())
}
}
@@ -349,7 +318,6 @@ impl ProviderAccountStateUpdate {
impl ProviderAccountObservation {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
if self.provider_quota_json.is_some() != self.quota_observed_at.is_some() {
return Err(invalid("quota JSON and quota_observed_at must agree"));
}
@@ -363,14 +331,14 @@ impl ProviderAccountObservation {
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,
availability_reason, cooldown_until, provider_quota_json,
provider_quota_json,
availability_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,
availability_reason, cooldown_until, provider_quota_json,
provider_quota_json,
availability_observed_at, quota_observed_at, created_at, updated_at
from provider_accounts
where id = any($1::text[]) and provider_kind = $2
@@ -420,11 +388,7 @@ pub(crate) fn core_account_from_summary(
summary.upstream_account_id,
summary.plan_type,
)
.with_runtime_state(
summary.enabled,
summary.availability,
summary.cooldown_until.map(Into::into),
)
.with_runtime_state(summary.enabled, summary.availability)
.with_refresh_schedule(
summary.has_refresh_token,
summary.next_refresh_at.map(Into::into),
@@ -474,8 +438,6 @@ pub(crate) fn account_summary_from_row(
next_refresh_at: get(&row, "next_refresh_at")?,
enabled: get(&row, "enabled")?,
availability: parse_availability(&availability)?,
availability_reason: get(&row, "availability_reason")?,
cooldown_until: get(&row, "cooldown_until")?,
availability_observed_at: get(&row, "availability_observed_at")?,
quota_observed_at: get(&row, "quota_observed_at")?,
created_at: get(&row, "created_at")?,
@@ -1,22 +1,13 @@
//! 请求调度用的可丢失账号 cooldown Redis 存储。
use std::{sync::Arc, time::SystemTime};
//! 请求调度用的可丢失 Provider 级 cooldown Redis 存储。
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_core::{
engine::credential::{
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision, LoadedCredential, NewProviderAccount, ProviderAccount,
ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, QuotaObservation,
QuotaWriteOutcome,
},
error::StoreError as CoreStoreError,
engine::credential::{CredentialRevision, ProviderAccountId},
provider_ports::{
ProviderCooldown, ProviderCooldownPort, ProviderCooldownScope, ProviderScopedCooldown,
ProviderStoreError, ProviderStoreErrorKind,
},
routing::ProviderKind,
};
use redis::{Script, aio::ConnectionManager};
@@ -352,216 +343,6 @@ impl ProviderCooldownPort for RedisCredentialCooldownRepository {
}
}
/// 账号状态写入到可丢失 cooldown cache 的桥接 adapter。
pub struct CooldownCachingProviderAccountStore {
authoritative: Arc<dyn ProviderAccountStore>,
cooldowns: Arc<dyn CredentialCooldownRepository>,
}
impl CooldownCachingProviderAccountStore {
#[must_use]
pub const fn new(
authoritative: Arc<dyn ProviderAccountStore>,
cooldowns: Arc<dyn CredentialCooldownRepository>,
) -> Self {
Self {
authoritative,
cooldowns,
}
}
/// 启动时从 PostgreSQL 尽力重建仍有效的 cooldown cache。
pub async fn hydrate(&self, now: SystemTime) -> u64 {
let Ok(accounts) = self.authoritative.list_accounts().await else {
return 0;
};
let mut hydrated = 0_u64;
for account in accounts {
let Some(cooldown) = projection(
account.id(),
account.revision(),
account.availability(),
account.cooldown_until(),
now,
) else {
continue;
};
if matches!(
self.cooldowns.cache_credential_cooldown(&cooldown).await,
Ok(true)
) {
hydrated = hydrated.saturating_add(1);
}
}
hydrated
}
async fn mirror(
&self,
account_id: &ProviderAccountId,
credential_revision: CredentialRevision,
availability: AccountAvailability,
cooldown_until: Option<SystemTime>,
) {
let Some(cooldown) = projection(
account_id,
credential_revision,
availability,
cooldown_until,
SystemTime::now(),
) else {
let Ok(revision) = Revision::new(credential_revision.get()) else {
return;
};
let _ = self
.cooldowns
.invalidate_credential_cooldown(account_id.as_str(), revision)
.await;
return;
};
let _ = self.cooldowns.cache_credential_cooldown(&cooldown).await;
}
}
#[async_trait]
impl ProviderAccountStore for CooldownCachingProviderAccountStore {
async fn create_account(&self, account: NewProviderAccount) -> Result<(), CoreStoreError> {
self.authoritative.create_account(account).await
}
async fn get_account(
&self,
account: &ProviderAccountId,
) -> Result<Option<ProviderAccount>, CoreStoreError> {
self.authoritative.get_account(account).await
}
async fn list_accounts(&self) -> Result<Vec<ProviderAccount>, CoreStoreError> {
self.authoritative.list_accounts().await
}
async fn list_for_provider(
&self,
provider: &ProviderKind,
) -> Result<Vec<ProviderAccount>, CoreStoreError> {
self.authoritative.list_for_provider(provider).await
}
async fn load_credential(
&self,
account: &ProviderAccountId,
expected_revision: CredentialRevision,
) -> Result<LoadedCredential, CoreStoreError> {
self.authoritative
.load_credential(account, expected_revision)
.await
}
async fn load_current_credential(
&self,
account: &ProviderAccountId,
) -> Result<LoadedCredential, CoreStoreError> {
self.authoritative.load_current_credential(account).await
}
async fn compare_and_swap_credential(
&self,
update: CredentialCasUpdate,
) -> Result<CredentialCasOutcome, CoreStoreError> {
let account_id = update.account_id().clone();
let outcome = self
.authoritative
.compare_and_swap_credential(update)
.await?;
if let CredentialCasOutcome::Updated(credential_revision) = outcome
&& let Ok(revision) = Revision::new(credential_revision.get())
{
let _ = self
.cooldowns
.invalidate_credential_cooldown(account_id.as_str(), revision)
.await;
}
Ok(outcome)
}
async fn get_quotas(
&self,
accounts: &[ProviderAccountId],
) -> Result<Vec<QuotaObservation>, CoreStoreError> {
self.authoritative.get_quotas(accounts).await
}
async fn compare_and_swap_quota(
&self,
observation: QuotaObservation,
) -> Result<QuotaWriteOutcome, CoreStoreError> {
self.authoritative.compare_and_swap_quota(observation).await
}
async fn apply_state_change(&self, change: AccountStateChange) -> Result<(), CoreStoreError> {
let account_id = change.account_id.clone();
let credential_revision = change.expected_revision;
let availability = change.availability;
let cooldown_until = change.cooldown_until;
self.authoritative.apply_state_change(change).await?;
let Some(account) = self.authoritative.get_account(&account_id).await? else {
return Ok(());
};
if !account.enabled() {
return Ok(());
}
self.mirror(
&account_id,
credential_revision,
availability,
cooldown_until,
)
.await;
Ok(())
}
async fn update_account(&self, update: ProviderAccountUpdate) -> Result<(), CoreStoreError> {
self.authoritative.update_account(update).await
}
async fn set_enabled(
&self,
account: &ProviderAccountId,
enabled: bool,
) -> Result<(), CoreStoreError> {
self.authoritative.set_enabled(account, enabled).await
}
async fn delete_account(&self, account: &ProviderAccountId) -> Result<(), CoreStoreError> {
self.authoritative.delete_account(account).await?;
if let Ok(through_revision) = Revision::new(u64::MAX) {
let _ = self
.cooldowns
.invalidate_credential_cooldown(account.as_str(), through_revision)
.await;
}
Ok(())
}
}
fn projection(
account_id: &ProviderAccountId,
credential_revision: CredentialRevision,
availability: AccountAvailability,
cooldown_until: Option<SystemTime>,
now: SystemTime,
) -> Option<CredentialCooldown> {
let cooldown_until = cooldown_until
.filter(|_| availability == AccountAvailability::Cooldown)
.filter(|until| *until > now)
.map(DateTime::<Utc>::from)?;
Some(CredentialCooldown {
provider_account_id: account_id.as_str().to_owned(),
credential_revision: Revision::new(credential_revision.get()).ok()?,
cooldown_until,
})
}
fn invalid(message: &str) -> StoreError {
StoreError::InvalidData {
entity: "credential cooldown",
@@ -512,6 +512,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,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
@@ -28,9 +28,9 @@ use gateway_store::{
postgres::{
AdminAuditActorKind, AdminAuditEvent, DeleteProviderAccounts, ImportProviderAccounts,
NewProviderAccount, PgAdminAccountStore, PgProviderAccountRepository,
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountObservation,
ProviderAccountRepository, ProviderCredentialUpdate, RotateProviderAccount,
SetProviderAccountEnabled, UpdateProviderAccount,
ProviderAccountAdminRepository, ProviderAccountAdminScope, ProviderAccountRepository,
ProviderCredentialUpdate, RotateProviderAccount, SetProviderAccountEnabled,
UpdateProviderAccount,
},
};
use serde_json::json;
@@ -161,8 +161,6 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
account_id: account_id.clone(),
expected_revision: revision,
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: Some(newer),
observed_at: newer,
})
.await
@@ -173,8 +171,6 @@ async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
account_id: account_id.clone(),
expected_revision: revision,
availability: AccountAvailability::Ready,
reason: None,
cooldown_until: None,
observed_at: older,
})
.await
@@ -747,7 +743,6 @@ async fn admin_import_updates_the_same_verified_identity_without_rebinding_it()
updated.name = "updated import".to_owned();
updated.provider_credentials_json = credential_json("updated-import-secret");
updated.availability = AccountAvailability::Banned;
updated.availability_reason = Some("upstream_account_deactivated".to_owned());
let imported = repository
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
@@ -819,8 +814,6 @@ async fn core_refresh_cas_updates_profile_and_credential_under_one_revision() {
next_refresh_at: None,
enabled: true,
availability: AccountAvailability::Ready,
availability_reason: None,
cooldown_until: None,
availability_observed_at: Utc::now(),
})
.await
@@ -942,8 +935,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
};
let mut cooldown_account = account("acct_admin_b", "user-admin-b");
cooldown_account.availability = AccountAvailability::Cooldown;
cooldown_account.cooldown_until = Some(Utc::now() + TimeDelta::minutes(10));
cooldown_account.availability = AccountAvailability::Ready;
let imported = repository
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
@@ -1158,7 +1150,6 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
};
let mut exhausted = account("acct_rotation_quota", "user-rotation-quota");
exhausted.availability = AccountAvailability::QuotaExhausted;
exhausted.availability_reason = Some("quota_exhausted".to_owned());
repository
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
@@ -1224,29 +1215,20 @@ async fn xai_free_quota_exhaustion_state_is_persisted() {
account_id: account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
availability: AccountAvailability::QuotaExhausted,
reason: Some("upstream_free_quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("persist xAI free quota exhaustion");
let current: (String, String, Option<String>) = sqlx::query_as(
"select provider_kind, availability, availability_reason
let current: (String, String) = sqlx::query_as(
"select provider_kind, availability
from provider_accounts where id = $1",
)
.bind(account_id.as_str())
.fetch_one(&database.pool)
.await
.expect("load persisted xAI quota state");
assert_eq!(
current,
(
"xai".to_owned(),
"quota_exhausted".to_owned(),
Some("upstream_free_quota_exhausted".to_owned()),
),
);
assert_eq!(current, ("xai".to_owned(), "quota_exhausted".to_owned()),);
database.close().await;
}
@@ -1266,38 +1248,25 @@ async fn xai_resettable_usage_limit_state_is_persisted() {
.expect("insert xAI account");
let observed_at = SystemTime::now();
let cooldown_until = observed_at
.checked_add(Duration::from_secs(24 * 60 * 60))
.expect("cooldown deadline");
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
availability: AccountAvailability::Cooldown,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: Some(cooldown_until),
availability: AccountAvailability::Ready,
observed_at,
})
.await
.expect("persist xAI resettable usage limit");
let current: (String, String, Option<String>, bool) = sqlx::query_as(
"select provider_kind, availability, availability_reason, cooldown_until is not null
let current: (String, String) = sqlx::query_as(
"select provider_kind, availability
from provider_accounts where id = $1",
)
.bind(account_id.as_str())
.fetch_one(&database.pool)
.await
.expect("load persisted xAI usage limit state");
assert_eq!(
current,
(
"xai".to_owned(),
"cooldown".to_owned(),
Some("usage_limit_exhausted".to_owned()),
true,
),
);
assert_eq!(current, ("xai".to_owned(), "ready".to_owned()));
database.close().await;
}
@@ -1315,7 +1284,6 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
let mut disabled = account(account_id.as_str(), "user-disabled-refresh");
disabled.enabled = false;
disabled.availability = AccountAvailability::Expired;
disabled.availability_reason = Some("credential_expired".to_owned());
repository
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
@@ -1334,8 +1302,6 @@ async fn disabled_account_preserves_user_state_during_refresh_writes() {
account_id: account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
availability: AccountAvailability::Ready,
reason: None,
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -1389,8 +1355,6 @@ fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
next_refresh_at: None,
enabled: true,
availability: AccountAvailability::Ready,
availability_reason: None,
cooldown_until: None,
availability_observed_at: Utc::now(),
}
}
@@ -1517,31 +1481,15 @@ fn provider_credentials_are_redacted_from_debug() {
assert!(!format!("{object:?}").contains(secret));
}
#[test]
fn cooldown_observation_requires_expiry() {
let observation = ProviderAccountObservation {
account_id: "account-1".to_owned(),
availability: AccountAvailability::Cooldown,
availability_reason: None,
cooldown_until: None,
provider_quota_json: None,
availability_observed_at: Utc::now(),
quota_observed_at: None,
};
assert!(observation.validate().is_err());
}
#[test]
fn quota_exhausted_account_accepts_cooldown_runtime_facts() {
let mut imported = account("acct_quota_cooldown", "user-quota-cooldown");
imported.availability = AccountAvailability::QuotaExhausted;
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
assert!(imported.validate().is_ok());
}
#[test]
fn imported_account_rejects_cooldown_until_for_ready_runtime_facts() {
let mut imported = account("acct_invalid_cooldown", "user-invalid-cooldown");
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
assert!(imported.validate().is_err());
fn imported_ready_account_passes_validation() {
let imported = account("acct_ready_import", "user-ready-import");
assert!(imported.validate().is_ok());
}
@@ -1,27 +1,13 @@
use std::{
sync::{Arc, Mutex},
time::{Duration as StdDuration, SystemTime},
};
use std::time::{Duration as StdDuration, SystemTime};
use async_trait::async_trait;
use chrono::{DateTime, Duration, Utc};
use gateway_core::{
engine::credential::{
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision, LoadedCredential, NewProviderAccount, PlaintextCredential,
ProviderAccount, ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate,
QuotaObservation, QuotaWriteOutcome,
},
error::{StoreError as CoreStoreError, StoreErrorKind},
engine::credential::{CredentialRevision, ProviderAccountId},
provider_ports::{ProviderCooldownPort, ProviderCooldownScope, ProviderScopedCooldown},
routing::ProviderKind,
};
use gateway_store::{
Revision, StoreBackend, StoreError, StoreResult,
redis::{
CooldownCachingProviderAccountStore, CredentialCooldown, CredentialCooldownRepository,
RedisCredentialCooldownRepository,
},
Revision,
redis::{CredentialCooldown, CredentialCooldownRepository, RedisCredentialCooldownRepository},
};
use redis::aio::ConnectionManager;
use uuid::Uuid;
@@ -36,409 +22,6 @@ fn credential_cooldown_is_revision_fenced() {
assert_eq!(cooldown.credential_revision.get(), 2);
}
struct CooldownTestAccountStore {
accounts: Vec<ProviderAccount>,
calls: Arc<Mutex<Vec<&'static str>>>,
fail_state_write: bool,
}
#[async_trait]
impl ProviderAccountStore for CooldownTestAccountStore {
async fn create_account(&self, _account: NewProviderAccount) -> Result<(), CoreStoreError> {
unreachable!("cooldown tests do not create accounts")
}
async fn get_account(
&self,
account: &ProviderAccountId,
) -> Result<Option<ProviderAccount>, CoreStoreError> {
Ok(self
.accounts
.iter()
.find(|candidate| candidate.id() == account)
.cloned())
}
async fn list_accounts(&self) -> Result<Vec<ProviderAccount>, CoreStoreError> {
Ok(self.accounts.clone())
}
async fn list_for_provider(
&self,
_provider: &ProviderKind,
) -> Result<Vec<ProviderAccount>, CoreStoreError> {
unreachable!("cooldown tests do not list one provider")
}
async fn load_credential(
&self,
_account: &ProviderAccountId,
_expected_revision: CredentialRevision,
) -> Result<LoadedCredential, CoreStoreError> {
unreachable!("cooldown tests do not load credentials")
}
async fn load_current_credential(
&self,
_account: &ProviderAccountId,
) -> Result<LoadedCredential, CoreStoreError> {
unreachable!("cooldown tests do not load credentials")
}
async fn compare_and_swap_credential(
&self,
_update: CredentialCasUpdate,
) -> Result<CredentialCasOutcome, CoreStoreError> {
self.calls.lock().expect("calls lock").push("postgres");
Ok(CredentialCasOutcome::Updated(
CredentialRevision::new(8).expect("positive revision"),
))
}
async fn get_quotas(
&self,
_accounts: &[ProviderAccountId],
) -> Result<Vec<QuotaObservation>, CoreStoreError> {
unreachable!("cooldown tests do not read quota")
}
async fn compare_and_swap_quota(
&self,
_observation: QuotaObservation,
) -> Result<QuotaWriteOutcome, CoreStoreError> {
unreachable!("cooldown tests do not write quota")
}
async fn apply_state_change(&self, _change: AccountStateChange) -> Result<(), CoreStoreError> {
self.calls.lock().expect("calls lock").push("postgres");
if self.fail_state_write {
Err(CoreStoreError::new(StoreErrorKind::Unavailable))
} else {
Ok(())
}
}
async fn update_account(&self, _update: ProviderAccountUpdate) -> Result<(), CoreStoreError> {
unreachable!("cooldown tests do not update account metadata")
}
async fn set_enabled(
&self,
_account: &ProviderAccountId,
_enabled: bool,
) -> Result<(), CoreStoreError> {
unreachable!("cooldown tests do not change account switches")
}
async fn delete_account(&self, _account: &ProviderAccountId) -> Result<(), CoreStoreError> {
self.calls.lock().expect("calls lock").push("postgres");
Ok(())
}
}
struct CooldownTestCache {
calls: Arc<Mutex<Vec<&'static str>>>,
cached: Mutex<Vec<CredentialCooldown>>,
invalidated: Mutex<Vec<(String, Revision)>>,
fail: bool,
}
#[async_trait]
impl CredentialCooldownRepository for CooldownTestCache {
async fn cache_credential_cooldown(&self, cooldown: &CredentialCooldown) -> StoreResult<bool> {
self.calls.lock().expect("calls lock").push("redis");
if self.fail {
return Err(StoreError::Unavailable {
backend: StoreBackend::Redis,
message: "test cache unavailable".to_owned(),
});
}
self.cached
.lock()
.expect("cached lock")
.push(cooldown.clone());
Ok(true)
}
async fn read_credential_cooldown(
&self,
_provider_account_id: &str,
) -> StoreResult<Option<CredentialCooldown>> {
unreachable!("write-through adapter does not consume cooldown cache")
}
async fn invalidate_credential_cooldown(
&self,
provider_account_id: &str,
through_revision: Revision,
) -> StoreResult<bool> {
self.calls.lock().expect("calls lock").push("redis");
if self.fail {
return Err(StoreError::Unavailable {
backend: StoreBackend::Redis,
message: "test cache unavailable".to_owned(),
});
}
self.invalidated
.lock()
.expect("invalidated lock")
.push((provider_account_id.to_owned(), through_revision));
Ok(true)
}
}
fn composition_fixture(
fail_state_write: bool,
fail_cache: bool,
) -> (
CooldownCachingProviderAccountStore,
Arc<CooldownTestCache>,
Arc<Mutex<Vec<&'static str>>>,
) {
composition_fixture_with_accounts(
vec![ready_test_account("acct_cooldown_composition")],
fail_state_write,
fail_cache,
)
}
fn composition_fixture_with_accounts(
accounts: Vec<ProviderAccount>,
fail_state_write: bool,
fail_cache: bool,
) -> (
CooldownCachingProviderAccountStore,
Arc<CooldownTestCache>,
Arc<Mutex<Vec<&'static str>>>,
) {
let calls = Arc::new(Mutex::new(Vec::new()));
let authoritative = Arc::new(CooldownTestAccountStore {
accounts,
calls: Arc::clone(&calls),
fail_state_write,
});
let cache = Arc::new(CooldownTestCache {
calls: Arc::clone(&calls),
cached: Mutex::new(Vec::new()),
invalidated: Mutex::new(Vec::new()),
fail: fail_cache,
});
let store = CooldownCachingProviderAccountStore::new(
authoritative,
Arc::clone(&cache) as Arc<dyn CredentialCooldownRepository>,
);
(store, cache, calls)
}
#[tokio::test]
async fn cooldown_composition_writes_cache_only_after_postgres_commit() {
let (store, cache, calls) = composition_fixture(false, false);
store
.apply_state_change(cooldown_state_change(60))
.await
.expect("persist cooldown");
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
assert_eq!(cache.cached.lock().expect("cached lock").len(), 1);
}
#[tokio::test]
async fn cooldown_composition_does_not_cache_failed_postgres_state() {
let (store, cache, calls) = composition_fixture(true, false);
assert!(
store
.apply_state_change(cooldown_state_change(60))
.await
.is_err()
);
assert_eq!(*calls.lock().expect("calls lock"), ["postgres"]);
assert!(cache.cached.lock().expect("cached lock").is_empty());
}
#[tokio::test]
async fn cooldown_composition_keeps_postgres_terminal_when_redis_fails() {
let (store, _cache, calls) = composition_fixture(false, true);
store
.apply_state_change(cooldown_state_change(60))
.await
.expect("Redis is only a cache");
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
}
#[tokio::test]
async fn disabled_account_state_write_does_not_change_runtime_cooldown() {
let disabled = test_provider_account("acct_cooldown_composition").with_runtime_state(
false,
AccountAvailability::Ready,
None,
);
let (store, cache, calls) = composition_fixture_with_accounts(vec![disabled], false, false);
store
.apply_state_change(cooldown_state_change(60))
.await
.expect("disabled state write is a no-op");
assert_eq!(*calls.lock().expect("calls lock"), ["postgres"]);
assert!(cache.cached.lock().expect("cached lock").is_empty());
assert!(
cache
.invalidated
.lock()
.expect("invalidated lock")
.is_empty()
);
}
#[tokio::test]
async fn ready_state_invalidates_same_revision_cooldown_after_postgres_commit() {
let (store, cache, calls) = composition_fixture(false, false);
let mut change = cooldown_state_change(60);
change.availability = AccountAvailability::Ready;
change.reason = None;
change.cooldown_until = None;
store
.apply_state_change(change)
.await
.expect("persist ready state");
assert_eq!(
cache.invalidated.lock().expect("invalidated lock")[0].1,
Revision::new(7).expect("positive revision")
);
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
}
#[tokio::test]
async fn credential_rotation_invalidates_cooldown_through_new_revision() {
let (store, cache, calls) = composition_fixture(false, false);
let outcome = store
.compare_and_swap_credential(credential_update())
.await
.expect("rotate credential");
assert_eq!(
outcome,
CredentialCasOutcome::Updated(CredentialRevision::new(8).expect("positive revision"))
);
assert_eq!(
cache.invalidated.lock().expect("invalidated lock")[0].1,
Revision::new(8).expect("positive revision")
);
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
}
#[tokio::test]
async fn account_delete_invalidates_all_stale_cooldown_revisions() {
let (store, cache, calls) = composition_fixture(false, false);
let account_id = ProviderAccountId::new("acct_cooldown_delete").expect("valid account ID");
store
.delete_account(&account_id)
.await
.expect("delete account");
assert_eq!(
cache.invalidated.lock().expect("invalidated lock")[0],
(
account_id.to_string(),
Revision::new(u64::MAX).expect("positive revision")
)
);
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
}
#[tokio::test]
async fn cooldown_startup_hydrates_only_active_postgres_cooldowns() {
let now = SystemTime::now();
let calls = Arc::new(Mutex::new(Vec::new()));
let authoritative = Arc::new(CooldownTestAccountStore {
accounts: vec![
cooldown_test_account("acct_hydrate_active", now + StdDuration::from_secs(60)),
cooldown_test_account("acct_hydrate_expired", now - StdDuration::from_secs(1)),
ready_test_account("acct_hydrate_ready"),
],
calls: Arc::clone(&calls),
fail_state_write: false,
});
let cache = Arc::new(CooldownTestCache {
calls,
cached: Mutex::new(Vec::new()),
invalidated: Mutex::new(Vec::new()),
fail: false,
});
let store = CooldownCachingProviderAccountStore::new(
authoritative,
Arc::clone(&cache) as Arc<dyn CredentialCooldownRepository>,
);
assert_eq!(store.hydrate(now).await, 1);
assert_eq!(
cache.cached.lock().expect("cached lock")[0].provider_account_id,
"acct_hydrate_active"
);
}
fn cooldown_state_change(seconds: u64) -> AccountStateChange {
AccountStateChange {
account_id: ProviderAccountId::new("acct_cooldown_composition").expect("valid account ID"),
expected_revision: CredentialRevision::new(7).expect("positive revision"),
availability: AccountAvailability::Cooldown,
reason: Some("rate_limited".to_owned()),
cooldown_until: Some(SystemTime::now() + StdDuration::from_secs(seconds)),
observed_at: SystemTime::now(),
}
}
fn credential_update() -> CredentialCasUpdate {
let account_id = ProviderAccountId::new("acct_cooldown_rotation").expect("valid account ID");
CredentialCasUpdate::new(
account_id.clone(),
CredentialRevision::new(7).expect("positive revision"),
ProviderAccountUpdate {
account_id,
name: "Rotated cooldown account".to_owned(),
email: None,
plan_type: None,
},
PlaintextCredential::new(serde_json::Map::new()),
false,
Some(SystemTime::now() + StdDuration::from_secs(3_600)),
None,
)
.expect("valid credential update")
}
fn cooldown_test_account(account_id: &str, cooldown_until: SystemTime) -> ProviderAccount {
test_provider_account(account_id).with_runtime_state(
true,
AccountAvailability::Cooldown,
Some(cooldown_until),
)
}
fn ready_test_account(account_id: &str) -> ProviderAccount {
test_provider_account(account_id).with_runtime_state(true, AccountAvailability::Ready, None)
}
fn test_provider_account(account_id: &str) -> ProviderAccount {
ProviderAccount::new(
ProviderAccountId::new(account_id).expect("valid account ID"),
ProviderKind::new("openai").expect("valid provider kind"),
"Cooldown test".to_owned(),
"upstream-user".to_owned(),
"oauth".to_owned(),
CredentialRevision::new(7).expect("positive revision"),
Some(SystemTime::now() + StdDuration::from_secs(3_600)),
)
}
#[tokio::test]
async fn credential_cooldown_round_trips_without_raw_account_id_in_key() {
let Some((repository, mut connection, namespace)) = repository().await else {
+31 -56
View File
@@ -579,8 +579,6 @@ fn prepared_create(
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.enabled(),
availability,
availability_reason: None,
cooldown_until: account.cooldown_until().map(DateTime::<Utc>::from),
availability_observed_at: observed_at,
})
}
@@ -682,11 +680,7 @@ 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,
account.cooldown_until.map(SystemTime::from),
)
.with_runtime_state(account.enabled, account.availability)
.with_refresh_schedule(
account.has_refresh_token,
account.next_refresh_at.map(SystemTime::from),
@@ -753,9 +747,12 @@ fn project_quota(
snapshot: Option<CodexAccountQuotaSnapshot>,
account: &ProviderAccount,
) -> ProviderQuota {
let mut quota = snapshot.map_or_else(empty_quota, project_quota_snapshot);
// 展示前先滚动已过期窗口,避免 reset 后的旧窗口值悬挂。
let mut quota = snapshot
.map(|snapshot| project_quota_snapshot(snapshot.roll_expired_windows(SystemTime::now())))
.unwrap_or_else(empty_quota);
if account.availability() == AccountAvailability::QuotaExhausted {
force_confirmed_exhaustion_projection(&mut quota, account.cooldown_until());
force_confirmed_exhaustion_projection(&mut quota);
}
quota
}
@@ -790,6 +787,7 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
key: window.key().to_owned(),
group: quota_group(window.kind()).to_owned(),
label: codex_quota_window_label(
window.limit_name(),
window.kind(),
window.source(),
window.window_seconds(),
@@ -798,6 +796,7 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
window_seconds: window.window_seconds(),
used_percent: window.used_percent(),
reset_at: window.reset_at(),
limit_reached: window.limit_reached(),
local_usage: None,
provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(data))),
}
@@ -813,12 +812,9 @@ fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota
}
}
fn force_confirmed_exhaustion_projection(
quota: &mut ProviderQuota,
cooldown_until: Option<SystemTime>,
) {
// 真实 429 对“当前已耗尽”的证明强于可能延迟结算的 usage 文档。这里只覆盖
// Admin 展示投影,原始 Provider JSON 仍作为后续判断窗口是否真正恢复的基线。
fn force_confirmed_exhaustion_projection(quota: &mut ProviderQuota) {
// 402 已确认耗尽:只覆盖 Admin 展示投影,原始 Provider JSON 仍作为
// 后续判断窗口是否真正恢复的基线。
quota.provider_data = Some(ProviderDocument::new(OpaqueProviderData::new(
Map::from_iter([
(
@@ -828,7 +824,6 @@ fn force_confirmed_exhaustion_projection(
("exhausted".to_owned(), Value::Bool(true)),
]),
)));
let reset_at = cooldown_until.map(DateTime::<Utc>::from);
if quota.windows.is_empty() {
quota.windows.push(ProviderQuotaWindow {
key: "confirmed-quota-exhaustion".to_owned(),
@@ -837,42 +832,28 @@ fn force_confirmed_exhaustion_projection(
source: None,
window_seconds: None,
used_percent: Some(100.0),
reset_at,
reset_at: None,
limit_reached: true,
local_usage: None,
provider_data: None,
});
return;
}
let matched_reset = reset_at.and_then(|reset_at| {
quota
.windows
.iter()
.enumerate()
.filter_map(|(index, window)| {
let distance = window.reset_at.map(|window_reset| {
(window_reset.timestamp() - reset_at.timestamp()).unsigned_abs()
})?;
Some((index, distance))
})
.min_by_key(|(_, distance)| *distance)
.map(|(index, _)| index)
});
let index = matched_reset
.or_else(|| {
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)
// 选已用比例最高的窗口强制 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 {
@@ -891,11 +872,17 @@ const fn quota_role(role: CodexQuotaWindowRole) -> &'static str {
}
}
/// 窗口展示名,优先使用上游 `limit_name`(原文,`_` 转空格);
/// 无 limit_name 时按窗口时长猜(±5% 容差,汉化输出),猜不中兜底「额度」。
fn codex_quota_window_label(
limit_name: Option<&str>,
kind: CodexQuotaWindowKind,
source: &str,
window_seconds: Option<u64>,
) -> String {
if let Some(limit_name) = limit_name {
return limit_name.replace('_', " ");
}
let base = match kind {
CodexQuotaWindowKind::Monthly => "月限额".to_owned(),
CodexQuotaWindowKind::Weekly => "周限额".to_owned(),
@@ -908,12 +895,9 @@ fn codex_quota_window_label(
}
CodexQuotaWindowKind::Other => custom_quota_window_label(window_seconds),
};
if matches!(source, "core" | "codex" | "spend_control" | "monthly_limit") {
if matches!(source, "core" | "codex" | "code_review" | "spend_control") {
return base;
}
if is_codex_review_limit(source) {
return format!("代码审查 · {base}");
}
format!("{source} · {base}")
}
@@ -930,15 +914,6 @@ fn custom_quota_window_label(window_seconds: Option<u64>) -> String {
}
}
fn is_codex_review_limit(value: &str) -> bool {
let normalized = value.trim().to_ascii_lowercase().replace(['-', ' '], "_");
matches!(
normalized.as_str(),
"review" | "code_review" | "codex_review" | "codex_code_review"
) || normalized.contains("code_review")
|| normalized.contains("codex_review")
}
/// 将 Provider-owned PKCE/OIDC 状态保存到 Store 提供的 Redis 原子端口。
pub(crate) struct OpenAiOAuthPendingStore {
port: Arc<dyn OAuthPendingFlowPort>,
@@ -409,7 +409,7 @@ impl CodexCredentialAdmin {
Some(input.verified_account.chatgpt_account_id),
input.verified_account.plan_type,
)
.with_runtime_state(input.enabled, AccountAvailability::Ready, None)
.with_runtime_state(input.enabled, AccountAvailability::Ready)
.with_refresh_schedule(
input.secret.refresh_token.is_some(),
optional_time(input.next_refresh_at),
@@ -806,10 +806,7 @@ impl CodexCredentialAdminService {
)) {
return Err(CodexCredentialAdminError::InvalidInput);
}
account.account =
account
.account
.with_runtime_state(enabled, availability, None);
account.account = account.account.with_runtime_state(enabled, availability);
account
}
};
@@ -933,7 +930,7 @@ impl CodexCredentialAdminService {
Some(upstream_account_id),
candidate.plan_type.clone(),
)
.with_runtime_state(enabled, AccountAvailability::Ready, None)
.with_runtime_state(enabled, AccountAvailability::Ready)
.with_refresh_schedule(false, None);
Ok(NewProviderAccount {
account,
@@ -1000,9 +997,7 @@ fn cpr_status(account: &ProviderAccount) -> &'static str {
AccountAvailability::QuotaExhausted => "quota_exhausted",
AccountAvailability::Expired | AccountAvailability::Invalid => "expired",
AccountAvailability::Banned => "banned",
AccountAvailability::Unknown
| AccountAvailability::Ready
| AccountAvailability::Cooldown => "active",
AccountAvailability::Unknown | AccountAvailability::Ready => "active",
}
}
@@ -783,7 +783,6 @@ fn eligible_catalog_account(account: &ProviderAccount, now: SystemTime) -> bool
&& match account.availability() {
AccountAvailability::Unknown
| AccountAvailability::Ready
| AccountAvailability::Cooldown
| AccountAvailability::QuotaExhausted => true,
AccountAvailability::Expired
| AccountAvailability::Banned
File diff suppressed because it is too large Load Diff
@@ -192,7 +192,6 @@ impl CodexCredentialRefreshService {
self.persist_terminal(
&due.account,
AccountAvailability::Expired,
"invalid_grant",
CodexCredentialRefreshOutcome::Invalidated { account_id },
)
.await
@@ -201,7 +200,6 @@ impl CodexCredentialRefreshService {
self.persist_terminal(
&due.account,
AccountAvailability::Banned,
"account_banned",
CodexCredentialRefreshOutcome::Banned { account_id },
)
.await
@@ -330,7 +328,6 @@ impl CodexCredentialRefreshService {
.persist_terminal(
&due.account,
AccountAvailability::Invalid,
"refreshed_identity_rejected",
CodexCredentialRefreshOutcome::Invalidated {
account_id: due.account.id().to_string(),
},
@@ -433,7 +430,6 @@ impl CodexCredentialRefreshService {
&self,
account: &ProviderAccount,
availability: AccountAvailability,
reason: &'static str,
outcome: CodexCredentialRefreshOutcome,
) -> Result<CodexCredentialRefreshOutcome, CodexCredentialRefreshError> {
match self.repository.load_runtime_credential(account).await {
@@ -447,13 +443,7 @@ impl CodexCredentialRefreshService {
}
match self
.repository
.apply_state(
account,
availability,
Some(reason.to_owned()),
None,
SystemTime::now(),
)
.apply_state(account, availability, SystemTime::now())
.await
{
Ok(()) => Ok(outcome),
@@ -233,8 +233,6 @@ impl CodexCredentialRepository {
&self,
account: &ProviderAccount,
availability: AccountAvailability,
reason: Option<String>,
cooldown_until: Option<SystemTime>,
observed_at: SystemTime,
) -> Result<(), CredentialRepositoryError> {
if !account.enabled() {
@@ -245,8 +243,6 @@ impl CodexCredentialRepository {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability,
reason,
cooldown_until,
observed_at,
})
.await?;
@@ -32,8 +32,6 @@ use super::types::{
RuntimeCodexCookie,
};
const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
const CLOUDFLARE_PATH_BLOCK_COOLDOWN: Duration = Duration::from_secs(30);
const CLOUDFLARE_RECOVERY_STALE_AFTER: Duration = Duration::from_secs(60 * 60);
const CLOUDFLARE_CHALLENGE_BACKOFF: [Duration; 4] = [
Duration::from_secs(10),
@@ -84,7 +82,6 @@ struct RiskRecoveryState {
#[derive(Debug, Clone, Copy)]
enum CookieRecovery {
None,
ExpireAt(SystemTime),
Clear,
}
@@ -321,6 +318,7 @@ impl CodexCredentialSelector {
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
failure_rate_basis_points: None,
first_output_latency_ms: None,
})
@@ -329,9 +327,16 @@ impl CodexCredentialSelector {
AccountCandidate { account, signals }
})
.collect::<Vec<_>>();
// 旧版开关只影响“候选资格”。保留原始账号供 revision-fenced 凭据加载,
// 用影子候选承载放宽后的资格,避免伪造已持久化账号的运行状态。
let quota_exhausted_candidates = (!self.skip_exhausted).then(|| {
// skip_exhausted 语义:true 时排除任一窗口触顶(limit_reached)的账号,
// 仅由上游 429 兜底;false 时保留全部候选(含 QuotaExhausted 的放宽投影)。
// 保留原始账号供 revision-fenced 凭据加载,用影子候选承载资格放宽。
let scheduling_candidates = if self.skip_exhausted {
candidates
.iter()
.filter(|candidate| !candidate.signals.quota_limit_reached)
.cloned()
.collect::<Vec<_>>()
} else {
candidates
.iter()
.cloned()
@@ -340,8 +345,8 @@ impl CodexCredentialSelector {
candidate
})
.collect::<Vec<_>>()
});
let scheduling_candidates = quota_exhausted_candidates.as_deref().unwrap_or(&candidates);
};
let scheduling_candidates = scheduling_candidates.as_slice();
let mut affinity = if diagnostic {
AffinitySelection::default()
} else {
@@ -682,118 +687,64 @@ impl CodexCredentialSelector {
failure: CodexAccountFailure,
) -> Result<(), CredentialSelectionError> {
let now = SystemTime::now();
let (availability, reason, cooldown_until, cookie_recovery) = match failure {
CodexAccountFailure::CredentialExpired => (
AccountAvailability::Expired,
Some("credential_expired".to_owned()),
None,
CookieRecovery::None,
),
CodexAccountFailure::IdentityVerificationRequired => (
AccountAvailability::Invalid,
Some("identity_verification_required".to_owned()),
None,
CookieRecovery::None,
),
CodexAccountFailure::Banned => (
AccountAvailability::Banned,
Some("account_banned".to_owned()),
None,
CookieRecovery::None,
),
CodexAccountFailure::QuotaExhausted => (
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
CookieRecovery::None,
),
CodexAccountFailure::UsageLimitExhausted { retry_after } => {
return self
.record_resettable_limit_failure(
account,
retry_after,
"usage_limit_exhausted",
now,
)
.await;
match failure {
CodexAccountFailure::CredentialExpired => {
self.apply_state(account, AccountAvailability::Expired, now)
.await
}
CodexAccountFailure::IdentityVerificationRequired => {
self.apply_state(account, AccountAvailability::Invalid, now)
.await
}
CodexAccountFailure::Banned => {
self.apply_state(account, AccountAvailability::Banned, now)
.await
}
// 402:真额度耗尽,由 quota worker 按窗口 reset 自动恢复。
CodexAccountFailure::QuotaExhausted
| CodexAccountFailure::UsageLimitExhausted { .. } => {
self.apply_state(account, AccountAvailability::QuotaExhausted, now)
.await
}
// 429:临时限流只写 quota,availability 不改变。
CodexAccountFailure::RateLimited { retry_after } => {
return self
.record_resettable_limit_failure(account, retry_after, "rate_limited", now)
.await;
self.quota
.apply_rate_limit_429(account, retry_after, now)
.await
.map_err(|_| CredentialSelectionError::Store)?;
Ok(())
}
// Cloudflare 挑战:内存退避表(记录风险计数),不写 availability。
CodexAccountFailure::CloudflareChallenge { retry_after } => {
let delay = self.cloudflare_challenge_delay(account.id(), now, retry_after);
let cooldown_until = now.checked_add(delay);
(
AccountAvailability::Cooldown,
Some("cloudflare_challenge".to_owned()),
cooldown_until,
cooldown_until.map_or(CookieRecovery::Clear, CookieRecovery::ExpireAt),
)
let recovery = now
.checked_add(delay)
.map_or(CookieRecovery::Clear, CookieRecovery::ExpireAt);
self.apply_cookie_recovery(account, recovery).await?;
Ok(())
}
// Cloudflare 路径被封:连续超阈值才置 Invalid;否则内存退避。
CodexAccountFailure::CloudflarePathBlocked => {
let blocked = self.record_cloudflare_path_block(account.id(), now);
if blocked >= CLOUDFLARE_PATH_BLOCK_THRESHOLD {
(
AccountAvailability::Invalid,
Some("cloudflare_path_blocked".to_owned()),
None,
CookieRecovery::Clear,
)
} else {
(
AccountAvailability::Cooldown,
Some("cloudflare_path_blocked".to_owned()),
now.checked_add(CLOUDFLARE_PATH_BLOCK_COOLDOWN),
CookieRecovery::Clear,
)
self.apply_state(account, AccountAvailability::Invalid, now)
.await?;
}
self.apply_cookie_recovery(account, CookieRecovery::Clear)
.await?;
Ok(())
}
};
self.repository
.apply_state(account, availability, reason, cooldown_until, now)
.await?;
self.apply_cookie_recovery(account, cookie_recovery).await?;
Ok(())
}
}
async fn record_resettable_limit_failure(
async fn apply_state(
&self,
account: &ProviderAccount,
retry_after: Option<Duration>,
reason: &'static str,
availability: AccountAvailability,
observed_at: SystemTime,
) -> Result<(), CredentialSelectionError> {
let current = self.current_account(account.id()).await?;
if current.revision() != account.revision() {
return Err(CredentialRepositoryError::RevisionConflict.into());
}
let cooldown_until =
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN));
let cooldown_until = match (current.availability(), current.cooldown_until()) {
(AccountAvailability::QuotaExhausted, Some(existing)) => {
cooldown_until.map_or(Some(existing), |until| Some(existing.max(until)))
}
_ => cooldown_until,
};
let reason = if current.availability() == AccountAvailability::QuotaExhausted
&& current.cooldown_until().is_none()
{
Some("quota_exhausted".to_owned())
} else {
Some(reason.to_owned())
};
// 一次真实 429 证明上游已经接受当前凭据,可纠正旧的认证/封禁终态。
// resettable 用量限制沿用 v2 的 quota lock 语义,并保留/延长复核时间。
self.repository
.apply_state(
&current,
AccountAvailability::QuotaExhausted,
reason,
cooldown_until,
observed_at,
)
.apply_state(account, availability, observed_at)
.await?;
Ok(())
}
@@ -857,7 +808,6 @@ impl CodexCredentialSelector {
|| !matches!(
current.availability(),
AccountAvailability::Unknown
| AccountAvailability::Cooldown
| AccountAvailability::Expired
| AccountAvailability::Invalid
| AccountAvailability::Banned
@@ -867,13 +817,7 @@ impl CodexCredentialSelector {
}
if let Err(error) = self
.repository
.apply_state(
&current,
AccountAvailability::Ready,
None,
None,
SystemTime::now(),
)
.apply_state(&current, AccountAvailability::Ready, SystemTime::now())
.await
{
tracing::warn!(
@@ -985,15 +929,11 @@ impl CodexCredentialSelector {
account: &ProviderAccount,
recovery: CookieRecovery,
) -> Result<(), CredentialSelectionError> {
if matches!(recovery, CookieRecovery::None) {
return Ok(());
}
let mut data = self.repository.load_complete_data(account).await?;
if data.cookies().is_empty() {
return Ok(());
}
match recovery {
CookieRecovery::None => return Ok(()),
CookieRecovery::ExpireAt(expires_at) => {
let expires_at = chrono::DateTime::<chrono::Utc>::from(expires_at);
for cookie in data.cookies_mut() {
@@ -1013,11 +953,9 @@ impl CodexCredentialSelector {
fn quota_exhausted_scheduling_projection(account: &ProviderAccount) -> ProviderAccount {
if account.availability() == AccountAvailability::QuotaExhausted {
return account.clone().with_runtime_state(
account.enabled(),
AccountAvailability::Ready,
account.cooldown_until(),
);
return account
.clone()
.with_runtime_state(account.enabled(), AccountAvailability::Ready);
}
account.clone()
}
@@ -1032,7 +970,6 @@ fn affinity_unavailable_reason(account: &ProviderAccount, now: SystemTime) -> Af
}
match account.availability() {
AccountAvailability::QuotaExhausted => AffinityEscapeReason::QuotaExhausted,
AccountAvailability::Cooldown => AffinityEscapeReason::Cooldown,
AccountAvailability::Unknown
| AccountAvailability::Ready
| AccountAvailability::Expired
@@ -2533,7 +2533,6 @@ impl ScheduledTask for OpenAiQuotaTask {
updated = summary.updated,
exhausted = summary.exhausted,
banned = summary.banned,
cooldown = summary.cooldown,
transient = summary.transient,
stale = summary.stale,
"OpenAI quota cycle contained operational failures"
@@ -498,7 +498,8 @@ async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_qu
.collect::<Vec<_>>();
assert_eq!(monthly.len(), 1);
assert_eq!(monthly[0].label, "月限额");
// limit_name 优先显示,不再猜成「月限额」。
assert_eq!(monthly[0].label, "custom codex label");
assert_eq!(monthly[0].source.as_deref(), Some("core"));
assert_eq!(monthly[0].used_percent, Some(2.0));
}
@@ -544,8 +545,6 @@ async fn openai_admin_projects_confirmed_quota_exhaustion_as_full_without_mutati
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: CoreAccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: Some(SystemTime::UNIX_EPOCH + Duration::from_secs(reset_at)),
observed_at: SystemTime::now(),
})
.await
@@ -586,8 +585,6 @@ async fn openai_admin_projects_confirmed_quota_exhaustion_as_full_without_mutati
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: CoreAccountAvailability::Ready,
reason: None,
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -645,7 +642,6 @@ async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit(
stale_record.next_refresh_at = None;
stale_record.enabled = false;
stale_record.availability = AdminAccountAvailability::Banned;
stale_record.cooldown_until = Some(Utc::now() + chrono::Duration::minutes(5));
let rotation_error = admin
.prepare_rotation(PrepareCredentialRotation {
account: stale_record,
@@ -712,8 +708,6 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.enabled(),
availability: AdminAccountAvailability::Ready,
availability_reason: None,
cooldown_until: account.cooldown_until().map(DateTime::<Utc>::from),
availability_observed_at: now,
quota_observed_at: None,
created_at: now,
@@ -194,7 +194,7 @@ fn agent_export_item(id: &str) -> (ExportManagedCodexCredential, String) {
Some("agent-account-export".to_owned()),
Some("pro".to_owned()),
)
.with_runtime_state(true, AccountAvailability::Ready, None)
.with_runtime_state(true, AccountAvailability::Ready)
.with_refresh_schedule(false, None);
(
ExportManagedCodexCredential {
@@ -66,7 +66,7 @@ fn agent_account(
Some("agent-account".to_owned()),
Some("pro".to_owned()),
)
.with_runtime_state(true, AccountAvailability::Ready, None)
.with_runtime_state(true, AccountAvailability::Ready)
.with_refresh_schedule(false, None);
NewProviderAccount {
account,
@@ -510,8 +510,6 @@ async fn selector_should_escape_a_quota_exhausted_affinity_account() {
account_id: first.id().clone(),
expected_revision: first.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -550,57 +548,6 @@ async fn selector_should_escape_a_quota_exhausted_affinity_account() {
);
}
#[tokio::test]
async fn selector_should_escape_an_affinity_account_during_cooldown() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first", "at-first");
create_account(&store, "acct_second", "at-second");
let first = store.account("acct_first").expect("first account");
store
.apply_state_change(AccountStateChange {
account_id: first.id().clone(),
expected_revision: first.revision(),
availability: AccountAvailability::Cooldown,
reason: Some("rate_limited".to_owned()),
cooldown_until: Some(SystemTime::now() + Duration::from_secs(60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark first account cooling down");
let affinity = Arc::new(MemorySessionAffinity::default());
let provider = ProviderKind::new("openai").expect("provider");
let key = ProviderSessionAffinityKey::try_new("cooldown-session").expect("affinity key");
affinity
.bind(&provider, &key, first.id(), Duration::from_secs(60))
.await
.expect("seed affinity");
let selector =
selector_with_affinity(&store, Arc::new(TestLeaseCoordinator::default()), affinity);
let request_url =
Url::parse("https://chatgpt.com/backend-api/codex/responses").expect("request URL");
let request_attempt = attempt(BTreeSet::new());
let selected = selector
.select(&SelectCodexCredential {
upstream_model: "gpt-5.4",
request_url: &request_url,
attempt: &request_attempt,
session_affinity_key: Some(&key),
})
.await
.expect("select fallback account");
assert_eq!(
(
selected.account_id().as_str(),
selected.affinity_hit(),
selected.escape_reason(),
selected.account_switch(),
),
("acct_second", false, Some("cooldown"), true)
);
}
#[test]
fn selector_honors_attempt_local_account_exclusion() {
let store = Arc::new(MemoryAccountStore::default());
@@ -728,7 +675,7 @@ fn credential_expired_failure_marks_unified_account_expired() {
}
#[test]
fn rate_limited_failure_marks_quota_exhausted_with_a_reset_deadline() {
fn rate_limited_failure_records_quota_window_without_changing_availability() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
@@ -753,12 +700,33 @@ fn rate_limited_failure_marks_quota_exhausted_with_a_reset_deadline() {
))
.expect("record rate-limit failure");
// 429 不改变 availability(限流不改变账号可用性),只把限流窗口合成进 quota JSON。
let account = store.account("acct_primary").expect("account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert_eq!(account.availability(), AccountAvailability::Ready);
let quota = store.quota_json("acct_primary").expect("quota");
assert_eq!(
quota
.pointer("/rate_limit/limit_reached")
.and_then(serde_json::Value::as_bool),
Some(true)
);
assert_eq!(
quota
.pointer("/rate_limit/primary_window/used_percent")
.and_then(serde_json::Value::as_f64),
Some(100.0)
);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
quota
.pointer("/rate_limit/primary_window/reset_at")
.and_then(serde_json::Value::as_i64)
.is_some_and(|reset_at| {
reset_at
> SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("clock")
.as_secs() as i64
})
);
}
@@ -781,11 +749,6 @@ async fn usage_limit_exhaustion_marks_quota_exhausted_without_usage_probe() {
let account = store.account("acct_primary").expect("persisted account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert_eq!(store.quota_reads(), 0);
}
@@ -798,8 +761,6 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
}))
.expect("mark exhausted");
@@ -813,6 +774,7 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
))
.expect("record rate-limit failure");
// 429 不写 availability,402 的 QuotaExhausted 保持不变。
assert_eq!(
store
.account("acct_primary")
@@ -820,47 +782,6 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
.availability(),
AccountAvailability::QuotaExhausted
);
assert!(
store
.account("acct_primary")
.expect("persisted account")
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
}
#[test]
fn rate_limited_failure_keeps_longer_existing_quota_deadline() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let existing_deadline = SystemTime::now() + Duration::from_secs(300);
block_on(store.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: Some(existing_deadline),
observed_at: SystemTime::now(),
}))
.expect("mark exhausted");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&account,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record shorter rate-limit failure");
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.cooldown_until(),
Some(existing_deadline)
);
}
#[test]
@@ -883,13 +804,14 @@ fn rate_limited_failure_does_not_consult_stale_quota_snapshot() {
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::QuotaExhausted
AccountAvailability::Ready
);
assert_eq!(store.quota_reads(), 0);
// 429 合成限流窗口需要读写 quota JSON。
assert!(store.quota_reads() >= 1);
}
#[test]
fn authenticated_rate_limit_recovers_stale_authentication_state() {
fn rate_limited_failure_does_not_overwrite_stale_authentication_state() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
@@ -897,8 +819,6 @@ fn authenticated_rate_limit_recovers_stale_authentication_state() {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Invalid,
reason: Some("identity_verification_required".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
}))
.expect("mark account invalid");
@@ -913,12 +833,13 @@ fn authenticated_rate_limit_recovers_stale_authentication_state() {
))
.expect("record rate-limit failure");
// 429 不写 availability(限流不改变账号可用性),认证终态保持不变。
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::QuotaExhausted
AccountAvailability::Invalid
);
}
@@ -936,8 +857,6 @@ fn successful_upstream_response_recovers_non_quota_terminal_states() {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: stale,
reason: Some("stale_observation".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
}))
.expect("seed stale state");
@@ -996,7 +915,7 @@ fn rate_limited_failures_for_distinct_accounts_do_not_conflict() {
.expect("persisted second account")
.availability(),
],
[AccountAvailability::QuotaExhausted; 2]
[AccountAvailability::Ready; 2]
);
}
@@ -1082,7 +1001,7 @@ fn identity_verification_failure_isolates_only_selected_account() {
}
#[test]
fn cloudflare_challenge_backoff_escalates_and_success_resets_it() {
fn cloudflare_challenge_does_not_change_availability() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
@@ -1099,48 +1018,21 @@ fn cloudflare_challenge_backoff_escalates_and_success_resets_it() {
)
.expect("select account");
// Cloudflare 挑战只走内存退避表,不写 availability(限流不改变账号可用性)。
block_on(selector.record_failure(
lease.account(),
CodexAccountFailure::CloudflareChallenge { retry_after: None },
))
.expect("record first challenge");
let first = store
.account("acct_primary")
.expect("account")
.cooldown_until()
.expect("first cooldown")
.duration_since(SystemTime::now())
.expect("future cooldown");
assert!(first >= Duration::from_secs(9) && first <= Duration::from_secs(10));
block_on(selector.record_failure(
lease.account(),
CodexAccountFailure::CloudflareChallenge { retry_after: None },
))
.expect("record second challenge");
let second = store
.account("acct_primary")
.expect("account")
.cooldown_until()
.expect("second cooldown")
.duration_since(SystemTime::now())
.expect("future cooldown");
assert!(second >= Duration::from_secs(29) && second <= Duration::from_secs(30));
.expect("record challenge");
assert_eq!(
store
.account("acct_primary")
.expect("account")
.availability(),
AccountAvailability::Ready
);
block_on(selector.record_success(lease.account(), None));
block_on(selector.record_failure(
lease.account(),
CodexAccountFailure::CloudflareChallenge { retry_after: None },
))
.expect("record reset challenge");
let reset = store
.account("acct_primary")
.expect("account")
.cooldown_until()
.expect("reset cooldown")
.duration_since(SystemTime::now())
.expect("future cooldown");
assert!(reset >= Duration::from_secs(9) && reset <= Duration::from_secs(10));
}
#[test]
@@ -1226,14 +1118,12 @@ fn cloudflare_challenge_expires_provider_owned_cookies_at_cooldown_boundary() {
.expect("record challenge");
let account = store.account("acct_primary").expect("account");
let cooldown_until = account.cooldown_until().expect("cooldown");
let data = block_on(store.repository().load_complete_data(&account)).expect("credential data");
assert_eq!(data.cookies().len(), 1);
assert!(
data.cookies()[0]
.expires_at
.is_some_and(|expires_at| SystemTime::from(expires_at) <= cooldown_until)
);
assert!(data.cookies()[0].expires_at.is_some_and(|expires_at| {
let expires_at = SystemTime::from(expires_at);
expires_at > SystemTime::now() && expires_at <= SystemTime::now() + Duration::from_secs(120)
}));
}
#[test]
@@ -291,14 +291,13 @@ async fn persisted_provider_quota_projects_dynamic_windows_without_network_io()
.expect("read quota")
.expect("quota snapshot");
assert_eq!(snapshot.windows().len(), 3);
assert_eq!(snapshot.windows().len(), 2);
assert_eq!(
snapshot.windows()[0].kind(),
CodexQuotaWindowKind::ShortTerm
);
assert_eq!(snapshot.windows()[0].window_seconds(), Some(18_000));
assert_eq!(snapshot.windows()[1].kind(), CodexQuotaWindowKind::Weekly);
assert_eq!(snapshot.windows()[2].kind(), CodexQuotaWindowKind::Monthly);
}
#[tokio::test]
@@ -358,7 +357,78 @@ async fn persisted_codex_additional_limit_replaces_the_top_level_rate_limit() {
}
#[tokio::test]
async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision_fence() {
async fn code_review_limit_projects_as_independent_window_with_limit_name() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_code_review").await;
let account = store.account("acct_code_review").expect("account");
let raw = json!({
"rate_limit": {
"primary_window": {
"used_percent": 37,
"reset_at": 1_900_000_000,
"limit_window_seconds": 18_000
}
},
"code_review_rate_limit": {
"primary_window": {
"used_percent": 80,
"reset_at": 1_900_604_800,
"limit_window_seconds": 604_800
}
},
"additional_rate_limits": [{
"limit_name": "code review",
"metered_feature": "code_review",
"rate_limit": {
"primary_window": {
"used_percent": 55,
"reset_at": 1_900_604_800,
"limit_window_seconds": 604_800
}
}
}],
"spend_control": {
"reached": false,
"individual_limit": {
"used_percent": 12,
"reset_at": 1_902_592_000
}
}
});
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
raw.as_object().expect("quota object").clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist quota");
let snapshot = quota_service(&store)
.read_account(account.id())
.await
.expect("read quota")
.expect("quota snapshot");
// spend_control 不再生成窗口(只作 exhaustion 信号),
// 顶层 code_review_rate_limit 与 additional 的 code_review 桶同源替换后保留一个。
let review = snapshot
.windows()
.iter()
.filter(|window| window.source() == "code_review")
.collect::<Vec<_>>();
assert_eq!(review.len(), 1);
assert_eq!(review[0].limit_name(), Some("code review"));
assert_eq!(review[0].used_percent(), Some(55.0));
assert_eq!(snapshot.windows().len(), 2);
assert!(!snapshot.fact().exhausted());
}
#[tokio::test]
async fn passive_rate_limit_headers_update_quota_with_revision_fence() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_passive_quota").await;
let account = store.account("acct_passive_quota").expect("account");
@@ -384,7 +454,8 @@ async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision
let current = store
.account("acct_passive_quota")
.expect("current account");
assert_eq!(current.availability(), AccountAvailability::QuotaExhausted);
// 被动头限流只写 quota,不改变 availability(限流不改变账号可用性)。
assert_eq!(current.availability(), AccountAvailability::Ready);
let snapshot = service
.read_account(account.id())
.await
@@ -434,11 +505,10 @@ async fn passive_full_percent_with_explicit_allowance_keeps_account_ready() {
assert_eq!(
(
current.availability(),
current.cooldown_until(),
snapshot.fact().remaining_percent(),
snapshot.fact().exhausted(),
),
(AccountAvailability::Ready, None, Some(0), false)
(AccountAvailability::Ready, Some(0), false)
);
}
@@ -453,8 +523,6 @@ async fn manual_quota_refresh_keeps_confirmed_exhaustion_until_the_recorded_rese
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Expired,
reason: Some("stale_expired_state".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -495,7 +563,8 @@ async fn manual_quota_refresh_keeps_confirmed_exhaustion_until_the_recorded_rese
.account(account_id)
.expect("exhausted account")
.availability(),
AccountAvailability::QuotaExhausted
AccountAvailability::Expired,
"manual refresh never changes non-QuotaExhausted availability"
);
server.reset().await;
@@ -526,8 +595,8 @@ async fn manual_quota_refresh_keeps_confirmed_exhaustion_until_the_recorded_rese
.account(account_id)
.expect("account after contradictory quota")
.availability(),
AccountAvailability::QuotaExhausted,
"an unchanged usage window cannot override a confirmed quota failure"
AccountAvailability::Expired,
"an unchanged usage window cannot override an expired credential state"
);
server.reset().await;
@@ -558,8 +627,8 @@ async fn manual_quota_refresh_keeps_confirmed_exhaustion_until_the_recorded_rese
.account(account_id)
.expect("account after lower usage")
.availability(),
AccountAvailability::QuotaExhausted,
"a lower used_percent before the recorded reset is not quota recovery"
AccountAvailability::Expired,
"a lower used_percent before the recorded reset is not quota recovery for an expired credential"
);
}
@@ -574,8 +643,6 @@ async fn manual_quota_refresh_preserves_disabled_account_state() {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Expired,
reason: Some("disabled_expired_state".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -629,8 +696,6 @@ async fn manual_quota_auth_rejection_does_not_invalidate_refreshable_credential(
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -782,7 +847,6 @@ async fn synchronize_without_accounts_is_a_noop_before_network_io() {
assert_eq!(summary.updated, 0);
assert_eq!(summary.exhausted, 0);
assert_eq!(summary.banned, 0);
assert_eq!(summary.cooldown, 0);
assert_eq!(summary.transient, 0);
assert_eq!(summary.stale, 0);
}
@@ -824,7 +888,7 @@ async fn periodic_quota_synchronization_does_not_scan_stale_ready_accounts() {
}
#[tokio::test]
async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_quota() {
async fn periodic_quota_synchronization_preserves_availability_when_usage_peaks() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_stuck_cooldown").await;
let account = store
@@ -835,9 +899,7 @@ async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_q
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Cooldown,
reason: Some("rate_limited".to_owned()),
cooldown_until: SystemTime::now().checked_add(Duration::from_secs(24 * 60 * 60)),
availability: AccountAvailability::Ready,
observed_at: SystemTime::now(),
})
.await
@@ -884,19 +946,11 @@ async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_q
service.synchronize().await.expect("quota synchronization");
// 429/触顶不再写入 availability(限流不改变账号可用性),worker 只更新 quota 数据。
let reconciled = store
.account("acct_periodic_stuck_cooldown")
.expect("reconciled account");
assert_eq!(
(reconciled.availability(), reconciled.cooldown_until()),
(
AccountAvailability::QuotaExhausted,
Some(
SystemTime::UNIX_EPOCH
+ Duration::from_secs(u64::try_from(reset_at).expect("positive reset")),
),
)
);
assert_eq!(reconciled.availability(), AccountAvailability::Ready);
}
#[tokio::test]
@@ -911,8 +965,6 @@ async fn periodic_quota_synchronization_does_not_downgrade_exhaustion_on_rate_li
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -938,11 +990,6 @@ async fn periodic_quota_synchronization_does_not_downgrade_exhaustion_on_rate_li
.account("acct_periodic_exhausted")
.expect("persisted account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
}
#[tokio::test]
@@ -957,8 +1004,6 @@ async fn periodic_quota_synchronization_respects_retry_after_on_service_unavaila
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -985,15 +1030,10 @@ async fn periodic_quota_synchronization_respects_retry_after_on_service_unavaila
.account("acct_periodic_503_retry_after")
.expect("persisted account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
}
#[tokio::test]
async fn periodic_quota_synchronization_skips_ordinary_cooldown_without_exhausted_quota() {
async fn periodic_quota_synchronization_skips_ready_accounts_without_quota_reached() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_temporary_cooldown").await;
let account = store
@@ -1003,9 +1043,7 @@ async fn periodic_quota_synchronization_skips_ordinary_cooldown_without_exhauste
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Cooldown,
reason: Some("rate_limited".to_owned()),
cooldown_until: SystemTime::now().checked_add(Duration::from_secs(60)),
availability: AccountAvailability::Ready,
observed_at: SystemTime::now(),
})
.await
@@ -1022,7 +1060,7 @@ async fn periodic_quota_synchronization_skips_ordinary_cooldown_without_exhauste
.account("acct_periodic_temporary_cooldown")
.expect("cooldown account")
.availability(),
AccountAvailability::Cooldown
AccountAvailability::Ready
);
}
@@ -1038,8 +1076,6 @@ async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -1065,8 +1101,6 @@ async fn periodic_quota_synchronization_rechecks_exhausted_accounts_without_wait
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: SystemTime::now().checked_add(Duration::from_secs(24 * 60 * 60)),
observed_at: SystemTime::now(),
})
.await
@@ -1131,8 +1165,6 @@ async fn periodic_quota_synchronization_does_not_recover_from_percent_drop_befor
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: SystemTime::now().checked_add(Duration::from_secs(24 * 60 * 60)),
observed_at: SystemTime::now(),
})
.await
@@ -1223,21 +1255,19 @@ async fn periodic_quota_synchronization_recovers_after_recorded_reset() {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: SystemTime::now().checked_sub(Duration::from_secs(60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted with an elapsed reset");
let server = MockServer::start().await;
let next_reset = Utc::now().timestamp() + 24 * 60 * 60;
let elapsed_reset = Utc::now().timestamp() - 1;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 1, "reset_at": next_reset}
"primary_window": {"used_percent": 1, "reset_at": elapsed_reset}
}
})))
.expect(1)
@@ -1260,7 +1290,6 @@ async fn periodic_quota_synchronization_recovers_after_recorded_reset() {
.account(account_id)
.expect("account after elapsed reset");
assert_eq!(current.availability(), AccountAvailability::Ready);
assert_eq!(current.cooldown_until(), None);
}
#[tokio::test]
@@ -1275,8 +1304,6 @@ async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -1295,15 +1322,11 @@ async fn periodic_quota_synchronization_does_not_bypass_throttle_for_a_new_reset
let account_id = "acct_periodic_new_reset";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
let initial_reset = SystemTime::UNIX_EPOCH + Duration::from_secs(1);
let confirmed_reset = SystemTime::UNIX_EPOCH + Duration::from_secs(2);
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: Some(initial_reset),
observed_at: SystemTime::now(),
})
.await
@@ -1330,10 +1353,6 @@ async fn periodic_quota_synchronization_does_not_bypass_throttle_for_a_new_reset
);
let first = service.synchronize().await.expect("first reset recheck");
let persisted_reset = store
.account(account_id)
.expect("account after first recheck")
.cooldown_until();
let second = service
.synchronize()
.await
@@ -1346,20 +1365,8 @@ async fn periodic_quota_synchronization_does_not_bypass_throttle_for_a_new_reset
.len();
assert_eq!(
(
first.exhausted,
persisted_reset,
second.exhausted,
repeated,
request_count,
),
(
1,
Some(confirmed_reset),
0,
CodexQuotaSyncSummary::default(),
1,
)
(first.exhausted, second.exhausted, repeated, request_count,),
(1, 0, CodexQuotaSyncSummary::default(), 1,)
);
}
@@ -1374,8 +1381,6 @@ async fn invalid_quota_json_for_one_account_does_not_abort_the_batch() {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -621,8 +621,6 @@ async fn expired_access_token_refreshes_even_when_next_refresh_is_deferred() {
.apply_state(
&deferred,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
SystemTime::now(),
)
.await
@@ -685,18 +683,11 @@ async fn expired_access_token_refreshes_during_active_cooldown() {
.await
.expect("seed expired access token with deferred refresh");
let deferred = store.account("acct_refresh").expect("deferred account");
let cooldown_until = SystemTime::now() + Duration::from_secs(12 * 60 * 60);
store
.repository()
.apply_state(
&deferred,
AccountAvailability::Cooldown,
Some("rate_limited".to_owned()),
Some(cooldown_until),
SystemTime::now(),
)
.apply_state(&deferred, AccountAvailability::Ready, SystemTime::now())
.await
.expect("mark account cooling down");
.expect("mark account ready");
let refresher = Arc::new(Refresher::new(Ok(success_tokens())));
let service = CodexCredentialRefreshService::new(
store.repository(),
@@ -724,8 +715,7 @@ async fn expired_access_token_refreshes_during_active_cooldown() {
["rt-old-access"]
);
let refreshed = store.account("acct_refresh").expect("refreshed account");
assert_eq!(refreshed.availability(), AccountAvailability::Cooldown);
assert_eq!(refreshed.cooldown_until(), Some(cooldown_until));
assert_eq!(refreshed.availability(), AccountAvailability::Ready);
assert!(
refreshed
.access_token_expires_at()
@@ -1263,12 +1263,8 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_quota_loc
vec!["response.created", "response.failed"]
);
let account = store.account(account_id).expect("rate-limited account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
// 429 只记录限流窗口,不改变 availability(限流不改变账号可用性)。
assert_eq!(account.availability(), AccountAvailability::Ready);
let _ = release.send(());
server.abort();
let _ = server.await;
@@ -1358,11 +1354,6 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
assert_eq!(failure.kind(), ProviderErrorKind::QuotaExhausted);
let account = store.account(account_id).expect("usage-limit account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
tokio::time::sleep(Duration::from_millis(200)).await;
let requests = server.received_requests().await.expect("received requests");
assert_eq!(
@@ -1799,8 +1790,6 @@ async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_st
.apply_state(
&account,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
SystemTime::now(),
)
.await
@@ -1907,7 +1896,7 @@ async fn successful_response_preserves_quota_exhaustion_from_rate_limit_headers(
.account(account_id)
.expect("account after successful response")
.availability(),
AccountAvailability::QuotaExhausted
AccountAvailability::Ready
);
assert!(store.has_quota(account_id));
}
@@ -1969,16 +1958,8 @@ async fn successful_http_sse_rate_limit_event_persists_structured_exhaustion() {
.account(account_id)
.expect("account after HTTP SSE response");
assert_eq!(
(
account.availability(),
account.cooldown_until(),
store.has_quota(account_id),
),
(
AccountAvailability::QuotaExhausted,
Some(SystemTime::UNIX_EPOCH + Duration::from_secs(reset_at)),
true,
)
(account.availability(), store.has_quota(account_id),),
(AccountAvailability::Ready, true,)
);
}
@@ -117,6 +117,17 @@ impl MemoryAccountStore {
.is_some_and(|stored| stored.quota.is_some())
}
pub(crate) fn quota_json(&self, id: &str) -> Option<serde_json::Value> {
let id = ProviderAccountId::new(id).ok()?;
self.accounts
.lock()
.expect("account store lock")
.get(&id)
.and_then(|stored| stored.quota.as_ref())
.and_then(|observation| observation.quota.as_ref())
.map(|data| serde_json::Value::Object(data.expose_to_provider().clone()))
}
pub(crate) fn fail_provider_listing(&self) {
self.fail_provider_listing.store(true, Ordering::SeqCst);
}
@@ -235,7 +246,6 @@ impl ProviderAccountStore for MemoryAccountStore {
revision: next,
enabled: stored.account.enabled(),
availability: stored.account.availability(),
cooldown_until: stored.account.cooldown_until(),
access_token_expires_at,
has_refresh_token,
next_refresh_at,
@@ -303,7 +313,6 @@ impl ProviderAccountStore for MemoryAccountStore {
revision: stored.account.revision(),
enabled: stored.account.enabled(),
availability: change.availability,
cooldown_until: change.cooldown_until,
access_token_expires_at: stored.account.access_token_expires_at(),
has_refresh_token: stored.account.has_refresh_token(),
next_refresh_at: stored.account.next_refresh_at(),
@@ -325,7 +334,6 @@ impl ProviderAccountStore for MemoryAccountStore {
revision: stored.account.revision(),
enabled: stored.account.enabled(),
availability: stored.account.availability(),
cooldown_until: stored.account.cooldown_until(),
access_token_expires_at: stored.account.access_token_expires_at(),
has_refresh_token: stored.account.has_refresh_token(),
next_refresh_at: stored.account.next_refresh_at(),
@@ -350,7 +358,6 @@ impl ProviderAccountStore for MemoryAccountStore {
revision: stored.account.revision(),
enabled,
availability: stored.account.availability(),
cooldown_until: stored.account.cooldown_until(),
access_token_expires_at: stored.account.access_token_expires_at(),
has_refresh_token: stored.account.has_refresh_token(),
next_refresh_at: stored.account.next_refresh_at(),
@@ -378,7 +385,6 @@ struct AccountRebuild {
revision: CredentialRevision,
enabled: bool,
availability: AccountAvailability,
cooldown_until: Option<SystemTime>,
access_token_expires_at: Option<SystemTime>,
has_refresh_token: bool,
next_refresh_at: Option<SystemTime>,
@@ -407,11 +413,7 @@ fn rebuild_account(current: &ProviderAccount, rebuild: AccountRebuild) -> Provid
current.upstream_account_id().map(str::to_owned),
plan_type,
)
.with_runtime_state(
rebuild.enabled,
rebuild.availability,
rebuild.cooldown_until,
)
.with_runtime_state(rebuild.enabled, rebuild.availability)
.with_refresh_schedule(rebuild.has_refresh_token, rebuild.next_refresh_at)
}
@@ -464,6 +466,7 @@ impl ProviderLeasePort for TestLeaseCoordinator {
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
@@ -765,6 +768,6 @@ pub(crate) fn codex_account(id: &str) -> ProviderAccount {
Some(format!("chatgpt-{id}")),
Some("pro".to_owned()),
)
.with_runtime_state(true, AccountAvailability::Ready, None)
.with_runtime_state(true, AccountAvailability::Ready)
.with_refresh_schedule(true, Some(SystemTime::now() + Duration::from_secs(2_700)))
}
+5 -7
View File
@@ -913,8 +913,6 @@ fn prepared_create(
next_refresh_at: account.next_refresh_at().map(Into::into),
enabled: account.enabled(),
availability: account.availability(),
availability_reason: None,
cooldown_until: account.cooldown_until().map(Into::into),
availability_observed_at: observed_at,
})
}
@@ -1018,11 +1016,7 @@ 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,
account.cooldown_until.map(Into::into),
)
.with_runtime_state(account.enabled, account.availability)
.with_refresh_schedule(
account.has_refresh_token,
account.next_refresh_at.map(Into::into),
@@ -1087,6 +1081,9 @@ fn project_quota(
window_seconds,
used_percent: billing.used_percent(),
reset_at: billing.period_end().and_then(parse_utc),
limit_reached: billing
.used_percent()
.is_some_and(|used| used.is_finite() && used >= 100.0),
local_usage: None,
provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(data))),
}
@@ -1099,6 +1096,7 @@ fn project_quota(
window_seconds: Some(crate::GROK_FREE_ROLLING_WINDOW_SECONDS),
used_percent: None,
reset_at: None,
limit_reached: false,
local_usage: rolling_usage,
provider_data: None,
}
@@ -482,7 +482,6 @@ impl GrokCredentialQuotaService {
GrokCredentialAvailability::QuotaExhausted
)
.then_some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at,
})
.await
@@ -1151,7 +1150,6 @@ fn eligible_catalog_candidate(candidate: &LoadedGrokCredential) -> bool {
&& match account.availability() {
AccountAvailability::Unknown
| AccountAvailability::Ready
| AccountAvailability::Cooldown
| AccountAvailability::QuotaExhausted => true,
AccountAvailability::Expired
| AccountAvailability::Banned
@@ -1294,8 +1292,15 @@ fn quota_scheduling_signals(billing: &GrokBillingPresentation) -> Option<Account
.and_then(|value| DateTime::parse_from_rfc3339(value).ok())
.map(|value| SystemTime::from(value.to_utc()))
.filter(|reset_at| *reset_at > now);
(remaining_rank.is_some() || reset_at.is_some())
.then(|| AccountQuotaSignals::new(reset_at, remaining_rank))
(remaining_rank.is_some() || reset_at.is_some()).then(|| {
AccountQuotaSignals::new(
reset_at,
remaining_rank,
billing
.used_percent()
.is_some_and(|used| used.is_finite() && used >= 100.0),
)
})
}
fn quota_is_exhausted(billing: &GrokBillingPresentation) -> bool {
@@ -1343,7 +1348,6 @@ fn quota_refresh_may_update_state(before: &ProviderAccount, current: &ProviderAc
.is_some_and(|expires_at| expires_at > SystemTime::now())
&& current.revision() == before.revision()
&& current.availability() == before.availability()
&& current.cooldown_until() == before.cooldown_until()
}
fn quota_refresh_availability(
@@ -1357,15 +1361,13 @@ fn quota_refresh_availability(
AccountAvailability::QuotaExhausted => {
matches!(exhausted, Some(false)).then_some(GrokCredentialAvailability::Ready)
}
AccountAvailability::Cooldown | AccountAvailability::Unknown => {
exhausted.map(|exhausted| {
if exhausted {
GrokCredentialAvailability::QuotaExhausted
} else {
GrokCredentialAvailability::Ready
}
})
}
AccountAvailability::Unknown => exhausted.map(|exhausted| {
if exhausted {
GrokCredentialAvailability::QuotaExhausted
} else {
GrokCredentialAvailability::Ready
}
}),
AccountAvailability::Ready => {
matches!(exhausted, Some(true)).then_some(GrokCredentialAvailability::QuotaExhausted)
}
@@ -705,7 +705,6 @@ impl GrokCredentialRefreshService {
expected_revision: credential.credential_revision,
availability,
availability_reason: Some(reason.to_owned()),
cooldown_until: None,
observed_at: Utc::now(),
})
.await
@@ -852,8 +851,6 @@ fn account_due(
| AccountAvailability::Banned
| AccountAvailability::Invalid
)
&& (availability != AccountAvailability::Cooldown
|| account.cooldown_until().is_some_and(|until| until <= now))
&& account.next_refresh_at().is_some_and(|next| next <= now)
}
@@ -168,11 +168,7 @@ impl GrokCredentialAdmin {
input.account.upstream_account_id.clone(),
input.account.plan_type.clone(),
)
.with_runtime_state(
input.enabled,
input.initial_availability.into(),
input.initial_cooldown_until.map(to_system_time),
)
.with_runtime_state(input.enabled, input.initial_availability.into())
.with_refresh_schedule(true, Some(to_system_time(input.next_refresh_at)));
Ok(NewProviderAccount {
account,
@@ -222,7 +218,6 @@ impl GrokCredentialAdmin {
enabled: input.enabled,
initial_availability: GrokCredentialAvailability::Ready,
initial_availability_reason: None,
initial_cooldown_until: None,
})
}
@@ -434,15 +429,12 @@ impl GrokCredentialRepository {
input: &UpdateGrokCredentialState,
) -> Result<(), GrokCredentialRepositoryError> {
validate_reason(input.availability_reason.as_deref())?;
validate_cooldown(input.availability, input.cooldown_until)?;
self.ensure_xai_account(&input.account_id).await?;
self.store
.apply_state_change(AccountStateChange {
account_id: input.account_id.clone(),
expected_revision: input.expected_revision,
availability: input.availability.into(),
reason: input.availability_reason.clone(),
cooldown_until: input.cooldown_until.map(to_system_time),
observed_at: to_system_time(input.observed_at),
})
.await
@@ -708,8 +700,7 @@ fn validate_create(input: &CreateGrokCredential) -> Result<(), GrokCredentialRep
validate_profile(&input.account)?;
validate_refresh_schedule(input.next_refresh_at, input.account.access_token_expires_at)?;
validate_secret(&input.secret)?;
validate_reason(input.initial_availability_reason.as_deref())?;
validate_cooldown(input.initial_availability, input.initial_cooldown_until)
validate_reason(input.initial_availability_reason.as_deref())
}
fn validate_profile(profile: &GrokAccountProfile) -> Result<(), GrokCredentialRepositoryError> {
@@ -816,18 +807,6 @@ fn validate_reason(reason: Option<&str>) -> Result<(), GrokCredentialRepositoryE
Ok(())
}
fn validate_cooldown(
availability: GrokCredentialAvailability,
cooldown_until: Option<DateTime<Utc>>,
) -> Result<(), GrokCredentialRepositoryError> {
if matches!(availability, GrokCredentialAvailability::Cooldown) != cooldown_until.is_some() {
return Err(GrokCredentialRepositoryError::InvalidInput(
"cooldown_until",
));
}
Ok(())
}
fn ensure_xai(account: &ProviderAccount) -> Result<(), GrokCredentialRepositoryError> {
if account.provider().as_str() == XAI_PROVIDER_KIND
&& account.authentication_kind() == XAI_AUTHENTICATION_KIND_OAUTH
@@ -4,7 +4,7 @@ use std::fmt;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use chrono::{DateTime, Utc};
use chrono::Utc;
use gateway_core::engine::credential::{
AccountAvailability, AccountAvailabilityPolicy, AccountCandidate, AccountFeedbackStats,
AccountSelectionContext, AccountSelector, ProviderAccount, ProviderAccountId,
@@ -300,36 +300,26 @@ impl GrokAccountSessionSelector {
for ((account, account_read), model_read) in
accounts.iter().zip(account_reads).zip(model_reads)
{
if let Ok(Some(cooldown)) = account_read {
if cooldown.credential_revision() != account.revision() {
if cooldown.credential_revision() < account.revision() {
let _ = self.cooldowns.clear(account.id(), account.revision()).await;
}
} else if cooldown.until() > now {
insert_runtime_cooldown(
&mut cooled,
account.id(),
cooldown.until(),
RuntimeCooldownScope::Account,
);
}
if let Ok(Some(cooldown)) = account_read
&& cooldown.until() > now
{
// 账号级限流冷却跨凭据轮换保留:上游限流不因本地凭据更换解除。
insert_runtime_cooldown(
&mut cooled,
account.id(),
cooldown.until(),
RuntimeCooldownScope::Account,
);
}
if let Ok(Some(cooldown)) = model_read {
if cooldown.credential_revision() != account.revision() {
if cooldown.credential_revision() < account.revision() {
let _ = self
.cooldowns
.clear_scoped(account.id(), &model_scope, account.revision())
.await;
}
} else if cooldown.until() > now {
insert_runtime_cooldown(
&mut cooled,
account.id(),
cooldown.until(),
RuntimeCooldownScope::Model,
);
}
if let Ok(Some(cooldown)) = model_read
&& cooldown.until() > now
{
insert_runtime_cooldown(
&mut cooled,
account.id(),
cooldown.until(),
RuntimeCooldownScope::Model,
);
}
}
cooled
@@ -374,7 +364,6 @@ impl GrokAccountSessionSelector {
session: &SelectedGrokSession,
availability: GrokCredentialAvailability,
reason: &'static str,
cooldown_until: Option<DateTime<Utc>>,
observed_at: chrono::DateTime<Utc>,
) {
let update = |expected_revision| UpdateGrokCredentialState {
@@ -382,7 +371,6 @@ impl GrokAccountSessionSelector {
expected_revision,
availability,
availability_reason: Some(reason.to_owned()),
cooldown_until,
observed_at,
};
match self
@@ -467,28 +455,22 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
GrokCredentialFailure::QuotaExhausted => (
GrokCredentialAvailability::QuotaExhausted,
"upstream_quota_exhausted",
None,
),
GrokCredentialFailure::FreeQuotaExhausted => (
GrokCredentialAvailability::QuotaExhausted,
"upstream_free_quota_exhausted",
None,
),
GrokCredentialFailure::PaymentRequired { .. } => (
GrokCredentialAvailability::QuotaExhausted,
"upstream_payment_quota_exhausted",
None,
),
// xAI 免费模型额度按账号滚动窗口恢复:限流只写运行时冷却
// (Redis 跨重启保留),不进入持久化账号状态。
GrokCredentialFailure::ModelQuotaExhausted { retry_after, .. } => {
let retry_after =
bounded_cooldown(retry_after, MODEL_QUOTA_COOLDOWN, MAX_MODEL_COOLDOWN);
// xAI 免费模型额度按账号滚动窗口恢复;持久化使页面和各调度节点收敛。
self.record_runtime_cooldown(session, retry_after).await;
(
GrokCredentialAvailability::Cooldown,
"usage_limit_exhausted",
persistent_cooldown_until(observed_at, retry_after),
)
return;
}
GrokCredentialFailure::ModelAccessDenied {
upstream_model,
@@ -508,31 +490,17 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)
.min(MAX_RATE_LIMIT_COOLDOWN);
self.record_runtime_cooldown(session, retry_after).await;
(
GrokCredentialAvailability::Cooldown,
"upstream_rate_limited",
persistent_cooldown_until(observed_at, retry_after),
)
return;
}
GrokCredentialFailure::StreamInterrupted => {
self.record_runtime_cooldown(session, STREAM_INTERRUPTION_COOLDOWN)
.await;
(
GrokCredentialAvailability::Cooldown,
"upstream_stream_interrupted",
persistent_cooldown_until(observed_at, STREAM_INTERRUPTION_COOLDOWN),
)
return;
}
};
let (availability, reason, cooldown_until) = persistent;
self.persist_account_failure(
session,
availability,
reason,
cooldown_until,
observed_at,
)
.await;
let (availability, reason) = persistent;
self.persist_account_failure(session, availability, reason, observed_at)
.await;
})
}
@@ -552,9 +520,7 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
|| !account.enabled()
|| !matches!(
account.availability(),
AccountAvailability::Unknown
| AccountAvailability::Cooldown
| AccountAvailability::Expired
AccountAvailability::Unknown | AccountAvailability::Expired
)
{
return;
@@ -566,7 +532,6 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
expected_revision: account.revision(),
availability: GrokCredentialAvailability::Ready,
availability_reason: None,
cooldown_until: None,
observed_at: Utc::now(),
})
.await
@@ -600,15 +565,6 @@ fn bounded_cooldown(
.min(maximum)
}
fn persistent_cooldown_until(
observed_at: DateTime<Utc>,
duration: Duration,
) -> Option<DateTime<Utc>> {
chrono::Duration::from_std(duration)
.ok()
.map(|duration| observed_at + duration)
}
fn insert_runtime_cooldown(
cooldowns: &mut std::collections::BTreeMap<ProviderAccountId, RuntimeCooldown>,
account_id: &ProviderAccountId,
@@ -64,11 +64,13 @@ impl fmt::Debug for GrokAccountProfile {
}
/// 与 `provider_accounts.availability` 一一对应的 xAI 状态。
///
/// 限流(429/滚动窗口)不进入持久化状态:由 `ProviderCooldownPort`(Redis 跨重启
/// 保留)驱动调度排除,账号状态保持 `Ready`。
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum GrokCredentialAvailability {
Unknown,
Ready,
Cooldown,
QuotaExhausted,
Expired,
Banned,
@@ -81,7 +83,6 @@ impl From<GrokCredentialAvailability> for AccountAvailability {
match value {
GrokCredentialAvailability::Unknown => Self::Unknown,
GrokCredentialAvailability::Ready => Self::Ready,
GrokCredentialAvailability::Cooldown => Self::Cooldown,
GrokCredentialAvailability::QuotaExhausted => Self::QuotaExhausted,
GrokCredentialAvailability::Expired => Self::Expired,
GrokCredentialAvailability::Banned => Self::Banned,
@@ -100,7 +101,6 @@ pub struct CreateGrokCredential {
pub enabled: bool,
pub initial_availability: GrokCredentialAvailability,
pub initial_availability_reason: Option<String>,
pub initial_cooldown_until: Option<DateTime<Utc>>,
}
impl fmt::Debug for CreateGrokCredential {
@@ -243,7 +243,6 @@ pub struct UpdateGrokCredentialState {
pub expected_revision: CredentialRevision,
pub availability: GrokCredentialAvailability,
pub availability_reason: Option<String>,
pub cooldown_until: Option<DateTime<Utc>>,
pub observed_at: DateTime<Utc>,
}
+1 -2
View File
@@ -472,8 +472,6 @@ fn account_record(account: &ProviderAccount) -> AccountRecord {
next_refresh_at: account.next_refresh_at().map(DateTime::<Utc>::from),
enabled: account.enabled(),
availability: account.availability(),
availability_reason: None,
cooldown_until: account.cooldown_until().map(DateTime::<Utc>::from),
availability_observed_at: now,
quota_observed_at: None,
created_at: now,
@@ -504,6 +502,7 @@ impl ProviderLeasePort for TestLeases {
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
@@ -201,15 +201,12 @@ async fn set_account_state(
store: &MemoryProviderAccountStore,
id: &ProviderAccountId,
availability: AccountAvailability,
cooldown_until: Option<SystemTime>,
) {
store
.apply_state_change(AccountStateChange {
account_id: id.clone(),
expected_revision: CredentialRevision::new(1).expect("revision"),
availability,
reason: Some("test isolation".to_owned()),
cooldown_until,
observed_at: SystemTime::now(),
})
.await
@@ -491,8 +488,6 @@ async fn quota_exhausted_account_remains_eligible_for_catalog_discovery() {
account_id: id,
expected_revision: CredentialRevision::new(1).expect("revision"),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
@@ -686,18 +681,14 @@ async fn quota_refresh_persists_dynamic_provider_document_and_projects_known_fie
}
#[tokio::test]
async fn recovered_quota_refresh_releases_existing_quota_and_cooldown_isolation() {
for (suffix, availability, cooldown_until) in [
("recover-quota", AccountAvailability::QuotaExhausted, None),
(
"recover-cooldown",
AccountAvailability::Cooldown,
SystemTime::now().checked_add(Duration::from_secs(300)),
),
async fn recovered_quota_refresh_releases_existing_quota_isolation() {
for (suffix, availability) in [
("recover-quota", AccountAvailability::QuotaExhausted),
("recover-ready", AccountAvailability::Ready),
] {
let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await;
let id = account_id(suffix);
set_account_state(&store, &id, availability, cooldown_until).await;
set_account_state(&store, &id, availability).await;
crate::support::grok_quota_service(
repository,
QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":25}}"#),
@@ -708,7 +699,6 @@ async fn recovered_quota_refresh_releases_existing_quota_and_cooldown_isolation(
let account = store.account(&id).expect("account");
assert_eq!(account.availability(), AccountAvailability::Ready);
assert_eq!(account.cooldown_until(), None);
}
}
@@ -717,7 +707,7 @@ async fn authoritative_exhaustion_preserves_quota_exhausted_state() {
let (store, repository) =
repository_with_accounts(&[("still-exhausted", "still-exhausted")]).await;
let id = account_id("still-exhausted");
set_account_state(&store, &id, AccountAvailability::QuotaExhausted, None).await;
set_account_state(&store, &id, AccountAvailability::QuotaExhausted).await;
crate::support::grok_quota_service(
repository,
@@ -742,7 +732,7 @@ async fn recovered_quota_does_not_clear_terminal_account_states() {
] {
let (store, repository) = repository_with_accounts(&[(suffix, suffix)]).await;
let id = account_id(suffix);
set_account_state(&store, &id, availability, None).await;
set_account_state(&store, &id, availability).await;
crate::support::grok_quota_service(
repository,
QueueBillingTransport::success(br#"{"config":{"creditUsagePercent":10}}"#),
@@ -759,20 +749,15 @@ async fn recovered_quota_does_not_clear_terminal_account_states() {
}
#[tokio::test]
async fn quota_refresh_does_not_overwrite_a_newer_cooldown() {
async fn quota_refresh_preserves_ready_availability_across_concurrent_state_write() {
let (store, repository) = repository_with_accounts(&[("new-cooldown", "new-cooldown")]).await;
let id = account_id("new-cooldown");
let new_cooldown = SystemTime::now()
.checked_add(Duration::from_secs(600))
.expect("cooldown time");
let transport = Arc::new(MutatingBillingTransport {
store: Arc::clone(&store),
mutation: Mutex::new(Some(BillingMutation::State(AccountStateChange {
account_id: id.clone(),
expected_revision: CredentialRevision::new(1).expect("revision"),
availability: AccountAvailability::Cooldown,
reason: Some("new rate limit".to_owned()),
cooldown_until: Some(new_cooldown),
availability: AccountAvailability::Ready,
observed_at: SystemTime::now(),
}))),
body: br#"{"config":{"creditUsagePercent":10}}"#.to_vec(),
@@ -781,11 +766,10 @@ async fn quota_refresh_does_not_overwrite_a_newer_cooldown() {
crate::support::grok_quota_service(repository, transport)
.refresh_account(&id)
.await
.expect("refresh around concurrent cooldown");
.expect("refresh around concurrent state write");
let account = store.account(&id).expect("account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert_eq!(account.cooldown_until(), Some(new_cooldown));
assert_eq!(account.availability(), AccountAvailability::Ready);
}
#[tokio::test]
@@ -878,7 +862,7 @@ async fn non_authoritative_quota_refresh_does_not_clear_existing_quota_exhaustio
let (store, repository) =
repository_with_accounts(&[("free-quota", "subject-free-quota")]).await;
let id = account_id("free-quota");
set_account_state(&store, &id, AccountAvailability::QuotaExhausted, None).await;
set_account_state(&store, &id, AccountAvailability::QuotaExhausted).await;
let transport = QueueBillingTransport::success(
br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"onDemandCap":{"val":0},"onDemandUsed":{"val":0},"prepaidBalance":{"val":0}}}"#,
);
@@ -219,44 +219,18 @@ async fn state_update_uses_credential_revision_fence() {
let (store, repository) = repository();
let input = create_input("state", "subject-state");
seed_input(&store, &input).await.expect("create");
let cooldown_until = Utc::now() + chrono::Duration::minutes(1);
repository
.update_state(&UpdateGrokCredentialState {
account_id: input.account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("revision"),
availability: GrokCredentialAvailability::Cooldown,
availability_reason: Some("rate_limited".to_owned()),
cooldown_until: Some(cooldown_until),
availability: GrokCredentialAvailability::Ready,
availability_reason: Some("upstream_rate_limited".to_owned()),
observed_at: Utc::now(),
})
.await
.expect("state update");
let account = store.account(&input.account_id).expect("account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(account.cooldown_until().is_some());
}
#[tokio::test]
async fn cooldown_requires_matching_deadline() {
let (store, repository) = repository();
let input = create_input("bad-cooldown", "subject-cooldown");
seed_input(&store, &input).await.expect("create");
let result = repository
.update_state(&UpdateGrokCredentialState {
account_id: input.account_id,
expected_revision: CredentialRevision::new(1).expect("revision"),
availability: GrokCredentialAvailability::Cooldown,
availability_reason: None,
cooldown_until: None,
observed_at: Utc::now(),
})
.await;
assert_eq!(
result,
Err(GrokCredentialRepositoryError::InvalidInput(
"cooldown_until"
))
);
assert_eq!(account.availability(), AccountAvailability::Ready);
}
#[tokio::test]
@@ -110,7 +110,6 @@ impl SelectorFixture {
expected_revision: CredentialRevision::new(1).expect("revision"),
availability: GrokCredentialAvailability::Ready,
availability_reason: None,
cooldown_until: None,
observed_at: Utc::now(),
})
.await
@@ -132,6 +131,7 @@ impl SelectorFixture {
last_started_at: None,
quota_reset_at: None,
quota_remaining_rank: None,
quota_limit_reached: false,
failure_rate_basis_points: None,
first_output_latency_ms: None,
},
@@ -369,21 +369,17 @@ async fn account_scoped_cooldown_survives_credential_rotation() {
CredentialCasOutcome::Updated(revision) if revision.get() == 2
));
// 账号级限流冷却跨凭据轮换保留:轮换后仍被冷却排除。
assert!(matches!(
fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await,
Err(GrokSessionSelectorError::NoEligibleSession)
Err(GrokSessionSelectorError::AccountCoolingDown { .. })
));
let account = fixture.store.account(&id).expect("rotated account");
assert_eq!(account.revision().get(), 2);
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert_eq!(account.availability(), AccountAvailability::Ready);
}
#[tokio::test]
@@ -405,12 +401,7 @@ async fn rate_limit_feedback_persists_the_grok2api_cooldown_state() {
)
.await;
let account = fixture.store.account(&selected).expect("account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert_eq!(account.availability(), AccountAvailability::Ready);
assert!(
fixture
.cooldowns
@@ -450,7 +441,6 @@ async fn successful_request_clears_the_persisted_cooldown_state() {
.account(session.account_id())
.expect("account");
assert_eq!(account.availability(), AccountAvailability::Ready);
assert!(account.cooldown_until().is_none());
}
#[tokio::test]
@@ -511,13 +501,8 @@ async fn free_model_quota_feedback_persists_resettable_account_exhaustion() {
.await;
let account = fixture.store.account(&selected).expect("selected account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert_eq!(account.availability(), AccountAvailability::Ready);
let minimum_until = SystemTime::now() + Duration::from_secs(23 * 60 * 60);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > minimum_until)
);
let scope = ProviderCooldownScope::upstream_model(failed_model);
assert!(
fixture
@@ -604,12 +589,7 @@ async fn interrupted_stream_feedback_persists_the_grok2api_cooldown_state() {
.store
.account(session.account_id())
.expect("account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert_eq!(account.availability(), AccountAvailability::Ready);
assert!(
fixture
.cooldowns
+1 -11
View File
@@ -328,7 +328,6 @@ impl ProviderAccountStore for MemoryProviderAccountStore {
revision: next,
access_token_expires_at,
availability: stored.account.availability(),
cooldown_until: stored.account.cooldown_until(),
enabled: stored.account.enabled(),
has_refresh_token,
next_refresh_at,
@@ -381,7 +380,6 @@ impl ProviderAccountStore for MemoryProviderAccountStore {
revision: stored.account.revision(),
access_token_expires_at: stored.account.access_token_expires_at(),
availability: change.availability,
cooldown_until: change.cooldown_until,
enabled: stored.account.enabled(),
has_refresh_token: stored.account.has_refresh_token(),
next_refresh_at: stored.account.next_refresh_at(),
@@ -402,7 +400,6 @@ impl ProviderAccountStore for MemoryProviderAccountStore {
revision: stored.account.revision(),
access_token_expires_at: stored.account.access_token_expires_at(),
availability: stored.account.availability(),
cooldown_until: stored.account.cooldown_until(),
enabled: stored.account.enabled(),
has_refresh_token: stored.account.has_refresh_token(),
next_refresh_at: stored.account.next_refresh_at(),
@@ -427,7 +424,6 @@ impl ProviderAccountStore for MemoryProviderAccountStore {
revision: stored.account.revision(),
access_token_expires_at: stored.account.access_token_expires_at(),
availability: stored.account.availability(),
cooldown_until: stored.account.cooldown_until(),
enabled,
has_refresh_token: stored.account.has_refresh_token(),
next_refresh_at: stored.account.next_refresh_at(),
@@ -449,7 +445,6 @@ struct AccountReplacement {
revision: CredentialRevision,
access_token_expires_at: Option<SystemTime>,
availability: AccountAvailability,
cooldown_until: Option<SystemTime>,
enabled: bool,
has_refresh_token: bool,
next_refresh_at: Option<SystemTime>,
@@ -473,11 +468,7 @@ fn rebuild_account(previous: &ProviderAccount, replacement: AccountReplacement)
previous.upstream_account_id().map(str::to_owned),
replacement.plan_type,
)
.with_runtime_state(
replacement.enabled,
replacement.availability,
replacement.cooldown_until,
)
.with_runtime_state(replacement.enabled, replacement.availability)
.with_refresh_schedule(replacement.has_refresh_token, replacement.next_refresh_at)
}
@@ -549,7 +540,6 @@ pub fn create_input(suffix: &str, subject: &str) -> CreateGrokCredential {
enabled: true,
initial_availability: GrokCredentialAvailability::Unknown,
initial_availability_reason: None,
initial_cooldown_until: None,
}
}
@@ -0,0 +1,43 @@
-- 移除账号级 cooldown 状态模型(照 v2 语义:429 限流不进 availability,由 quota 数据驱动)。
--
-- 存量迁移:
-- - availability='cooldown' 的账号:429/Cloudflare 等临时冷却 → 'ready'
-- (限流事实留在 provider_quota_json,由投影滚动兜底);
-- 402 语义的 usage_limit_exhausted 且未到期 → 'quota_exhausted'。
-- 1. 迁移存量 cooldown 账号。
update provider_accounts
set availability = case
when availability_reason = 'usage_limit_exhausted'
and cooldown_until is not null
and cooldown_until > now()
then 'quota_exhausted'
else 'ready'
end
where availability = 'cooldown';
-- 2. 删除依赖 cooldown_until 的约束与索引。
alter table provider_accounts drop constraint provider_accounts_cooldown_ck;
drop index provider_accounts_availability_idx;
-- 3. availability CHECK 去掉 'cooldown'。
alter table provider_accounts drop constraint provider_accounts_availability_ck;
alter table provider_accounts
add constraint provider_accounts_availability_ck check (
availability in (
'unknown',
'ready',
'quota_exhausted',
'expired',
'banned',
'invalid'
)
);
-- 4. 删除列。
alter table provider_accounts drop column availability_reason;
alter table provider_accounts drop column cooldown_until;
-- 5. 重建可用性索引(不含 cooldown_until)。
create index provider_accounts_availability_idx
on provider_accounts (availability, id);
+1 -1
View File
@@ -7,6 +7,7 @@ export interface AccountQuotaWindow {
labelDisplay: string
usedPercent: number | null
usedPercentDisplay: string
limitReached: boolean
localUsage?: unknown
resetAtDisplay: string
}
@@ -102,7 +103,6 @@ export interface Account {
hasRefreshToken: boolean
status: string
displayStatus: string
tokenRefreshing: boolean
availability: string
enabled: boolean
accessTokenExpiresAt: string | null
+1
View File
@@ -144,6 +144,7 @@ export interface DashboardAccountUsage {
export interface DashboardPoolSummary {
total: number
active: number
rateLimited: number
expired: number
quotaExhausted: number
refreshing: number | null
@@ -230,7 +230,7 @@ export function useAccountConnectionTest(options: AccountConnectionTestOptions)
testingAccount.value = {
...account,
status,
displayStatus: account.tokenRefreshing ? account.displayStatus : status,
displayStatus: status,
}
options.onAccountStatus(account.id, status)
}
@@ -221,7 +221,7 @@ export function useAccountMutations(options: {
? {
...account,
status,
displayStatus: account.tokenRefreshing ? 'refreshing' : status,
displayStatus: status,
}
: account,
)
+18 -3
View File
@@ -128,7 +128,7 @@ export const statusLabels: Record<string, string> = {
disabled: '已禁用',
banned: '已封禁',
quota_exhausted: '配额耗尽',
refreshing: '刷新中',
rate_limited: '限流中',
}
export const statusTones: Record<string, 'success' | 'danger' | 'warning' | 'info' | 'normal'> = {
@@ -137,18 +137,33 @@ export const statusTones: Record<string, 'success' | 'danger' | 'warning' | 'inf
disabled: 'normal',
banned: 'danger',
quota_exhausted: 'warning',
refreshing: 'info',
rate_limited: 'warning',
}
export const accountStatusFilterOptions = [
{ label: '全部状态', value: '' },
{ label: statusLabels.active, value: 'active' },
{ label: statusLabels.expired, value: 'expired' },
{ label: statusLabels.quota_exhausted, value: 'quota_exhausted' },
{ label: statusLabels.expired, value: 'expired' },
{ label: statusLabels.disabled, value: 'disabled' },
{ label: statusLabels.banned, value: 'banned' },
]
/**
* 照 v2:`rate_limited` 不是后端状态,而是由「任一 quota 窗口 limit_reached」派生。
* 账号 `active` 但任一窗口触顶 → 显示「限流中」。
*/
export function derivedAccountStatus(row: AccountRow): string {
if (row.status === 'active' && isQuotaWindowReached(row)) {
return 'rate_limited'
}
return row.status
}
function isQuotaWindowReached(row: AccountRow): boolean {
return row.quota.windows.some(window => window.limitReached)
}
export function visibleSummaryQuotaWindows(windows: AccountQuotaWindow[]) {
const known = [...windows]
.filter(window => summaryGroupOrder.has(window.group))
+3 -3
View File
@@ -23,7 +23,7 @@ import { useAccountConnectionTest } from './composables/useAccountConnectionTest
import { useAccountMutations } from './composables/useAccountMutations'
import { useAccountsQuery } from './composables/useAccountsQuery'
import { useAccountsTable } from './composables/useAccountsTable'
import { accountColumns } from './constants'
import { accountColumns, derivedAccountStatus } from './constants'
const selectedIds = ref<Set<string>>(new Set())
const {
@@ -204,7 +204,7 @@ const {
</template>
<template #status="{ row }">
<AccountStatusBadge :status="row.tokenRefreshing ? 'refreshing' : row.status" />
<AccountStatusBadge :status="derivedAccountStatus(row)" />
</template>
<template #planType="{ row }">
@@ -219,7 +219,7 @@ const {
<AccountTableActions
:account="row"
:deleting="deletingAccount"
:refreshing="refreshingAccountIds.has(row.id) || row.tokenRefreshing"
:refreshing="refreshingAccountIds.has(row.id)"
:schedule-label="scheduleActionLabel(row)"
:testing="testingConnectionIds.has(row.id)"
:updating-status="updatingStatusAccountIds.has(row.id)"
@@ -112,8 +112,15 @@ const statusRows = computed(() => {
icon: RefreshCw,
},
{
label: '额度受限',
description: '已触发额度阈值',
label: '限流中',
description: '任一额度窗口已触顶',
value: String(p?.rateLimited ?? 0),
tone: 'warning',
icon: TriangleAlert,
},
{
label: '配额耗尽',
description: '已确认 402 耗尽',
value: String(quota),
tone: 'warning',
icon: TriangleAlert,
@@ -141,11 +148,13 @@ const statusBars = computed(() => {
return []
const active = (p.active / p.total) * 100
const refreshing = ((p.refreshing ?? 0) / p.total) * 100
const rateLimited = ((p.rateLimited ?? 0) / p.total) * 100
const quota = (p.quotaExhausted / p.total) * 100
const unavailable = ((p.expired + p.disabled + p.banned) / p.total) * 100
return [
{ pct: active, cls: 'bg-(--cp-success)' },
{ pct: refreshing, cls: 'bg-(--cp-normal)' },
{ pct: rateLimited, cls: 'bg-(--cp-warning)' },
{ pct: quota, cls: 'bg-(--cp-warning)' },
{ pct: unavailable, cls: 'bg-(--cp-danger)' },
].filter(b => b.pct > 0)
@@ -367,6 +367,7 @@ function dashboardUsageWindow(item: DashboardSummary['accountUsage'][number]) {
labelDisplay: '额度',
usedPercent,
usedPercentDisplay: usedPercent === null ? '—' : `${usedPercent}%`,
limitReached: (usedPercent ?? 0) >= 100,
windowSeconds: null,
resetAtDisplay: '—',
}
@@ -388,6 +389,7 @@ function dashboardUsageWindow(item: DashboardSummary['accountUsage'][number]) {
labelDisplay: '日请求',
usedPercent: null,
usedPercentDisplay: '—',
limitReached: false,
windowSeconds: 86_400,
resetAtDisplay: '—',
localUsage: {