fix(openai): harden exhausted quota recovery

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-30 21:07:53 +08:00
co-authored by Codex
parent 36321a9af4
commit f5dca4848a
9 changed files with 591 additions and 22 deletions
@@ -630,7 +630,8 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
else availability_observed_at
end,
updated_at = case when enabled then greatest(now(), $6) else updated_at end
where id = $1 and credential_revision = $2",
where id = $1 and credential_revision = $2
and (availability_observed_at is null or availability_observed_at <= $6)",
)
.bind(update.account_id)
.bind(to_i64(update.expected_revision.get())?)
@@ -675,7 +676,8 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
"update provider_accounts
set provider_quota_json = $3, quota_observed_at = $4,
updated_at = greatest(now(), coalesce($4, now()))
where id = $1 and credential_revision = $2",
where id = $1 and credential_revision = $2
and ($4 is null or quota_observed_at is null or quota_observed_at <= $4)",
)
.bind(account_id)
.bind(to_i64(expected_revision.get())?)
@@ -111,6 +111,97 @@ async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() {
database.close().await;
}
#[tokio::test]
async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
let Some(database) = TestDatabase::create("provider_account_observation_fence").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let account_id = ProviderAccountId::new("acct_observation_fence").expect("account id");
repository
.insert_provider_account(account(account_id.as_str(), "user-observation-fence"))
.await
.expect("insert observation fence fixture");
let revision = CredentialRevision::new(1).expect("revision");
let newer = SystemTime::now()
.checked_add(Duration::from_secs(60))
.expect("newer observation time");
let older = newer
.checked_sub(Duration::from_secs(30))
.expect("older observation time");
let quota = |marker: &str, observed_at| QuotaObservation {
account_id: account_id.clone(),
expected_revision: revision,
quota: Some(OpaqueProviderData::new(
json!({"marker": marker})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(observed_at),
};
assert_eq!(
repository
.compare_and_swap_quota(quota("newer", newer))
.await
.expect("persist newer quota"),
QuotaWriteOutcome::Updated
);
assert_eq!(
repository
.compare_and_swap_quota(quota("older", older))
.await
.expect("reject older quota"),
QuotaWriteOutcome::Conflict
);
repository
.apply_state_change(AccountStateChange {
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
.expect("persist newer account state");
assert!(
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: revision,
availability: AccountAvailability::Ready,
reason: None,
cooldown_until: None,
observed_at: older,
})
.await
.is_err(),
"a delayed state observation must conflict"
);
let current = repository
.get_account(&account_id)
.await
.expect("load fenced account")
.expect("fenced account");
assert_eq!(current.availability(), AccountAvailability::QuotaExhausted);
let observed = repository
.get_quotas(std::slice::from_ref(&account_id))
.await
.expect("load fenced quota")
.pop()
.expect("fenced quota");
assert_eq!(
observed.quota.expect("quota document").expose_to_provider()["marker"],
"newer"
);
database.close().await;
}
#[tokio::test]
async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_usage() {
let Some(database) = TestDatabase::create("provider_account_terminal_list").await else {
+83 -7
View File
@@ -25,8 +25,9 @@ use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
use gateway_core::accounting::Money;
use gateway_core::engine::credential::{
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
AccountAvailability, CredentialRevision, LoadedCredential, NewProviderAccount,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountStore,
};
use gateway_core::error::StoreErrorKind;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
@@ -458,21 +459,23 @@ impl ProviderAdmin for OpenAiAdminProvider {
refresh,
rolling_usage: _,
} = request;
self.account(&account_id).await?;
let mut account = self.account(&account_id).await?;
let snapshot = if refresh {
Some(
let snapshot = Some(
self.quota
.refresh_account(&account_id)
.await
.map_err(map_quota_error)?,
)
);
account = self.account(&account_id).await?;
snapshot
} else {
self.quota
.read_account(&account_id)
.await
.map_err(map_quota_error)?
};
Ok(snapshot.map_or_else(empty_quota, project_quota))
Ok(project_quota(snapshot, &account))
}
async fn models(
@@ -746,7 +749,18 @@ fn empty_quota() -> ProviderQuota {
}
}
fn project_quota(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota {
fn project_quota(
snapshot: Option<CodexAccountQuotaSnapshot>,
account: &ProviderAccount,
) -> ProviderQuota {
let mut quota = snapshot.map_or_else(empty_quota, project_quota_snapshot);
if account.availability() == AccountAvailability::QuotaExhausted {
force_confirmed_exhaustion_projection(&mut quota, account.cooldown_until());
}
quota
}
fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota {
let mut provider_data = Map::new();
provider_data.insert(
"remaining_percent".to_owned(),
@@ -799,6 +813,68 @@ fn project_quota(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota {
}
}
fn force_confirmed_exhaustion_projection(
quota: &mut ProviderQuota,
cooldown_until: Option<SystemTime>,
) {
// 真实 429 对“当前已耗尽”的证明强于可能延迟结算的 usage 文档。这里只覆盖
// Admin 展示投影,原始 Provider JSON 仍作为后续判断窗口是否真正恢复的基线。
quota.provider_data = Some(ProviderDocument::new(OpaqueProviderData::new(
Map::from_iter([
(
"remaining_percent".to_owned(),
Value::Number(Number::from(0)),
),
("exhausted".to_owned(), Value::Bool(true)),
]),
)));
let reset_at = cooldown_until.map(DateTime::<Utc>::from);
if quota.windows.is_empty() {
quota.windows.push(ProviderQuotaWindow {
key: "confirmed-quota-exhaustion".to_owned(),
group: "other".to_owned(),
label: "额度".to_owned(),
source: None,
window_seconds: None,
used_percent: Some(100.0),
reset_at,
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)
})
.unwrap_or(0);
quota.windows[index].used_percent = Some(100.0);
}
const fn quota_group(kind: CodexQuotaWindowKind) -> &'static str {
match kind {
CodexQuotaWindowKind::Monthly => "monthly",
@@ -194,6 +194,19 @@ pub struct CodexQuotaSyncSummary {
pub stale: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum QuotaRecoveryEvidence {
SuccessfulResponse,
RefreshedSnapshot { confirms_recovery: bool },
PostFailureSnapshot,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum QuotaRefreshMode {
Reconcile,
AfterConfirmedFailure,
}
impl CodexQuotaSyncSummary {
#[must_use]
pub const fn has_operational_failures(self) -> bool {
@@ -636,9 +649,20 @@ impl CodexCredentialQuotaService {
let observed_at = SystemTime::now();
match self.fetch_usage_with_recovery(&client, &account).await {
Ok(FetchedCodexQuota { account, value }) => {
let previous_snapshot = observations
.as_ref()
.and_then(|observations| observations.get(account.id()))
.filter(|observation| observation.expected_revision == account.revision())
.and_then(quota_snapshot_from_observation);
// 单账号解析或落库失败只影响该账号;其余账号继续同步。
if let Err(error) = self
.apply_fetched_quota(&account, &value, observed_at, &mut summary)
.apply_fetched_quota(
&account,
&value,
observed_at,
previous_snapshot.as_ref(),
&mut summary,
)
.await
{
summary.transient += 1;
@@ -753,9 +777,16 @@ impl CodexCredentialQuotaService {
account: &ProviderAccount,
value: &Value,
observed_at: SystemTime,
previous_snapshot: Option<&CodexAccountQuotaSnapshot>,
summary: &mut CodexQuotaSyncSummary,
) -> Result<(), CodexCredentialQuotaError> {
let fact = parse_codex_quota_usage(value)?;
let snapshot = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
value,
)?;
let fact = snapshot.fact();
let object = value
.as_object()
.cloned()
@@ -793,6 +824,9 @@ impl CodexCredentialQuotaService {
current.cooldown_until(),
fact,
SystemTime::now(),
QuotaRecoveryEvidence::RefreshedSnapshot {
confirms_recovery: quota_snapshot_confirms_recovery(previous_snapshot, &snapshot),
},
) {
let _ = self
.repository
@@ -859,6 +893,7 @@ impl CodexCredentialQuotaService {
current.cooldown_until(),
fact,
observed_at,
QuotaRecoveryEvidence::SuccessfulResponse,
) {
self.repository
.apply_state(
@@ -934,6 +969,24 @@ impl CodexCredentialQuotaService {
pub async fn refresh_account(
&self,
account_id: &ProviderAccountId,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
self.refresh_account_with_mode(account_id, QuotaRefreshMode::Reconcile)
.await
}
/// 真实限额失败后的异步刷新只补齐展示快照,不允许同一轮陈旧 usage 撤销失败状态。
pub(crate) async fn refresh_account_after_failure(
&self,
account_id: &ProviderAccountId,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
self.refresh_account_with_mode(account_id, QuotaRefreshMode::AfterConfirmedFailure)
.await
}
async fn refresh_account_with_mode(
&self,
account_id: &ProviderAccountId,
mode: QuotaRefreshMode,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
let account = self
.store
@@ -941,6 +994,7 @@ impl CodexCredentialQuotaService {
.await?
.filter(|account| account.provider().as_str() == "openai")
.ok_or(CodexCredentialQuotaError::NotFound)?;
let previous_snapshot = self.read_snapshot_for(&account).await?;
let observed_at = SystemTime::now();
if !access_token_is_current(&account, observed_at) {
return Err(CodexCredentialQuotaError::CredentialRefreshRequired);
@@ -1017,6 +1071,17 @@ impl CodexCredentialQuotaService {
current.cooldown_until(),
snapshot.fact(),
SystemTime::now(),
match mode {
QuotaRefreshMode::AfterConfirmedFailure => {
QuotaRecoveryEvidence::PostFailureSnapshot
}
QuotaRefreshMode::Reconcile => QuotaRecoveryEvidence::RefreshedSnapshot {
confirms_recovery: quota_snapshot_confirms_recovery(
previous_snapshot.as_ref(),
&snapshot,
),
},
},
) {
self.repository
.apply_state(
@@ -1239,9 +1304,13 @@ fn quota_success_state(
cooldown_until: Option<SystemTime>,
fact: CodexQuotaFact,
now: SystemTime,
recovery_evidence: QuotaRecoveryEvidence,
) -> Option<(AccountAvailability, Option<SystemTime>)> {
if fact.exhausted() {
let observed_cooldown = fact.resets_at().map(SystemTime::from).or(cooldown_until);
let observed_cooldown = match (fact.resets_at().map(SystemTime::from), cooldown_until) {
(Some(observed), Some(confirmed)) => Some(observed.max(confirmed)),
(observed, confirmed) => observed.or(confirmed),
};
return (current != AccountAvailability::QuotaExhausted
|| cooldown_until != observed_cooldown)
.then_some((AccountAvailability::QuotaExhausted, observed_cooldown));
@@ -1251,8 +1320,13 @@ fn quota_success_state(
AccountAvailability::Invalid
| AccountAvailability::Expired
| AccountAvailability::Banned
| AccountAvailability::Unknown
| AccountAvailability::QuotaExhausted => Some(AccountAvailability::Ready),
| AccountAvailability::Unknown => Some(AccountAvailability::Ready),
AccountAvailability::QuotaExhausted
if !quota_recovery_confirmed(recovery_evidence, cooldown_until, now) =>
{
None
}
AccountAvailability::QuotaExhausted => Some(AccountAvailability::Ready),
AccountAvailability::Ready => None,
AccountAvailability::Cooldown
if cooldown_until.is_some_and(|cooldown_until| cooldown_until <= now) =>
@@ -1264,6 +1338,59 @@ fn quota_success_state(
availability.map(|availability| (availability, None))
}
fn quota_recovery_confirmed(
evidence: QuotaRecoveryEvidence,
cooldown_until: Option<SystemTime>,
now: SystemTime,
) -> bool {
if cooldown_until.is_some_and(|cooldown_until| cooldown_until <= now) {
return true;
}
match evidence {
QuotaRecoveryEvidence::SuccessfulResponse => true,
QuotaRecoveryEvidence::RefreshedSnapshot { confirms_recovery } => confirms_recovery,
QuotaRecoveryEvidence::PostFailureSnapshot => false,
}
}
fn quota_snapshot_confirms_recovery(
previous: Option<&CodexAccountQuotaSnapshot>,
current: &CodexAccountQuotaSnapshot,
) -> bool {
if current.fact().exhausted() {
return false;
}
let Some(previous) = previous else {
return false;
};
if previous
.fact()
.remaining_percent()
.zip(current.fact().remaining_percent())
.is_some_and(|(previous, current)| current > previous)
{
return true;
}
// 单窗口账号可以用同一窗口的用量回落或 reset 前移证明上游提前翻新;多窗口
// 无法从账号级 429 确认具体受限窗口,保守等待整体剩余额度改善或真实请求成功。
let ([previous], [current]) = (previous.windows(), current.windows()) else {
return false;
};
if previous.key() != current.key() {
return false;
}
let usage_recovered = previous
.used_percent()
.zip(current.used_percent())
.is_some_and(|(previous, current)| current < previous);
let window_advanced = previous
.reset_at()
.zip(current.reset_at())
.is_some_and(|(previous, current)| current > previous);
usage_recovered || window_advanced
}
fn quota_state_transition(
error: &CodexClientError,
observed_at: SystemTime,
@@ -2009,15 +2009,15 @@ fn schedule_authoritative_quota_refresh_after_failure(
account: &ProviderAccount,
) {
// 限额错误可以确认账号状态,但 Responses 事件中的 used_percent 可能仍停在上一结算点。
// usage 快照在后台补齐,不能阻塞原始失败响应。先等待 2 秒让上游结算,
// 再查询;5 秒仅限制实际查询自身。
// usage 快照在后台补齐展示基线,不得撤销同一轮真实失败,也不能阻塞原始响应。
// 先等待 2 秒让上游结算,再查询;5 秒仅限制实际查询自身。
let quota = Arc::clone(quota);
let account_id = account.id().clone();
drop(tokio::spawn(async move {
tokio::time::sleep(QUOTA_FAILURE_REFRESH_DELAY).await;
match tokio::time::timeout(
QUOTA_FAILURE_REFRESH_TIMEOUT,
quota.refresh_account(&account_id),
quota.refresh_account_after_failure(&account_id),
)
.await
{
+104 -2
View File
@@ -20,8 +20,8 @@ use gateway_admin::model::provider_credentials::{
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::ProviderAdminErrorKind;
use gateway_core::engine::credential::{
CredentialRevision, OpaqueProviderData, ProviderAccount, ProviderAccountId,
ProviderAccountStore, QuotaObservation,
AccountAvailability as CoreAccountAvailability, AccountStateChange, CredentialRevision,
OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore, QuotaObservation,
};
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::provider_ports::{
@@ -503,6 +503,108 @@ async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_qu
assert_eq!(monthly[0].used_percent, Some(2.0));
}
#[tokio::test]
async fn openai_admin_projects_confirmed_quota_exhaustion_as_full_without_mutating_raw_usage() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_admin_confirmed_exhaustion";
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
name: "admin confirmed exhaustion".to_owned(),
secret: secret("admin-confirmed-exhaustion-access"),
verified_account: profile("chatgpt-admin-confirmed-exhaustion"),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
let account = store.account(account_id).expect("stored account");
let reset_at = 1_900_000_000_u64;
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 86, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist raw quota");
store
.apply_state_change(AccountStateChange {
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
.expect("mark account exhausted");
let config = valid_config();
let bundle = provider_openai::initialize(
config.config.clone(),
provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())),
)
.await
.expect("OpenAI bundle");
let exhausted = bundle
.admin_provider()
.quota(ProviderQuotaRequest {
account_id: account.id().clone(),
refresh: false,
rolling_usage: None,
})
.await
.expect("project exhausted quota");
assert_eq!(exhausted.windows.len(), 1);
assert_eq!(exhausted.windows[0].used_percent, Some(100.0));
let raw = store
.get_quotas(std::slice::from_ref(account.id()))
.await
.expect("read raw quota")
.pop()
.expect("raw quota");
assert_eq!(
raw.quota.expect("raw document").expose_to_provider()["rate_limit"]["primary_window"]["used_percent"],
86
);
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: CoreAccountAvailability::Ready,
reason: None,
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("recover account");
let recovered = bundle
.admin_provider()
.quota(ProviderQuotaRequest {
account_id: account.id().clone(),
refresh: false,
rolling_usage: None,
})
.await
.expect("project recovered quota");
assert_eq!(recovered.windows[0].used_percent, Some(86.0));
}
#[tokio::test]
async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit() {
let store = Arc::new(MemoryAccountStore::default());
@@ -498,6 +498,38 @@ async fn manual_quota_refresh_updates_account_state_from_provider_response() {
AccountAvailability::QuotaExhausted
);
server.reset().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header(
"authorization",
format!("Bearer token-{account_id}"),
))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
})))
.mount(&server)
.await;
let stale = service
.refresh_account(account.id())
.await
.expect("refresh contradictory full quota");
assert!(!stale.fact().exhausted());
assert_eq!(
store
.account(account_id)
.expect("account after contradictory quota")
.availability(),
AccountAvailability::QuotaExhausted,
"an unchanged usage window cannot override a confirmed quota failure"
);
server.reset().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
@@ -1018,6 +1050,118 @@ async fn periodic_quota_synchronization_rechecks_exhausted_accounts_without_wait
assert_eq!(repeated, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn periodic_quota_synchronization_requires_a_real_snapshot_improvement_to_recover() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_periodic_recovery_evidence";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
let reset_at = 1_900_000_000;
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 86, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("seed post-failure quota snapshot");
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: SystemTime::now().checked_add(Duration::from_secs(24 * 60 * 60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted");
let server = MockServer::start().await;
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": 86, "reset_at": reset_at}
}
})))
.expect(1)
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let unchanged = service
.synchronize()
.await
.expect("unchanged quota recheck");
assert_eq!(unchanged.updated, 1);
assert_eq!(
store
.account(account_id)
.expect("account after unchanged snapshot")
.availability(),
AccountAvailability::QuotaExhausted
);
server.reset().await;
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": 20, "reset_at": reset_at}
}
})))
.expect(1)
.mount(&server)
.await;
// 新 worker 实例模拟下一次 30 分钟复核,避免测试等待真实节流周期。
let next_cycle = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let improved = next_cycle
.synchronize()
.await
.expect("improved quota recheck");
assert_eq!(improved.updated, 1);
assert_eq!(
store
.account(account_id)
.expect("account after improved snapshot")
.availability(),
AccountAvailability::Ready
);
}
#[tokio::test]
async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
let store = Arc::new(MemoryAccountStore::default());
@@ -1036,15 +1036,16 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header("authorization", format!("Bearer at-{account_id}")))
// 限额错误本身只能确认状态;这里必须覆盖上一帧被动观察到的 99%。
// 生产环境会短暂返回互相矛盾的快照:数字已到 100%,但布尔位仍声称可用。
// 后台同步应覆盖展示数据,不能据此撤销真实请求已经确认的额度耗尽。
.respond_with(
ResponseTemplate::new(200)
// 验证后台 usage 同步不能把原始的额度错误响应拖到查询完成之后。
.set_delay(Duration::from_millis(750))
.set_body_json(json!({
"rate_limit": {
"allowed": false,
"limit_reached": true,
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
})),
@@ -1142,6 +1143,14 @@ async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_qu
.count(),
1
);
assert_eq!(
store
.account(account_id)
.expect("account after stale full usage refresh")
.availability(),
AccountAvailability::QuotaExhausted,
"a contradictory full usage snapshot must not unlock a confirmed exhausted account"
);
}
#[tokio::test]
@@ -36,6 +36,7 @@ struct StoredAccount {
account: ProviderAccount,
credential: gateway_core::engine::credential::PlaintextCredential,
quota: Option<QuotaObservation>,
state_observed_at: Option<SystemTime>,
}
#[derive(Default)]
@@ -129,6 +130,7 @@ impl ProviderAccountStore for MemoryAccountStore {
account: input.account,
credential: input.credential,
quota: None,
state_observed_at: None,
},
);
Ok(())
@@ -260,6 +262,15 @@ impl ProviderAccountStore for MemoryAccountStore {
if stored.account.revision() != observation.expected_revision {
return Ok(QuotaWriteOutcome::Conflict);
}
if stored
.quota
.as_ref()
.and_then(|quota| quota.observed_at)
.zip(observation.observed_at)
.is_some_and(|(current, incoming)| current > incoming)
{
return Ok(QuotaWriteOutcome::Conflict);
}
stored.quota = Some(observation);
Ok(QuotaWriteOutcome::Updated)
}
@@ -272,6 +283,12 @@ impl ProviderAccountStore for MemoryAccountStore {
if stored.account.revision() != change.expected_revision {
return Err(store_error(StoreErrorKind::Conflict));
}
if stored
.state_observed_at
.is_some_and(|current| current > change.observed_at)
{
return Err(store_error(StoreErrorKind::Conflict));
}
stored.account = rebuild_account(
&stored.account,
AccountRebuild {
@@ -285,6 +302,7 @@ impl ProviderAccountStore for MemoryAccountStore {
profile: None,
},
);
stored.state_observed_at = Some(change.observed_at);
Ok(())
}