fix(openai): restore periodic quota checks and timely reset recovery

Recheck exhausted accounts every 30 minutes even before their recorded reset, so early provider resets can be discovered. Allow each pending reset boundary to trigger an earlier check after a two-minute grace period, then return to periodic retries if exhaustion persists.

Separate quota checks from model catalog refreshes. Scan quota eligibility every 30 seconds while retaining the configured catalog interval and independent per-window recovery rules. Keep quota timing regressions in the mirrored integration-test directory and exercise the public service.

Validation: quota integration tests (46 passed), gateway application and architecture tests (36 passed), strict OpenAI Clippy with all targets and features, workspace rustfmt, and diff whitespace checks passed.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-09-09 15:39:41 +08:00
co-authored by Codex
parent 35988d0eb7
commit 18a04f1213
8 changed files with 315 additions and 32 deletions
@@ -213,6 +213,7 @@ impl CodexWebSocketPoolSettings {
#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct CodexQuotaSettings {
/// 保留模型目录的刷新周期;额度独立每 30 秒检查周期复核和 reset 到期条件。
pub refresh_interval_minutes: u64,
}
@@ -57,8 +57,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_FALLBACK_RECHECK_INTERVAL: Duration = Duration::from_secs(30 * 60);
const EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL: Duration = Duration::from_secs(30 * 60);
const EXHAUSTED_QUOTA_RESET_GRACE: Duration = Duration::from_secs(2 * 60);
/// 首次 OAuth 异步观察失败时,由既有 quota worker 兜底重试的单轮上限。
const INITIAL_QUOTA_SYNC_BATCH: usize = 100;
// 5xx 上游拒绝的短退避重试预算;指数退避 1s/2s,吞掉瞬时抖动。
@@ -68,8 +68,8 @@ const QUOTA_FETCH_5XX_BASE_DELAY: Duration = Duration::from_secs(1);
/// OpenAI Provider 主动额度刷新的调度策略。
///
/// 正常账号依赖请求响应的被动额度同步;周期 worker 仅复核已耗尽账号。
/// 该策略只控制 Provider 自己的 quota worker 频率,避免把 OpenAI
/// 额度语义泄漏到公共调度层。
/// 该策略保留模型目录的周期刷新频率;额度到期检查使用独立的短周期,
/// 避免到达 reset 后还要等待完整的目录刷新周期。
#[derive(Debug, Clone, Copy)]
pub struct CodexQuotaRefreshPolicy {
interval: Duration,
@@ -204,7 +204,12 @@ struct CodexQuotaSchedulingProjection {
struct CodexQuotaProjectionState {
next_version: u64,
entries: BTreeMap<ProviderAccountId, CodexQuotaSchedulingEntry>,
last_periodic_refresh_at: BTreeMap<ProviderAccountId, Instant>,
last_periodic_refresh_at: BTreeMap<ProviderAccountId, CodexQuotaRefreshAttempt>,
}
struct CodexQuotaRefreshAttempt {
monotonic_at: Instant,
wall_at: SystemTime,
}
#[derive(Debug, Clone, Copy)]
@@ -398,16 +403,27 @@ impl CodexQuotaSchedulingProjection {
.last_periodic_refresh_at
.retain(|account_id, _| candidate_ids.contains(account_id));
// 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。定时器只复核
// 已耗尽且已到 reset/fallback 时刻的账号;重复复核至少间隔 30 分钟。
// 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。
// 已耗尽账号每 30 分钟复核,不能等待旧 reset:官方活动可能提前重置额度。
// reset + 2 分钟可提前触发一次复核,给上游重置留出传播时间。
let mut reserved = Vec::new();
for account in candidates {
if !periodic_quota_refresh_due(&state, account.id(), refreshed_at) {
if !periodic_quota_refresh_due(
&state,
account.id(),
account.quota().reset_at(),
now,
refreshed_at,
) {
continue;
}
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
state.last_periodic_refresh_at.insert(
account.id().clone(),
CodexQuotaRefreshAttempt {
monotonic_at: refreshed_at,
wall_at: now,
},
);
reserved.push(account);
}
reserved
@@ -415,24 +431,27 @@ impl CodexQuotaSchedulingProjection {
}
fn quota_refresh_candidate(account: ProviderAccount, now: SystemTime) -> Option<ProviderAccount> {
// reset 到期是“应刷新”的信号,不是“已恢复”的证据。
(eligible_periodic_quota_refresh(&account, now)
&& account
.quota()
.exhaustion_refresh_due(now, EXHAUSTED_QUOTA_FALLBACK_RECHECK_INTERVAL))
.then_some(account)
(eligible_periodic_quota_refresh(&account, now) && account.quota().is_exhausted())
.then_some(account)
}
fn periodic_quota_refresh_due(
state: &CodexQuotaProjectionState,
account_id: &ProviderAccountId,
now: Instant,
reset_at: Option<SystemTime>,
now: SystemTime,
monotonic_now: Instant,
) -> bool {
state
.last_periodic_refresh_at
.get(account_id)
.is_none_or(|last| {
now.saturating_duration_since(*last) >= EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL
monotonic_now.saturating_duration_since(last.monotonic_at)
>= EXHAUSTED_QUOTA_REFRESH_RETRY_INTERVAL
|| reset_at
.and_then(|reset| reset.checked_add(EXHAUSTED_QUOTA_RESET_GRACE))
// 已在该边界之后复核过时回到周期重试,避免过期 reset 每轮触发。
.is_some_and(|due_at| last.wall_at < due_at && due_at <= now)
})
}
@@ -7,8 +7,10 @@ pub(super) const WORKER_MAXIMUM_BACKOFF: Duration = Duration::from_secs(60);
pub(super) const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60);
pub(super) const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60);
pub(super) const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30);
pub(super) const QUOTA_CHECK_INTERVAL: Duration = Duration::from_secs(30);
pub(super) const DESKTOP_RELEASE_WORKER_OWNER: &str = "openai-desktop-release";
pub(super) const MODEL_ETAG_WORKER_OWNER: &str = "openai-model-etag";
pub(super) const MODEL_CATALOG_WORKER_OWNER: &str = "openai-model-catalog";
pub(crate) fn worker_contributions(
refresh: Arc<CodexCredentialRefreshService>,
@@ -20,6 +22,7 @@ pub(crate) fn worker_contributions(
) -> Result<Vec<WorkerContribution>, WorkerDefinitionError> {
let refresh_id = WorkerId::try_new(WorkerKind::OAuthRefresh, PROVIDER_NAME)?;
let quota_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, PROVIDER_NAME)?;
let catalog_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, MODEL_CATALOG_WORKER_OWNER)?;
let etag_id = WorkerId::try_new(WorkerKind::QuotaCatalogHealth, MODEL_ETAG_WORKER_OWNER)?;
let desktop_release_id =
WorkerId::try_new(WorkerKind::QuotaCatalogHealth, DESKTOP_RELEASE_WORKER_OWNER)?;
@@ -34,9 +37,13 @@ pub(crate) fn worker_contributions(
contributions.extend([
WorkerContribution::Registration(scheduled_registration(
quota_id,
QUOTA_CHECK_INTERVAL,
Box::new(OpenAiQuotaTask { quota }),
)?),
WorkerContribution::Registration(scheduled_registration(
catalog_id,
quota_refresh_policy.interval(),
Box::new(OpenAiQuotaTask {
quota,
Box::new(OpenAiCatalogTask {
catalog: Arc::clone(&catalog),
}),
)?),
@@ -159,6 +166,9 @@ impl ScheduledTask for OpenAiOAuthRefreshTask {
pub(super) struct OpenAiQuotaTask {
quota: Arc<CodexCredentialQuotaService>,
}
pub(super) struct OpenAiCatalogTask {
catalog: Arc<CodexCredentialCatalogService>,
}
@@ -195,7 +205,6 @@ impl ScheduledTask for OpenAiQuotaTask {
if context.cancellation().is_cancelled() {
return Ok(());
}
let mut failures = false;
match self.quota.synchronize().await {
Ok(summary) if summary.has_operational_failures() => {
tracing::warn!(
@@ -209,30 +218,33 @@ impl ScheduledTask for OpenAiQuotaTask {
}
Ok(_) => {}
Err(error) => {
failures = true;
tracing::warn!(
error = %error,
"OpenAI quota synchronization failed"
);
return Err(WorkerTaskError::safe("OpenAI quota synchronization failed"));
}
}
Ok(())
})
}
}
impl ScheduledTask for OpenAiCatalogTask {
fn run_cycle(&self, context: WorkerCycleContext) -> BoxFuture<'_, Result<(), WorkerTaskError>> {
Box::pin(async move {
if context.cancellation().is_cancelled() {
return Ok(());
}
match self.catalog.refresh_catalogs().await {
Ok(_) | Err(CodexCredentialCatalogError::NoEligibleCredential) => {}
Ok(_) | Err(CodexCredentialCatalogError::NoEligibleCredential) => Ok(()),
Err(error) => {
failures = true;
tracing::warn!(error = %error, "OpenAI model catalog refresh failed");
Err(WorkerTaskError::safe(
"OpenAI model catalog synchronization failed",
))
}
}
if failures {
Err(WorkerTaskError::safe(
"OpenAI quota or catalog synchronization failed",
))
} else {
Ok(())
}
})
}
}
+25 -1
View File
@@ -74,7 +74,7 @@ async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions
assert_eq!(bundle.core_provider().name(), "openai");
assert_eq!(bundle.admin_provider().provider_kind().as_str(), "openai");
let contributions = bundle.take_worker_contributions();
assert_eq!(contributions.len(), 4);
assert_eq!(contributions.len(), 5);
assert!(
contributions
.iter()
@@ -101,6 +101,30 @@ async fn openai_bundle_exposes_one_core_provider_and_drains_worker_contributions
panic!("Desktop release worker must be scheduled");
};
assert_eq!(schedule.interval(), APPCAST_POLL_INTERVAL);
for (owner, interval) in [
("openai", Duration::from_secs(30)),
(
"openai-model-catalog",
config.config.quota_refresh_policy().interval(),
),
] {
let schedule = contributions
.iter()
.find_map(|contribution| match contribution {
WorkerContribution::Registration(registration)
if registration.id.kind() == WorkerKind::QuotaCatalogHealth
&& registration.id.owner() == owner =>
{
match &registration.runnable {
WorkerRunnable::Scheduled { schedule, .. } => Some(schedule),
WorkerRunnable::Daemon { .. } => None,
}
}
_ => None,
})
.expect("quota/catalog scheduled worker");
assert_eq!(schedule.interval(), interval, "{owner}");
}
assert!(contributions.iter().any(|contribution| {
matches!(
contribution,
@@ -1,6 +1,7 @@
//! OpenAI 额度事实边界与展示快照回归。
mod recovery;
mod refresh_timing;
mod scheduling;
mod slots;
mod snapshot;
@@ -0,0 +1,130 @@
//! 通过 quota 服务验证周期复核与 reset 宽限期,不依赖内部调度状态。
use std::time::Duration;
use super::*;
async fn mount_usage(server: &MockServer, value: serde_json::Value) {
server.reset().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(value))
.mount(server)
.await;
}
async fn wait_for_reset_grace(reset: i64) {
let due = SystemTime::UNIX_EPOCH + Duration::from_secs((reset + 120) as u64);
if let Ok(remaining) = due.duration_since(SystemTime::now()) {
tokio::time::sleep(remaining + Duration::from_millis(50)).await;
}
}
#[tokio::test]
async fn reset_grace_bypasses_periodic_throttle_once_then_allows_the_next_window() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_quota_timing").await;
let account = store.account("acct_quota_timing").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() - 115;
let week_reset = short_reset + 3;
let usage = |short_used, short_reset, week_used, week_reset| {
json!({"rate_limit": {
"allowed": false,
"primary_window": {
"used_percent": short_used, "reset_at": short_reset, "limit_window_seconds": 18_000,
},
"secondary_window": {
"used_percent": week_used, "reset_at": week_reset, "limit_window_seconds": 604_800,
},
}})
};
mount_usage(&server, usage(100, short_reset, 100, week_reset)).await;
service
.refresh_account(account.id())
.await
.expect("seed exhaustion");
service.synchronize().await.expect("initial periodic check");
let requests = server.received_requests().await.expect("requests").len();
assert!(
Utc::now().timestamp() < short_reset + 120,
"fixture must precede grace deadline"
);
service.synchronize().await.expect("before grace deadline");
assert_eq!(
server.received_requests().await.expect("requests").len(),
requests
);
mount_usage(&server, usage(0, short_reset + 18_000, 100, week_reset)).await;
wait_for_reset_grace(short_reset).await;
let partial = service.synchronize().await.expect("short reset check");
assert_eq!(partial.exhausted, 1);
assert_eq!(
store
.account("acct_quota_timing")
.expect("account")
.quota()
.reset_at(),
Some(SystemTime::UNIX_EPOCH + Duration::from_secs(week_reset as u64))
);
service
.synchronize()
.await
.expect("next scan before weekly grace deadline");
assert_eq!(server.received_requests().await.expect("requests").len(), 1);
// 周窗口的到期复核仍未恢复时,也不能每轮扫描重复请求。
wait_for_reset_grace(week_reset).await;
assert_eq!(
service
.synchronize()
.await
.expect("weekly reset check")
.exhausted,
1
);
service
.synchronize()
.await
.expect("repeat scan after weekly check");
assert_eq!(server.received_requests().await.expect("requests").len(), 2);
}
#[tokio::test]
async fn periodic_checks_continue_when_reset_is_unknown_or_far_in_the_future() {
for reset in [None, Some(Utc::now().timestamp() + 604_800)] {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_quota_periodic").await;
let account = store.account("acct_quota_periodic").expect("account");
let server = MockServer::start().await;
let service = quota_service_with_base_url(&store, reqwest::Client::new(), server.uri());
let mut usage = json!({"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "limit_window_seconds": 604_800},
}});
if let Some(reset) = reset {
usage["rate_limit"]["primary_window"]["reset_at"] = json!(reset);
}
mount_usage(&server, usage.clone()).await;
service
.refresh_account(account.id())
.await
.expect("seed exhaustion");
mount_usage(&server, usage).await;
assert_eq!(
service
.synchronize()
.await
.expect("periodic check")
.exhausted,
1
);
service.synchronize().await.expect("throttled repeat check");
assert_eq!(server.received_requests().await.expect("requests").len(), 1);
}
}
@@ -172,6 +172,99 @@ async fn recovery_checks_each_exhausted_slot_in_worker_and_manual_refresh() {
}
}
#[tokio::test]
async fn worker_detects_early_resets_without_unlocking_other_exhausted_windows() {
let short_reset = Utc::now().timestamp() + 18_000;
let week_reset = Utc::now().timestamp() + 604_800;
for (name, old_short, old_week, new_short, new_week, recovered) in [
("short_only", Some(100), None, Some(0), None, true),
("weekly_only", None, Some(100), None, Some(0), true),
(
"short_with_weekly_available",
Some(100),
Some(45),
Some(0),
Some(45),
true,
),
(
"weekly_still_exhausted",
Some(100),
Some(100),
Some(0),
Some(100),
false,
),
] {
let store = Arc::new(MemoryAccountStore::default());
let account_id = format!("acct_early_{name}");
create_account(&store, &account_id).await;
let account = store.account(&account_id).expect("account");
let server = MockServer::start().await;
let service = quota_service_with_base_url(&store, reqwest::Client::new(), server.uri());
let window = |used: Option<u64>, reset, seconds| {
used.map(|used| {
json!({
"used_percent": used,
"reset_at": reset,
"limit_window_seconds": seconds,
})
})
};
mount_usage(
&server,
json!({"rate_limit": {
"allowed": false,
"primary_window": window(old_short, short_reset, 18_000),
"secondary_window": window(old_week, week_reset, 604_800),
}}),
)
.await;
let baseline = service
.refresh_account(account.id())
.await
.expect("seed exhaustion");
assert!(baseline.quota().reset_at().expect("future reset") > SystemTime::now());
mount_usage(
&server,
json!({"rate_limit": {
"allowed": recovered,
"primary_window": window(new_short, short_reset + 18_000, 18_000),
"secondary_window": window(new_week, if new_week == Some(0) {
week_reset + 604_800
} else {
week_reset
}, 604_800),
}}),
)
.await;
let summary = service
.synchronize()
.await
.expect("periodic early reset check");
assert_eq!(summary.updated, u64::from(recovered), "{name}");
assert_eq!(summary.exhausted, u64::from(!recovered), "{name}");
assert_eq!(
store
.account(&account_id)
.expect("updated account")
.quota()
.is_exhausted(),
!recovered,
"{name}"
);
// 未恢复的周窗口仍在未来,也必须保留复核节流;恢复后则退出定时复核。
service.synchronize().await.expect("next worker cycle");
assert_eq!(
server.received_requests().await.expect("requests").len(),
1,
"{name}"
);
}
}
#[tokio::test]
async fn partial_recovery_survives_restart_missing_windows_and_changed_roles() {
let store = Arc::new(MemoryAccountStore::default());
+3
View File
@@ -306,6 +306,9 @@ OAuth start 使用:
但上游 authorization code 本身通常只能交换一次;已完成过 token exchange 时应重新创建 OAuth flow。
- `GET /accounts/quota` 只读取最后一次落库快照;`POST /accounts/quota/refresh` 才访问上游。access token
已过期时,额度刷新要求先走 credential 刷新或重新授权,不会拿过期 token 探测额度。
- OpenAI 已耗尽账号每 30 分钟主动复核一次,也会在最早未恢复窗口的 `resetAt + 2 分钟` 到期后
提前复核。后台每 30 秒检查触发条件;同一重置边界复核后仍未恢复时回到 30 分钟重试,
避免旧 reset 持续触发请求。各窗口独立确认恢复,时间到期本身不会直接解除账号耗尽。
- `POST /accounts/recover` 是管理员对本地事实的强制恢复:它清除 Redis cooldown 和已保存的额度/错误,
把账号重新启用并恢复为可调度 credential;它不验证上游账号是否已经恢复,下一次真实请求仍可重新写入
失败事实。