mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 16:06:08 +00:00
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 <noreply@openai.com>
This commit is contained in:
@@ -806,11 +806,12 @@ pub struct ProviderQuotaWindow {
|
||||
pub provider_data: Option<ProviderDocument>,
|
||||
}
|
||||
|
||||
/// 账号用量统计周期,按周优先、月次之选择;短期限流窗口不参与统计面板
|
||||
/// 账号用量统计周期,完整周/月额度优先,随后使用 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<Item = (&ProviderQuotaWindow, AccountUsagePeriod)> {
|
||||
|
||||
@@ -62,14 +62,15 @@ pub fn account_quota_forecasts(
|
||||
now: DateTime<Utc>,
|
||||
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("缺少周/月额度窗口,请刷新账号额度"),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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![
|
||||
|
||||
@@ -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());
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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::<AdminTestState>()
|
||||
.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!(
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -1192,20 +1192,9 @@ fn project_quota(
|
||||
refresh_token_expires_at: Option<DateTime<Utc>>,
|
||||
rolling_usage: Option<gateway_admin::model::accounts::AccountUsage>,
|
||||
) -> 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,
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<Item = Result<GrokRefreshTokens, GrokRefreshFailure>>,
|
||||
@@ -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]
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
|
||||
+2
-1
@@ -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 列表。例如:
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user