From 4c57a0405df277ba59bffc482df8f6359cd5667c Mon Sep 17 00:00:00 2001 From: Zyy Date: Wed, 4 Nov 2026 12:39:10 +0800 Subject: [PATCH] fix(xai): preserve quota snapshots and restore daily usage Preserve account profiles and quota observations during OAuth refresh and retry deferral. Project rolling daily usage independently of upstream snapshots and expose its period through the existing account usage API. Keep quota capacity forecasts scoped to weekly and monthly windows. Cover refresh persistence, daily statistics, API projection, and PostgreSQL CAS behavior. Co-authored-by: Codex --- .../src/model/provider_credentials.rs | 24 +++++-- .../gateway-admin/src/model/quota_forecast.rs | 11 +-- backend/crates/gateway-admin/tests/main.rs | 52 -------------- .../tests/model/provider_credentials.rs | 30 ++++++++ .../tests/model/quota_forecast.rs | 5 +- .../src/admin/accounts/presenter.rs | 2 + .../tests/admin/accounts/handlers.rs | 70 +++++++++++++++++-- .../tests/postgres/provider_accounts/mod.rs | 20 ++++++ backend/crates/providers/xai/src/admin.rs | 24 +++---- .../providers/xai/src/credential/refresh.rs | 3 +- .../xai/src/credential/repository.rs | 5 +- backend/crates/providers/xai/tests/admin.rs | 36 ++++++++-- .../providers/xai/tests/credential/refresh.rs | 40 ++++++++++- backend/crates/providers/xai/tests/support.rs | 8 ++- docs/api.md | 3 +- docs/architecture.md | 2 + 16 files changed, 238 insertions(+), 97 deletions(-) diff --git a/backend/crates/gateway-admin/src/model/provider_credentials.rs b/backend/crates/gateway-admin/src/model/provider_credentials.rs index fee68c9b1..c6e896bbd 100644 --- a/backend/crates/gateway-admin/src/model/provider_credentials.rs +++ b/backend/crates/gateway-admin/src/model/provider_credentials.rs @@ -806,11 +806,12 @@ pub struct ProviderQuotaWindow { pub provider_data: Option, } -/// 账号用量统计周期，按周优先、月次之选择；短期限流窗口不参与统计面板 +/// 账号用量统计周期，完整周/月额度优先，随后使用 Provider 提供的本地滚动日统计 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] pub enum AccountUsagePeriod { Weekly, Monthly, + Daily, } /// Provider 解释 quota 所需的公共请求事实 @@ -998,16 +999,27 @@ impl ProviderQuota { } } - /// 选择账号用量统计的周/月窗口，同时返回展示所需的周期事实 + /// 选择账号用量统计窗口，同时返回展示所需的周期事实 /// - /// 只使用边界完整、可归属到整个账号的窗口；周优先于月，同周期保持 - /// Provider 投影顺序 - /// 没有符合条件的窗口时不回退到短期或历史累计 + /// 周/月额度使用边界完整、可归属到整个账号的窗口 + /// 周优先于月，同周期保持 Provider 投影顺序 + /// Provider 提供的无固定重置点的滚动日统计可独立展示 #[must_use] pub fn usage_window(&self) -> Option<(&ProviderQuotaWindow, AccountUsagePeriod)> { self.usage_windows() .filter(|(window, _)| window.local_usage.is_some()) .min_by_key(|(_, period)| *period) + .or_else(|| { + self.windows + .iter() + .find(|window| { + window.local_usage_attribution == QuotaLocalUsageAttribution::AccountWide + && window.window_seconds == Some(86_400) + && window.reset_at.is_none() + && window.local_usage.is_some() + }) + .map(|window| (window, AccountUsagePeriod::Daily)) + }) } /// 互联账号按最长的 Key 消费周期统计本地用量，池内账号窗口不代表此 Key 的周期 @@ -1025,7 +1037,7 @@ impl ProviderQuota { .max_by_key(|window| window.window_seconds) } - /// 统计与预测共用窗口归属规则，避免把模型专属桶或短期限流当作账号容量 + /// 周/月额度的归属规则，供额度周期统计与容量预测共用 pub(crate) fn usage_windows( &self, ) -> impl Iterator { diff --git a/backend/crates/gateway-admin/src/model/quota_forecast.rs b/backend/crates/gateway-admin/src/model/quota_forecast.rs index 4c943bf8d..0be0bbfa7 100644 --- a/backend/crates/gateway-admin/src/model/quota_forecast.rs +++ b/backend/crates/gateway-admin/src/model/quota_forecast.rs @@ -62,14 +62,15 @@ pub fn account_quota_forecasts( now: DateTime, samples: &[QuotaForecastSample], ) -> [AccountQuotaForecast; 2] { - [AccountUsagePeriod::Weekly, AccountUsagePeriod::Monthly].map(|period| { + [ + (AccountUsagePeriod::Weekly, 7 * DAY_SECONDS), + (AccountUsagePeriod::Monthly, 30 * DAY_SECONDS), + ] + .map(|(period, target_seconds)| { let selected = quota_forecast_source_window(quota, period); let mut forecast = AccountQuotaForecast { period, - target_seconds: match period { - AccountUsagePeriod::Weekly => 7 * DAY_SECONDS, - AccountUsagePeriod::Monthly => 30 * DAY_SECONDS, - }, + target_seconds, extrapolated: false, source: None, unavailable_reason: Some("缺少周/月额度窗口，请刷新账号额度"), diff --git a/backend/crates/gateway-admin/tests/main.rs b/backend/crates/gateway-admin/tests/main.rs index 791a1fc08..2ebb9dfc2 100644 --- a/backend/crates/gateway-admin/tests/main.rs +++ b/backend/crates/gateway-admin/tests/main.rs @@ -7,7 +7,6 @@ use gateway_core::account::OpaqueProviderData; use gateway_admin::model::{ PageSize, Revision, - accounts::AccountUsage, auth::LoginCommand, client_keys::ClientKeyPageSize, observability::{DecimalAmount, RequestOutcome, TimeRange}, @@ -146,57 +145,6 @@ fn representative_quota_should_prefer_account_wide_window_over_model_specific_wi assert_eq!(quota.representative_used_percent(), Some(5.0)); } -#[test] -fn usage_window_should_not_use_daily_rolling_usage_for_weekly_statistics() { - let usage = AccountUsage { - account_id: "acct_xai".to_owned(), - request_count: 5, - success_count: 5, - input_tokens: Some(100), - output_tokens: Some(20), - cached_input_tokens: Some(80), - cache_write_input_tokens: Some(0), - reasoning_output_tokens: Some(0), - image_input_tokens: Some(0), - image_output_tokens: Some(0), - image_request_count: 0, - image_request_failed_count: 0, - total_tokens: Some(120), - cost_coverage: Default::default(), - costs: Vec::new(), - last_used_at: Some(Utc::now()), - request_buckets: Vec::new(), - models: Vec::new(), - }; - let quota = ProviderQuota { - upstream: None, - credits: None, - plan_type: None, - observed_at: None, - refresh_token_expires_at: None, - windows: vec![ProviderQuotaWindow { - upstream: None, - key: "free-rolling-24h".to_owned(), - group: "shortTerm".to_owned(), - label: "日限额".to_owned(), - limit_id: None, - limit_name: None, - role: None, - local_usage_attribution: QuotaLocalUsageAttribution::AccountWide, - window_seconds: Some(86_400), - used_percent: None, - reset_at: None, - limit_reached: false, - local_usage: Some(usage), - provider_data: None, - }], - limit_reached: false, - provider_data: None, - }; - - assert_eq!(quota.usage_window(), None); -} - #[test] fn exhausted_quota_should_project_full_usage_to_only_the_representative_window() { let mut quota = ProviderQuota { diff --git a/backend/crates/gateway-admin/tests/model/provider_credentials.rs b/backend/crates/gateway-admin/tests/model/provider_credentials.rs index 07c200a52..743b85e20 100644 --- a/backend/crates/gateway-admin/tests/model/provider_credentials.rs +++ b/backend/crates/gateway-admin/tests/model/provider_credentials.rs @@ -65,6 +65,36 @@ fn usage_window_should_not_fallback_to_short_daily_or_incomplete_windows() { assert_eq!(quota.usage_window(), None); } +#[test] +fn usage_window_should_preserve_provider_rolling_daily_statistics() { + let mut rolling = usage_window("rolling-day", "shortTerm", 86_400); + rolling.reset_at = None; + rolling.used_percent = None; + let mut quota = usage_quota(vec![rolling.clone()]); + assert_eq!( + selected_usage_window("a), + Some(("rolling-day", AccountUsagePeriod::Daily)) + ); + + quota.limit_reached = true; + quota.apply_limit_reached_display(); + assert_eq!( + selected_usage_window("a), + Some(("rolling-day", AccountUsagePeriod::Daily)) + ); + + quota + .windows + .push(usage_window("week", "shortTerm", 7 * 86_400)); + assert_eq!( + selected_usage_window("a), + Some(("week", AccountUsagePeriod::Weekly)) + ); + + rolling.local_usage_attribution = QuotaLocalUsageAttribution::Unavailable; + assert_eq!(usage_quota(vec![rolling]).usage_window(), None); +} + #[test] fn usage_window_should_preserve_provider_order_for_equal_periods() { let quota = usage_quota(vec![ diff --git a/backend/crates/gateway-admin/tests/model/quota_forecast.rs b/backend/crates/gateway-admin/tests/model/quota_forecast.rs index 81f03d9be..6284af338 100644 --- a/backend/crates/gateway-admin/tests/model/quota_forecast.rs +++ b/backend/crates/gateway-admin/tests/model/quota_forecast.rs @@ -162,7 +162,10 @@ fn source_selection_keeps_provider_order_with_the_maximum_window_count() { fn short_or_model_specific_windows_are_not_account_capacity() { let mut model = window("model", 7); model.local_usage_attribution = QuotaLocalUsageAttribution::Unavailable; - for item in forecast("a(vec![window("day", 1), model])) { + let mut rolling = window("rolling-day", 1); + rolling.reset_at = None; + rolling.used_percent = None; + for item in forecast("a(vec![window("day", 1), rolling, model])) { assert!(item.source.is_none()); assert!(item.unavailable_reason.is_some()); assert!(item.estimated_tokens.is_none()); diff --git a/backend/crates/gateway-api/src/admin/accounts/presenter.rs b/backend/crates/gateway-api/src/admin/accounts/presenter.rs index 622dbfac8..9ef0f8fd6 100644 --- a/backend/crates/gateway-api/src/admin/accounts/presenter.rs +++ b/backend/crates/gateway-api/src/admin/accounts/presenter.rs @@ -193,6 +193,7 @@ fn quota_forecast_view( period: match forecast.period { AccountUsagePeriod::Weekly => "weekly", AccountUsagePeriod::Monthly => "monthly", + AccountUsagePeriod::Daily => "daily", }, target_days: forecast.target_seconds as f64 / 86_400.0, extrapolated: forecast.extrapolated, @@ -389,6 +390,7 @@ pub(super) fn account_usage_view( }; AccountUsageView { window_kind: match period { + Some(AccountUsagePeriod::Daily) => AccountUsageWindowKind::Daily, Some(AccountUsagePeriod::Weekly) => AccountUsageWindowKind::Weekly, Some(AccountUsagePeriod::Monthly) => AccountUsageWindowKind::Monthly, None => AccountUsageWindowKind::Quota, diff --git a/backend/crates/gateway-api/tests/admin/accounts/handlers.rs b/backend/crates/gateway-api/tests/admin/accounts/handlers.rs index 769dca432..9a9ca855a 100644 --- a/backend/crates/gateway-api/tests/admin/accounts/handlers.rs +++ b/backend/crates/gateway-api/tests/admin/accounts/handlers.rs @@ -444,9 +444,12 @@ fn account_page_item( } #[tokio::test] -async fn account_usage_preserves_cache_read_and_write_counts_and_missing_values() { +async fn account_usage_preserves_periods_cache_counts_and_missing_values() { use chrono::Utc; use gateway_admin::model::accounts::AccountUsage; + use gateway_admin::model::provider_credentials::{ + ProviderQuota, ProviderQuotaWindow, QuotaLocalUsageAttribution, + }; use gateway_core::account::{AccountStatusFacts, CredentialState, QuotaState}; let fixture = AdminTestFixture::new().await; @@ -461,9 +464,19 @@ async fn account_usage_preserves_cache_read_and_write_counts_and_missing_values( last_error_reason: None, last_error_message: None, }; - *fixture.account.lock().unwrap() = Some(account_page_item(facts, now, "api_key")); - for (read, write) in [(Some(13), Some(7)), (None, None)] { - *fixture.account_usage.lock().unwrap() = vec![AccountUsage { + for (authentication_kind, read, write) in [ + ("api_key", Some(13), Some(7)), + ("api_key", None, None), + ("oauth", Some(13), Some(7)), + ("oauth", None, None), + ] { + let mut account = account_page_item(facts.clone(), now, authentication_kind); + if authentication_kind == "oauth" { + account.account.provider_kind = + gateway_core::routing::ProviderKind::new("xai").unwrap(); + } + *fixture.account.lock().unwrap() = Some(account); + let local_usage = AccountUsage { account_id: "acct_cooldown".to_owned(), request_count: 1, success_count: 1, @@ -480,9 +493,33 @@ async fn account_usage_preserves_cache_read_and_write_counts_and_missing_values( cost_coverage: Default::default(), costs: Vec::new(), last_used_at: Some(now), - request_buckets: Vec::new(), + request_buckets: vec![gateway_admin::model::accounts::AccountRequestBucket { + bucket_start: now, + request_count: 1, + }], models: Vec::new(), - }]; + }; + *fixture.account_usage.lock().unwrap() = vec![local_usage.clone()]; + *fixture.provider_quota.lock().unwrap() = + (authentication_kind == "oauth").then(|| ProviderQuota { + windows: vec![ProviderQuotaWindow { + upstream: None, + key: "rolling-day".to_owned(), + group: "shortTerm".to_owned(), + label: "日限额".to_owned(), + limit_id: None, + limit_name: None, + role: None, + local_usage_attribution: QuotaLocalUsageAttribution::AccountWide, + window_seconds: Some(86_400), + used_percent: None, + reset_at: None, + limit_reached: false, + local_usage: Some(local_usage), + provider_data: None, + }], + ..Default::default() + }); let response = admin::accounts::router::() .with_state(fixture.state()) .oneshot( @@ -500,6 +537,27 @@ async fn account_usage_preserves_cache_read_and_write_counts_and_missing_values( let value: serde_json::Value = serde_json::from_slice(&body).unwrap(); assert_eq!(status, StatusCode::OK, "{value}"); let usage = &value["data"]["items"][0]["usage"]; + assert_eq!( + usage["windowKind"], + if authentication_kind == "oauth" { + "daily" + } else { + "general" + } + ); + assert_eq!(usage["requestCount"], 1); + assert_eq!(usage["totalTokens"], 120); + if authentication_kind == "oauth" { + let window = &value["data"]["items"][0]["quota"]["windows"][0]; + assert_eq!(window["localUsage"]["requestCount"], 1); + assert_eq!( + window["localUsage"]["requestBuckets"] + .as_array() + .unwrap() + .len(), + 1 + ); + } assert_eq!(usage["cachedInputTokens"], serde_json::json!(read)); assert_eq!(usage["cacheWriteInputTokens"], serde_json::json!(write)); assert_eq!( diff --git a/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs b/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs index 3da93a05a..7ab301814 100644 --- a/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs +++ b/backend/crates/gateway-store/tests/postgres/provider_accounts/mod.rs @@ -508,6 +508,10 @@ async fn profile_changes_survive_inflight_background_and_manual_token_refresh() assert_eq!(upgraded.plan_type(), Some("pro")); assert_eq!(upgraded.revision(), original.revision()); assert_eq!(current_revision(&database.pool).await, config_revision); + let mut saved_quota = repository + .get_quotas(std::slice::from_ref(&id)) + .await + .unwrap(); // 刷新已准备好旧资料，CAS 提交仍须保留数据库中的新名称和套餐 let refresh = CredentialCasUpdate::new( @@ -541,6 +545,14 @@ async fn profile_changes_survive_inflight_background_and_manual_token_refresh() let current = repository.get_account(&id).await.unwrap().unwrap(); assert_eq!(current.plan_type(), Some("pro")); assert_eq!(current.name(), "renamed team"); + saved_quota[0].expected_revision = current.revision(); + assert_eq!( + repository + .get_quotas(std::slice::from_ref(&id)) + .await + .unwrap(), + saved_quota + ); let stale = QuotaObservation { expected_connection: None, plan_type: Some("free".to_owned()), @@ -571,6 +583,14 @@ async fn profile_changes_survive_inflight_background_and_manual_token_refresh() assert_eq!(refreshed.plan_type(), Some("pro")); assert_eq!(refreshed.name(), "renamed team"); assert_eq!(refreshed.email(), original.email()); + saved_quota[0].expected_revision = refreshed.revision(); + assert_eq!( + repository + .get_quotas(std::slice::from_ref(&id)) + .await + .unwrap(), + saved_quota + ); let observed_at = SystemTime::now(); repository .compare_and_swap_quota(QuotaObservation { diff --git a/backend/crates/providers/xai/src/admin.rs b/backend/crates/providers/xai/src/admin.rs index d09c8176b..005b15979 100644 --- a/backend/crates/providers/xai/src/admin.rs +++ b/backend/crates/providers/xai/src/admin.rs @@ -1192,20 +1192,9 @@ fn project_quota( refresh_token_expires_at: Option>, rolling_usage: Option, ) -> ProviderQuota { - let Some(snapshot) = snapshot else { - return ProviderQuota { - upstream: None, - credits: None, - plan_type: None, - observed_at: None, - refresh_token_expires_at, - windows: Vec::new(), - limit_reached: quota_exhausted, - provider_data: None, - }; - }; - let billing = snapshot.billing(); - let window = if billing.has_authoritative_quota() { + let billing = snapshot.as_ref().map(GrokQuotaSnapshot::billing); + let window = if let Some(billing) = billing.filter(|billing| billing.has_authoritative_quota()) + { let period_kind = billing.period_kind(); let window_seconds = quota_window_seconds(billing.period_start(), billing.period_end()); let mut data = Map::new(); @@ -1250,6 +1239,7 @@ fn project_quota( provider_data: Some(ProviderDocument::new(OpaqueProviderData::new(data))), } } else { + // 本地滚动统计独立于上游额度观察，导入后尚无快照时也保留请求图与日用量 ProviderQuotaWindow { upstream: None, key: "free-rolling-24h".to_owned(), @@ -1270,8 +1260,10 @@ fn project_quota( ProviderQuota { upstream: None, credits: None, - plan_type: billing.plan_type().map(str::to_owned), - observed_at: Some(snapshot.observed_at()), + plan_type: billing + .and_then(|billing| billing.plan_type()) + .map(str::to_owned), + observed_at: snapshot.as_ref().map(GrokQuotaSnapshot::observed_at), refresh_token_expires_at, windows: vec![window], limit_reached: quota_exhausted, diff --git a/backend/crates/providers/xai/src/credential/refresh.rs b/backend/crates/providers/xai/src/credential/refresh.rs index fde6a45ee..a580b01ff 100644 --- a/backend/crates/providers/xai/src/credential/refresh.rs +++ b/backend/crates/providers/xai/src/credential/refresh.rs @@ -489,7 +489,7 @@ impl GrokCredentialRefreshService { .upstream_user_id() .ok_or(GrokCredentialRepositoryError::InvalidCredentialData(None))? .to_owned(); - let prepared = GrokCredentialAdmin + let mut prepared = GrokCredentialAdmin .prepare_rotation(&RotateManagedGrokCredential { current, secret: GrokOAuthSecret { @@ -512,6 +512,7 @@ impl GrokCredentialRefreshService { }, }) .map_err(GrokCredentialRefreshError::from)?; + prepared.credential = prepared.credential.preserving_profile(); Ok(prepared.with_refresh_guards(capacity_guard, account_guard)) } diff --git a/backend/crates/providers/xai/src/credential/repository.rs b/backend/crates/providers/xai/src/credential/repository.rs index e16915905..c8bcd2f75 100644 --- a/backend/crates/providers/xai/src/credential/repository.rs +++ b/backend/crates/providers/xai/src/credential/repository.rs @@ -394,7 +394,7 @@ impl GrokCredentialRepository { )?; let outcome = self .store - .compare_and_swap_credential(prepared.credential) + .compare_and_swap_credential(prepared.credential.preserving_profile()) .await .map_err(map_store_error)?; let CredentialCasOutcome::Updated(revision) = outcome else { @@ -439,7 +439,8 @@ impl GrokCredentialRepository { GrokCredentialRepositoryError::InvalidCredentialData(Some( gateway_core::error::ErrorSource::new(source), )) - })?; + })? + .preserving_profile(); match self .store .compare_and_swap_credential(update) diff --git a/backend/crates/providers/xai/tests/admin.rs b/backend/crates/providers/xai/tests/admin.rs index 47170e861..c9d7f0218 100644 --- a/backend/crates/providers/xai/tests/admin.rs +++ b/backend/crates/providers/xai/tests/admin.rs @@ -8,12 +8,12 @@ use std::time::{Duration, SystemTime}; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use chrono::{DateTime, Utc}; use futures::future::BoxFuture; -use gateway_admin::model::accounts::AccountRecord; +use gateway_admin::model::accounts::{AccountRecord, AccountUsage}; use gateway_admin::model::observability::{ CurrencyCost, DesktopReleaseStatus, ProviderBillingInput, }; use gateway_admin::model::provider_credentials::{ - CompleteAuthorization, PrepareCredentialImport, PrepareCredentialRefresh, + AccountUsagePeriod, CompleteAuthorization, PrepareCredentialImport, PrepareCredentialRefresh, PrepareCredentialRotation, ProviderDocument, ProviderExportCredentialInput, ProviderQuotaRequest, }; @@ -529,15 +529,43 @@ async fn xai_admin_provider_projects_cached_quota_models_and_canonical_export() Some(true) ); + let rolling_usage = AccountUsage { + account_id: account.id().to_string(), + request_count: 5, + success_count: 5, + input_tokens: Some(100), + output_tokens: Some(20), + cached_input_tokens: Some(80), + cache_write_input_tokens: None, + reasoning_output_tokens: None, + image_input_tokens: None, + image_output_tokens: None, + image_request_count: 0, + image_request_failed_count: 0, + total_tokens: Some(120), + cost_coverage: Default::default(), + costs: Vec::new(), + last_used_at: Some(Utc::now()), + request_buckets: Vec::new(), + models: Vec::new(), + }; let quota = admin .quota(ProviderQuotaRequest { account_id: account.id().clone(), refresh: false, - rolling_usage: None, + rolling_usage: Some(rolling_usage.clone()), }) .await .expect("cached quota"); - assert!(quota.windows.is_empty()); + assert_eq!(quota.windows.len(), 1); + assert!(quota.observed_at.is_none()); + assert!(quota.plan_type.is_none()); + assert_eq!(quota.windows[0].window_seconds, Some(86_400)); + assert_eq!(quota.windows[0].local_usage.as_ref(), Some(&rolling_usage)); + assert_eq!( + quota.usage_window().map(|(_, period)| period), + Some(AccountUsagePeriod::Daily) + ); store .compare_and_swap_quota(gateway_core::account::QuotaObservation { expected_connection: None, diff --git a/backend/crates/providers/xai/tests/credential/refresh.rs b/backend/crates/providers/xai/tests/credential/refresh.rs index 09cb6513f..4fb835420 100644 --- a/backend/crates/providers/xai/tests/credential/refresh.rs +++ b/backend/crates/providers/xai/tests/credential/refresh.rs @@ -9,8 +9,9 @@ use std::time::{Duration, SystemTime}; use async_trait::async_trait; use chrono::Utc; use gateway_core::account::{ - CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState, OutboundProxy, - ProviderAccountId, ProviderAccountStore, ProviderAccountUpdate, + CredentialCasOutcome, CredentialCasUpdate, CredentialRevision, CredentialState, + OpaqueProviderData, OutboundProxy, ProviderAccountId, ProviderAccountStore, + ProviderAccountUpdate, QuotaObservation, }; use gateway_core::provider_ports::{ ProviderCredentialState, ProviderCredentialStatePort, ProviderLeaseAcquisition, @@ -316,6 +317,29 @@ fn due_input(suffix: &str) -> provider_xai::CreateGrokCredential { input } +async fn seed_quota_snapshot( + store: &MemoryProviderAccountStore, + id: &ProviderAccountId, +) -> QuotaObservation { + let account = store.account(id).unwrap(); + let quota = QuotaObservation { + account_id: id.clone(), + expected_revision: account.revision(), + expected_connection: None, + plan_type: None, + quota: OpaqueProviderData::new( + serde_json::json!({"subscriptionTier": "free", "config": null}) + .as_object() + .unwrap() + .clone(), + ), + observed_at: SystemTime::now(), + state: account.quota(), + }; + store.compare_and_swap_quota(quota.clone()).await.unwrap(); + quota +} + async fn fixture( input: provider_xai::CreateGrokCredential, responses: impl IntoIterator>, @@ -620,6 +644,7 @@ async fn successful_refresh_rotates_plaintext_tokens_once() { let id = input.account_id.clone(); let (store, _, refresher, service) = fixture(input, [Ok(success_tokens(Some("new-refresh")))], true).await; + let mut quota = seed_quota_snapshot(&store, &id).await; let outcomes = service.refresh_due().await.expect("refresh cycle"); assert!(matches!( @@ -630,6 +655,11 @@ async fn successful_refresh_rotates_plaintext_tokens_once() { }] if account_id == &id && credential_revision.get() == 2 )); assert_eq!(refresher.prepare_calls.load(Ordering::SeqCst), 1); + quota.expected_revision = CredentialRevision::new(2).unwrap(); + assert_eq!( + store.get_quotas(std::slice::from_ref(&id)).await.unwrap(), + vec![quota] + ); let credential = store.credential(&id).expect("credential"); assert_eq!( credential_object(&credential) @@ -747,6 +777,7 @@ async fn manual_refresh_returns_prepared_rotation_without_writing_store() { let id = input.account_id.clone(); let (store, _, refresher, service) = fixture(input, [Ok(success_tokens(Some("manual-refresh")))], true).await; + let mut quota = seed_quota_snapshot(&store, &id).await; let current = store .load_current_credential(&id) .await @@ -767,6 +798,8 @@ async fn manual_refresh_returns_prepared_rotation_without_writing_store() { .expect("persist prepared rotation"), CredentialCasOutcome::Updated(revision) if revision.get() == 2 )); + quota.expected_revision = CredentialRevision::new(2).unwrap(); + assert_eq!(store.get_quotas(&[id]).await.unwrap(), vec![quota]); drop(guard); } @@ -952,6 +985,7 @@ async fn pre_send_transient_failure_applies_bounded_oauth_backoff() { let input = due_input("transient"); let id = input.account_id.clone(); let (store, _, _, service) = fixture(input, [Err(GrokRefreshFailure::Transient)], true).await; + let mut quota = seed_quota_snapshot(&store, &id).await; service.refresh_due().await.expect("refresh"); let account = store.account(&id).expect("account"); assert_eq!(account.credential_state(), CredentialState::Unknown); @@ -961,6 +995,8 @@ async fn pre_send_transient_failure_applies_bounded_oauth_backoff() { .next_refresh_at() .is_some_and(|retry| retry > std::time::SystemTime::now()) ); + quota.expected_revision = CredentialRevision::new(2).unwrap(); + assert_eq!(store.get_quotas(&[id]).await.unwrap(), vec![quota]); } #[tokio::test] diff --git a/backend/crates/providers/xai/tests/support.rs b/backend/crates/providers/xai/tests/support.rs index eebe98c97..ad10bd052 100644 --- a/backend/crates/providers/xai/tests/support.rs +++ b/backend/crates/providers/xai/tests/support.rs @@ -463,7 +463,13 @@ impl ProviderAccountStore for MemoryProviderAccountStore { }, ); stored.credential = credential; - stored.quota = None; + if preserve_profile { + if let Some(quota) = &mut stored.quota { + quota.expected_revision = next; + } + } else { + stored.quota = None; + } Ok(CredentialCasOutcome::Updated(next)) } diff --git a/docs/api.md b/docs/api.md index cc967a4cd..b31e29882 100644 --- a/docs/api.md +++ b/docs/api.md @@ -1440,7 +1440,8 @@ API Key 账号与 OAuth 共用现有 Responses、Images 和 standalone Search 互联账号按上游 Key 已提供的最长消费窗口统计本机 Token、费用和模型用量。 日、周周期的 `usage.windowKind` 分别为 `daily`、`weekly`，其他周期为 `quota`。查询边界同账号额度窗口，为 `[resetAt - windowSeconds, resetAt)`； 未开启或缺少重置边界时显示无数据，不回退到历史累计，也不采用池内其他账号的周期。 -OAuth 账号仍按实际周/月额度窗口统计 +OAuth 账号优先按实际周/月额度窗口统计。xAI 缺少可展示的上游额度时，使用过去 24 小时的本地请求统计， +`usage.windowKind` 为 `daily`，请求图和用量统计共用该滚动范围。该统计可在上游额度快照缺失时展示 批量导入 AT / RT 使用 `accounts` JSON 数组，最多 200 项，不接受纯文本 token 列表。例如： diff --git a/docs/architecture.md b/docs/architecture.md index 516c45154..740e94c8e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -991,6 +991,8 @@ credential 与 quota 是两组独立事实：credential refresh 不等于 quota - xAI 使用 OAuth session；API Key 不是受支持的账号 credential。刷新额度时同步查询官方实时订阅， 只把套餐事实写入现有 quota JSON。明确无付费订阅的个人账号显示 Free；查询失败、缺失字段或 团队身份不推断为 Free，订阅查询失败不影响额度观测 + 普通 Token 续期与刷新退避保留提交时的账号资料及额度快照，显式凭据轮换仍按原子提交合同失效旧观测。 + xAI Provider 在缺少可展示上游额度时投影本地 24 小时滚动统计，Admin 按该窗口展示日用量，额度预测使用完整周/月窗口 - 账号导入和 OAuth complete（包括重新授权）在 credential 提交、Provider 事实失效及快照发布后， 由 Admin 共用流程后台读取一次额度；不等待观测完成才返回管理请求，失败记录告警但不回滚账号事务。 手工和后台 credential refresh 仍不隐式刷新 quota