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

fix(openai): refresh expired quota windows after reset grace (#170)

Browse Source 
Refresh non-exhausted accounts after the quota reset grace period while preserving observed usage and exhausted-window attribution. Keep periodic checks for exhausted accounts to detect early upstream resets, and cover refresh timing, throttling, and display regressions.

Co-authored-by: Zyy <zyycn@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
2han9wen71an
2026-09-19 18:13:33 +08:00
committed by GitHub
co-authored by Zyy Codex
parent 4669a29b46
commit d0180b74b3
4 changed files with 282 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/providers/openai/src/credential/quota/mod.rs
+69 -36
View File
@@ -60,8 +60,8 @@ use snapshot::{
const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
pub(crate) const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
const EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL: Duration = Duration::from_secs(30 * 60);
const EXHAUSTED_QUOTA_RESET_GRACE: Duration = Duration::from_secs(2 * 60);
const PERIODIC_QUOTA_REFRESH_RETRY_INTERVAL: Duration = Duration::from_secs(30 * 60);
const QUOTA_RESET_GRACE: Duration = Duration::from_secs(2 * 60);
/// 首次 OAuth 异步观察失败时，由既有 quota worker 兜底重试的单轮上限。
const INITIAL_QUOTA_SYNC_BATCH: usize = 100;
// 5xx 上游拒绝的短退避重试预算；指数退避 1s/2s，吞掉瞬时抖动。
@@ -394,15 +394,19 @@ impl CodexQuotaSchedulingProjection {
fn reserve_periodic_refreshes(
&self,
accounts: Vec<ProviderAccount>,
observed_snapshots: &BTreeMap<ProviderAccountId, CodexAccountQuotaSnapshot>,
now: SystemTime,
) -> Vec<ProviderAccount> {
let candidates = accounts
.into_iter()
.filter_map(|account| quota_refresh_candidate(account, now))
.filter_map(|account| {
let snapshot = observed_snapshots.get(account.id());
quota_refresh_candidate(account, snapshot, now)
})
.collect::<Vec<_>>();
let candidate_ids = candidates
.iter()
.map(|account| account.id().clone())
.map(|(account, _)| account.id().clone())
.collect::<BTreeSet<_>>();
let refreshed_at = Instant::now();
let mut state = self
@@ -413,18 +417,12 @@ impl CodexQuotaSchedulingProjection {
.last_periodic_refresh_at
.retain(|account_id, _| candidate_ids.contains(account_id));
// 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。
// 已耗尽账号每 30 分钟复核，不能等待旧 reset：官方活动可能提前重置额度。
// reset + 2 分钟可提前触发一次复核，给上游重置留出传播时间。
// 已耗尽账号首次立即复核，之后每 30 分钟复核，以发现官方提前重置。
// reset + 2 分钟额外触发一次复核，给上游重置留出传播时间。
// 正常账号仅在非零用量窗口经过宽限期后参与，未更新时复用周期节流。
let mut reserved = Vec::new();
for account in candidates {
if !periodic_quota_refresh_due(
&state,
account.id(),
account.quota().reset_at(),
now,
refreshed_at,
) {
for (account, target_reset) in candidates {
if !periodic_quota_refresh_due(&state, account.id(), target_reset, now, refreshed_at) {
continue;
}
state.last_periodic_refresh_at.insert(
@@ -440,9 +438,35 @@ impl CodexQuotaSchedulingProjection {
}
}
fn quota_refresh_candidate(account: ProviderAccount, now: SystemTime) -> Option<ProviderAccount> {
(eligible_periodic_quota_refresh(&account, now) && account.quota().is_exhausted())
.then_some(account)
fn quota_refresh_candidate(
account: ProviderAccount,
snapshot: Option<&CodexAccountQuotaSnapshot>,
now: SystemTime,
) -> Option<(ProviderAccount, Option<SystemTime>)> {
if !eligible_periodic_quota_refresh(&account, now) {
return None;
}
if account.quota().is_exhausted() {
let reset_at = account.quota().reset_at();
return Some((account, reset_at));
}
// 正常账号首次进入周期候选也要等待宽限期，不能由“无刷新历史”绕过。
let snapshot = snapshot?;
let expired_window_reset = snapshot
.windows()
.iter()
.filter(|window| {
window.used_percent().is_some_and(|used| used > 0.0) || window.limit_reached()
})
.filter_map(CodexQuotaWindow::reset_at)
.map(SystemTime::from)
.filter(|reset| {
reset
.checked_add(QUOTA_RESET_GRACE)
.is_some_and(|due_at| due_at <= now)
})
.min();
expired_window_reset.map(|reset_at| (account, Some(reset_at)))
}
fn periodic_quota_refresh_due(
@@ -457,9 +481,9 @@ fn periodic_quota_refresh_due(
.get(account_id)
.is_none_or(|last| {
monotonic_now.saturating_duration_since(last.monotonic_at)
>= EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL
>= PERIODIC_QUOTA_REFRESH_RETRY_INTERVAL
|| reset_at
.and_then(|reset| reset.checked_add(EXHAUSTED_QUOTA_RESET_GRACE))
.and_then(|reset| reset.checked_add(QUOTA_RESET_GRACE))
// 已在该边界之后复核过时回到周期重试，避免过期 reset 每轮触发。
.is_some_and(|due_at| last.wall_at < due_at && due_at <= now)
})
@@ -781,8 +805,26 @@ impl CodexCredentialQuotaService {
});
let mut summary = CodexQuotaSyncSummary::default();
let now = SystemTime::now();
let initial = self.initial_quota_sync_accounts(&accounts, now).await?;
let periodic = self.scheduling.reserve_periodic_refreshes(accounts, now);
let account_ids = accounts
.iter()
.map(|account| account.id().clone())
.collect::<Vec<_>>();
let observed = self.store.get_quotas(&account_ids).await?;
let observed_snapshots = observed
.iter()
.filter_map(|obs| {
quota_snapshot_from_observation(obs)
.map(|snapshot| (obs.account_id.clone(), snapshot))
})
.collect::<BTreeMap<_, _>>();
let observed_ids = observed
.into_iter()
.map(|observation| observation.account_id)
.collect::<BTreeSet<_>>();
let initial = Self::initial_quota_sync_accounts(&accounts, &observed_ids, now);
let periodic =
self.scheduling
.reserve_periodic_refreshes(accounts, &observed_snapshots, now);
let accounts = initial
.into_iter()
.chain(periodic)
@@ -857,28 +899,19 @@ impl CodexCredentialQuotaService {
}
/// `quota_observed_at` 为空代表首次异步观察尚未成功；不另建同步状态表。
async fn initial_quota_sync_accounts(
&self,
fn initial_quota_sync_accounts(
accounts: &[ProviderAccount],
observed_ids: &BTreeSet<ProviderAccountId>,
now: SystemTime,
) -> Result<Vec<ProviderAccount>, CodexCredentialQuotaError> {
let account_ids = accounts
.iter()
.map(|account| account.id().clone())
.collect::<Vec<_>>();
let observed = self.store.get_quotas(&account_ids).await?;
let observed_ids = observed
.into_iter()
.map(|observation| observation.account_id)
.collect::<BTreeSet<_>>();
Ok(accounts
) -> Vec<ProviderAccount> {
accounts
.iter()
.filter(|account| {
!observed_ids.contains(account.id()) && eligible_initial_quota_sync(account, now)
})
.take(INITIAL_QUOTA_SYNC_BATCH)
.cloned()
.collect())
.collect()
}
/// 解析并 revision-fenced 落库单账号的 Provider quota JSON。
backend/crates/providers/openai/tests/admin.rs
+115
View File
@@ -1142,6 +1142,121 @@ async fn openai_admin_keeps_confirmed_exhaustion_separate_from_raw_usage_display
assert_eq!(recovered.windows[0].used_percent, Some(86.0));
}
#[tokio::test]
async fn openai_admin_preserves_expired_window_usage_and_exhaustion_attribution() {
for exhausted in [false, true] {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_admin_expired_window";
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
name: "admin expired window".to_owned(),
secret: secret("admin-expired-window-access"),
verified_account: profile("chatgpt-admin-expired-window"),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
let account = store.account(account_id).expect("stored account");
let past_reset_at = Utc::now().timestamp() - 60;
let observed_at = SystemTime::now() - Duration::from_secs(300);
let weekly_used = if exhausted { 100.0 } else { 74.0 };
let state = if exhausted {
QuotaState::exhausted(
QuotaEvidence::ProviderDenied,
observed_at,
Some(SystemTime::UNIX_EPOCH + Duration::from_secs(past_reset_at as u64)),
)
} else {
QuotaState::allowed(observed_at)
};
store
.compare_and_swap_quota(QuotaObservation {
plan_type: None,
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": !exhausted,
"limit_reached": exhausted,
"primary_window": {
"used_percent": 15,
"reset_at": past_reset_at + 18_000,
"limit_window_seconds": 18_000,
},
"secondary_window": {
"used_percent": weekly_used,
"reset_at": past_reset_at,
"limit_window_seconds": 604_800,
}
}
})
.as_object()
.expect("quota object")
.clone(),
),
observed_at,
state,
})
.await
.expect("persist raw quota");
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 mut projected = bundle
.admin_provider()
.quota(ProviderQuotaRequest {
account_id: account.id().clone(),
refresh: false,
rolling_usage: None,
})
.await
.expect("project quota");
assert_eq!(projected.limit_reached, exhausted);
// 账号接口还会归一化耗尽展示；过期周窗口不能把触顶错误转移到短期窗口。
projected.apply_limit_reached_display();
let primary = projected
.windows
.iter()
.find(|w| w.window_seconds == Some(18_000))
.expect("primary");
let weekly = projected
.windows
.iter()
.find(|w| w.window_seconds == Some(604_800))
.expect("weekly");
assert_eq!(
(primary.used_percent, primary.limit_reached),
(Some(15.0), false)
);
assert_eq!(
(weekly.used_percent, weekly.limit_reached),
(Some(weekly_used), exhausted)
);
assert_eq!(
weekly.reset_at.map(|reset| reset.timestamp()),
Some(past_reset_at)
);
let raw = store
.get_quotas(std::slice::from_ref(account.id()))
.await
.expect("raw quota")
.pop()
.expect("observation");
assert_eq!(raw.observed_at, observed_at);
assert_eq!(
raw.quota.expose_to_provider()["rate_limit"]["secondary_window"]["used_percent"],
weekly_used
);
}
}
#[tokio::test]
async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit() {
let store = Arc::new(MemoryAccountStore::default());
backend/crates/providers/openai/tests/credential/quota/refresh_timing.rs
+94
View File
@@ -128,3 +128,97 @@ async fn periodic_checks_continue_when_reset_is_unknown_or_far_in_the_future() {
assert_eq!(server.received_requests().await.expect("requests").len(), 1);
}
}
#[tokio::test]
async fn allowed_account_with_expired_window_synchronizes_at_reset_grace() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_allowed_expired_window").await;
let account = store
.account("acct_allowed_expired_window")
.expect("account");
let server = MockServer::start().await;
let service = quota_service_with_base_url(&store, reqwest::Client::new(), server.uri());
let short_reset = Utc::now().timestamp() - 114;
let usage = |used: u32, reset: i64| {
json!({"rate_limit": {
"allowed": true,
"primary_window": {
"used_percent": used,
"reset_at": reset,
"limit_window_seconds": 18_000,
}
}})
};
mount_usage(&server, usage(100, short_reset)).await;
service
.refresh_account(account.id())
.await
.expect("seed quota with 100% usage");
assert!(
Utc::now().timestamp() < short_reset + 120,
"fixture must precede grace deadline"
);
let requests = server.received_requests().await.expect("requests").len();
service.synchronize().await.expect("initial periodic check");
service
.synchronize()
.await
.expect("repeat check before grace");
assert_eq!(
server.received_requests().await.expect("requests").len(),
requests,
"正常账号首次复核也必须等待 reset 宽限期"
);
// 当到达 reset + 120s 宽限期后，账号虽然处于 allowed 状态，但包含已到期的非零用量窗口，被调度主动同步。
mount_usage(&server, usage(0, short_reset + 18_000)).await;
wait_for_reset_grace(short_reset).await;
let summary = service
.synchronize()
.await
.expect("synchronize expired window");
assert_eq!(summary.updated, 1);
assert_eq!(server.received_requests().await.expect("requests").len(), 1);
let snapshot = service
.read_account(account.id())
.await
.expect("read quota")
.expect("snapshot");
assert_eq!(snapshot.windows()[0].used_percent(), Some(0.0));
service.synchronize().await.expect("skip fresh window");
assert_eq!(server.received_requests().await.expect("requests").len(), 1);
}
#[tokio::test]
async fn allowed_expired_window_refresh_is_throttled_when_observation_is_unchanged() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_allowed_expired_retry").await;
let account = store
.account("acct_allowed_expired_retry")
.expect("account");
let server = MockServer::start().await;
let service = quota_service_with_base_url(&store, reqwest::Client::new(), server.uri());
let usage = json!({"rate_limit": {
"allowed": true,
"primary_window": {
"used_percent": 74,
"reset_at": Utc::now().timestamp() - 180,
"limit_window_seconds": 18_000,
}
}});
mount_usage(&server, usage.clone()).await;
service
.refresh_account(account.id())
.await
.expect("seed quota");
mount_usage(&server, usage).await;
assert_eq!(service.synchronize().await.expect("first check").updated, 1);
service.synchronize().await.expect("throttled repeat check");
assert_eq!(server.received_requests().await.expect("requests").len(), 1);
}
docs/api.md
+4 -2
View File
@@ -694,8 +694,10 @@ OAuth start 使用：
- `GET /accounts/quota` 只读取最后一次落库快照；`POST /accounts/quota/refresh` 才访问上游。access token
已过期时，额度刷新要求先走 credential 刷新或重新授权，不会拿过期 token 探测额度。
- OpenAI 已耗尽账号每 30 分钟主动复核一次，也会在最早未恢复窗口的 `resetAt + 2 分钟` 到期后
提前复核。后台每 30 秒检查触发条件；同一重置边界复核后仍未恢复时回到 30 分钟重试，
避免旧 reset 持续触发请求。各窗口独立确认恢复，时间到期本身不会直接解除账号耗尽。
提前复核；周期复核不等待旧重置时间，因此也能发现官方提前重置。正常账号有非零用量或触顶窗口时，
在该窗口的 `resetAt + 2 分钟` 后主动复核。后台每 30 秒检查触发条件；同一重置边界复核后仍未更新时
回到 30 分钟重试，避免旧 reset 持续触发请求。各窗口独立确认恢复，时间到期本身不会直接解除账号
耗尽或将展示用量归零。
- `POST /accounts/recover` 是管理员对本地事实的强制恢复：它清除 Redis cooldown 和已保存的额度/错误，
把账号重新启用并恢复为可调度 credential；它不验证上游账号是否已经恢复，下一次真实请求仍可重新写入
失败事实。