starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 16:06:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(openai): preserve exhausted quotas until reset

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-31 02:20:37 +08:00
co-authored by Codex
parent ab4bf75d0e
commit 756c9c7158
2 changed files with 80 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/providers/openai/src/credential/quota.rs
+11 -83
View File
@@ -197,14 +197,7 @@ pub struct CodexQuotaSyncSummary {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum QuotaRecoveryEvidence {
SuccessfulResponse,
RefreshedSnapshot { confirms_recovery: bool },
PostFailureSnapshot,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum QuotaRefreshMode {
Reconcile,
AfterConfirmedFailure,
RefreshedSnapshot,
}
impl CodexQuotaSyncSummary {
@@ -649,20 +642,9 @@ 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,
previous_snapshot.as_ref(),
&mut summary,
)
.apply_fetched_quota(&account, &value, observed_at, &mut summary)
.await
{
summary.transient += 1;
@@ -777,7 +759,6 @@ impl CodexCredentialQuotaService {
account: &ProviderAccount,
value: &Value,
observed_at: SystemTime,
previous_snapshot: Option<&CodexAccountQuotaSnapshot>,
summary: &mut CodexQuotaSyncSummary,
) -> Result<(), CodexCredentialQuotaError> {
let snapshot = parse_account_quota_snapshot(
@@ -824,9 +805,7 @@ impl CodexCredentialQuotaService {
current.cooldown_until(),
fact,
SystemTime::now(),
QuotaRecoveryEvidence::RefreshedSnapshot {
confirms_recovery: quota_snapshot_confirms_recovery(previous_snapshot, &snapshot),
},
QuotaRecoveryEvidence::RefreshedSnapshot,
) {
let _ = self
.repository
@@ -970,23 +949,20 @@ impl CodexCredentialQuotaService {
&self,
account_id: &ProviderAccountId,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
self.refresh_account_with_mode(account_id, QuotaRefreshMode::Reconcile)
.await
self.refresh_account_snapshot(account_id).await
}
/// 真实限额失败后的异步刷新只补齐展示快照，不允许同一轮陈旧 usage 撤销失败状态。
/// 真实限额失败后的异步刷新只补齐展示快照，不允许 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
self.refresh_account_snapshot(account_id).await
}
async fn refresh_account_with_mode(
async fn refresh_account_snapshot(
&self,
account_id: &ProviderAccountId,
mode: QuotaRefreshMode,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
let account = self
.store
@@ -994,7 +970,6 @@ 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);
@@ -1071,17 +1046,7 @@ 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,
),
},
},
QuotaRecoveryEvidence::RefreshedSnapshot,
) {
self.repository
.apply_state(
@@ -1348,49 +1313,12 @@ fn quota_recovery_confirmed(
}
match evidence {
QuotaRecoveryEvidence::SuccessfulResponse => true,
QuotaRecoveryEvidence::RefreshedSnapshot { confirms_recovery } => confirms_recovery,
QuotaRecoveryEvidence::PostFailureSnapshot => false,
// used_percent 是可滞后、会取整的观测值；确认 429 后的 100% 也是
// 管理端投影。它们的回落不能覆盖仍未到期的上游 reset 时间。
QuotaRecoveryEvidence::RefreshedSnapshot => 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,
backend/crates/providers/openai/tests/credential/quota.rs
+69 -16
View File
@@ -443,7 +443,7 @@ async fn passive_full_percent_with_explicit_allowance_keeps_account_ready() {
}
#[tokio::test]
async fn manual_quota_refresh_updates_account_state_from_provider_response() {
async fn manual_quota_refresh_keeps_confirmed_exhaustion_until_the_recorded_reset() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_manual_quota_state";
create_account(&store, account_id).await;
@@ -541,24 +541,25 @@ async fn manual_quota_refresh_updates_account_state_from_provider_response() {
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 20, "reset_at": 1_900_000_000}
"primary_window": {"used_percent": 99, "reset_at": 1_900_000_000}
}
})))
.mount(&server)
.await;
let recovered = service
let unchanged_reset = service
.refresh_account(account.id())
.await
.expect("refresh recovered quota");
.expect("refresh lower usage quota");
assert!(!recovered.fact().exhausted());
assert!(!unchanged_reset.fact().exhausted());
assert_eq!(
store
.account(account_id)
.expect("recovered account")
.expect("account after lower usage")
.availability(),
AccountAvailability::Ready
AccountAvailability::QuotaExhausted,
"a lower used_percent before the recorded reset is not quota recovery"
);
}
@@ -1051,7 +1052,7 @@ async fn periodic_quota_synchronization_rechecks_exhausted_accounts_without_wait
}
#[tokio::test]
async fn periodic_quota_synchronization_requires_a_real_snapshot_improvement_to_recover() {
async fn periodic_quota_synchronization_does_not_recover_from_percent_drop_before_reset() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_periodic_recovery_evidence";
create_account(&store, account_id).await;
@@ -1066,7 +1067,7 @@ async fn periodic_quota_synchronization_requires_a_real_snapshot_improvement_to_
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 86, "reset_at": reset_at}
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
@@ -1095,7 +1096,7 @@ async fn periodic_quota_synchronization_requires_a_real_snapshot_improvement_to_
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 86, "reset_at": reset_at}
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})))
.expect(1)
@@ -1131,7 +1132,7 @@ async fn periodic_quota_synchronization_requires_a_real_snapshot_improvement_to_
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 20, "reset_at": reset_at}
"primary_window": {"used_percent": 99, "reset_at": reset_at}
}
})))
.expect(1)
@@ -1147,21 +1148,73 @@ async fn periodic_quota_synchronization_requires_a_real_snapshot_improvement_to_
server.uri(),
);
let improved = next_cycle
let lower_percent = next_cycle
.synchronize()
.await
.expect("improved quota recheck");
.expect("lower percent quota recheck");
assert_eq!(improved.updated, 1);
assert_eq!(lower_percent.updated, 1);
assert_eq!(
store
.account(account_id)
.expect("account after improved snapshot")
.expect("account after lower percent snapshot")
.availability(),
AccountAvailability::Ready
AccountAvailability::QuotaExhausted,
"a provider percentage change is not sufficient to release a confirmed exhaustion"
);
}
#[tokio::test]
async fn periodic_quota_synchronization_recovers_after_recorded_reset() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_periodic_reset_elapsed";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
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_sub(Duration::from_secs(60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted with an elapsed reset");
let server = MockServer::start().await;
let next_reset = Utc::now().timestamp() + 24 * 60 * 60;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 1, "reset_at": next_reset}
}
})))
.expect(1)
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
// 定时复核可提前执行，但只有已记录的 reset 到期后才释放额度耗尽状态。
let summary = service.synchronize().await.expect("quota synchronization");
assert_eq!(summary.updated, 1);
let current = store
.account(account_id)
.expect("account after elapsed reset");
assert_eq!(current.availability(), AccountAvailability::Ready);
assert_eq!(current.cooldown_until(), None);
}
#[tokio::test]
async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
let store = Arc::new(MemoryAccountStore::default());