mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-10 00:06:00 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a321ccf737 | ||
|
|
4533ae8850 | ||
|
|
7542ed365e | ||
|
|
bf4da7a3e7 | ||
|
|
0fe354b083 |
@@ -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);
|
||||
|
||||
@@ -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(¤t.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(¤t, observed_at).await,
|
||||
Ok(true)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
let snapshot = match self.read_snapshot_for(¤t).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(
|
||||
¤t,
|
||||
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 {
|
||||
¤t,
|
||||
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 {
|
||||
¤t,
|
||||
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(¤t, 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(¤t, availability, reason, cooldown_until, observed_at)
|
||||
.apply_state(
|
||||
¤t,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
reason,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 @@
|
||||
version: "3.0.22"
|
||||
version: "3.0.24"
|
||||
|
||||
Reference in New Issue
Block a user