Compare commits

...
5 Commits
Author SHA1 Message Date
ZyyandCodex a321ccf737 chore: bump version to 3.0.24
Co-authored-by: Codex <noreply@openai.com>
2026-07-29 14:49:27 +08:00
CodexandCodex 4533ae8850 fix(openai): preserve resettable quota exhaustion state
Co-authored-by: Codex <noreply@openai.com>
2026-07-29 14:48:23 +08:00
zyycnandCodex 7542ed365e chore: bump version to 3.0.23
Co-authored-by: Codex <noreply@openai.com>
2026-07-29 03:52:17 +08:00
zyycnandCodex bf4da7a3e7 fix(accounts): unify credential and quota state
Apply OpenAI reauthorization to the selected local account ID while atomically replacing its verified upstream identity. Converge successful quota observations across refresh paths and preserve explicitly disabled account state during connection checks, quota refreshes, and credential rotation.

Co-authored-by: Codex <noreply@openai.com>
2026-07-29 03:51:54 +08:00
zyycnandCodex 0fe354b083 fix(openai): preserve requested service tier
Restore 2.x semantics: a request-specified tier drives persisted display and pricing, while the upstream response tier remains separately observable. Cover HTTP SSE and WebSocket conflicts plus the response-only fallback.

Co-authored-by: Codex <noreply@openai.com>
2026-07-29 03:09:22 +08:00
35 changed files with 1460 additions and 349 deletions
@@ -5,7 +5,7 @@ use std::fmt;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId},
engine::credential::{OpaqueProviderData, ProviderAccountId, ProviderAccountIdentity},
routing::{ProviderKind, UpstreamModelId},
};
@@ -496,6 +496,7 @@ pub struct PreparedCredentialRotationFacts {
pub account_id: ProviderAccountId,
pub provider_kind: ProviderKind,
pub expected_credential_revision: Revision,
pub replacement_identity: Option<ProviderAccountIdentity>,
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
@@ -155,6 +155,7 @@ impl FakeProviderAdmin {
.current_credential_revision
.lock()
.expect("current credential revision"),
replacement_identity: None,
name: account.name.clone(),
email: account.email.clone(),
plan_type: account.plan_type.clone(),
@@ -801,6 +802,35 @@ async fn connection_test_should_probe_unavailable_account() {
));
}
#[tokio::test]
async fn connection_test_should_preserve_disabled_account_status() {
let provider = FakeProviderAdmin::new("xai", events());
let mut account = account_record("xai");
account.enabled = false;
let store = FakeAccountStore::with_account(account, events());
let services =
accounts_service_with_probe(provider, store, Arc::new(FailingAccountProbe)).await;
let events = services
.accounts()
.test_connection(
ProviderAccountId::new("acct_test").expect("account ID"),
gateway_core::routing::UpstreamModelId::new("grok-4.5").expect("model"),
)
.await
.expect("connection test stream")
.collect::<Vec<_>>()
.await;
assert!(matches!(
events.last(),
Some(AccountConnectionTestEvent::Failed {
account_status: gateway_admin::model::accounts::AccountStatus::Disabled,
..
})
));
}
#[tokio::test]
async fn accounts_export_should_pass_store_loaded_timestamps_and_material_to_provider() {
let provider = FakeProviderAdmin::new("openai", events());
@@ -44,6 +44,46 @@ impl fmt::Display for ProviderAccountId {
}
}
/// credential 轮换时可选替换的上游账号身份。
#[derive(Clone, PartialEq, Eq)]
pub struct ProviderAccountIdentity {
upstream_user_id: String,
upstream_account_id: Option<String>,
}
impl fmt::Debug for ProviderAccountIdentity {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ProviderAccountIdentity")
.field("upstream_user_id", &"<redacted>")
.field(
"upstream_account_id",
&self.upstream_account_id.as_ref().map(|_| "<redacted>"),
)
.finish()
}
}
impl ProviderAccountIdentity {
#[must_use]
pub const fn new(upstream_user_id: String, upstream_account_id: Option<String>) -> Self {
Self {
upstream_user_id,
upstream_account_id,
}
}
#[must_use]
pub fn upstream_user_id(&self) -> &str {
&self.upstream_user_id
}
#[must_use]
pub fn upstream_account_id(&self) -> Option<&str> {
self.upstream_account_id.as_deref()
}
}
/// `provider_accounts.credential_revision` 的正数 CAS revision。
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CredentialRevision(NonZeroU64);
@@ -9,7 +9,8 @@ use gateway_core::engine::credential::{
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker,
AccountSelectionContext, AccountSelectionPolicy, AccountSelector, CredentialCasUpdate,
CredentialRevision, OpaqueProviderData, PlaintextCredential, PreferredAccountSelection,
ProviderAccount, ProviderAccountId, ProviderAccountUpdate, RotationStrategy,
ProviderAccount, ProviderAccountId, ProviderAccountIdentity, ProviderAccountUpdate,
RotationStrategy,
};
use gateway_core::routing::ProviderKind;
@@ -73,6 +74,18 @@ fn opaque_provider_data_should_not_expose_quota_values_in_debug() {
assert!(!format!("{quota:?}").contains("private-window"));
}
#[test]
fn provider_account_identity_debug_should_redact_upstream_ids() {
let identity = ProviderAccountIdentity::new(
"private-user-id".to_owned(),
Some("private-account-id".to_owned()),
);
let debug = format!("{identity:?}");
assert!(!debug.contains("private-user-id"));
assert!(!debug.contains("private-account-id"));
}
#[test]
fn availability_should_round_trip_all_database_values() {
assert_eq!(
@@ -392,6 +392,12 @@ pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(Strin
&format!("{prefix}-secondary"),
details.secondary,
);
if let Some(allowed) = details.allowed {
headers.push((format!("{prefix}-allowed"), allowed.to_string()));
}
if let Some(limit_reached) = details.limit_reached {
headers.push((format!("{prefix}-limit-reached"), limit_reached.to_string()));
}
if let Some(limit_name) = &details.limit_name {
headers.push((format!("{prefix}-limit-name"), limit_name.clone()));
}
@@ -440,6 +440,8 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
"metered_limit_name": "codex-other",
"limit_name": "Other requests",
"rate_limits": {
"allowed": false,
"limit_reached": true,
"primary": {
"used_percent": 33,
"window_minutes": 60,
@@ -450,7 +452,8 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
"has_credits": true,
"unlimited": false,
"balance": "9.25"
}
},
"rate_limit_reached_type": "usage_limit_reached"
});
let parsed = parse_rate_limits_event(&event).expect("rate-limit event");
@@ -460,16 +463,22 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
(
parsed.active_limit.as_deref(),
details.limit_name.as_deref(),
details.allowed,
details.limit_reached,
details.primary.expect("primary window").window_minutes,
parsed.plan_type.as_deref(),
parsed.credits.expect("credits").balance,
parsed.rate_limit_reached_type.as_deref(),
),
(
Some("codex_other"),
Some("Other requests"),
Some(false),
Some(true),
Some(60),
Some("team"),
Some("9.25".to_owned()),
Some("usage_limit_reached"),
)
);
}
@@ -498,9 +507,15 @@ fn rate_limit_header_round_trip_should_preserve_wire_facts() {
"300".to_owned(),
),
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-allowed".to_owned(), "false".to_owned()),
("x-codex-limit-reached".to_owned(), "true".to_owned()),
("x-codex-credits-has-credits".to_owned(), "true".to_owned()),
("x-codex-credits-unlimited".to_owned(), "false".to_owned()),
("x-codex-plan-type".to_owned(), "plus".to_owned()),
(
"x-codex-rate-limit-reached-type".to_owned(),
"usage_limit_reached".to_owned(),
),
];
let parsed = parse_rate_limit_headers(&source).expect("initial parse");
let encoded = rate_limits_to_header_pairs(&parsed);
@@ -1992,27 +1992,35 @@ async fn provider_account_metrics(
observed_at: DateTime<Utc>,
) -> StoreResult<ProviderAccountMetrics> {
let row = sqlx::query(
"select count(*)::bigint as total,
"with normalized as (
select enabled,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
when availability = 'quota_exhausted'
or (
availability = 'cooldown'
and availability_reason = 'usage_limit_exhausted'
and (cooldown_until is null or cooldown_until > $1)
)
then 'quota_exhausted'
when availability in ('expired', 'invalid')
or (access_token_expires_at is not null and access_token_expires_at <= $1)
then 'expired'
else 'active'
end as status
from provider_accounts
)
select count(*)::bigint as total,
count(*) filter (where enabled)::bigint as enabled,
count(*) filter (
where not (
enabled and availability = 'ready'
and (access_token_expires_at is null or access_token_expires_at > $1)
)
)::bigint as unavailable,
count(*) filter (
where enabled and availability = 'ready'
and (access_token_expires_at is null or access_token_expires_at > $1)
)::bigint as active,
count(*) filter (
where availability = 'expired'
or (access_token_expires_at is not null and access_token_expires_at <= $1)
)::bigint as expired,
count(*) filter (where availability = 'quota_exhausted')::bigint
as quota_exhausted,
count(*) filter (where not enabled)::bigint as disabled,
count(*) filter (where availability = 'banned')::bigint as banned
from provider_accounts",
count(*) filter (where status in ('expired', 'disabled', 'banned'))::bigint
as unavailable,
count(*) filter (where status = 'active')::bigint as active,
count(*) filter (where status = 'expired')::bigint as expired,
count(*) filter (where status = 'quota_exhausted')::bigint as quota_exhausted,
count(*) filter (where status = 'disabled')::bigint as disabled,
count(*) filter (where status = 'banned')::bigint as banned
from normalized",
)
.bind(observed_at)
.fetch_one(pool)
@@ -2035,11 +2043,29 @@ async fn active_provider_account_ids(
observed_at: DateTime<Utc>,
) -> StoreResult<Vec<String>> {
sqlx::query(
"select id
from provider_accounts
where enabled and availability = 'ready'
and (access_token_expires_at is null or access_token_expires_at > $1)
order by id",
"with normalized as (
select id,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
when availability = 'quota_exhausted'
or (
availability = 'cooldown'
and availability_reason = 'usage_limit_exhausted'
and (cooldown_until is null or cooldown_until > $1)
)
then 'quota_exhausted'
when availability in ('expired', 'invalid')
or (access_token_expires_at is not null and access_token_expires_at <= $1)
then 'expired'
else 'active'
end as status
from provider_accounts
)
select id
from normalized
where status = 'active'
order by id",
)
.bind(observed_at)
.fetch_all(pool)
@@ -38,8 +38,8 @@ use gateway_core::engine::credential::{
CredentialRevision as CoreCredentialRevision, LoadedCredential,
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
ProviderAccountStore, ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation,
QuotaWriteOutcome,
ProviderAccountIdentity, ProviderAccountStore,
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome,
};
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
use gateway_core::routing::ProviderKind;
@@ -125,6 +125,26 @@ fn parse_availability(value: &str) -> StoreResult<AccountAvailability> {
AccountAvailability::parse(value).ok_or_else(|| invalid("unknown availability value"))
}
fn validate_availability_cooldown(
availability: AccountAvailability,
has_cooldown_until: bool,
) -> StoreResult<()> {
if availability == AccountAvailability::Cooldown && !has_cooldown_until {
return Err(invalid("cooldown availability requires cooldown_until"));
}
if has_cooldown_until
&& !matches!(
availability,
AccountAvailability::Cooldown | AccountAvailability::QuotaExhausted
)
{
return Err(invalid(
"cooldown_until is only allowed for cooldown or quota_exhausted availability",
));
}
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountSummary {
pub id: String,
@@ -224,11 +244,7 @@ impl NewProviderAccount {
if !self.has_refresh_token && self.next_refresh_at.is_some() {
return Err(invalid("next_refresh_at requires a refresh token"));
}
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown_until must be present exactly for cooldown availability",
));
}
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
validate_object_size(
"provider_credentials_json",
&self.provider_credentials_json,
@@ -296,6 +312,7 @@ impl ImportProviderAccounts {
pub struct RotateProviderAccount {
pub scope: ProviderAccountAdminScope,
pub profile: UpdateProviderAccount,
pub replacement_identity: Option<ProviderAccountIdentity>,
pub credential: ProviderCredentialUpdate,
pub audit: AdminAuditEvent,
}
@@ -306,6 +323,7 @@ impl fmt::Debug for RotateProviderAccount {
.debug_struct("RotateProviderAccount")
.field("scope", &self.scope)
.field("profile", &self.profile)
.field("replacement_identity", &self.replacement_identity)
.field("credential", &self.credential)
.field("audit", &self.audit)
.finish()
@@ -378,11 +396,7 @@ pub struct ProviderAccountStateUpdate {
impl ProviderAccountStateUpdate {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown availability and cooldown_until must agree",
));
}
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
Ok(())
}
}
@@ -390,11 +404,7 @@ impl ProviderAccountStateUpdate {
impl ProviderAccountObservation {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown availability and cooldown_until must agree",
));
}
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
if self.provider_quota_json.is_some() != self.quota_observed_at.is_some() {
return Err(invalid("quota JSON and quota_observed_at must agree"));
}
@@ -612,8 +622,14 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
update.validate()?;
let result = sqlx::query(
"update provider_accounts
set availability = $3, availability_reason = $4, cooldown_until = $5,
availability_observed_at = $6, updated_at = greatest(now(), $6)
set availability = case when enabled then $3 else availability end,
availability_reason = case when enabled then $4 else availability_reason end,
cooldown_until = case when enabled then $5 else cooldown_until end,
availability_observed_at = case
when enabled then $6
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",
)
.bind(update.account_id)
@@ -755,6 +771,12 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
command.scope.validate()?;
require_nonempty(ENTITY, "account_id", &command.profile.id)?;
require_nonempty(ENTITY, "name", &command.profile.name)?;
if let Some(identity) = &command.replacement_identity {
require_nonempty(ENTITY, "upstream_user_id", identity.upstream_user_id())?;
if let Some(account_id) = identity.upstream_account_id() {
require_nonempty(ENTITY, "upstream_account_id", account_id)?;
}
}
if command.profile.id != command.credential.account_id {
return Err(invalid("rotated profile and credential account IDs differ"));
}
@@ -770,6 +792,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
&mut transaction,
&command.scope,
&command.profile,
command.replacement_identity.as_ref(),
&command.credential,
)
.await?;
@@ -1040,6 +1063,7 @@ impl PgAdminAccountStore {
email: prepared.email,
plan_type: prepared.plan_type,
},
replacement_identity: prepared.replacement_identity,
credential: ProviderCredentialUpdate {
account_id: account_id.as_str().to_owned(),
expected_revision: store_revision(prepared.expected_credential_revision)?,
@@ -1975,8 +1999,13 @@ async fn rotate_provider_account_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
scope: &ProviderAccountAdminScope,
profile: &UpdateProviderAccount,
replacement_identity: Option<&ProviderAccountIdentity>,
update: &ProviderCredentialUpdate,
) -> StoreResult<Revision> {
let replace_identity = replacement_identity.is_some();
let upstream_user_id = replacement_identity.map(ProviderAccountIdentity::upstream_user_id);
let upstream_account_id =
replacement_identity.and_then(ProviderAccountIdentity::upstream_account_id);
let next = sqlx::query_scalar::<_, i64>(
"update provider_accounts
set name = $4,
@@ -1987,19 +2016,25 @@ async fn rotate_provider_account_in_transaction(
has_refresh_token = $8,
access_token_expires_at = $9,
next_refresh_at = $10,
upstream_user_id = case when $11::boolean then $12::text else upstream_user_id end,
upstream_account_id = case when $11::boolean then $13::text else upstream_account_id end,
availability = case
when not enabled then availability
when availability <> 'quota_exhausted' then 'ready'
else availability
end,
availability_reason = case
when not enabled then availability_reason
when availability <> 'quota_exhausted' then null
else availability_reason
end,
cooldown_until = case
when not enabled then cooldown_until
when availability <> 'quota_exhausted' then null
else cooldown_until
end,
availability_observed_at = case
when not enabled then availability_observed_at
when availability <> 'quota_exhausted' then now()
else availability_observed_at
end,
@@ -2018,9 +2053,25 @@ async fn rotate_provider_account_in_transaction(
.bind(update.has_refresh_token)
.bind(update.access_token_expires_at)
.bind(update.next_refresh_at)
.bind(replace_identity)
.bind(upstream_user_id)
.bind(upstream_account_id)
.fetch_optional(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("rotate provider account in admin transaction"))?
.map_err(|error| {
if error
.as_database_error()
.is_some_and(sqlx::error::DatabaseError::is_unique_violation)
{
StoreError::Conflict {
entity: ENTITY,
id: update.account_id.clone(),
kind: ConflictKind::InvalidTransition,
}
} else {
postgres_unavailable("rotate provider account in admin transaction")
}
})?
.ok_or_else(|| StoreError::Conflict {
entity: ENTITY,
id: update.account_id.clone(),
@@ -504,6 +504,12 @@ impl ProviderAccountStore for CooldownCachingProviderAccountStore {
let availability = change.availability;
let cooldown_until = change.cooldown_until;
self.authoritative.apply_state_change(change).await?;
let Some(account) = self.authoritative.get_account(&account_id).await? else {
return Ok(());
};
if !account.enabled() {
return Ok(());
}
self.mirror(
&account_id,
credential_revision,
@@ -76,6 +76,12 @@ impl TestDatabase {
.execute(&mut *migration)
.await
.expect("apply model request service tier migration");
sqlx::raw_sql(include_str!(
"../../../../migrations/0005_allow_quota_exhausted_cooldown.sql"
))
.execute(&mut *migration)
.await
.expect("apply quota exhausted cooldown migration");
migration.commit().await.expect("commit terminal migration");
Some(Self {
admin,
@@ -171,7 +177,7 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
.expect("drop migration test database");
admin.close().await;
assert_eq!((first_table_count, migration_count), (8, 4));
assert_eq!((first_table_count, migration_count), (8, 5));
assert_eq!(response_id_types, ["bytea", "bytea"]);
assert!(!raw_response_id_index_exists);
}
@@ -183,6 +189,7 @@ fn migrations_should_leave_transaction_ownership_to_sqlx() {
include_str!("../../../../migrations/0002_drop_opaque_response_id_index.sql"),
include_str!("../../../../migrations/0003_store_opaque_response_ids_as_bytes.sql"),
include_str!("../../../../migrations/0004_add_model_request_service_tier.sql"),
include_str!("../../../../migrations/0005_allow_quota_exhausted_cooldown.sql"),
]
.into_iter()
.flat_map(str::lines)
@@ -122,16 +122,27 @@ async fn dashboard_account_metrics_should_partition_available_and_unavailable_ac
id, provider_kind, name, upstream_user_id, authentication_kind,
provider_credentials_json, credential_revision, has_refresh_token,
access_token_expires_at, enabled, availability,
availability_reason, cooldown_until,
availability_observed_at, created_at, updated_at
) values
('acct_available', 'openai', 'available', 'user-available', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', $1, $1, $1),
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready',
null, null, $1, $1, $1),
('acct_expired', 'openai', 'expired', 'user-expired', 'oauth',
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready', $1, $1, $1),
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready',
null, null, $1, $1, $1),
('acct_rate_cooldown', 'openai', 'rate-cooldown', 'user-rate-cooldown', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'cooldown',
'rate_limited', $1 + interval '2 hours', $1, $1, $1),
('acct_usage_limit', 'openai', 'usage-limit', 'user-usage-limit', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'cooldown',
'usage_limit_exhausted', $1 + interval '2 hours', $1, $1, $1),
('acct_banned', 'xai', 'banned', 'user-banned', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned', $1, $1, $1),
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned',
null, null, $1, $1, $1),
('acct_disabled', 'xai', 'disabled', 'user-disabled', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', $1, $1, $1)",
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready',
null, null, $1, $1, $1)",
)
.bind(now)
.execute(&database.pool)
@@ -154,10 +165,11 @@ async fn dashboard_account_metrics_should_partition_available_and_unavailable_ac
metrics.active,
metrics.unavailable,
metrics.expired,
metrics.quota_exhausted,
metrics.disabled,
metrics.banned,
),
(4, 3, 1, 3, 1, 1, 1),
(6, 5, 2, 3, 1, 1, 1, 1),
);
database.close().await;
}
@@ -17,9 +17,10 @@ use gateway_admin::{
ports::store::AccountStore,
};
use gateway_core::engine::credential::{
AccountAvailability, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision,
OpaqueProviderData, PlaintextCredential, ProviderAccountId, ProviderAccountStore,
ProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome,
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision, OpaqueProviderData, PlaintextCredential, ProviderAccountId,
ProviderAccountIdentity, ProviderAccountStore, ProviderAccountUpdate, QuotaObservation,
QuotaWriteOutcome,
};
use gateway_core::routing::ProviderKind;
use gateway_store::{
@@ -898,6 +899,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
.rotate_provider_account(RotateProviderAccount {
scope: wrong_scope,
profile: profile("acct_admin_a", "wrong scope"),
replacement_identity: None,
credential: credential_update("acct_admin_a", 1, "wrong-scope-secret"),
audit: audit("audit_wrong_scope", "rotate", "acct_admin_a"),
})
@@ -926,6 +928,10 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
.rotate_provider_account(RotateProviderAccount {
scope: scope.clone(),
profile: profile("acct_admin_a", "rotated account"),
replacement_identity: Some(ProviderAccountIdentity::new(
"user-admin-rebound".to_owned(),
Some("workspace-admin-rebound".to_owned()),
)),
credential: credential_update("acct_admin_a", 1, "rotated-secret"),
audit: audit("audit_account_rotate", "rotate", "acct_admin_a"),
})
@@ -933,15 +939,71 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
.expect("rotate provider account");
assert_eq!(rotation.config_revision.get(), 3);
assert_eq!(rotation.credential_revision.get(), 2);
let restored: (String, Option<String>, Option<chrono::DateTime<Utc>>) = sqlx::query_as(
"select availability, availability_reason, cooldown_until
let restored: (
String,
Option<String>,
Option<chrono::DateTime<Utc>>,
String,
Option<String>,
) = sqlx::query_as(
"select availability, availability_reason, cooldown_until,
upstream_user_id, upstream_account_id
from provider_accounts where id = $1",
)
.bind("acct_admin_a")
.fetch_one(&database.pool)
.await
.expect("load restored account state");
assert_eq!(restored, ("ready".to_owned(), None, None));
assert_eq!(
restored,
(
"ready".to_owned(),
None,
None,
"user-admin-rebound".to_owned(),
Some("workspace-admin-rebound".to_owned()),
)
);
let identity_conflict = repository
.rotate_provider_account(RotateProviderAccount {
scope: scope.clone(),
profile: profile("acct_admin_a", "must roll back"),
replacement_identity: Some(ProviderAccountIdentity::new(
"user-admin-b".to_owned(),
None,
)),
credential: credential_update("acct_admin_a", 2, "must-not-persist"),
audit: audit(
"audit_account_identity_conflict",
"reauthorize",
"acct_admin_a",
),
})
.await
.expect_err("upstream identity collision must roll back rotation");
assert!(matches!(
identity_conflict,
StoreError::Conflict {
kind: ConflictKind::InvalidTransition,
..
}
));
assert_eq!(current_revision(&database.pool).await, 3);
let unchanged: (String, serde_json::Value, i64, String, Option<String>) = sqlx::query_as(
"select name, provider_credentials_json, credential_revision,
upstream_user_id, upstream_account_id
from provider_accounts where id = $1",
)
.bind("acct_admin_a")
.fetch_one(&database.pool)
.await
.expect("load account after identity collision");
assert_eq!(unchanged.0, "rotated account");
assert_eq!(unchanged.1["access_token"], "rotated-secret");
assert_eq!(unchanged.2, 2);
assert_eq!(unchanged.3, "user-admin-rebound");
assert_eq!(unchanged.4.as_deref(), Some("workspace-admin-rebound"));
let revision = repository
.set_provider_account_enabled_admin(SetProviderAccountEnabled {
@@ -1013,6 +1075,7 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
.rotate_provider_account(RotateProviderAccount {
scope,
profile: profile("acct_rotation_quota", "reauthorized account"),
replacement_identity: None,
credential: credential_update("acct_rotation_quota", 1, "reauthorized-secret"),
audit: audit(
"audit_rotation_quota_rotate",
@@ -1041,6 +1104,77 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
database.close().await;
}
#[tokio::test]
async fn disabled_account_preserves_user_state_during_refresh_writes() {
let Some(database) = TestDatabase::create("provider_account_disabled_refresh").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let scope = ProviderAccountAdminScope {
provider_kind: "openai".to_owned(),
};
let account_id = ProviderAccountId::new("acct_disabled_refresh").expect("account ID");
let mut disabled = account(account_id.as_str(), "user-disabled-refresh");
disabled.enabled = false;
disabled.availability = AccountAvailability::Expired;
disabled.availability_reason = Some("credential_expired".to_owned());
repository
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![disabled],
audit: audit(
"audit_disabled_refresh_import",
"import",
account_id.as_str(),
),
})
.await
.expect("import disabled account");
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
availability: AccountAvailability::Ready,
reason: None,
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("disabled state write is a no-op");
repository
.rotate_provider_account(RotateProviderAccount {
scope,
profile: profile(account_id.as_str(), "refreshed disabled account"),
replacement_identity: None,
credential: credential_update(account_id.as_str(), 1, "disabled-refreshed-secret"),
audit: audit(
"audit_disabled_refresh_rotate",
"refresh",
account_id.as_str(),
),
})
.await
.expect("refresh disabled account credential");
let current: (bool, String, Option<String>, serde_json::Value, i64) = sqlx::query_as(
"select enabled, availability, availability_reason,
provider_credentials_json, credential_revision
from provider_accounts where id = $1",
)
.bind(account_id.as_str())
.fetch_one(&database.pool)
.await
.expect("load disabled account after refresh writes");
assert!(!current.0);
assert_eq!(current.1, "expired");
assert_eq!(current.2.as_deref(), Some("credential_expired"));
assert_eq!(current.3["access_token"], "disabled-refreshed-secret");
assert_eq!(current.4, 2);
database.close().await;
}
fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
NewProviderAccount {
id: id.to_owned(),
@@ -1200,7 +1334,15 @@ fn cooldown_observation_requires_expiry() {
}
#[test]
fn imported_account_rejects_mismatched_cooldown_runtime_facts() {
fn quota_exhausted_account_accepts_cooldown_runtime_facts() {
let mut imported = account("acct_quota_cooldown", "user-quota-cooldown");
imported.availability = AccountAvailability::QuotaExhausted;
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
assert!(imported.validate().is_ok());
}
#[test]
fn imported_account_rejects_cooldown_until_for_ready_runtime_facts() {
let mut imported = account("acct_invalid_cooldown", "user-invalid-cooldown");
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
assert!(imported.validate().is_err());
@@ -50,9 +50,13 @@ impl ProviderAccountStore for CooldownTestAccountStore {
async fn get_account(
&self,
_account: &ProviderAccountId,
account: &ProviderAccountId,
) -> Result<Option<ProviderAccount>, CoreStoreError> {
unreachable!("cooldown tests do not get one account")
Ok(self
.accounts
.iter()
.find(|candidate| candidate.id() == account)
.cloned())
}
async fn list_accounts(&self) -> Result<Vec<ProviderAccount>, CoreStoreError> {
@@ -190,10 +194,26 @@ fn composition_fixture(
CooldownCachingProviderAccountStore,
Arc<CooldownTestCache>,
Arc<Mutex<Vec<&'static str>>>,
) {
composition_fixture_with_accounts(
vec![ready_test_account("acct_cooldown_composition")],
fail_state_write,
fail_cache,
)
}
fn composition_fixture_with_accounts(
accounts: Vec<ProviderAccount>,
fail_state_write: bool,
fail_cache: bool,
) -> (
CooldownCachingProviderAccountStore,
Arc<CooldownTestCache>,
Arc<Mutex<Vec<&'static str>>>,
) {
let calls = Arc::new(Mutex::new(Vec::new()));
let authoritative = Arc::new(CooldownTestAccountStore {
accounts: Vec::new(),
accounts,
calls: Arc::clone(&calls),
fail_state_write,
});
@@ -249,6 +269,31 @@ async fn cooldown_composition_keeps_postgres_terminal_when_redis_fails() {
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
}
#[tokio::test]
async fn disabled_account_state_write_does_not_change_runtime_cooldown() {
let disabled = test_provider_account("acct_cooldown_composition").with_runtime_state(
false,
AccountAvailability::Ready,
None,
);
let (store, cache, calls) = composition_fixture_with_accounts(vec![disabled], false, false);
store
.apply_state_change(cooldown_state_change(60))
.await
.expect("disabled state write is a no-op");
assert_eq!(*calls.lock().expect("calls lock"), ["postgres"]);
assert!(cache.cached.lock().expect("cached lock").is_empty());
assert!(
cache
.invalidated
.lock()
.expect("invalidated lock")
.is_empty()
);
}
#[tokio::test]
async fn ready_state_invalidates_same_revision_cooldown_after_postgres_commit() {
let (store, cache, calls) = composition_fixture(false, false);
+2 -1
View File
@@ -586,7 +586,7 @@ fn prepared_rotation(
prepared: crate::credential::PreparedCodexCredentialRotation,
provider_kind: ProviderKind,
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
let (profile, credential, guard) = prepared.into_parts();
let (profile, credential, replacement_identity, guard) = prepared.into_parts();
let (
account_id,
expected_revision,
@@ -606,6 +606,7 @@ fn prepared_rotation(
account_id,
provider_kind,
expected_credential_revision,
replacement_identity,
name: profile.name,
email: profile.email,
plan_type: profile.plan_type,
@@ -11,7 +11,8 @@ use std::time::SystemTime;
use chrono::{DateTime, FixedOffset, Utc};
use gateway_core::engine::credential::{
AccountAvailability, CredentialCasUpdate, CredentialRevision, LoadedCredential,
NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountUpdate,
NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountIdentity,
ProviderAccountUpdate,
};
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest,
@@ -286,6 +287,7 @@ impl std::fmt::Debug for RotateManagedCodexCredential {
pub struct PreparedCodexCredentialRotation {
pub profile: ProviderAccountUpdate,
pub credential: CredentialCasUpdate,
pub replacement_identity: Option<ProviderAccountIdentity>,
refresh_guards: Option<ProviderRefreshGuards>,
}
@@ -300,6 +302,7 @@ impl std::fmt::Debug for PreparedCodexCredentialRotation {
.debug_struct("PreparedCodexCredentialRotation")
.field("profile", &self.profile)
.field("credential", &self.credential)
.field("replacement_identity", &self.replacement_identity)
.field(
"refresh_guards",
&self.refresh_guards.as_ref().map(|_| "<held>"),
@@ -316,11 +319,13 @@ impl PreparedCodexCredentialRotation {
) -> (
ProviderAccountUpdate,
CredentialCasUpdate,
Option<ProviderAccountIdentity>,
PreparedCodexCredentialRotationGuard,
) {
(
self.profile,
self.credential,
self.replacement_identity,
PreparedCodexCredentialRotationGuard(self.refresh_guards),
)
}
@@ -494,7 +499,7 @@ impl CodexCredentialAdmin {
self.prepare_oauth_rotation(input, false)
}
/// 完整 OAuth 授权允许认证主体投影更新,但不允许目标账号或用户发生变化。
/// 完整 OAuth 授权按目标数据 ID 替换凭据与上游账号身份。
pub(crate) fn prepare_reauthorization(
&self,
input: RotateManagedCodexCredential,
@@ -505,7 +510,7 @@ impl CodexCredentialAdmin {
fn prepare_oauth_rotation(
&self,
input: RotateManagedCodexCredential,
replace_principal: bool,
replace_identity: bool,
) -> Result<PreparedCodexCredentialRotation, CodexCredentialAdminError> {
let access_token_expires_at =
required_time(input.verified_account.access_token_expires_at)?;
@@ -516,16 +521,23 @@ impl CodexCredentialAdmin {
.ok_or(CodexCredentialAdminError::IdentityMismatch)?;
if input.current.account.provider().as_str() != PROVIDER_NAME
|| input.current.account.authentication_kind() != CODEX_AUTHENTICATION_KIND_OAUTH
|| input.current.account.upstream_account_id()
!= Some(input.verified_account.chatgpt_account_id.as_str())
|| input.current.account.upstream_user_id() != input.verified_account.chatgpt_user_id
|| (!replace_principal
&& (oauth.principal.oauth_subject != input.verified_account.oauth_subject
|| (!replace_identity
&& (input.current.account.upstream_account_id()
!= Some(input.verified_account.chatgpt_account_id.as_str())
|| input.current.account.upstream_user_id()
!= input.verified_account.chatgpt_user_id
|| oauth.principal.oauth_subject != input.verified_account.oauth_subject
|| oauth.principal.poid != input.verified_account.poid))
{
return Err(CodexCredentialAdminError::IdentityMismatch);
}
if replace_principal {
let replacement_identity = replace_identity.then(|| {
ProviderAccountIdentity::new(
input.verified_account.chatgpt_user_id.clone(),
Some(input.verified_account.chatgpt_account_id.clone()),
)
});
if replace_identity {
oauth
.principal
.oauth_subject
@@ -567,6 +579,7 @@ impl CodexCredentialAdmin {
Ok(PreparedCodexCredentialRotation {
profile,
credential,
replacement_identity,
refresh_guards: None,
})
}
@@ -29,7 +29,6 @@ use super::admin::{
use super::identity::{
CodexAccountIdentityVerifier, CodexIdentityExpectation, CodexIdentityVerificationError,
};
use super::security::CodexCredentialCodec;
use super::token_client::{
AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant,
OFFICIAL_CODEX_OAUTH_CLIENT_ID, OFFICIAL_CODEX_REDIRECT_URI,
@@ -500,11 +499,7 @@ impl CodexOAuthAdminService {
} else {
None
};
let expectation = current
.as_ref()
.map(oauth_identity_expectation)
.transpose()?
.unwrap_or_default();
let expectation = CodexIdentityExpectation::default();
let profile = self
.verifier
.verify_authorization(&secret, &id_token, &pending.nonce, &expectation)
@@ -679,26 +674,6 @@ impl CodexOAuthAdminService {
}
}
fn oauth_identity_expectation(
current: &gateway_core::engine::credential::LoadedCredential,
) -> Result<CodexIdentityExpectation, CodexOAuthAdminError> {
let runtime = CodexCredentialCodec::decode(&current.credential)
.map_err(|_| CodexOAuthAdminError::Credential)?;
if runtime.authentication.oauth().is_none() {
return Err(CodexOAuthAdminError::Credential);
}
let account_id = current
.account
.upstream_account_id()
.ok_or(CodexOAuthAdminError::Credential)?;
CodexIdentityExpectation::reauthorization(
account_id.to_owned(),
current.account.upstream_user_id().to_owned(),
runtime.installation_id,
)
.map_err(|_| CodexOAuthAdminError::Credential)
}
fn authorization_url(
pending: &CodexPendingAuthorization,
oauth_client_id: &str,
@@ -30,6 +30,7 @@ const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60);
const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
const QUOTA_RESET_RECHECK_GRACE: Duration = Duration::from_secs(30);
// ready 账号主要依赖真实请求的被动观察;这里只提供低频、定额的状态兜底。
const QUOTA_READY_REFRESH_MAX_AGE: Duration = Duration::from_secs(2 * 60 * 60);
const QUOTA_READY_REFRESH_BATCH_SIZE: usize = 8;
@@ -307,7 +308,7 @@ enum CodexQuotaFetchError {
enum CodexQuotaStateTransition {
Banned { reason: &'static str },
Exhausted,
Cooldown { until: Option<SystemTime> },
ResettableLimit { until: Option<SystemTime> },
}
impl CodexQuotaStateTransition {
@@ -315,7 +316,7 @@ impl CodexQuotaStateTransition {
match self {
Self::Banned { .. } => AccountAvailability::Banned,
Self::Exhausted => AccountAvailability::QuotaExhausted,
Self::Cooldown { .. } => AccountAvailability::Cooldown,
Self::ResettableLimit { .. } => AccountAvailability::QuotaExhausted,
}
}
@@ -323,13 +324,13 @@ impl CodexQuotaStateTransition {
match self {
Self::Banned { reason } => reason,
Self::Exhausted => "quota_exhausted",
Self::Cooldown { .. } => "quota_rate_limited",
Self::ResettableLimit { .. } => "quota_rate_limited",
}
}
const fn cooldown_until(self) -> Option<SystemTime> {
match self {
Self::Cooldown { until } => until,
Self::ResettableLimit { until } => until,
Self::Banned { .. } | Self::Exhausted => None,
}
}
@@ -547,7 +548,9 @@ fn quota_refresh_candidate(
let snapshot = quota_snapshot_for_account(&account, observation);
let kind = match account.availability() {
AccountAvailability::QuotaExhausted => CodexQuotaRefreshKind::Exhausted {
reset_at: quota_reset_at_from_snapshot(snapshot.as_ref()),
reset_at: account
.cooldown_until()
.or_else(|| quota_reset_at_from_snapshot(snapshot.as_ref())),
},
AccountAvailability::Cooldown
if snapshot
@@ -618,22 +621,24 @@ fn reserve_exhausted_quota_refresh(
now: SystemTime,
refreshed_at: Instant,
) -> bool {
let reset_recheck_due = reset_at.is_some_and(|reset_at| {
reset_at <= now
&& state
.last_exhausted_reset_recheck
.get(account.id())
.is_none_or(|marker| {
marker.revision != account.revision() || marker.reset_at != reset_at
})
});
if !reset_recheck_due && !periodic_quota_refresh_due(state, account.id(), refreshed_at) {
return false;
}
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
if let Some(reset_at) = reset_at.filter(|reset_at| *reset_at <= now) {
if let Some(reset_at) = reset_at {
if !quota_reset_recheck_ready(reset_at, now) {
return false;
}
let same_reset_already_checked = state
.last_exhausted_reset_recheck
.get(account.id())
.is_some_and(|marker| {
marker.revision == account.revision() && marker.reset_at == reset_at
});
if same_reset_already_checked
&& !periodic_quota_refresh_due(state, account.id(), refreshed_at)
{
return false;
}
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
state.last_exhausted_reset_recheck.insert(
account.id().clone(),
CodexQuotaResetMarker {
@@ -641,10 +646,23 @@ fn reserve_exhausted_quota_refresh(
reset_at,
},
);
return true;
}
if !periodic_quota_refresh_due(state, account.id(), refreshed_at) {
return false;
}
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
true
}
fn quota_reset_recheck_ready(reset_at: SystemTime, now: SystemTime) -> bool {
reset_at
.checked_add(QUOTA_RESET_RECHECK_GRACE)
.is_some_and(|ready_at| ready_at <= now)
}
fn periodic_quota_refresh_due(
state: &CodexQuotaProjectionState,
account_id: &ProviderAccountId,
@@ -802,16 +820,12 @@ impl CodexCredentialQuotaService {
summary.transient += 1;
continue;
};
let preserve_exhaustion = account.availability()
== AccountAvailability::QuotaExhausted
&& 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;
CodexQuotaStateTransition::Exhausted
| CodexQuotaStateTransition::ResettableLimit { .. } => {
summary.exhausted += 1;
}
}
let _ = self
.repository
@@ -856,19 +870,23 @@ impl CodexCredentialQuotaService {
{
continue;
}
if !matches!(
self.has_active_exhaustion(&current, observed_at).await,
Ok(true)
) {
continue;
}
let snapshot = match self.read_snapshot_for(&current).await {
Ok(Some(snapshot))
if quota_fact_is_active_exhaustion(snapshot.fact(), observed_at) =>
{
snapshot
}
_ => continue,
};
let cooldown_until =
quota_reset_at_from_snapshot(Some(&snapshot)).or_else(|| current.cooldown_until());
if let Err(error) = self
.repository
.apply_state(
&current,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await
@@ -881,8 +899,11 @@ impl CodexCredentialQuotaService {
continue;
}
let enabled = current.enabled();
*account =
current.with_runtime_state(enabled, AccountAvailability::QuotaExhausted, None);
*account = current.with_runtime_state(
enabled,
AccountAvailability::QuotaExhausted,
cooldown_until,
);
}
}
@@ -927,10 +948,10 @@ impl CodexCredentialQuotaService {
} else {
summary.updated += 1;
}
if let Some(availability) = quota_success_availability(
if let Some((availability, cooldown_until)) = quota_success_state(
current.availability(),
current.cooldown_until(),
fact.exhausted(),
fact,
SystemTime::now(),
) {
let _ = self
@@ -939,7 +960,7 @@ impl CodexCredentialQuotaService {
&current,
availability,
fact.exhausted().then_some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await;
@@ -993,10 +1014,10 @@ impl CodexCredentialQuotaService {
}
self.scheduling
.observe(current.id().clone(), current.revision(), observed_at, fact);
if let Some(availability) = quota_success_availability(
if let Some((availability, cooldown_until)) = quota_success_state(
current.availability(),
current.cooldown_until(),
fact.exhausted(),
fact,
observed_at,
) {
self.repository
@@ -1004,7 +1025,7 @@ impl CodexCredentialQuotaService {
&current,
availability,
fact.exhausted().then_some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await?;
@@ -1026,19 +1047,6 @@ impl CodexCredentialQuotaService {
self.read_snapshot_for(&account).await
}
pub(super) async fn has_active_exhaustion(
&self,
account: &ProviderAccount,
now: SystemTime,
) -> Result<bool, CodexCredentialQuotaError> {
match self.read_snapshot_for(account).await {
Ok(snapshot) => Ok(snapshot
.is_some_and(|snapshot| quota_fact_is_active_exhaustion(snapshot.fact(), now))),
Err(CodexCredentialQuotaError::RevisionConflict) => Ok(false),
Err(error) => Err(error),
}
}
async fn read_snapshot_for(
&self,
account: &ProviderAccount,
@@ -1164,10 +1172,10 @@ impl CodexCredentialQuotaService {
snapshot.observed_at(),
snapshot.fact(),
);
if let Some(availability) = quota_success_availability(
if let Some((availability, cooldown_until)) = quota_success_state(
current.availability(),
current.cooldown_until(),
snapshot.fact().exhausted(),
snapshot.fact(),
SystemTime::now(),
) {
self.repository
@@ -1178,7 +1186,7 @@ impl CodexCredentialQuotaService {
.fact()
.exhausted()
.then_some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await?;
@@ -1386,31 +1394,34 @@ fn passive_rate_limit_window(window: RateLimitWindow) -> Map<String, Value> {
snapshot
}
fn quota_success_availability(
fn quota_success_state(
current: AccountAvailability,
cooldown_until: Option<SystemTime>,
exhausted: bool,
fact: CodexQuotaFact,
now: SystemTime,
) -> Option<AccountAvailability> {
match current {
) -> Option<(AccountAvailability, Option<SystemTime>)> {
if fact.exhausted() {
let observed_cooldown = fact.resets_at().map(SystemTime::from).or(cooldown_until);
return (current != AccountAvailability::QuotaExhausted
|| cooldown_until != observed_cooldown)
.then_some((AccountAvailability::QuotaExhausted, observed_cooldown));
}
let availability = match current {
AccountAvailability::Invalid
| AccountAvailability::Expired
| AccountAvailability::Banned => None,
AccountAvailability::QuotaExhausted => (!exhausted).then_some(AccountAvailability::Ready),
AccountAvailability::Ready => exhausted.then_some(AccountAvailability::QuotaExhausted),
AccountAvailability::Cooldown if exhausted => Some(AccountAvailability::QuotaExhausted),
| AccountAvailability::Banned
| AccountAvailability::Unknown
| AccountAvailability::QuotaExhausted => Some(AccountAvailability::Ready),
AccountAvailability::Ready => None,
AccountAvailability::Cooldown
if cooldown_until.is_some_and(|cooldown_until| cooldown_until <= now) =>
{
Some(AccountAvailability::Ready)
}
AccountAvailability::Cooldown => None,
AccountAvailability::Unknown => Some(if exhausted {
AccountAvailability::QuotaExhausted
} else {
AccountAvailability::Ready
}),
}
};
availability.map(|availability| (availability, None))
}
fn quota_state_transition(
@@ -1438,7 +1449,7 @@ fn quota_state_transition(
retry_after_seconds.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN.as_secs()),
)
.min(MAX_RATE_LIMIT_COOLDOWN);
Some(CodexQuotaStateTransition::Cooldown {
Some(CodexQuotaStateTransition::ResettableLimit {
until: observed_at.checked_add(duration),
})
}
@@ -1876,7 +1887,6 @@ impl QuotaAggregate {
self.remaining_percent
.map_or(remaining, |current| current.min(remaining)),
);
self.exhausted |= used >= 100.0;
}
if let Some(reset) = object.get("reset_at") {
let seconds = reset
@@ -256,8 +256,6 @@ impl CodexCredentialRefreshService {
| AccountAvailability::Banned
| AccountAvailability::Invalid
)
&& (account.availability() != AccountAvailability::Cooldown
|| account.cooldown_until().is_some_and(|until| until <= now))
&& refresh_due_at(account, now).is_some()
});
accounts.sort_by_key(|account| (refresh_due_at(account, now), account.id().clone()));
@@ -237,6 +237,9 @@ impl CodexCredentialRepository {
cooldown_until: Option<SystemTime>,
observed_at: SystemTime,
) -> Result<(), CredentialRepositoryError> {
if !account.enabled() {
return Ok(());
}
self.store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
@@ -769,31 +769,31 @@ impl CodexCredentialSelector {
if current.revision() != account.revision() {
return Err(CredentialRepositoryError::RevisionConflict.into());
}
// 一次真实 429 证明上游已经接受当前凭据,可纠正旧的认证/封禁终态。
// 结构化额度耗尽仍拥有更高优先级,不能被通用限流降级。
if current.availability() == AccountAvailability::QuotaExhausted {
return Ok(());
}
let quota_exhausted = self
.quota
.has_active_exhaustion(&current, observed_at)
.await
.map_err(|_| CredentialSelectionError::Store)?;
let (availability, reason, cooldown_until) = if quota_exhausted {
(
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
)
} else {
(
AccountAvailability::Cooldown,
Some(reason.to_owned()),
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)),
)
let cooldown_until =
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN));
let cooldown_until = match (current.availability(), current.cooldown_until()) {
(AccountAvailability::QuotaExhausted, Some(existing)) => {
cooldown_until.map_or(Some(existing), |until| Some(existing.max(until)))
}
_ => cooldown_until,
};
let reason = if current.availability() == AccountAvailability::QuotaExhausted
&& current.cooldown_until().is_none()
{
Some("quota_exhausted".to_owned())
} else {
Some(reason.to_owned())
};
// 一次真实 429 证明上游已经接受当前凭据,可纠正旧的认证/封禁终态。
// resettable 用量限制沿用 v2 的 quota lock 语义,并保留/延长复核时间。
self.repository
.apply_state(&current, availability, reason, cooldown_until, observed_at)
.apply_state(
&current,
AccountAvailability::QuotaExhausted,
reason,
cooldown_until,
observed_at,
)
.await?;
Ok(())
}
+35 -14
View File
@@ -68,6 +68,7 @@ use crate::transport::request::{
scope_request_to_account,
};
use crate::transport::session::CodexSessionIdentity;
use crate::transport::usage::normalize_service_tier;
use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason};
use crate::transport::{
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
@@ -822,8 +823,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
// 下游按字节转发上游原文,避免 serde 往返改写数值/精度(大整数→f64、logprobs 等)。
// WS 帧由 reducer 以 encode_sse_event(&event, raw) 逐字节内嵌上游原始 JSON
// (transport/protocol/websocket.rs),push_frames 抽出的 data 即上游原文。
let mut decoder =
CodexCanonicalDecoder::new(upstream_model.as_str()).with_raw_sse_passthrough();
let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str())
.with_requested_service_tier(request.service_tier())
.with_raw_sse_passthrough();
let mut pre_commit_events = PreCommitClientEvents::new(stream_commit_policy);
loop {
let Some(stream_deadline) = remaining(context.deadline()) else {
@@ -895,7 +897,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
};
pre_commit_events.observe_chunk(chunk_len);
let service_tier_changed = observation_state
.observe_service_tier(decoder.response_service_tier());
.observe_upstream_service_tier(decoder.response_service_tier());
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
(
@@ -990,8 +992,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
)
});
let timing_signals = decoder.take_timing_signals();
let service_tier_changed =
observation_state.observe_service_tier(decoder.response_service_tier());
let service_tier_changed = observation_state
.observe_upstream_service_tier(decoder.response_service_tier());
let timing_changed = observation_state
.observe_timing_signals(timing_signals, context.timing_started_at());
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
@@ -1070,7 +1072,8 @@ struct OpenAiResponseObservationState {
requested_model: String,
stream: bool,
compact: bool,
service_tier: Option<String>,
requested_service_tier: Option<String>,
upstream_service_tier: Option<String>,
rate_limit_headers: Vec<(String, String)>,
account_id: String,
attempt_index: u32,
@@ -1102,8 +1105,8 @@ impl OpenAiResponseObservationState {
requested_model: request.model().to_owned(),
stream: request.stream(),
compact: semantics.compact,
// 请求档位保留在 requestSummary;一等观测只接受响应确认的实际档位。
service_tier: None,
requested_service_tier: normalize_service_tier(request.service_tier()),
upstream_service_tier: None,
rate_limit_headers: selected_observation_headers(&response.rate_limit_headers),
account_id: account_id.to_owned(),
attempt_index,
@@ -1127,7 +1130,7 @@ impl OpenAiResponseObservationState {
if let Some(metadata) = self.provider_metadata() {
observation = observation.with_provider_metadata(metadata);
}
if let Some(service_tier) = self.service_tier.as_deref() {
if let Some(service_tier) = self.effective_service_tier() {
observation = observation.with_service_tier_if_valid(service_tier.to_owned());
}
Some(observation)
@@ -1158,17 +1161,23 @@ impl OpenAiResponseObservationState {
changed
}
fn observe_service_tier(&mut self, service_tier: Option<&str>) -> bool {
let Some(service_tier) = service_tier else {
fn observe_upstream_service_tier(&mut self, service_tier: Option<&str>) -> bool {
let Some(service_tier) = normalize_service_tier(service_tier) else {
return false;
};
if self.service_tier.as_deref() == Some(service_tier) {
if self.upstream_service_tier.as_deref() == Some(service_tier.as_str()) {
return false;
}
self.service_tier = Some(service_tier.to_owned());
self.upstream_service_tier = Some(service_tier);
true
}
fn effective_service_tier(&self) -> Option<&str> {
self.requested_service_tier
.as_deref()
.or(self.upstream_service_tier.as_deref())
}
fn merge_rate_limit_headers(&mut self, updates: &[(String, String)]) -> bool {
let mut changed = false;
for (name, value) in selected_observation_headers(updates) {
@@ -1245,9 +1254,21 @@ impl OpenAiResponseObservationState {
)),
);
metadata.insert("requestSummary".to_owned(), self.request_summary.clone());
if let Some(service_tier) = &self.service_tier {
if let Some(service_tier) = self.effective_service_tier() {
metadata.insert(
"serviceTier".to_owned(),
Value::String(service_tier.to_owned()),
);
}
if let Some(service_tier) = &self.requested_service_tier {
metadata.insert(
"requestedServiceTier".to_owned(),
Value::String(service_tier.clone()),
);
}
if let Some(service_tier) = &self.upstream_service_tier {
metadata.insert(
"upstreamServiceTier".to_owned(),
Value::String(service_tier.clone()),
);
}
@@ -21,10 +21,9 @@ use thiserror::Error;
use super::protocol::responses::{
ResponseEventSignals, ResponsesSseFailure, response_event_signals,
};
use super::usage::openai_billing_breakdown;
use super::usage::{normalize_service_tier, openai_billing_breakdown};
const CONTENTS_PER_OUTPUT: u32 = 1_024;
const MAX_SERVICE_TIER_BYTES: usize = 64;
/// 单 attempt 的增量 Responses decoder。
///
@@ -42,6 +41,7 @@ pub struct CodexCanonicalDecoder {
reasoning_output_seen: BTreeSet<u32>,
usage_emitted: bool,
semantic_output_seen: bool,
requested_service_tier: Option<String>,
response_service_tier: Option<String>,
timing_signals: ResponseEventSignals,
raw_sse_passthrough: bool,
@@ -137,6 +137,7 @@ impl CodexCanonicalDecoder {
reasoning_output_seen: BTreeSet::new(),
usage_emitted: false,
semantic_output_seen: false,
requested_service_tier: None,
response_service_tier: None,
timing_signals: ResponseEventSignals::default(),
raw_sse_passthrough: false,
@@ -152,6 +153,13 @@ impl CodexCanonicalDecoder {
self
}
/// 请求显式指定的档位优先用于计费;响应档位只在请求未指定时兜底。
#[must_use]
pub fn with_requested_service_tier(mut self, service_tier: Option<&str>) -> Self {
self.requested_service_tier = normalize_service_tier(service_tier);
self
}
pub fn push(&mut self, chunk: &[u8]) -> CodexCanonicalOutcome {
self.timing_signals = ResponseEventSignals::default();
if self.raw_sse_passthrough {
@@ -308,19 +316,14 @@ impl CodexCanonicalDecoder {
}
fn observe_response_service_tier(&mut self, value: &Value) {
let Some(service_tier) = response_object(value)
.and_then(|response| response.get("service_tier"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| {
!value.is_empty()
&& value.len() <= MAX_SERVICE_TIER_BYTES
&& !value.chars().any(char::is_control)
})
else {
let Some(service_tier) = normalize_service_tier(
response_object(value)
.and_then(|response| response.get("service_tier"))
.and_then(Value::as_str),
) else {
return;
};
self.response_service_tier = Some(service_tier.to_ascii_lowercase());
self.response_service_tier = Some(service_tier);
}
/// 构造下发的 wire event。
@@ -836,9 +839,9 @@ impl CodexCanonicalDecoder {
.filter(|model| !model.is_empty())
.unwrap_or(&self.fallback_model)
.to_owned();
let service_tier = response
.get("service_tier")
.and_then(Value::as_str)
let service_tier = self
.requested_service_tier
.as_deref()
.or(self.response_service_tier.as_deref());
if let Some(breakdown) = usage
.filter(|usage| billable_usage_is_complete(response, *usage))
@@ -3,7 +3,8 @@ use std::{sync::Arc, time::Instant};
use futures::{StreamExt, TryStreamExt};
use gateway_protocol::openai::{
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
events::retry_after_seconds_from_body,
events::{self, retry_after_seconds_from_body},
sse::{SseEventDecoder, SseFrame},
};
use reqwest::{
Client, Response as ReqwestResponse,
@@ -130,13 +131,14 @@ impl CodexBackendClient {
});
}
let rate_limit_header_updates = Arc::new(tokio::sync::Mutex::new(Vec::new()));
Ok(CodexBackendStreamingResponse {
body: http_sse_stream(response),
body: http_sse_stream(response, Arc::clone(&rate_limit_header_updates)),
transport: CodexBackendTransport::HttpSse,
turn_state,
set_cookie_headers,
rate_limit_headers,
rate_limit_header_updates: None,
rate_limit_header_updates: Some(rate_limit_header_updates),
turn_state_update: None,
websocket_pool_decision: None,
diagnostics,
@@ -594,20 +596,74 @@ fn websocket_connection_profile(headers: &HeaderMap) -> String {
.join("\0")
}
fn http_sse_stream(response: ReqwestResponse) -> CodexBackendSseStream {
fn http_sse_stream(
response: ReqwestResponse,
rate_limit_updates: CodexRateLimitHeaderUpdates,
) -> CodexBackendSseStream {
let stream: CodexBackendSseStream =
Box::pin(response.bytes_stream().map_err(CodexClientError::Http));
Box::pin(futures::stream::unfold(Some(stream), |stream| async move {
let mut stream = stream?;
match tokio::time::timeout(UPSTREAM_STREAM_IDLE_TIMEOUT, stream.next()).await {
Ok(Some(chunk)) => Some((chunk, Some(stream))),
Ok(None) => None,
Err(_) => Some((
Err(CodexClientError::StreamIdleTimeout {
timeout: UPSTREAM_STREAM_IDLE_TIMEOUT,
}),
None,
)),
let stream: CodexBackendSseStream =
Box::pin(futures::stream::unfold(Some(stream), |stream| async move {
let mut stream = stream?;
match tokio::time::timeout(UPSTREAM_STREAM_IDLE_TIMEOUT, stream.next()).await {
Ok(Some(chunk)) => Some((chunk, Some(stream))),
Ok(None) => None,
Err(_) => Some((
Err(CodexClientError::StreamIdleTimeout {
timeout: UPSTREAM_STREAM_IDLE_TIMEOUT,
}),
None,
)),
}
}));
observe_http_sse_rate_limits(stream, rate_limit_updates)
}
fn observe_http_sse_rate_limits(
stream: CodexBackendSseStream,
updates: CodexRateLimitHeaderUpdates,
) -> CodexBackendSseStream {
Box::pin(futures::stream::unfold(
(stream, SseEventDecoder::default(), updates),
|(mut stream, mut decoder, updates)| async move {
match stream.next().await {
Some(chunk) => {
if let Ok(bytes) = &chunk {
append_http_sse_rate_limit_updates(decoder.push_frames(bytes), &updates)
.await;
}
Some((chunk, (stream, decoder, updates)))
}
None => {
append_http_sse_rate_limit_updates(decoder.finish_frames(), &updates).await;
None
}
}
},
))
}
async fn append_http_sse_rate_limit_updates(
frames: Vec<SseFrame>,
updates: &CodexRateLimitHeaderUpdates,
) {
let mut headers = Vec::new();
for frame in frames {
for event in frame.events() {
if event
.event
.as_deref()
.is_some_and(|event| event != "codex.rate_limits")
{
continue;
}
let Some(rate_limits) = events::parse_rate_limits_event_raw(&event.data) else {
continue;
};
headers.extend(events::rate_limits_to_header_pairs(&rate_limits));
}
}))
}
if !headers.is_empty() {
updates.lock().await.extend(headers);
}
}
@@ -298,10 +298,13 @@ fn model_matches_rule(model: &str, rule: &str) -> bool {
.is_some_and(|suffix| matches!(suffix.as_bytes().first(), Some(b'-' | b'.' | b':')))
}
fn normalize_service_tier(service_tier: Option<&str>) -> Option<String> {
/// 规范化请求或响应携带的服务档位,供观测与计费共用。
pub(crate) fn normalize_service_tier(service_tier: Option<&str>) -> Option<String> {
service_tier
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| {
!value.is_empty() && value.len() <= 64 && !value.chars().any(char::is_control)
})
.map(str::to_ascii_lowercase)
}
@@ -3,12 +3,10 @@ use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use chrono::Utc;
use futures::executor::block_on;
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountAvailability, AccountFeedbackStats, AccountSelectionPolicy,
AccountStateChange, OpaqueProviderData, ProviderAccountId, ProviderAccountStore as _,
QuotaObservation, RotationStrategy,
AccountStateChange, ProviderAccountId, ProviderAccountStore as _, RotationStrategy,
};
use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, RequestAttemptContext,
@@ -24,7 +22,6 @@ use provider_openai::credential::{
};
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use secrecy::ExposeSecret;
use serde_json::json;
use url::Url;
use crate::support::{
@@ -121,7 +118,7 @@ fn selector_with_runtime(
os_version: "6.8".to_owned(),
arch: "x86_64".to_owned(),
terminal: "selector-contract".to_owned(),
verified_at: Utc::now(),
verified_at: chrono::Utc::now(),
});
let http = reqwest::Client::builder().build().expect("HTTP client");
let agent_identity = agent_identity_service(store);
@@ -731,7 +728,7 @@ fn credential_expired_failure_marks_unified_account_expired() {
}
#[test]
fn rate_limited_failure_marks_account_cooldown_with_a_cooldown_deadline() {
fn rate_limited_failure_marks_quota_exhausted_with_a_reset_deadline() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
@@ -757,7 +754,7 @@ fn rate_limited_failure_marks_account_cooldown_with_a_cooldown_deadline() {
.expect("record rate-limit failure");
let account = store.account("acct_primary").expect("account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
@@ -765,28 +762,31 @@ fn rate_limited_failure_marks_account_cooldown_with_a_cooldown_deadline() {
);
}
#[test]
fn usage_limit_exhaustion_is_a_resettable_account_cooldown() {
#[tokio::test]
async fn usage_limit_exhaustion_marks_quota_exhausted_without_usage_probe() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&account,
CodexAccountFailure::UsageLimitExhausted {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record usage-limit exhaustion");
selector
.record_failure(
&account,
CodexAccountFailure::UsageLimitExhausted {
retry_after: Some(Duration::from_secs(30)),
},
)
.await
.expect("record usage-limit exhaustion");
let account = store.account("acct_primary").expect("persisted account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert_eq!(store.quota_reads(), 0);
}
#[test]
@@ -794,26 +794,15 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let reset_at = Utc::now().timestamp() + 24 * 60 * 60;
block_on(
store.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
}),
)
.expect("persist exhausted quota");
block_on(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(),
}))
.expect("mark exhausted");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
@@ -831,33 +820,54 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
.availability(),
AccountAvailability::QuotaExhausted
);
assert!(
store
.account("acct_primary")
.expect("persisted account")
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
}
#[test]
fn rate_limited_failure_ignores_exhaustion_after_its_reset_time() {
fn rate_limited_failure_keeps_longer_existing_quota_deadline() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let existing_deadline = SystemTime::now() + Duration::from_secs(300);
block_on(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: Some(existing_deadline),
observed_at: SystemTime::now(),
}))
.expect("mark exhausted");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&account,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record shorter rate-limit failure");
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.cooldown_until(),
Some(existing_deadline)
);
}
#[test]
fn rate_limited_failure_does_not_consult_stale_quota_snapshot() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let reset_at = Utc::now().timestamp() - 1;
block_on(
store.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
}),
)
.expect("persist expired quota exhaustion");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
@@ -873,8 +883,9 @@ fn rate_limited_failure_ignores_exhaustion_after_its_reset_time() {
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::Cooldown
AccountAvailability::QuotaExhausted
);
assert_eq!(store.quota_reads(), 0);
}
#[test]
@@ -907,7 +918,7 @@ fn authenticated_rate_limit_recovers_stale_authentication_state() {
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::Cooldown
AccountAvailability::QuotaExhausted
);
}
@@ -985,7 +996,7 @@ fn rate_limited_failures_for_distinct_accounts_do_not_conflict() {
.expect("persisted second account")
.availability(),
],
[AccountAvailability::Cooldown; 2]
[AccountAvailability::QuotaExhausted; 2]
);
}
@@ -600,10 +600,10 @@ async fn reauthorization_uses_revision_advanced_after_start_for_prepared_rotatio
);
}
struct RebindingVerifier;
struct ReplacementIdentityVerifier;
#[async_trait]
impl CodexAccountIdentityVerifier for RebindingVerifier {
impl CodexAccountIdentityVerifier for ReplacementIdentityVerifier {
async fn verify(
&self,
_: &CodexOAuthSecret,
@@ -617,8 +617,9 @@ impl CodexAccountIdentityVerifier for RebindingVerifier {
_: &CodexOAuthSecret,
_: &SecretString,
_: &SecretString,
_: &CodexIdentityExpectation,
expectation: &CodexIdentityExpectation,
) -> Result<CodexIdentityVerification, CodexIdentityVerificationError> {
assert_eq!(expectation, &CodexIdentityExpectation::default());
Ok(CodexIdentityVerification::Complete(profile(
"chatgpt-different-owner",
)))
@@ -699,7 +700,7 @@ async fn reauthorization_updates_changed_free_principal_for_same_account() {
}
#[tokio::test]
async fn reauthorization_rejects_identity_rebinding() {
async fn reauthorization_rebinds_target_record_to_authorized_identity() {
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
@@ -711,7 +712,7 @@ async fn reauthorization_rejects_identity_rebinding() {
enabled: true,
})
.await;
let service = service_with_store(store, Arc::new(RebindingVerifier));
let service = service_with_store(store, Arc::new(ReplacementIdentityVerifier));
let started = service
.start_authorization(StartCodexOAuthAuthorization {
mutation: reauthorization_mutation("request-rebind", "acct_oauth_rebind"),
@@ -723,7 +724,7 @@ async fn reauthorization_rejects_identity_rebinding() {
.query_pairs()
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
.expect("state");
let error = service
let completed = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
@@ -732,6 +733,24 @@ async fn reauthorization_rejects_identity_rebinding() {
)),
})
.await
.expect_err("identity rebind");
assert_eq!(error, CodexOAuthAdminError::Credential);
.expect("complete identity replacement");
let CompletedCodexOAuthCredential::Reauthorize(prepared) = completed.credential else {
panic!("expected prepared reauthorization");
};
assert_eq!(
prepared.credential.account_id().as_str(),
"acct_oauth_rebind"
);
let replacement = prepared
.replacement_identity
.as_ref()
.expect("replacement identity");
assert_eq!(
replacement.upstream_account_id(),
Some("chatgpt-different-owner")
);
assert_eq!(
replacement.upstream_user_id(),
"user-chatgpt-different-owner"
);
}
@@ -55,6 +55,36 @@ fn parser_treats_any_confirmed_provider_limit_as_exhausted() {
assert!(fact.exhausted());
}
#[test]
fn parser_keeps_full_percent_as_display_without_provider_exhaustion_signal() {
let fact = parse_codex_quota_usage(&json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
}))
.expect("valid full-percent quota");
assert_eq!(fact.remaining_percent(), Some(0));
assert!(!fact.exhausted());
}
#[test]
fn parser_treats_confirmed_limit_reached_as_exhausted_even_below_full_percent() {
let fact = parse_codex_quota_usage(&json!({
"rate_limit": {
"allowed": false,
"limit_reached": true,
"primary_window": {"used_percent": 98, "reset_at": 1_900_000_000}
}
}))
.expect("valid provider-confirmed quota");
assert_eq!(fact.remaining_percent(), Some(2));
assert!(fact.exhausted());
}
#[test]
fn parser_does_not_infer_exhaustion_from_unknown_credit_fields() {
let fact = parse_codex_quota_usage(&json!({
@@ -181,6 +211,14 @@ fn blocked_network_quota_service(store: &Arc<MemoryAccountStore>) -> CodexCreden
}
async fn create_account(store: &Arc<MemoryAccountStore>, account_id: &str) {
create_account_with_enabled(store, account_id, true).await;
}
async fn create_account_with_enabled(
store: &Arc<MemoryAccountStore>,
account_id: &str,
enabled: bool,
) {
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
@@ -188,7 +226,7 @@ async fn create_account(store: &Arc<MemoryAccountStore>, account_id: &str) {
secret: secret(&format!("token-{account_id}")),
verified_account: profile(&format!("chatgpt-{account_id}")),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
enabled,
})
.await;
}
@@ -361,12 +399,66 @@ async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision
);
}
#[tokio::test]
async fn passive_full_percent_with_explicit_allowance_keeps_account_ready() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_passive_full_allowed").await;
let account = store.account("acct_passive_full_allowed").expect("account");
let service = quota_service(&store);
let headers = vec![
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-primary-used-percent".to_owned(), "100".to_owned()),
(
"x-codex-primary-reset-at".to_owned(),
"1900000000".to_owned(),
),
("x-codex-allowed".to_owned(), "true".to_owned()),
("x-codex-limit-reached".to_owned(), "false".to_owned()),
];
assert!(
service
.synchronize_passive_headers(&account, &headers)
.await
.expect("passive quota")
);
let current = store
.account("acct_passive_full_allowed")
.expect("current account");
let snapshot = service
.read_account(account.id())
.await
.expect("read quota")
.expect("quota snapshot");
assert_eq!(
(
current.availability(),
current.cooldown_until(),
snapshot.fact().remaining_percent(),
snapshot.fact().exhausted(),
),
(AccountAvailability::Ready, None, Some(0), false)
);
}
#[tokio::test]
async fn manual_quota_refresh_updates_account_state_from_provider_response() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_manual_quota_state";
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::Expired,
reason: Some("stale_expired_state".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("seed stale expired state");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
@@ -438,6 +530,61 @@ async fn manual_quota_refresh_updates_account_state_from_provider_response() {
);
}
#[tokio::test]
async fn manual_quota_refresh_preserves_disabled_account_state() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_disabled_quota_state";
create_account_with_enabled(&store, account_id, false).await;
let account = store.account(account_id).expect("created disabled account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Expired,
reason: Some("disabled_expired_state".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("seed disabled account state");
let server = MockServer::start().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": 20, "reset_at": 1_900_000_000}
}
})))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
service
.refresh_account(account.id())
.await
.expect("refresh disabled account quota");
let current = store
.account(account_id)
.expect("disabled account after quota refresh");
assert!(!current.enabled());
assert_eq!(current.availability(), AccountAvailability::Expired);
assert!(store.has_quota(account_id));
}
#[tokio::test]
async fn manual_quota_auth_rejection_does_not_invalidate_refreshable_credential() {
let store = Arc::new(MemoryAccountStore::default());
@@ -697,12 +844,18 @@ async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_q
service.synchronize().await.expect("quota synchronization");
let reconciled = store
.account("acct_periodic_stuck_cooldown")
.expect("reconciled account");
assert_eq!(
store
.account("acct_periodic_stuck_cooldown")
.expect("reconciled account")
.availability(),
AccountAvailability::QuotaExhausted
(reconciled.availability(), reconciled.cooldown_until()),
(
AccountAvailability::QuotaExhausted,
Some(
SystemTime::UNIX_EPOCH
+ Duration::from_secs(u64::try_from(reset_at).expect("positive reset")),
),
)
);
}
@@ -741,12 +894,14 @@ async fn periodic_quota_synchronization_does_not_downgrade_exhaustion_on_rate_li
service.synchronize().await.expect("quota synchronization");
assert_eq!(
store
.account("acct_periodic_exhausted")
.expect("persisted account")
.availability(),
AccountAvailability::QuotaExhausted
let account = store
.account("acct_periodic_exhausted")
.expect("persisted account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
}
@@ -877,6 +1032,60 @@ async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() {
assert_eq!(summary.transient, 1);
}
#[tokio::test]
async fn periodic_quota_synchronization_waits_for_quota_cooldown_deadline_with_grace() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_exhausted_cooldown").await;
let account = store
.account("acct_periodic_exhausted_cooldown")
.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(29)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted with recent cooldown");
let service = blocked_network_quota_service(&store);
let waiting = service
.synchronize()
.await
.expect("quota synchronization before grace");
assert_eq!(waiting, CodexQuotaSyncSummary::default());
let current = store
.account("acct_periodic_exhausted_cooldown")
.expect("current account");
store
.apply_state_change(AccountStateChange {
account_id: current.id().clone(),
expected_revision: current.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: SystemTime::now().checked_sub(Duration::from_secs(31)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted past grace");
let due = service
.synchronize()
.await
.expect("quota synchronization after grace");
let repeated = service
.synchronize()
.await
.expect("quota synchronization repeated");
assert_eq!(due.transient, 1);
assert_eq!(repeated, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
let store = Arc::new(MemoryAccountStore::default());
@@ -955,8 +1164,8 @@ async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_
server.uri(),
);
let first = service.synchronize().await.expect("initial refresh");
assert_eq!(first.transient, 1);
let waiting = service.synchronize().await.expect("initial cooldown wait");
assert_eq!(waiting, CodexQuotaSyncSummary::default());
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
@@ -987,6 +1196,77 @@ async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_
assert_eq!(repeated, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn periodic_quota_synchronization_reschedules_a_confirmed_new_reset() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_periodic_new_reset";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
let initial_reset = SystemTime::UNIX_EPOCH + Duration::from_secs(1);
let confirmed_reset = SystemTime::UNIX_EPOCH + Duration::from_secs(2);
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: Some(initial_reset),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted with due reset");
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": false,
"limit_reached": true,
"primary_window": {"used_percent": 100, "reset_at": 2}
}
})))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let first = service.synchronize().await.expect("first reset recheck");
let persisted_reset = store
.account(account_id)
.expect("account after first recheck")
.cooldown_until();
let second = service.synchronize().await.expect("new reset recheck");
let repeated = service.synchronize().await.expect("same reset throttled");
let request_count = server
.received_requests()
.await
.expect("received requests")
.len();
assert_eq!(
(
first.exhausted,
persisted_reset,
second.exhausted,
repeated,
request_count,
),
(
1,
Some(confirmed_reset),
1,
CodexQuotaSyncSummary::default(),
2,
)
);
}
#[tokio::test]
async fn invalid_quota_json_for_one_account_does_not_abort_the_batch() {
let store = Arc::new(MemoryAccountStore::default());
@@ -665,6 +665,74 @@ async fn expired_access_token_refreshes_even_when_next_refresh_is_deferred() {
);
}
#[tokio::test]
async fn expired_access_token_refreshes_during_active_cooldown() {
let store = Arc::new(MemoryAccountStore::default());
seed_due_account(&store).await;
let original = store.account("acct_refresh").expect("seeded account");
let mut expired_profile = profile("chatgpt-acct_refresh");
expired_profile.access_token_expires_at =
Some(chrono::Utc::now() - chrono::Duration::minutes(5));
store
.repository()
.rotate_oauth_secret(RotateCodexCredential {
account_id: "acct_refresh".to_owned(),
expected_credential_revision: original.revision().get(),
secret: secret("old-access"),
verified_account: expired_profile,
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::hours(12)),
})
.await
.expect("seed expired access token with deferred refresh");
let deferred = store.account("acct_refresh").expect("deferred account");
let cooldown_until = SystemTime::now() + Duration::from_secs(12 * 60 * 60);
store
.repository()
.apply_state(
&deferred,
AccountAvailability::Cooldown,
Some("rate_limited".to_owned()),
Some(cooldown_until),
SystemTime::now(),
)
.await
.expect("mark account cooling down");
let refresher = Arc::new(Refresher::new(Ok(success_tokens())));
let service = CodexCredentialRefreshService::new(
store.repository(),
refresher.clone(),
Arc::new(VerifiedIdentity),
Arc::new(RefreshLeases {
available: true,
requests: Mutex::new(Vec::new()),
}),
Arc::new(CountingCredentialState::default()),
runtime_policy(),
);
let outcomes = service.refresh_due().await.expect("refresh due");
assert!(matches!(
outcomes.as_slice(),
[CodexCredentialRefreshOutcome::Refreshed {
account_id,
credential_revision: 3,
}] if account_id == "acct_refresh"
));
assert_eq!(
refresher.seen.lock().expect("seen tokens lock").as_slice(),
["rt-old-access"]
);
let refreshed = store.account("acct_refresh").expect("refreshed account");
assert_eq!(refreshed.availability(), AccountAvailability::Cooldown);
assert_eq!(refreshed.cooldown_until(), Some(cooldown_until));
assert!(
refreshed
.access_token_expires_at()
.is_some_and(|expires_at| expires_at > SystemTime::now())
);
}
#[tokio::test]
async fn unavailable_refresh_lease_prevents_duplicate_token_exchange() {
let (store, refresher, service) = setup(Ok(success_tokens()), false).await;
@@ -7,7 +7,7 @@ use base64::{Engine as _, engine::general_purpose::STANDARD};
use chrono::Utc;
use futures::StreamExt;
use gateway_core::engine::credential::{
AccountAvailability, AccountFeedbackStats, ProviderAccountId,
AccountAvailability, AccountFeedbackStats, ProviderAccountId, ProviderAccountStore as _,
};
use gateway_core::engine::provider::{Provider as _, ProviderRequest};
use gateway_core::engine::{
@@ -29,14 +29,14 @@ use provider_openai::credential::{
use provider_openai::transport::CodexWebSocketPool;
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use provider_openai::{CodexProvider, OFFICIAL_CODEX_BASE_URL};
use serde_json::{Map, json};
use serde_json::{Map, Value, json};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
sync::oneshot,
time::timeout,
};
use wiremock::matchers::{method, path};
use wiremock::matchers::{header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use crate::support::{
@@ -827,13 +827,23 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
.expect("prepare provider stream");
let mut observed_service_tier = None;
let mut upstream_service_tier = None;
while let Some(event) = stream.next().await {
let event = event.expect("provider event");
if let Some(service_tier) = event
.response_observation()
.and_then(|observation| observation.service_tier())
{
observed_service_tier = Some(service_tier.to_owned());
if let Some(observation) = event.response_observation() {
if let Some(service_tier) = observation.service_tier() {
observed_service_tier = Some(service_tier.to_owned());
}
upstream_service_tier = observation
.provider_metadata()
.and_then(|metadata| serde_json::from_str::<Value>(metadata.as_json()).ok())
.and_then(|metadata| {
metadata
.get("upstreamServiceTier")
.and_then(Value::as_str)
.map(str::to_owned)
})
.or(upstream_service_tier);
}
if event
.canonical_facts()
@@ -846,11 +856,12 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
drop(stream);
assert_eq!(affinity.binding_count(), 1);
assert_eq!(observed_service_tier.as_deref(), Some("default"));
assert_eq!(observed_service_tier.as_deref(), Some("priority"));
assert_eq!(upstream_service_tier.as_deref(), Some("default"));
}
#[tokio::test]
async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown() {
async fn response_failed_before_semantic_output_is_atomic_and_persists_quota_lock() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_atomic_failure";
create_account(&store, account_id).await;
@@ -895,7 +906,7 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown(
vec!["response.created", "response.failed"]
);
let account = store.account(account_id).expect("rate-limited account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
@@ -906,6 +917,68 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown(
let _ = server.await;
}
#[tokio::test]
async fn usage_limit_failure_marks_quota_exhausted_without_usage_probe() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_usage_limit_request_path";
create_account(&store, account_id).await;
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header("authorization", format!("Bearer at-{account_id}")))
.respond_with(ResponseTemplate::new(500))
.expect(0)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_usage_limit_confirm\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"retry_after_seconds\":86400,\"response\":{\"id\":\"resp_usage_limit_confirm\",\"status\":\"failed\",\"error\":{\"code\":\"usage_limit_reached\",\"message\":\"usage limit reached\"}}}\n\n"
)),
)
.mount(&server)
.await;
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", http_generate_operation()),
context("req_usage_limit_confirm", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let failure = loop {
match stream.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("usage-limit failure must surface a typed failure"),
}
};
assert_eq!(failure.kind(), ProviderErrorKind::QuotaExhausted);
let account = store.account(account_id).expect("usage-limit account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert!(!store.has_quota(account_id));
let requests = server.received_requests().await.expect("received requests");
assert_eq!(
requests
.iter()
.filter(|request| request.url.path() == "/api/codex/usage")
.count(),
0
);
}
#[tokio::test]
async fn ordinary_request_should_forward_created_before_later_failed_chunk() {
let store = Arc::new(MemoryAccountStore::default());
@@ -1207,8 +1280,8 @@ async fn response_failed_after_semantic_output_is_exposed_and_not_replay_safe()
async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_state() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_disabled_diagnostic";
create_account_with_enabled(&store, account_id, false).await;
let account = store.account(account_id).expect("disabled test account");
create_account(&store, account_id).await;
let account = store.account(account_id).expect("test account");
store
.repository()
.apply_state(
@@ -1220,6 +1293,10 @@ async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_st
)
.await
.expect("seed quota-exhausted state");
store
.set_enabled(account.id(), false)
.await
.expect("disable test account");
let affinity = Arc::new(MemorySessionAffinity::default());
let server = MockServer::start().await;
@@ -1323,6 +1400,76 @@ async fn successful_response_preserves_quota_exhaustion_from_rate_limit_headers(
assert!(store.has_quota(account_id));
}
#[tokio::test]
async fn successful_http_sse_rate_limit_event_persists_structured_exhaustion() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_http_sse_exhausted";
create_account(&store, account_id).await;
let reset_at = 1_900_000_000_u64;
let rate_limit_event = json!({
"type": "codex.rate_limits",
"rate_limits": {
"allowed": false,
"limit_reached": true,
"primary": {
"used_percent": 42,
"window_minutes": 300,
"reset_at": reset_at,
},
},
});
let completed_event = json!({
"type": "response.completed",
"response": {
"id": "resp_http_sse_exhausted",
"model": "gpt-5.4",
"status": "completed",
"output": [],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
},
});
let body = format!(
"event: codex.rate_limits\ndata: {rate_limit_event}\n\nevent: response.completed\ndata: {completed_event}\n\n"
);
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(body),
)
.mount(&server)
.await;
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", http_generate_operation()),
context("req_http_sse_exhausted", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
while let Some(event) = stream.next().await {
event.expect("successful upstream response");
}
let account = store
.account(account_id)
.expect("account after HTTP SSE response");
assert_eq!(
(
account.availability(),
account.cooldown_until(),
store.has_quota(account_id),
),
(
AccountAvailability::QuotaExhausted,
Some(SystemTime::UNIX_EPOCH + Duration::from_secs(reset_at)),
true,
)
);
}
#[test]
fn request_observation_reads_openai_metadata_without_changing_the_operation() {
let store = Arc::new(MemoryAccountStore::default());
@@ -203,18 +203,40 @@ fn decoder_should_emit_calculated_cost_for_complete_known_model_usage() {
}
#[test]
fn decoder_should_keep_response_service_tier_when_terminal_event_omits_it() {
fn decoder_should_prefer_requested_service_tier_over_response_tier_for_billing() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"priority\"}}\n\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"default\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
);
let mut decoder = CodexCanonicalDecoder::new("fallback");
let mut decoder =
CodexCanonicalDecoder::new("fallback").with_requested_service_tier(Some("priority"));
let events = decoder
.push(body.as_bytes())
.expect("canonical priority response");
assert_eq!(decoder.response_service_tier(), Some("default"));
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
event,
GatewayEvent::CalculatedCost(cost)
if cost.total().amount().scaled() == 6_875_000
)));
}
#[test]
fn decoder_should_use_response_service_tier_when_request_omits_it() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_response_tier_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"priority\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_response_tier_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
);
let mut decoder = CodexCanonicalDecoder::new("fallback");
let events = decoder
.push(body.as_bytes())
.expect("canonical response-tier priority response");
assert_eq!(decoder.response_service_tier(), Some("priority"));
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
event,
@@ -223,6 +245,31 @@ fn decoder_should_keep_response_service_tier_when_terminal_event_omits_it() {
)));
}
#[test]
fn websocket_decoder_should_prefer_requested_service_tier_over_response_tier_for_billing() {
let created = websocket_event_to_sse_frame(
r#"{"type":"response.created","response":{"id":"resp_ws_fast_cost","model":"gpt-5.4","service_tier":"default"}}"#,
)
.expect("created frame");
let completed = websocket_event_to_sse_frame(
r#"{"type":"response.completed","response":{"id":"resp_ws_fast_cost","model":"gpt-5.4","status":"completed","usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":25,"cache_write_tokens":0},"total_tokens":110}}}"#,
)
.expect("completed frame");
let mut decoder = CodexCanonicalDecoder::new("fallback")
.with_requested_service_tier(Some("priority"))
.with_raw_sse_passthrough();
let events = decoder
.push(format!("{created}{completed}").as_bytes())
.expect("canonical priority WebSocket response");
assert_eq!(decoder.response_service_tier(), Some("default"));
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
event,
GatewayEvent::CalculatedCost(cost)
if cost.total().amount().scaled() == 6_875_000
)));
}
#[test]
fn decoder_should_normalize_text_usage_and_completion() {
let body = concat!(
@@ -269,6 +269,54 @@ async fn codex_backend_client_should_parse_retry_after_from_rate_limit_error_bod
assert_eq!(retry_after_seconds, Some(12));
}
#[tokio::test]
async fn codex_backend_http_sse_should_capture_structured_rate_limit_event_updates() {
let server = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("POST"))
.and(wiremock::matchers::path("/codex/responses"))
.respond_with(
wiremock::ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: codex.rate_limits\n",
"data: {\"type\":\"codex.rate_limits\",\"rate_limits\":{\"allowed\":true,\"limit_reached\":false,\"primary\":{\"used_percent\":100,\"window_minutes\":300,\"reset_at\":1893456300}}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_http_rate_limits\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)),
)
.mount(&server)
.await;
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
test_wire_profile(),
);
let mut request = codex_request("gpt-5.5", "", Vec::new());
request.force_http_sse = true;
let response = client
.create_response(
&request,
request_context("req_http_rate_limit_event", Some("acct")),
)
.await
.expect("HTTP SSE response");
let structured = response
.rate_limit_headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-codex-allowed" | "x-codex-limit-reached"))
.cloned()
.collect::<Vec<_>>();
assert_eq!(
structured,
vec![
("x-codex-allowed".to_owned(), "true".to_owned()),
("x-codex-limit-reached".to_owned(), "false".to_owned()),
]
);
}
#[tokio::test(start_paused = true)]
async fn codex_backend_http_sse_should_fail_after_five_minutes_without_stream_data() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -942,6 +942,7 @@ fn prepared_rotation(
provider_kind,
expected_credential_revision: Revision::new(expected_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
replacement_identity: None,
name: profile.name,
email: profile.email,
plan_type: profile.plan_type,
+1
View File
@@ -2,3 +2,4 @@
0ae3945c02461347ab8dcfc64821027cd2366505778e0450d71bd82772c4e99d 0002_drop_opaque_response_id_index.sql
e608ed89991f9d6317491714447a00a61ede90caca4fb030eb30695654afb664 0003_store_opaque_response_ids_as_bytes.sql
df4e02cf2335733326c12429ced4cbd0b35b8a22b3fc24981c3c0549477d7c43 0004_add_model_request_service_tier.sql
fe8f36913c6db7053148db450f05e008dab6cdb9ca8cdcfcb22d97ad9a87bf49 0005_allow_quota_exhausted_cooldown.sql
@@ -0,0 +1,12 @@
alter table provider_accounts
drop constraint provider_accounts_cooldown_ck,
add constraint provider_accounts_cooldown_ck check (
(
availability != 'cooldown'
or cooldown_until is not null
)
and (
cooldown_until is null
or availability in ('cooldown', 'quota_exhausted')
)
);
+1 -1
View File
@@ -1 +1 @@
version: "3.0.22"
version: "3.0.24"