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

fix(openai): separate quota and credential auth state

Browse Source 
Keep quota probes from terminally invalidating OAuth accounts on 401 or 403, and stop manual quota refreshes from querying upstream with an expired access token. Model quota-owned transitions explicitly and cover both boundaries with regression tests.

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-29 01:09:07 +08:00
co-authored by Codex
parent 428e744710
commit 5544b0fb29
3 changed files with 167 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/providers/openai/src/admin.rs
+4 -3
View File
@@ -1360,9 +1360,10 @@ fn map_quota_error(error: CodexCredentialQuotaError) -> ProviderAdminError {
Error::InvalidCredentialData => ProviderAdminErrorKind::Invalid,
Error::NotFound => ProviderAdminErrorKind::NotFound,
Error::RevisionConflict => ProviderAdminErrorKind::Conflict,
Error::Repository(_) | Error::Store | Error::Upstream => {
ProviderAdminErrorKind::Unavailable
}
Error::CredentialRefreshRequired
| Error::Repository(_)
| Error::Store
| Error::Upstream => ProviderAdminErrorKind::Unavailable,
})
}
backend/crates/providers/openai/src/credential/quota.rs
+75 -54
View File
@@ -190,7 +190,7 @@ fn quota_fact_is_active_exhaustion(fact: CodexQuotaFact, now: SystemTime) -> boo
pub struct CodexQuotaSyncSummary {
pub updated: u64,
pub exhausted: u64,
pub invalid: u64,
pub banned: u64,
pub cooldown: u64,
pub transient: u64,
pub stale: u64,
@@ -207,6 +207,8 @@ impl CodexQuotaSyncSummary {
pub enum CodexCredentialQuotaError {
#[error("Codex quota response is invalid")]
InvalidCredentialData,
#[error("Codex OAuth access token must be refreshed before querying quota")]
CredentialRefreshRequired,
#[error(transparent)]
Repository(#[from] CredentialRepositoryError),
#[error("provider account store is unavailable")]
@@ -301,6 +303,38 @@ enum CodexQuotaFetchError {
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CodexQuotaStateTransition {
Banned { reason: &'static str },
Exhausted,
Cooldown { until: Option<SystemTime> },
}
impl CodexQuotaStateTransition {
const fn availability(self) -> AccountAvailability {
match self {
Self::Banned { .. } => AccountAvailability::Banned,
Self::Exhausted => AccountAvailability::QuotaExhausted,
Self::Cooldown { .. } => AccountAvailability::Cooldown,
}
}
const fn reason(self) -> &'static str {
match self {
Self::Banned { reason } => reason,
Self::Exhausted => "quota_exhausted",
Self::Cooldown { .. } => "quota_rate_limited",
}
}
const fn cooldown_until(self) -> Option<SystemTime> {
match self {
Self::Cooldown { until } => until,
Self::Banned { .. } | Self::Exhausted => None,
}
}
}
impl CodexQuotaSchedulingProjection {
fn hydration_targets(&self, accounts: &[ProviderAccount]) -> Vec<CodexQuotaHydrationTarget> {
let state = self
@@ -764,20 +798,17 @@ impl CodexCredentialQuotaService {
summary.transient += 1;
}
Err(CodexQuotaFetchError::Upstream { account, error }) => {
let (availability, reason, cooldown) = classify_error(&error, observed_at);
let Some(transition) = quota_state_transition(&error, observed_at) else {
summary.transient += 1;
continue;
};
let preserve_exhaustion = account.availability()
== AccountAvailability::QuotaExhausted
&& availability == Some(AccountAvailability::Cooldown);
match availability {
Some(AccountAvailability::Invalid | AccountAvailability::Banned) => {
summary.invalid += 1;
}
Some(AccountAvailability::QuotaExhausted) => summary.exhausted += 1,
Some(AccountAvailability::Cooldown) => summary.cooldown += 1,
_ => {
summary.transient += 1;
continue;
}
&& transition.availability() == AccountAvailability::Cooldown;
match transition {
CodexQuotaStateTransition::Banned { .. } => summary.banned += 1,
CodexQuotaStateTransition::Exhausted => summary.exhausted += 1,
CodexQuotaStateTransition::Cooldown { .. } => summary.cooldown += 1,
}
if preserve_exhaustion {
continue;
@@ -786,9 +817,9 @@ impl CodexCredentialQuotaService {
.repository
.apply_state(
&account,
availability.expect("classified availability"),
reason.map(str::to_owned),
cooldown,
transition.availability(),
Some(transition.reason().to_owned()),
transition.cooldown_until(),
observed_at,
)
.await;
@@ -1062,12 +1093,15 @@ impl CodexCredentialQuotaService {
.await?
.filter(|account| account.provider().as_str() == "openai")
.ok_or(CodexCredentialQuotaError::NotFound)?;
let observed_at = SystemTime::now();
if !access_token_is_current(&account, observed_at) {
return Err(CodexCredentialQuotaError::CredentialRefreshRequired);
}
let client = CodexBackendClient::new(
self.http.clone(),
self.base_url.clone(),
self.profile.clone(),
);
let observed_at = SystemTime::now();
let FetchedCodexQuota { account, value } =
match self.fetch_usage_with_recovery(&client, &account).await {
Ok(fetched) => fetched,
@@ -1078,15 +1112,14 @@ impl CodexCredentialQuotaService {
return Err(CodexCredentialQuotaError::Upstream);
}
Err(CodexQuotaFetchError::Upstream { account, error }) => {
let (availability, reason, cooldown) = classify_error(&error, observed_at);
if let Some(availability) = availability {
if let Some(transition) = quota_state_transition(&error, observed_at) {
let _ = self
.repository
.apply_state(
&account,
availability,
reason.map(str::to_owned),
cooldown,
transition.availability(),
Some(transition.reason().to_owned()),
transition.cooldown_until(),
observed_at,
)
.await;
@@ -1218,10 +1251,13 @@ async fn fetch_usage_once(
}
fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool {
account.enabled()
&& account
.access_token_expires_at()
.is_none_or(|expires_at| expires_at > now)
account.enabled() && access_token_is_current(account, now)
}
fn access_token_is_current(account: &ProviderAccount, now: SystemTime) -> bool {
account
.access_token_expires_at()
.is_none_or(|expires_at| expires_at > now)
}
fn merge_passive_quota(
@@ -1377,14 +1413,10 @@ fn quota_success_availability(
}
}
fn classify_error(
fn quota_state_transition(
error: &CodexClientError,
observed_at: SystemTime,
) -> (
Option<AccountAvailability>,
Option<&'static str>,
Option<SystemTime>,
) {
) -> Option<CodexQuotaStateTransition> {
let CodexClientError::Upstream {
status,
retry_after_seconds,
@@ -1392,38 +1424,27 @@ fn classify_error(
..
} = error
else {
return (None, None, None);
return None;
};
if *status == StatusCode::PAYMENT_REQUIRED && is_deactivated_workspace(body) {
return (
Some(AccountAvailability::Banned),
Some("quota_deactivated_workspace"),
None,
);
return Some(CodexQuotaStateTransition::Banned {
reason: "quota_deactivated_workspace",
});
}
match *status {
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => (
Some(AccountAvailability::Invalid),
Some("quota_auth_rejected"),
None,
),
StatusCode::PAYMENT_REQUIRED => (
Some(AccountAvailability::QuotaExhausted),
Some("quota_exhausted"),
None,
),
StatusCode::PAYMENT_REQUIRED => Some(CodexQuotaStateTransition::Exhausted),
StatusCode::TOO_MANY_REQUESTS => {
let duration = Duration::from_secs(
retry_after_seconds.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN.as_secs()),
)
.min(MAX_RATE_LIMIT_COOLDOWN);
(
Some(AccountAvailability::Cooldown),
Some("quota_rate_limited"),
observed_at.checked_add(duration),
)
Some(CodexQuotaStateTransition::Cooldown {
until: observed_at.checked_add(duration),
})
}
_ => (None, None, None),
// quota/usage 接口只能产生额度域状态；鉴权拒绝不具备判定 RT
// 永久失效的证据，终态只能由 OAuth refresh/身份校验状态机写入。
_ => None,
}
}
backend/crates/providers/openai/tests/credential/quota.rs
+88 -3
View File
@@ -9,8 +9,8 @@ use gateway_core::engine::credential::{
};
use provider_openai::OFFICIAL_CODEX_BASE_URL;
use provider_openai::credential::{
CodexCredentialQuotaService, CodexQuotaSyncSummary, CodexQuotaWindowKind,
ImportCodexOAuthCredential, parse_codex_quota_usage,
CodexCredentialQuotaError, CodexCredentialQuotaService, CodexQuotaSyncSummary,
CodexQuotaWindowKind, ImportCodexOAuthCredential, parse_codex_quota_usage,
};
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use serde_json::json;
@@ -438,6 +438,91 @@ async fn manual_quota_refresh_updates_account_state_from_provider_response() {
);
}
#[tokio::test]
async fn manual_quota_auth_rejection_does_not_invalidate_refreshable_credential() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_manual_quota_auth_rejected";
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("quota_exhausted".to_owned()),
cooldown_until: None,
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(401))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
assert!(matches!(
service.refresh_account(account.id()).await,
Err(CodexCredentialQuotaError::Upstream)
));
assert_eq!(
store
.account(account_id)
.expect("preserved account")
.availability(),
AccountAvailability::QuotaExhausted
);
}
#[tokio::test]
async fn manual_quota_refresh_waits_for_expired_oauth_token_to_be_refreshed() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_manual_quota_expired_token";
let mut expired_profile = profile(&format!("chatgpt-{account_id}"));
expired_profile.access_token_expires_at = Some(Utc::now() - chrono::Duration::minutes(1));
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
name: account_id.to_owned(),
secret: secret(&format!("token-{account_id}")),
verified_account: expired_profile,
next_refresh_at: Some(Utc::now() - chrono::Duration::minutes(5)),
enabled: true,
})
.await;
let account = store.account(account_id).expect("created account");
let server = MockServer::start().await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
assert!(matches!(
service.refresh_account(account.id()).await,
Err(CodexCredentialQuotaError::CredentialRefreshRequired)
));
assert!(
server
.received_requests()
.await
.expect("received requests")
.is_empty()
);
}
#[tokio::test]
async fn passive_rate_limit_headers_replace_stale_secondary_window() {
let store = Arc::new(MemoryAccountStore::default());
@@ -516,7 +601,7 @@ async fn synchronize_without_accounts_is_a_noop_before_network_io() {
assert_eq!(summary.updated, 0);
assert_eq!(summary.exhausted, 0);
assert_eq!(summary.invalid, 0);
assert_eq!(summary.banned, 0);
assert_eq!(summary.cooldown, 0);
assert_eq!(summary.transient, 0);
assert_eq!(summary.stale, 0);