mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-07 00:05:49 +00:00
fix(openai): harden exhausted quota recovery
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
{
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user