Compare commits

...
9 Commits
Author SHA1 Message Date
ZyyandCodex 5671a57dc8 chore: bump version to 3.0.26
Co-authored-by: Codex <noreply@openai.com>
2026-07-29 19:44:44 +08:00
Zyy 166e8754a2 fix(xai): restore tool replay and billing details 2026-07-29 18:28:22 +08:00
ZyyandCodex fbaf6c1bdd chore: bump version to 3.0.25
Co-authored-by: Codex <noreply@openai.com>
2026-07-29 17:43:34 +08:00
Zyy 7cdc2dba6f fix(xai): align account state and session replay 2026-07-29 17:36:09 +08:00
ZyyandCodex a321ccf737 chore: bump version to 3.0.24
Co-authored-by: Codex <noreply@openai.com>
2026-07-29 14:49:27 +08:00
CodexandCodex 4533ae8850 fix(openai): preserve resettable quota exhaustion state
Co-authored-by: Codex <noreply@openai.com>
2026-07-29 14:48:23 +08:00
zyycnandCodex 7542ed365e chore: bump version to 3.0.23
Co-authored-by: Codex <noreply@openai.com>
2026-07-29 03:52:17 +08:00
zyycnandCodex bf4da7a3e7 fix(accounts): unify credential and quota state
Apply OpenAI reauthorization to the selected local account ID while atomically replacing its verified upstream identity. Converge successful quota observations across refresh paths and preserve explicitly disabled account state during connection checks, quota refreshes, and credential rotation.

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

Co-authored-by: Codex <noreply@openai.com>
2026-07-29 03:09:22 +08:00
49 changed files with 3386 additions and 480 deletions
@@ -5,7 +5,7 @@ use std::fmt;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use gateway_core::{
engine::credential::{OpaqueProviderData, ProviderAccountId},
engine::credential::{OpaqueProviderData, ProviderAccountId, ProviderAccountIdentity},
routing::{ProviderKind, UpstreamModelId},
};
@@ -496,6 +496,7 @@ pub struct PreparedCredentialRotationFacts {
pub account_id: ProviderAccountId,
pub provider_kind: ProviderKind,
pub expected_credential_revision: Revision,
pub replacement_identity: Option<ProviderAccountIdentity>,
pub name: String,
pub email: Option<String>,
pub plan_type: Option<String>,
@@ -492,14 +492,14 @@ impl DefaultObservabilityService {
}
}
/// 逐条尽力把 calculated 总额升级为完整分解;单条脏数据(非法 Provider kind
/// 或费用规则失败)只保留该条已存的总额,不影响整页返回。
/// 逐条尽力把可校验的总额升级为完整分解;单条脏数据(非法 Provider kind、
/// 不支持的来源或费用规则失败)只保留该条已存的总额,不影响整页返回。
fn enrich_billing(&self, records: &mut [crate::model::observability::UsageRecord]) {
for record in records {
let Some(UsageBilling::Total { source, total }) = record.billing.as_ref() else {
continue;
};
if source != "calculated" {
if !matches!(source.as_str(), "calculated" | "provider_reported") {
continue;
}
let (Some(provider), Some(upstream_model_id)) = (
@@ -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());
@@ -327,9 +327,19 @@ async fn usage_records_should_tolerate_records_that_fail_billing_enrichment() {
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
let invalid_kind = "x".repeat(65);
store.replace_usage_records(vec![
calculated_total_record("request_invalid_kind", Some(&invalid_kind), now),
calculated_total_record("request_unregistered_kind", Some("anthropic"), now),
calculated_total_record("request_enriched", Some("openai"), now),
total_record(
"request_invalid_kind",
Some(&invalid_kind),
"calculated",
now,
),
total_record(
"request_unregistered_kind",
Some("anthropic"),
"calculated",
now,
),
total_record("request_enriched", Some("openai"), "calculated", now),
]);
let services = observability_services_with_calculated_billing(store).await;
@@ -354,6 +364,30 @@ async fn usage_records_should_tolerate_records_that_fail_billing_enrichment() {
);
}
#[tokio::test]
async fn usage_records_should_enrich_provider_reported_totals_when_pricing_matches() {
let now = Utc::now();
let store = Arc::new(FixtureObservabilityStore::new(observation_range(now)));
store.replace_usage_records(vec![total_record(
"request_provider_reported",
Some("openai"),
"provider_reported",
now,
)]);
let services = observability_services_with_calculated_billing(store).await;
let page = services
.observability()
.usage_records(usage_query(now))
.await
.expect("usage records");
assert!(matches!(
page.items[0].billing,
Some(UsageBilling::Calculated(_))
));
}
struct FixtureObservabilityStore {
trend: Mutex<Vec<RequestMetricPoint>>,
overview: Mutex<UsageOverview>,
@@ -567,9 +601,10 @@ fn usage_query(now: DateTime<Utc>) -> UsageQuery {
}
}
fn calculated_total_record(
fn total_record(
id: &str,
provider_kind: Option<&str>,
source: &str,
now: DateTime<Utc>,
) -> UsageRecord {
UsageRecord {
@@ -613,11 +648,11 @@ fn calculated_total_record(
image_input_tokens: None,
image_output_tokens: None,
total_tokens: Some(1_000),
cost_source: "calculated".to_owned(),
cost_source: source.to_owned(),
cost_amount: None,
cost_currency: None,
billing: Some(UsageBilling::Total {
source: "calculated".to_owned(),
source: source.to_owned(),
total: CurrencyCost {
currency: "USD".to_owned(),
amount: gateway_admin::model::observability::DecimalAmount::from_str("1.25")
@@ -44,6 +44,46 @@ impl fmt::Display for ProviderAccountId {
}
}
/// credential 轮换时可选替换的上游账号身份。
#[derive(Clone, PartialEq, Eq)]
pub struct ProviderAccountIdentity {
upstream_user_id: String,
upstream_account_id: Option<String>,
}
impl fmt::Debug for ProviderAccountIdentity {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ProviderAccountIdentity")
.field("upstream_user_id", &"<redacted>")
.field(
"upstream_account_id",
&self.upstream_account_id.as_ref().map(|_| "<redacted>"),
)
.finish()
}
}
impl ProviderAccountIdentity {
#[must_use]
pub const fn new(upstream_user_id: String, upstream_account_id: Option<String>) -> Self {
Self {
upstream_user_id,
upstream_account_id,
}
}
#[must_use]
pub fn upstream_user_id(&self) -> &str {
&self.upstream_user_id
}
#[must_use]
pub fn upstream_account_id(&self) -> Option<&str> {
self.upstream_account_id.as_deref()
}
}
/// `provider_accounts.credential_revision` 的正数 CAS revision。
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct CredentialRevision(NonZeroU64);
@@ -9,7 +9,8 @@ use gateway_core::engine::credential::{
AccountFeedbackStats, AccountQuotaSignals, AccountRuntimeSignals, AccountSchedulingBlocker,
AccountSelectionContext, AccountSelectionPolicy, AccountSelector, CredentialCasUpdate,
CredentialRevision, OpaqueProviderData, PlaintextCredential, PreferredAccountSelection,
ProviderAccount, ProviderAccountId, ProviderAccountUpdate, RotationStrategy,
ProviderAccount, ProviderAccountId, ProviderAccountIdentity, ProviderAccountUpdate,
RotationStrategy,
};
use gateway_core::routing::ProviderKind;
@@ -73,6 +74,18 @@ fn opaque_provider_data_should_not_expose_quota_values_in_debug() {
assert!(!format!("{quota:?}").contains("private-window"));
}
#[test]
fn provider_account_identity_debug_should_redact_upstream_ids() {
let identity = ProviderAccountIdentity::new(
"private-user-id".to_owned(),
Some("private-account-id".to_owned()),
);
let debug = format!("{identity:?}");
assert!(!debug.contains("private-user-id"));
assert!(!debug.contains("private-account-id"));
}
#[test]
fn availability_should_round_trip_all_database_values() {
assert_eq!(
@@ -392,6 +392,12 @@ pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(Strin
&format!("{prefix}-secondary"),
details.secondary,
);
if let Some(allowed) = details.allowed {
headers.push((format!("{prefix}-allowed"), allowed.to_string()));
}
if let Some(limit_reached) = details.limit_reached {
headers.push((format!("{prefix}-limit-reached"), limit_reached.to_string()));
}
if let Some(limit_name) = &details.limit_name {
headers.push((format!("{prefix}-limit-name"), limit_name.clone()));
}
@@ -440,6 +440,8 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
"metered_limit_name": "codex-other",
"limit_name": "Other requests",
"rate_limits": {
"allowed": false,
"limit_reached": true,
"primary": {
"used_percent": 33,
"window_minutes": 60,
@@ -450,7 +452,8 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
"has_credits": true,
"unlimited": false,
"balance": "9.25"
}
},
"rate_limit_reached_type": "usage_limit_reached"
});
let parsed = parse_rate_limits_event(&event).expect("rate-limit event");
@@ -460,16 +463,22 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
(
parsed.active_limit.as_deref(),
details.limit_name.as_deref(),
details.allowed,
details.limit_reached,
details.primary.expect("primary window").window_minutes,
parsed.plan_type.as_deref(),
parsed.credits.expect("credits").balance,
parsed.rate_limit_reached_type.as_deref(),
),
(
Some("codex_other"),
Some("Other requests"),
Some(false),
Some(true),
Some(60),
Some("team"),
Some("9.25".to_owned()),
Some("usage_limit_reached"),
)
);
}
@@ -498,9 +507,15 @@ fn rate_limit_header_round_trip_should_preserve_wire_facts() {
"300".to_owned(),
),
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-allowed".to_owned(), "false".to_owned()),
("x-codex-limit-reached".to_owned(), "true".to_owned()),
("x-codex-credits-has-credits".to_owned(), "true".to_owned()),
("x-codex-credits-unlimited".to_owned(), "false".to_owned()),
("x-codex-plan-type".to_owned(), "plus".to_owned()),
(
"x-codex-rate-limit-reached-type".to_owned(),
"usage_limit_reached".to_owned(),
),
];
let parsed = parse_rate_limit_headers(&source).expect("initial parse");
let encoded = rate_limits_to_header_pairs(&parsed);
@@ -1992,27 +1992,35 @@ async fn provider_account_metrics(
observed_at: DateTime<Utc>,
) -> StoreResult<ProviderAccountMetrics> {
let row = sqlx::query(
"select count(*)::bigint as total,
"with normalized as (
select enabled,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
when availability = 'quota_exhausted'
or (
availability = 'cooldown'
and availability_reason = 'usage_limit_exhausted'
and (cooldown_until is null or cooldown_until > $1)
)
then 'quota_exhausted'
when availability in ('expired', 'invalid')
or (access_token_expires_at is not null and access_token_expires_at <= $1)
then 'expired'
else 'active'
end as status
from provider_accounts
)
select count(*)::bigint as total,
count(*) filter (where enabled)::bigint as enabled,
count(*) filter (
where not (
enabled and availability = 'ready'
and (access_token_expires_at is null or access_token_expires_at > $1)
)
)::bigint as unavailable,
count(*) filter (
where enabled and availability = 'ready'
and (access_token_expires_at is null or access_token_expires_at > $1)
)::bigint as active,
count(*) filter (
where availability = 'expired'
or (access_token_expires_at is not null and access_token_expires_at <= $1)
)::bigint as expired,
count(*) filter (where availability = 'quota_exhausted')::bigint
as quota_exhausted,
count(*) filter (where not enabled)::bigint as disabled,
count(*) filter (where availability = 'banned')::bigint as banned
from provider_accounts",
count(*) filter (where status in ('expired', 'disabled', 'banned'))::bigint
as unavailable,
count(*) filter (where status = 'active')::bigint as active,
count(*) filter (where status = 'expired')::bigint as expired,
count(*) filter (where status = 'quota_exhausted')::bigint as quota_exhausted,
count(*) filter (where status = 'disabled')::bigint as disabled,
count(*) filter (where status = 'banned')::bigint as banned
from normalized",
)
.bind(observed_at)
.fetch_one(pool)
@@ -2035,11 +2043,29 @@ async fn active_provider_account_ids(
observed_at: DateTime<Utc>,
) -> StoreResult<Vec<String>> {
sqlx::query(
"select id
from provider_accounts
where enabled and availability = 'ready'
and (access_token_expires_at is null or access_token_expires_at > $1)
order by id",
"with normalized as (
select id,
case
when not enabled then 'disabled'
when availability = 'banned' then 'banned'
when availability = 'quota_exhausted'
or (
availability = 'cooldown'
and availability_reason = 'usage_limit_exhausted'
and (cooldown_until is null or cooldown_until > $1)
)
then 'quota_exhausted'
when availability in ('expired', 'invalid')
or (access_token_expires_at is not null and access_token_expires_at <= $1)
then 'expired'
else 'active'
end as status
from provider_accounts
)
select id
from normalized
where status = 'active'
order by id",
)
.bind(observed_at)
.fetch_all(pool)
@@ -38,8 +38,8 @@ use gateway_core::engine::credential::{
CredentialRevision as CoreCredentialRevision, LoadedCredential,
NewProviderAccount as CoreNewProviderAccount, OpaqueProviderData, PlaintextCredential,
ProviderAccount as CoreProviderAccount, ProviderAccountId as CoreProviderAccountId,
ProviderAccountStore, ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation,
QuotaWriteOutcome,
ProviderAccountIdentity, ProviderAccountStore,
ProviderAccountUpdate as CoreProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome,
};
use gateway_core::error::{StoreError as CoreStoreError, StoreErrorKind as CoreStoreErrorKind};
use gateway_core::routing::ProviderKind;
@@ -125,6 +125,26 @@ fn parse_availability(value: &str) -> StoreResult<AccountAvailability> {
AccountAvailability::parse(value).ok_or_else(|| invalid("unknown availability value"))
}
fn validate_availability_cooldown(
availability: AccountAvailability,
has_cooldown_until: bool,
) -> StoreResult<()> {
if availability == AccountAvailability::Cooldown && !has_cooldown_until {
return Err(invalid("cooldown availability requires cooldown_until"));
}
if has_cooldown_until
&& !matches!(
availability,
AccountAvailability::Cooldown | AccountAvailability::QuotaExhausted
)
{
return Err(invalid(
"cooldown_until is only allowed for cooldown or quota_exhausted availability",
));
}
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderAccountSummary {
pub id: String,
@@ -224,11 +244,7 @@ impl NewProviderAccount {
if !self.has_refresh_token && self.next_refresh_at.is_some() {
return Err(invalid("next_refresh_at requires a refresh token"));
}
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown_until must be present exactly for cooldown availability",
));
}
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
validate_object_size(
"provider_credentials_json",
&self.provider_credentials_json,
@@ -296,6 +312,7 @@ impl ImportProviderAccounts {
pub struct RotateProviderAccount {
pub scope: ProviderAccountAdminScope,
pub profile: UpdateProviderAccount,
pub replacement_identity: Option<ProviderAccountIdentity>,
pub credential: ProviderCredentialUpdate,
pub audit: AdminAuditEvent,
}
@@ -306,6 +323,7 @@ impl fmt::Debug for RotateProviderAccount {
.debug_struct("RotateProviderAccount")
.field("scope", &self.scope)
.field("profile", &self.profile)
.field("replacement_identity", &self.replacement_identity)
.field("credential", &self.credential)
.field("audit", &self.audit)
.finish()
@@ -378,11 +396,7 @@ pub struct ProviderAccountStateUpdate {
impl ProviderAccountStateUpdate {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown availability and cooldown_until must agree",
));
}
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
Ok(())
}
}
@@ -390,11 +404,7 @@ impl ProviderAccountStateUpdate {
impl ProviderAccountObservation {
pub fn validate(&self) -> StoreResult<()> {
require_nonempty(ENTITY, "account_id", &self.account_id)?;
if (self.availability == AccountAvailability::Cooldown) != self.cooldown_until.is_some() {
return Err(invalid(
"cooldown availability and cooldown_until must agree",
));
}
validate_availability_cooldown(self.availability, self.cooldown_until.is_some())?;
if self.provider_quota_json.is_some() != self.quota_observed_at.is_some() {
return Err(invalid("quota JSON and quota_observed_at must agree"));
}
@@ -612,8 +622,14 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
update.validate()?;
let result = sqlx::query(
"update provider_accounts
set availability = $3, availability_reason = $4, cooldown_until = $5,
availability_observed_at = $6, updated_at = greatest(now(), $6)
set availability = case when enabled then $3 else availability end,
availability_reason = case when enabled then $4 else availability_reason end,
cooldown_until = case when enabled then $5 else cooldown_until end,
availability_observed_at = case
when enabled then $6
else availability_observed_at
end,
updated_at = case when enabled then greatest(now(), $6) else updated_at end
where id = $1 and credential_revision = $2",
)
.bind(update.account_id)
@@ -755,6 +771,12 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
command.scope.validate()?;
require_nonempty(ENTITY, "account_id", &command.profile.id)?;
require_nonempty(ENTITY, "name", &command.profile.name)?;
if let Some(identity) = &command.replacement_identity {
require_nonempty(ENTITY, "upstream_user_id", identity.upstream_user_id())?;
if let Some(account_id) = identity.upstream_account_id() {
require_nonempty(ENTITY, "upstream_account_id", account_id)?;
}
}
if command.profile.id != command.credential.account_id {
return Err(invalid("rotated profile and credential account IDs differ"));
}
@@ -770,6 +792,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
&mut transaction,
&command.scope,
&command.profile,
command.replacement_identity.as_ref(),
&command.credential,
)
.await?;
@@ -1040,6 +1063,7 @@ impl PgAdminAccountStore {
email: prepared.email,
plan_type: prepared.plan_type,
},
replacement_identity: prepared.replacement_identity,
credential: ProviderCredentialUpdate {
account_id: account_id.as_str().to_owned(),
expected_revision: store_revision(prepared.expected_credential_revision)?,
@@ -1975,8 +1999,13 @@ async fn rotate_provider_account_in_transaction(
transaction: &mut Transaction<'_, Postgres>,
scope: &ProviderAccountAdminScope,
profile: &UpdateProviderAccount,
replacement_identity: Option<&ProviderAccountIdentity>,
update: &ProviderCredentialUpdate,
) -> StoreResult<Revision> {
let replace_identity = replacement_identity.is_some();
let upstream_user_id = replacement_identity.map(ProviderAccountIdentity::upstream_user_id);
let upstream_account_id =
replacement_identity.and_then(ProviderAccountIdentity::upstream_account_id);
let next = sqlx::query_scalar::<_, i64>(
"update provider_accounts
set name = $4,
@@ -1987,19 +2016,25 @@ async fn rotate_provider_account_in_transaction(
has_refresh_token = $8,
access_token_expires_at = $9,
next_refresh_at = $10,
upstream_user_id = case when $11::boolean then $12::text else upstream_user_id end,
upstream_account_id = case when $11::boolean then $13::text else upstream_account_id end,
availability = case
when not enabled then availability
when availability <> 'quota_exhausted' then 'ready'
else availability
end,
availability_reason = case
when not enabled then availability_reason
when availability <> 'quota_exhausted' then null
else availability_reason
end,
cooldown_until = case
when not enabled then cooldown_until
when availability <> 'quota_exhausted' then null
else cooldown_until
end,
availability_observed_at = case
when not enabled then availability_observed_at
when availability <> 'quota_exhausted' then now()
else availability_observed_at
end,
@@ -2018,9 +2053,25 @@ async fn rotate_provider_account_in_transaction(
.bind(update.has_refresh_token)
.bind(update.access_token_expires_at)
.bind(update.next_refresh_at)
.bind(replace_identity)
.bind(upstream_user_id)
.bind(upstream_account_id)
.fetch_optional(&mut **transaction)
.await
.map_err(|_| postgres_unavailable("rotate provider account in admin transaction"))?
.map_err(|error| {
if error
.as_database_error()
.is_some_and(sqlx::error::DatabaseError::is_unique_violation)
{
StoreError::Conflict {
entity: ENTITY,
id: update.account_id.clone(),
kind: ConflictKind::InvalidTransition,
}
} else {
postgres_unavailable("rotate provider account in admin transaction")
}
})?
.ok_or_else(|| StoreError::Conflict {
entity: ENTITY,
id: update.account_id.clone(),
@@ -504,6 +504,12 @@ impl ProviderAccountStore for CooldownCachingProviderAccountStore {
let availability = change.availability;
let cooldown_until = change.cooldown_until;
self.authoritative.apply_state_change(change).await?;
let Some(account) = self.authoritative.get_account(&account_id).await? else {
return Ok(());
};
if !account.enabled() {
return Ok(());
}
self.mirror(
&account_id,
credential_revision,
@@ -76,6 +76,12 @@ impl TestDatabase {
.execute(&mut *migration)
.await
.expect("apply model request service tier migration");
sqlx::raw_sql(include_str!(
"../../../../migrations/0005_allow_quota_exhausted_cooldown.sql"
))
.execute(&mut *migration)
.await
.expect("apply quota exhausted cooldown migration");
migration.commit().await.expect("commit terminal migration");
Some(Self {
admin,
@@ -171,7 +177,7 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
.expect("drop migration test database");
admin.close().await;
assert_eq!((first_table_count, migration_count), (8, 4));
assert_eq!((first_table_count, migration_count), (8, 5));
assert_eq!(response_id_types, ["bytea", "bytea"]);
assert!(!raw_response_id_index_exists);
}
@@ -183,6 +189,7 @@ fn migrations_should_leave_transaction_ownership_to_sqlx() {
include_str!("../../../../migrations/0002_drop_opaque_response_id_index.sql"),
include_str!("../../../../migrations/0003_store_opaque_response_ids_as_bytes.sql"),
include_str!("../../../../migrations/0004_add_model_request_service_tier.sql"),
include_str!("../../../../migrations/0005_allow_quota_exhausted_cooldown.sql"),
]
.into_iter()
.flat_map(str::lines)
@@ -122,16 +122,27 @@ async fn dashboard_account_metrics_should_partition_available_and_unavailable_ac
id, provider_kind, name, upstream_user_id, authentication_kind,
provider_credentials_json, credential_revision, has_refresh_token,
access_token_expires_at, enabled, availability,
availability_reason, cooldown_until,
availability_observed_at, created_at, updated_at
) values
('acct_available', 'openai', 'available', 'user-available', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', $1, $1, $1),
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready',
null, null, $1, $1, $1),
('acct_expired', 'openai', 'expired', 'user-expired', 'oauth',
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready', $1, $1, $1),
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready',
null, null, $1, $1, $1),
('acct_rate_cooldown', 'openai', 'rate-cooldown', 'user-rate-cooldown', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'cooldown',
'rate_limited', $1 + interval '2 hours', $1, $1, $1),
('acct_usage_limit', 'openai', 'usage-limit', 'user-usage-limit', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'cooldown',
'usage_limit_exhausted', $1 + interval '2 hours', $1, $1, $1),
('acct_banned', 'xai', 'banned', 'user-banned', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned', $1, $1, $1),
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned',
null, null, $1, $1, $1),
('acct_disabled', 'xai', 'disabled', 'user-disabled', 'oauth',
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', $1, $1, $1)",
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready',
null, null, $1, $1, $1)",
)
.bind(now)
.execute(&database.pool)
@@ -154,10 +165,11 @@ async fn dashboard_account_metrics_should_partition_available_and_unavailable_ac
metrics.active,
metrics.unavailable,
metrics.expired,
metrics.quota_exhausted,
metrics.disabled,
metrics.banned,
),
(4, 3, 1, 3, 1, 1, 1),
(6, 5, 2, 3, 1, 1, 1, 1),
);
database.close().await;
}
@@ -17,9 +17,10 @@ use gateway_admin::{
ports::store::AccountStore,
};
use gateway_core::engine::credential::{
AccountAvailability, CredentialCasOutcome, CredentialCasUpdate, CredentialRevision,
OpaqueProviderData, PlaintextCredential, ProviderAccountId, ProviderAccountStore,
ProviderAccountUpdate, QuotaObservation, QuotaWriteOutcome,
AccountAvailability, AccountStateChange, CredentialCasOutcome, CredentialCasUpdate,
CredentialRevision, OpaqueProviderData, PlaintextCredential, ProviderAccountId,
ProviderAccountIdentity, ProviderAccountStore, ProviderAccountUpdate, QuotaObservation,
QuotaWriteOutcome,
};
use gateway_core::routing::ProviderKind;
use gateway_store::{
@@ -898,6 +899,7 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
.rotate_provider_account(RotateProviderAccount {
scope: wrong_scope,
profile: profile("acct_admin_a", "wrong scope"),
replacement_identity: None,
credential: credential_update("acct_admin_a", 1, "wrong-scope-secret"),
audit: audit("audit_wrong_scope", "rotate", "acct_admin_a"),
})
@@ -926,6 +928,10 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
.rotate_provider_account(RotateProviderAccount {
scope: scope.clone(),
profile: profile("acct_admin_a", "rotated account"),
replacement_identity: Some(ProviderAccountIdentity::new(
"user-admin-rebound".to_owned(),
Some("workspace-admin-rebound".to_owned()),
)),
credential: credential_update("acct_admin_a", 1, "rotated-secret"),
audit: audit("audit_account_rotate", "rotate", "acct_admin_a"),
})
@@ -933,15 +939,71 @@ async fn provider_account_admin_mutations_are_scoped_audited_and_atomic() {
.expect("rotate provider account");
assert_eq!(rotation.config_revision.get(), 3);
assert_eq!(rotation.credential_revision.get(), 2);
let restored: (String, Option<String>, Option<chrono::DateTime<Utc>>) = sqlx::query_as(
"select availability, availability_reason, cooldown_until
let restored: (
String,
Option<String>,
Option<chrono::DateTime<Utc>>,
String,
Option<String>,
) = sqlx::query_as(
"select availability, availability_reason, cooldown_until,
upstream_user_id, upstream_account_id
from provider_accounts where id = $1",
)
.bind("acct_admin_a")
.fetch_one(&database.pool)
.await
.expect("load restored account state");
assert_eq!(restored, ("ready".to_owned(), None, None));
assert_eq!(
restored,
(
"ready".to_owned(),
None,
None,
"user-admin-rebound".to_owned(),
Some("workspace-admin-rebound".to_owned()),
)
);
let identity_conflict = repository
.rotate_provider_account(RotateProviderAccount {
scope: scope.clone(),
profile: profile("acct_admin_a", "must roll back"),
replacement_identity: Some(ProviderAccountIdentity::new(
"user-admin-b".to_owned(),
None,
)),
credential: credential_update("acct_admin_a", 2, "must-not-persist"),
audit: audit(
"audit_account_identity_conflict",
"reauthorize",
"acct_admin_a",
),
})
.await
.expect_err("upstream identity collision must roll back rotation");
assert!(matches!(
identity_conflict,
StoreError::Conflict {
kind: ConflictKind::InvalidTransition,
..
}
));
assert_eq!(current_revision(&database.pool).await, 3);
let unchanged: (String, serde_json::Value, i64, String, Option<String>) = sqlx::query_as(
"select name, provider_credentials_json, credential_revision,
upstream_user_id, upstream_account_id
from provider_accounts where id = $1",
)
.bind("acct_admin_a")
.fetch_one(&database.pool)
.await
.expect("load account after identity collision");
assert_eq!(unchanged.0, "rotated account");
assert_eq!(unchanged.1["access_token"], "rotated-secret");
assert_eq!(unchanged.2, 2);
assert_eq!(unchanged.3, "user-admin-rebound");
assert_eq!(unchanged.4.as_deref(), Some("workspace-admin-rebound"));
let revision = repository
.set_provider_account_enabled_admin(SetProviderAccountEnabled {
@@ -1013,6 +1075,7 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
.rotate_provider_account(RotateProviderAccount {
scope,
profile: profile("acct_rotation_quota", "reauthorized account"),
replacement_identity: None,
credential: credential_update("acct_rotation_quota", 1, "reauthorized-secret"),
audit: audit(
"audit_rotation_quota_rotate",
@@ -1041,6 +1104,77 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
database.close().await;
}
#[tokio::test]
async fn disabled_account_preserves_user_state_during_refresh_writes() {
let Some(database) = TestDatabase::create("provider_account_disabled_refresh").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let scope = ProviderAccountAdminScope {
provider_kind: "openai".to_owned(),
};
let account_id = ProviderAccountId::new("acct_disabled_refresh").expect("account ID");
let mut disabled = account(account_id.as_str(), "user-disabled-refresh");
disabled.enabled = false;
disabled.availability = AccountAvailability::Expired;
disabled.availability_reason = Some("credential_expired".to_owned());
repository
.import_provider_accounts(ImportProviderAccounts {
scope: scope.clone(),
accounts: vec![disabled],
audit: audit(
"audit_disabled_refresh_import",
"import",
account_id.as_str(),
),
})
.await
.expect("import disabled account");
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
availability: AccountAvailability::Ready,
reason: None,
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("disabled state write is a no-op");
repository
.rotate_provider_account(RotateProviderAccount {
scope,
profile: profile(account_id.as_str(), "refreshed disabled account"),
replacement_identity: None,
credential: credential_update(account_id.as_str(), 1, "disabled-refreshed-secret"),
audit: audit(
"audit_disabled_refresh_rotate",
"refresh",
account_id.as_str(),
),
})
.await
.expect("refresh disabled account credential");
let current: (bool, String, Option<String>, serde_json::Value, i64) = sqlx::query_as(
"select enabled, availability, availability_reason,
provider_credentials_json, credential_revision
from provider_accounts where id = $1",
)
.bind(account_id.as_str())
.fetch_one(&database.pool)
.await
.expect("load disabled account after refresh writes");
assert!(!current.0);
assert_eq!(current.1, "expired");
assert_eq!(current.2.as_deref(), Some("credential_expired"));
assert_eq!(current.3["access_token"], "disabled-refreshed-secret");
assert_eq!(current.4, 2);
database.close().await;
}
fn account(id: &str, upstream_user_id: &str) -> NewProviderAccount {
NewProviderAccount {
id: id.to_owned(),
@@ -1200,7 +1334,15 @@ fn cooldown_observation_requires_expiry() {
}
#[test]
fn imported_account_rejects_mismatched_cooldown_runtime_facts() {
fn quota_exhausted_account_accepts_cooldown_runtime_facts() {
let mut imported = account("acct_quota_cooldown", "user-quota-cooldown");
imported.availability = AccountAvailability::QuotaExhausted;
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
assert!(imported.validate().is_ok());
}
#[test]
fn imported_account_rejects_cooldown_until_for_ready_runtime_facts() {
let mut imported = account("acct_invalid_cooldown", "user-invalid-cooldown");
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
assert!(imported.validate().is_err());
@@ -50,9 +50,13 @@ impl ProviderAccountStore for CooldownTestAccountStore {
async fn get_account(
&self,
_account: &ProviderAccountId,
account: &ProviderAccountId,
) -> Result<Option<ProviderAccount>, CoreStoreError> {
unreachable!("cooldown tests do not get one account")
Ok(self
.accounts
.iter()
.find(|candidate| candidate.id() == account)
.cloned())
}
async fn list_accounts(&self) -> Result<Vec<ProviderAccount>, CoreStoreError> {
@@ -190,10 +194,26 @@ fn composition_fixture(
CooldownCachingProviderAccountStore,
Arc<CooldownTestCache>,
Arc<Mutex<Vec<&'static str>>>,
) {
composition_fixture_with_accounts(
vec![ready_test_account("acct_cooldown_composition")],
fail_state_write,
fail_cache,
)
}
fn composition_fixture_with_accounts(
accounts: Vec<ProviderAccount>,
fail_state_write: bool,
fail_cache: bool,
) -> (
CooldownCachingProviderAccountStore,
Arc<CooldownTestCache>,
Arc<Mutex<Vec<&'static str>>>,
) {
let calls = Arc::new(Mutex::new(Vec::new()));
let authoritative = Arc::new(CooldownTestAccountStore {
accounts: Vec::new(),
accounts,
calls: Arc::clone(&calls),
fail_state_write,
});
@@ -249,6 +269,31 @@ async fn cooldown_composition_keeps_postgres_terminal_when_redis_fails() {
assert_eq!(*calls.lock().expect("calls lock"), ["postgres", "redis"]);
}
#[tokio::test]
async fn disabled_account_state_write_does_not_change_runtime_cooldown() {
let disabled = test_provider_account("acct_cooldown_composition").with_runtime_state(
false,
AccountAvailability::Ready,
None,
);
let (store, cache, calls) = composition_fixture_with_accounts(vec![disabled], false, false);
store
.apply_state_change(cooldown_state_change(60))
.await
.expect("disabled state write is a no-op");
assert_eq!(*calls.lock().expect("calls lock"), ["postgres"]);
assert!(cache.cached.lock().expect("cached lock").is_empty());
assert!(
cache
.invalidated
.lock()
.expect("invalidated lock")
.is_empty()
);
}
#[tokio::test]
async fn ready_state_invalidates_same_revision_cooldown_after_postgres_commit() {
let (store, cache, calls) = composition_fixture(false, false);
+2 -1
View File
@@ -586,7 +586,7 @@ fn prepared_rotation(
prepared: crate::credential::PreparedCodexCredentialRotation,
provider_kind: ProviderKind,
) -> Result<PreparedCredentialRotation, ProviderAdminError> {
let (profile, credential, guard) = prepared.into_parts();
let (profile, credential, replacement_identity, guard) = prepared.into_parts();
let (
account_id,
expected_revision,
@@ -606,6 +606,7 @@ fn prepared_rotation(
account_id,
provider_kind,
expected_credential_revision,
replacement_identity,
name: profile.name,
email: profile.email,
plan_type: profile.plan_type,
@@ -11,7 +11,8 @@ use std::time::SystemTime;
use chrono::{DateTime, FixedOffset, Utc};
use gateway_core::engine::credential::{
AccountAvailability, CredentialCasUpdate, CredentialRevision, LoadedCredential,
NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountUpdate,
NewProviderAccount, ProviderAccount, ProviderAccountId, ProviderAccountIdentity,
ProviderAccountUpdate,
};
use gateway_core::provider_ports::{
ProviderLeaseAcquisition, ProviderLeaseGuard, ProviderLeasePort, ProviderLeaseRequest,
@@ -286,6 +287,7 @@ impl std::fmt::Debug for RotateManagedCodexCredential {
pub struct PreparedCodexCredentialRotation {
pub profile: ProviderAccountUpdate,
pub credential: CredentialCasUpdate,
pub replacement_identity: Option<ProviderAccountIdentity>,
refresh_guards: Option<ProviderRefreshGuards>,
}
@@ -300,6 +302,7 @@ impl std::fmt::Debug for PreparedCodexCredentialRotation {
.debug_struct("PreparedCodexCredentialRotation")
.field("profile", &self.profile)
.field("credential", &self.credential)
.field("replacement_identity", &self.replacement_identity)
.field(
"refresh_guards",
&self.refresh_guards.as_ref().map(|_| "<held>"),
@@ -316,11 +319,13 @@ impl PreparedCodexCredentialRotation {
) -> (
ProviderAccountUpdate,
CredentialCasUpdate,
Option<ProviderAccountIdentity>,
PreparedCodexCredentialRotationGuard,
) {
(
self.profile,
self.credential,
self.replacement_identity,
PreparedCodexCredentialRotationGuard(self.refresh_guards),
)
}
@@ -494,7 +499,7 @@ impl CodexCredentialAdmin {
self.prepare_oauth_rotation(input, false)
}
/// 完整 OAuth 授权允许认证主体投影更新,但不允许目标账号或用户发生变化。
/// 完整 OAuth 授权按目标数据 ID 替换凭据与上游账号身份。
pub(crate) fn prepare_reauthorization(
&self,
input: RotateManagedCodexCredential,
@@ -505,7 +510,7 @@ impl CodexCredentialAdmin {
fn prepare_oauth_rotation(
&self,
input: RotateManagedCodexCredential,
replace_principal: bool,
replace_identity: bool,
) -> Result<PreparedCodexCredentialRotation, CodexCredentialAdminError> {
let access_token_expires_at =
required_time(input.verified_account.access_token_expires_at)?;
@@ -516,16 +521,23 @@ impl CodexCredentialAdmin {
.ok_or(CodexCredentialAdminError::IdentityMismatch)?;
if input.current.account.provider().as_str() != PROVIDER_NAME
|| input.current.account.authentication_kind() != CODEX_AUTHENTICATION_KIND_OAUTH
|| input.current.account.upstream_account_id()
!= Some(input.verified_account.chatgpt_account_id.as_str())
|| input.current.account.upstream_user_id() != input.verified_account.chatgpt_user_id
|| (!replace_principal
&& (oauth.principal.oauth_subject != input.verified_account.oauth_subject
|| (!replace_identity
&& (input.current.account.upstream_account_id()
!= Some(input.verified_account.chatgpt_account_id.as_str())
|| input.current.account.upstream_user_id()
!= input.verified_account.chatgpt_user_id
|| oauth.principal.oauth_subject != input.verified_account.oauth_subject
|| oauth.principal.poid != input.verified_account.poid))
{
return Err(CodexCredentialAdminError::IdentityMismatch);
}
if replace_principal {
let replacement_identity = replace_identity.then(|| {
ProviderAccountIdentity::new(
input.verified_account.chatgpt_user_id.clone(),
Some(input.verified_account.chatgpt_account_id.clone()),
)
});
if replace_identity {
oauth
.principal
.oauth_subject
@@ -567,6 +579,7 @@ impl CodexCredentialAdmin {
Ok(PreparedCodexCredentialRotation {
profile,
credential,
replacement_identity,
refresh_guards: None,
})
}
@@ -29,7 +29,6 @@ use super::admin::{
use super::identity::{
CodexAccountIdentityVerifier, CodexIdentityExpectation, CodexIdentityVerificationError,
};
use super::security::CodexCredentialCodec;
use super::token_client::{
AuthorizationCodeExchangeError, AuthorizationCodeExchanger, AuthorizationCodeGrant,
OFFICIAL_CODEX_OAUTH_CLIENT_ID, OFFICIAL_CODEX_REDIRECT_URI,
@@ -500,11 +499,7 @@ impl CodexOAuthAdminService {
} else {
None
};
let expectation = current
.as_ref()
.map(oauth_identity_expectation)
.transpose()?
.unwrap_or_default();
let expectation = CodexIdentityExpectation::default();
let profile = self
.verifier
.verify_authorization(&secret, &id_token, &pending.nonce, &expectation)
@@ -679,26 +674,6 @@ impl CodexOAuthAdminService {
}
}
fn oauth_identity_expectation(
current: &gateway_core::engine::credential::LoadedCredential,
) -> Result<CodexIdentityExpectation, CodexOAuthAdminError> {
let runtime = CodexCredentialCodec::decode(&current.credential)
.map_err(|_| CodexOAuthAdminError::Credential)?;
if runtime.authentication.oauth().is_none() {
return Err(CodexOAuthAdminError::Credential);
}
let account_id = current
.account
.upstream_account_id()
.ok_or(CodexOAuthAdminError::Credential)?;
CodexIdentityExpectation::reauthorization(
account_id.to_owned(),
current.account.upstream_user_id().to_owned(),
runtime.installation_id,
)
.map_err(|_| CodexOAuthAdminError::Credential)
}
fn authorization_url(
pending: &CodexPendingAuthorization,
oauth_client_id: &str,
@@ -30,6 +30,7 @@ const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60);
const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
const QUOTA_RESET_RECHECK_GRACE: Duration = Duration::from_secs(30);
// ready 账号主要依赖真实请求的被动观察;这里只提供低频、定额的状态兜底。
const QUOTA_READY_REFRESH_MAX_AGE: Duration = Duration::from_secs(2 * 60 * 60);
const QUOTA_READY_REFRESH_BATCH_SIZE: usize = 8;
@@ -307,7 +308,7 @@ enum CodexQuotaFetchError {
enum CodexQuotaStateTransition {
Banned { reason: &'static str },
Exhausted,
Cooldown { until: Option<SystemTime> },
ResettableLimit { until: Option<SystemTime> },
}
impl CodexQuotaStateTransition {
@@ -315,7 +316,7 @@ impl CodexQuotaStateTransition {
match self {
Self::Banned { .. } => AccountAvailability::Banned,
Self::Exhausted => AccountAvailability::QuotaExhausted,
Self::Cooldown { .. } => AccountAvailability::Cooldown,
Self::ResettableLimit { .. } => AccountAvailability::QuotaExhausted,
}
}
@@ -323,13 +324,13 @@ impl CodexQuotaStateTransition {
match self {
Self::Banned { reason } => reason,
Self::Exhausted => "quota_exhausted",
Self::Cooldown { .. } => "quota_rate_limited",
Self::ResettableLimit { .. } => "quota_rate_limited",
}
}
const fn cooldown_until(self) -> Option<SystemTime> {
match self {
Self::Cooldown { until } => until,
Self::ResettableLimit { until } => until,
Self::Banned { .. } | Self::Exhausted => None,
}
}
@@ -547,7 +548,9 @@ fn quota_refresh_candidate(
let snapshot = quota_snapshot_for_account(&account, observation);
let kind = match account.availability() {
AccountAvailability::QuotaExhausted => CodexQuotaRefreshKind::Exhausted {
reset_at: quota_reset_at_from_snapshot(snapshot.as_ref()),
reset_at: account
.cooldown_until()
.or_else(|| quota_reset_at_from_snapshot(snapshot.as_ref())),
},
AccountAvailability::Cooldown
if snapshot
@@ -618,22 +621,24 @@ fn reserve_exhausted_quota_refresh(
now: SystemTime,
refreshed_at: Instant,
) -> bool {
let reset_recheck_due = reset_at.is_some_and(|reset_at| {
reset_at <= now
&& state
.last_exhausted_reset_recheck
.get(account.id())
.is_none_or(|marker| {
marker.revision != account.revision() || marker.reset_at != reset_at
})
});
if !reset_recheck_due && !periodic_quota_refresh_due(state, account.id(), refreshed_at) {
return false;
}
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
if let Some(reset_at) = reset_at.filter(|reset_at| *reset_at <= now) {
if let Some(reset_at) = reset_at {
if !quota_reset_recheck_ready(reset_at, now) {
return false;
}
let same_reset_already_checked = state
.last_exhausted_reset_recheck
.get(account.id())
.is_some_and(|marker| {
marker.revision == account.revision() && marker.reset_at == reset_at
});
if same_reset_already_checked
&& !periodic_quota_refresh_due(state, account.id(), refreshed_at)
{
return false;
}
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
state.last_exhausted_reset_recheck.insert(
account.id().clone(),
CodexQuotaResetMarker {
@@ -641,10 +646,23 @@ fn reserve_exhausted_quota_refresh(
reset_at,
},
);
return true;
}
if !periodic_quota_refresh_due(state, account.id(), refreshed_at) {
return false;
}
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
true
}
fn quota_reset_recheck_ready(reset_at: SystemTime, now: SystemTime) -> bool {
reset_at
.checked_add(QUOTA_RESET_RECHECK_GRACE)
.is_some_and(|ready_at| ready_at <= now)
}
fn periodic_quota_refresh_due(
state: &CodexQuotaProjectionState,
account_id: &ProviderAccountId,
@@ -802,16 +820,12 @@ impl CodexCredentialQuotaService {
summary.transient += 1;
continue;
};
let preserve_exhaustion = account.availability()
== AccountAvailability::QuotaExhausted
&& transition.availability() == AccountAvailability::Cooldown;
match transition {
CodexQuotaStateTransition::Banned { .. } => summary.banned += 1,
CodexQuotaStateTransition::Exhausted => summary.exhausted += 1,
CodexQuotaStateTransition::Cooldown { .. } => summary.cooldown += 1,
}
if preserve_exhaustion {
continue;
CodexQuotaStateTransition::Exhausted
| CodexQuotaStateTransition::ResettableLimit { .. } => {
summary.exhausted += 1;
}
}
let _ = self
.repository
@@ -856,19 +870,23 @@ impl CodexCredentialQuotaService {
{
continue;
}
if !matches!(
self.has_active_exhaustion(&current, observed_at).await,
Ok(true)
) {
continue;
}
let snapshot = match self.read_snapshot_for(&current).await {
Ok(Some(snapshot))
if quota_fact_is_active_exhaustion(snapshot.fact(), observed_at) =>
{
snapshot
}
_ => continue,
};
let cooldown_until =
quota_reset_at_from_snapshot(Some(&snapshot)).or_else(|| current.cooldown_until());
if let Err(error) = self
.repository
.apply_state(
&current,
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await
@@ -881,8 +899,11 @@ impl CodexCredentialQuotaService {
continue;
}
let enabled = current.enabled();
*account =
current.with_runtime_state(enabled, AccountAvailability::QuotaExhausted, None);
*account = current.with_runtime_state(
enabled,
AccountAvailability::QuotaExhausted,
cooldown_until,
);
}
}
@@ -927,10 +948,10 @@ impl CodexCredentialQuotaService {
} else {
summary.updated += 1;
}
if let Some(availability) = quota_success_availability(
if let Some((availability, cooldown_until)) = quota_success_state(
current.availability(),
current.cooldown_until(),
fact.exhausted(),
fact,
SystemTime::now(),
) {
let _ = self
@@ -939,7 +960,7 @@ impl CodexCredentialQuotaService {
&current,
availability,
fact.exhausted().then_some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await;
@@ -993,10 +1014,10 @@ impl CodexCredentialQuotaService {
}
self.scheduling
.observe(current.id().clone(), current.revision(), observed_at, fact);
if let Some(availability) = quota_success_availability(
if let Some((availability, cooldown_until)) = quota_success_state(
current.availability(),
current.cooldown_until(),
fact.exhausted(),
fact,
observed_at,
) {
self.repository
@@ -1004,7 +1025,7 @@ impl CodexCredentialQuotaService {
&current,
availability,
fact.exhausted().then_some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await?;
@@ -1026,19 +1047,6 @@ impl CodexCredentialQuotaService {
self.read_snapshot_for(&account).await
}
pub(super) async fn has_active_exhaustion(
&self,
account: &ProviderAccount,
now: SystemTime,
) -> Result<bool, CodexCredentialQuotaError> {
match self.read_snapshot_for(account).await {
Ok(snapshot) => Ok(snapshot
.is_some_and(|snapshot| quota_fact_is_active_exhaustion(snapshot.fact(), now))),
Err(CodexCredentialQuotaError::RevisionConflict) => Ok(false),
Err(error) => Err(error),
}
}
async fn read_snapshot_for(
&self,
account: &ProviderAccount,
@@ -1164,10 +1172,10 @@ impl CodexCredentialQuotaService {
snapshot.observed_at(),
snapshot.fact(),
);
if let Some(availability) = quota_success_availability(
if let Some((availability, cooldown_until)) = quota_success_state(
current.availability(),
current.cooldown_until(),
snapshot.fact().exhausted(),
snapshot.fact(),
SystemTime::now(),
) {
self.repository
@@ -1178,7 +1186,7 @@ impl CodexCredentialQuotaService {
.fact()
.exhausted()
.then_some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await?;
@@ -1386,31 +1394,34 @@ fn passive_rate_limit_window(window: RateLimitWindow) -> Map<String, Value> {
snapshot
}
fn quota_success_availability(
fn quota_success_state(
current: AccountAvailability,
cooldown_until: Option<SystemTime>,
exhausted: bool,
fact: CodexQuotaFact,
now: SystemTime,
) -> Option<AccountAvailability> {
match current {
) -> Option<(AccountAvailability, Option<SystemTime>)> {
if fact.exhausted() {
let observed_cooldown = fact.resets_at().map(SystemTime::from).or(cooldown_until);
return (current != AccountAvailability::QuotaExhausted
|| cooldown_until != observed_cooldown)
.then_some((AccountAvailability::QuotaExhausted, observed_cooldown));
}
let availability = match current {
AccountAvailability::Invalid
| AccountAvailability::Expired
| AccountAvailability::Banned => None,
AccountAvailability::QuotaExhausted => (!exhausted).then_some(AccountAvailability::Ready),
AccountAvailability::Ready => exhausted.then_some(AccountAvailability::QuotaExhausted),
AccountAvailability::Cooldown if exhausted => Some(AccountAvailability::QuotaExhausted),
| AccountAvailability::Banned
| AccountAvailability::Unknown
| AccountAvailability::QuotaExhausted => Some(AccountAvailability::Ready),
AccountAvailability::Ready => None,
AccountAvailability::Cooldown
if cooldown_until.is_some_and(|cooldown_until| cooldown_until <= now) =>
{
Some(AccountAvailability::Ready)
}
AccountAvailability::Cooldown => None,
AccountAvailability::Unknown => Some(if exhausted {
AccountAvailability::QuotaExhausted
} else {
AccountAvailability::Ready
}),
}
};
availability.map(|availability| (availability, None))
}
fn quota_state_transition(
@@ -1438,7 +1449,7 @@ fn quota_state_transition(
retry_after_seconds.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN.as_secs()),
)
.min(MAX_RATE_LIMIT_COOLDOWN);
Some(CodexQuotaStateTransition::Cooldown {
Some(CodexQuotaStateTransition::ResettableLimit {
until: observed_at.checked_add(duration),
})
}
@@ -1876,7 +1887,6 @@ impl QuotaAggregate {
self.remaining_percent
.map_or(remaining, |current| current.min(remaining)),
);
self.exhausted |= used >= 100.0;
}
if let Some(reset) = object.get("reset_at") {
let seconds = reset
@@ -256,8 +256,6 @@ impl CodexCredentialRefreshService {
| AccountAvailability::Banned
| AccountAvailability::Invalid
)
&& (account.availability() != AccountAvailability::Cooldown
|| account.cooldown_until().is_some_and(|until| until <= now))
&& refresh_due_at(account, now).is_some()
});
accounts.sort_by_key(|account| (refresh_due_at(account, now), account.id().clone()));
@@ -237,6 +237,9 @@ impl CodexCredentialRepository {
cooldown_until: Option<SystemTime>,
observed_at: SystemTime,
) -> Result<(), CredentialRepositoryError> {
if !account.enabled() {
return Ok(());
}
self.store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
@@ -769,31 +769,31 @@ impl CodexCredentialSelector {
if current.revision() != account.revision() {
return Err(CredentialRepositoryError::RevisionConflict.into());
}
// 一次真实 429 证明上游已经接受当前凭据,可纠正旧的认证/封禁终态。
// 结构化额度耗尽仍拥有更高优先级,不能被通用限流降级。
if current.availability() == AccountAvailability::QuotaExhausted {
return Ok(());
}
let quota_exhausted = self
.quota
.has_active_exhaustion(&current, observed_at)
.await
.map_err(|_| CredentialSelectionError::Store)?;
let (availability, reason, cooldown_until) = if quota_exhausted {
(
AccountAvailability::QuotaExhausted,
Some("quota_exhausted".to_owned()),
None,
)
} else {
(
AccountAvailability::Cooldown,
Some(reason.to_owned()),
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)),
)
let cooldown_until =
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN));
let cooldown_until = match (current.availability(), current.cooldown_until()) {
(AccountAvailability::QuotaExhausted, Some(existing)) => {
cooldown_until.map_or(Some(existing), |until| Some(existing.max(until)))
}
_ => cooldown_until,
};
let reason = if current.availability() == AccountAvailability::QuotaExhausted
&& current.cooldown_until().is_none()
{
Some("quota_exhausted".to_owned())
} else {
Some(reason.to_owned())
};
// 一次真实 429 证明上游已经接受当前凭据,可纠正旧的认证/封禁终态。
// resettable 用量限制沿用 v2 的 quota lock 语义,并保留/延长复核时间。
self.repository
.apply_state(&current, availability, reason, cooldown_until, observed_at)
.apply_state(
&current,
AccountAvailability::QuotaExhausted,
reason,
cooldown_until,
observed_at,
)
.await?;
Ok(())
}
+35 -14
View File
@@ -68,6 +68,7 @@ use crate::transport::request::{
scope_request_to_account,
};
use crate::transport::session::CodexSessionIdentity;
use crate::transport::usage::normalize_service_tier;
use crate::transport::websocket::{CodexWebSocketExchangeError, PreviousResponseUnavailableReason};
use crate::transport::{
CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient,
@@ -822,8 +823,9 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
// 下游按字节转发上游原文,避免 serde 往返改写数值/精度(大整数→f64、logprobs 等)。
// WS 帧由 reducer 以 encode_sse_event(&event, raw) 逐字节内嵌上游原始 JSON
// (transport/protocol/websocket.rs),push_frames 抽出的 data 即上游原文。
let mut decoder =
CodexCanonicalDecoder::new(upstream_model.as_str()).with_raw_sse_passthrough();
let mut decoder = CodexCanonicalDecoder::new(upstream_model.as_str())
.with_requested_service_tier(request.service_tier())
.with_raw_sse_passthrough();
let mut pre_commit_events = PreCommitClientEvents::new(stream_commit_policy);
loop {
let Some(stream_deadline) = remaining(context.deadline()) else {
@@ -895,7 +897,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
};
pre_commit_events.observe_chunk(chunk_len);
let service_tier_changed = observation_state
.observe_service_tier(decoder.response_service_tier());
.observe_upstream_service_tier(decoder.response_service_tier());
let terminal_failure = canonical_failure.map(|(error, semantic_output_seen)| {
let atomic_upstream_failure = matches!(&error, CodexCanonicalError::Upstream(_));
(
@@ -990,8 +992,8 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
)
});
let timing_signals = decoder.take_timing_signals();
let service_tier_changed =
observation_state.observe_service_tier(decoder.response_service_tier());
let service_tier_changed = observation_state
.observe_upstream_service_tier(decoder.response_service_tier());
let timing_changed = observation_state
.observe_timing_signals(timing_signals, context.timing_started_at());
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
@@ -1070,7 +1072,8 @@ struct OpenAiResponseObservationState {
requested_model: String,
stream: bool,
compact: bool,
service_tier: Option<String>,
requested_service_tier: Option<String>,
upstream_service_tier: Option<String>,
rate_limit_headers: Vec<(String, String)>,
account_id: String,
attempt_index: u32,
@@ -1102,8 +1105,8 @@ impl OpenAiResponseObservationState {
requested_model: request.model().to_owned(),
stream: request.stream(),
compact: semantics.compact,
// 请求档位保留在 requestSummary;一等观测只接受响应确认的实际档位。
service_tier: None,
requested_service_tier: normalize_service_tier(request.service_tier()),
upstream_service_tier: None,
rate_limit_headers: selected_observation_headers(&response.rate_limit_headers),
account_id: account_id.to_owned(),
attempt_index,
@@ -1127,7 +1130,7 @@ impl OpenAiResponseObservationState {
if let Some(metadata) = self.provider_metadata() {
observation = observation.with_provider_metadata(metadata);
}
if let Some(service_tier) = self.service_tier.as_deref() {
if let Some(service_tier) = self.effective_service_tier() {
observation = observation.with_service_tier_if_valid(service_tier.to_owned());
}
Some(observation)
@@ -1158,17 +1161,23 @@ impl OpenAiResponseObservationState {
changed
}
fn observe_service_tier(&mut self, service_tier: Option<&str>) -> bool {
let Some(service_tier) = service_tier else {
fn observe_upstream_service_tier(&mut self, service_tier: Option<&str>) -> bool {
let Some(service_tier) = normalize_service_tier(service_tier) else {
return false;
};
if self.service_tier.as_deref() == Some(service_tier) {
if self.upstream_service_tier.as_deref() == Some(service_tier.as_str()) {
return false;
}
self.service_tier = Some(service_tier.to_owned());
self.upstream_service_tier = Some(service_tier);
true
}
fn effective_service_tier(&self) -> Option<&str> {
self.requested_service_tier
.as_deref()
.or(self.upstream_service_tier.as_deref())
}
fn merge_rate_limit_headers(&mut self, updates: &[(String, String)]) -> bool {
let mut changed = false;
for (name, value) in selected_observation_headers(updates) {
@@ -1245,9 +1254,21 @@ impl OpenAiResponseObservationState {
)),
);
metadata.insert("requestSummary".to_owned(), self.request_summary.clone());
if let Some(service_tier) = &self.service_tier {
if let Some(service_tier) = self.effective_service_tier() {
metadata.insert(
"serviceTier".to_owned(),
Value::String(service_tier.to_owned()),
);
}
if let Some(service_tier) = &self.requested_service_tier {
metadata.insert(
"requestedServiceTier".to_owned(),
Value::String(service_tier.clone()),
);
}
if let Some(service_tier) = &self.upstream_service_tier {
metadata.insert(
"upstreamServiceTier".to_owned(),
Value::String(service_tier.clone()),
);
}
@@ -21,10 +21,9 @@ use thiserror::Error;
use super::protocol::responses::{
ResponseEventSignals, ResponsesSseFailure, response_event_signals,
};
use super::usage::openai_billing_breakdown;
use super::usage::{normalize_service_tier, openai_billing_breakdown};
const CONTENTS_PER_OUTPUT: u32 = 1_024;
const MAX_SERVICE_TIER_BYTES: usize = 64;
/// 单 attempt 的增量 Responses decoder。
///
@@ -42,6 +41,7 @@ pub struct CodexCanonicalDecoder {
reasoning_output_seen: BTreeSet<u32>,
usage_emitted: bool,
semantic_output_seen: bool,
requested_service_tier: Option<String>,
response_service_tier: Option<String>,
timing_signals: ResponseEventSignals,
raw_sse_passthrough: bool,
@@ -137,6 +137,7 @@ impl CodexCanonicalDecoder {
reasoning_output_seen: BTreeSet::new(),
usage_emitted: false,
semantic_output_seen: false,
requested_service_tier: None,
response_service_tier: None,
timing_signals: ResponseEventSignals::default(),
raw_sse_passthrough: false,
@@ -152,6 +153,13 @@ impl CodexCanonicalDecoder {
self
}
/// 请求显式指定的档位优先用于计费;响应档位只在请求未指定时兜底。
#[must_use]
pub fn with_requested_service_tier(mut self, service_tier: Option<&str>) -> Self {
self.requested_service_tier = normalize_service_tier(service_tier);
self
}
pub fn push(&mut self, chunk: &[u8]) -> CodexCanonicalOutcome {
self.timing_signals = ResponseEventSignals::default();
if self.raw_sse_passthrough {
@@ -308,19 +316,14 @@ impl CodexCanonicalDecoder {
}
fn observe_response_service_tier(&mut self, value: &Value) {
let Some(service_tier) = response_object(value)
.and_then(|response| response.get("service_tier"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| {
!value.is_empty()
&& value.len() <= MAX_SERVICE_TIER_BYTES
&& !value.chars().any(char::is_control)
})
else {
let Some(service_tier) = normalize_service_tier(
response_object(value)
.and_then(|response| response.get("service_tier"))
.and_then(Value::as_str),
) else {
return;
};
self.response_service_tier = Some(service_tier.to_ascii_lowercase());
self.response_service_tier = Some(service_tier);
}
/// 构造下发的 wire event。
@@ -836,9 +839,9 @@ impl CodexCanonicalDecoder {
.filter(|model| !model.is_empty())
.unwrap_or(&self.fallback_model)
.to_owned();
let service_tier = response
.get("service_tier")
.and_then(Value::as_str)
let service_tier = self
.requested_service_tier
.as_deref()
.or(self.response_service_tier.as_deref());
if let Some(breakdown) = usage
.filter(|usage| billable_usage_is_complete(response, *usage))
@@ -3,7 +3,8 @@ use std::{sync::Arc, time::Instant};
use futures::{StreamExt, TryStreamExt};
use gateway_protocol::openai::{
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
events::retry_after_seconds_from_body,
events::{self, retry_after_seconds_from_body},
sse::{SseEventDecoder, SseFrame},
};
use reqwest::{
Client, Response as ReqwestResponse,
@@ -130,13 +131,14 @@ impl CodexBackendClient {
});
}
let rate_limit_header_updates = Arc::new(tokio::sync::Mutex::new(Vec::new()));
Ok(CodexBackendStreamingResponse {
body: http_sse_stream(response),
body: http_sse_stream(response, Arc::clone(&rate_limit_header_updates)),
transport: CodexBackendTransport::HttpSse,
turn_state,
set_cookie_headers,
rate_limit_headers,
rate_limit_header_updates: None,
rate_limit_header_updates: Some(rate_limit_header_updates),
turn_state_update: None,
websocket_pool_decision: None,
diagnostics,
@@ -594,20 +596,74 @@ fn websocket_connection_profile(headers: &HeaderMap) -> String {
.join("\0")
}
fn http_sse_stream(response: ReqwestResponse) -> CodexBackendSseStream {
fn http_sse_stream(
response: ReqwestResponse,
rate_limit_updates: CodexRateLimitHeaderUpdates,
) -> CodexBackendSseStream {
let stream: CodexBackendSseStream =
Box::pin(response.bytes_stream().map_err(CodexClientError::Http));
Box::pin(futures::stream::unfold(Some(stream), |stream| async move {
let mut stream = stream?;
match tokio::time::timeout(UPSTREAM_STREAM_IDLE_TIMEOUT, stream.next()).await {
Ok(Some(chunk)) => Some((chunk, Some(stream))),
Ok(None) => None,
Err(_) => Some((
Err(CodexClientError::StreamIdleTimeout {
timeout: UPSTREAM_STREAM_IDLE_TIMEOUT,
}),
None,
)),
let stream: CodexBackendSseStream =
Box::pin(futures::stream::unfold(Some(stream), |stream| async move {
let mut stream = stream?;
match tokio::time::timeout(UPSTREAM_STREAM_IDLE_TIMEOUT, stream.next()).await {
Ok(Some(chunk)) => Some((chunk, Some(stream))),
Ok(None) => None,
Err(_) => Some((
Err(CodexClientError::StreamIdleTimeout {
timeout: UPSTREAM_STREAM_IDLE_TIMEOUT,
}),
None,
)),
}
}));
observe_http_sse_rate_limits(stream, rate_limit_updates)
}
fn observe_http_sse_rate_limits(
stream: CodexBackendSseStream,
updates: CodexRateLimitHeaderUpdates,
) -> CodexBackendSseStream {
Box::pin(futures::stream::unfold(
(stream, SseEventDecoder::default(), updates),
|(mut stream, mut decoder, updates)| async move {
match stream.next().await {
Some(chunk) => {
if let Ok(bytes) = &chunk {
append_http_sse_rate_limit_updates(decoder.push_frames(bytes), &updates)
.await;
}
Some((chunk, (stream, decoder, updates)))
}
None => {
append_http_sse_rate_limit_updates(decoder.finish_frames(), &updates).await;
None
}
}
},
))
}
async fn append_http_sse_rate_limit_updates(
frames: Vec<SseFrame>,
updates: &CodexRateLimitHeaderUpdates,
) {
let mut headers = Vec::new();
for frame in frames {
for event in frame.events() {
if event
.event
.as_deref()
.is_some_and(|event| event != "codex.rate_limits")
{
continue;
}
let Some(rate_limits) = events::parse_rate_limits_event_raw(&event.data) else {
continue;
};
headers.extend(events::rate_limits_to_header_pairs(&rate_limits));
}
}))
}
if !headers.is_empty() {
updates.lock().await.extend(headers);
}
}
@@ -298,10 +298,13 @@ fn model_matches_rule(model: &str, rule: &str) -> bool {
.is_some_and(|suffix| matches!(suffix.as_bytes().first(), Some(b'-' | b'.' | b':')))
}
fn normalize_service_tier(service_tier: Option<&str>) -> Option<String> {
/// 规范化请求或响应携带的服务档位,供观测与计费共用。
pub(crate) fn normalize_service_tier(service_tier: Option<&str>) -> Option<String> {
service_tier
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| {
!value.is_empty() && value.len() <= 64 && !value.chars().any(char::is_control)
})
.map(str::to_ascii_lowercase)
}
@@ -3,12 +3,10 @@ use std::num::NonZeroU32;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use chrono::Utc;
use futures::executor::block_on;
use gateway_core::engine::credential::{
AccountAttemptFeedback, AccountAvailability, AccountFeedbackStats, AccountSelectionPolicy,
AccountStateChange, OpaqueProviderData, ProviderAccountId, ProviderAccountStore as _,
QuotaObservation, RotationStrategy,
AccountStateChange, ProviderAccountId, ProviderAccountStore as _, RotationStrategy,
};
use gateway_core::engine::{
AccountAttemptContext, AttemptContext, CancellationToken, ModelRequestId, RequestAttemptContext,
@@ -24,7 +22,6 @@ use provider_openai::credential::{
};
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use secrecy::ExposeSecret;
use serde_json::json;
use url::Url;
use crate::support::{
@@ -121,7 +118,7 @@ fn selector_with_runtime(
os_version: "6.8".to_owned(),
arch: "x86_64".to_owned(),
terminal: "selector-contract".to_owned(),
verified_at: Utc::now(),
verified_at: chrono::Utc::now(),
});
let http = reqwest::Client::builder().build().expect("HTTP client");
let agent_identity = agent_identity_service(store);
@@ -731,7 +728,7 @@ fn credential_expired_failure_marks_unified_account_expired() {
}
#[test]
fn rate_limited_failure_marks_account_cooldown_with_a_cooldown_deadline() {
fn rate_limited_failure_marks_quota_exhausted_with_a_reset_deadline() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
@@ -757,7 +754,7 @@ fn rate_limited_failure_marks_account_cooldown_with_a_cooldown_deadline() {
.expect("record rate-limit failure");
let account = store.account("acct_primary").expect("account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
@@ -765,28 +762,31 @@ fn rate_limited_failure_marks_account_cooldown_with_a_cooldown_deadline() {
);
}
#[test]
fn usage_limit_exhaustion_is_a_resettable_account_cooldown() {
#[tokio::test]
async fn usage_limit_exhaustion_marks_quota_exhausted_without_usage_probe() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&account,
CodexAccountFailure::UsageLimitExhausted {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record usage-limit exhaustion");
selector
.record_failure(
&account,
CodexAccountFailure::UsageLimitExhausted {
retry_after: Some(Duration::from_secs(30)),
},
)
.await
.expect("record usage-limit exhaustion");
let account = store.account("acct_primary").expect("persisted account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert_eq!(store.quota_reads(), 0);
}
#[test]
@@ -794,26 +794,15 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let reset_at = Utc::now().timestamp() + 24 * 60 * 60;
block_on(
store.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
}),
)
.expect("persist exhausted quota");
block_on(store.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
}))
.expect("mark exhausted");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
@@ -831,33 +820,54 @@ fn rate_limited_failure_does_not_downgrade_persisted_quota_exhaustion() {
.availability(),
AccountAvailability::QuotaExhausted
);
assert!(
store
.account("acct_primary")
.expect("persisted account")
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
}
#[test]
fn rate_limited_failure_ignores_exhaustion_after_its_reset_time() {
fn rate_limited_failure_keeps_longer_existing_quota_deadline() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let existing_deadline = SystemTime::now() + Duration::from_secs(300);
block_on(store.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: Some(existing_deadline),
observed_at: SystemTime::now(),
}))
.expect("mark exhausted");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
&account,
CodexAccountFailure::RateLimited {
retry_after: Some(Duration::from_secs(30)),
},
))
.expect("record shorter rate-limit failure");
assert_eq!(
store
.account("acct_primary")
.expect("persisted account")
.cooldown_until(),
Some(existing_deadline)
);
}
#[test]
fn rate_limited_failure_does_not_consult_stale_quota_snapshot() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_primary", "at-primary");
let account = store.account("acct_primary").expect("account");
let reset_at = Utc::now().timestamp() - 1;
block_on(
store.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
}),
)
.expect("persist expired quota exhaustion");
let selector = selector(&store, Arc::new(TestLeaseCoordinator::default()));
block_on(selector.record_failure(
@@ -873,8 +883,9 @@ fn rate_limited_failure_ignores_exhaustion_after_its_reset_time() {
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::Cooldown
AccountAvailability::QuotaExhausted
);
assert_eq!(store.quota_reads(), 0);
}
#[test]
@@ -907,7 +918,7 @@ fn authenticated_rate_limit_recovers_stale_authentication_state() {
.account("acct_primary")
.expect("persisted account")
.availability(),
AccountAvailability::Cooldown
AccountAvailability::QuotaExhausted
);
}
@@ -985,7 +996,7 @@ fn rate_limited_failures_for_distinct_accounts_do_not_conflict() {
.expect("persisted second account")
.availability(),
],
[AccountAvailability::Cooldown; 2]
[AccountAvailability::QuotaExhausted; 2]
);
}
@@ -600,10 +600,10 @@ async fn reauthorization_uses_revision_advanced_after_start_for_prepared_rotatio
);
}
struct RebindingVerifier;
struct ReplacementIdentityVerifier;
#[async_trait]
impl CodexAccountIdentityVerifier for RebindingVerifier {
impl CodexAccountIdentityVerifier for ReplacementIdentityVerifier {
async fn verify(
&self,
_: &CodexOAuthSecret,
@@ -617,8 +617,9 @@ impl CodexAccountIdentityVerifier for RebindingVerifier {
_: &CodexOAuthSecret,
_: &SecretString,
_: &SecretString,
_: &CodexIdentityExpectation,
expectation: &CodexIdentityExpectation,
) -> Result<CodexIdentityVerification, CodexIdentityVerificationError> {
assert_eq!(expectation, &CodexIdentityExpectation::default());
Ok(CodexIdentityVerification::Complete(profile(
"chatgpt-different-owner",
)))
@@ -699,7 +700,7 @@ async fn reauthorization_updates_changed_free_principal_for_same_account() {
}
#[tokio::test]
async fn reauthorization_rejects_identity_rebinding() {
async fn reauthorization_rebinds_target_record_to_authorized_identity() {
let store = Arc::new(MemoryAccountStore::default());
store
.seed_oauth_credential(ImportCodexOAuthCredential {
@@ -711,7 +712,7 @@ async fn reauthorization_rejects_identity_rebinding() {
enabled: true,
})
.await;
let service = service_with_store(store, Arc::new(RebindingVerifier));
let service = service_with_store(store, Arc::new(ReplacementIdentityVerifier));
let started = service
.start_authorization(StartCodexOAuthAuthorization {
mutation: reauthorization_mutation("request-rebind", "acct_oauth_rebind"),
@@ -723,7 +724,7 @@ async fn reauthorization_rejects_identity_rebinding() {
.query_pairs()
.find_map(|(key, value)| (key == "state").then(|| value.into_owned()))
.expect("state");
let error = service
let completed = service
.complete_authorization(CompleteCodexOAuthAuthorization {
owner_ref: owner_ref(),
flow_id: started.flow_id,
@@ -732,6 +733,24 @@ async fn reauthorization_rejects_identity_rebinding() {
)),
})
.await
.expect_err("identity rebind");
assert_eq!(error, CodexOAuthAdminError::Credential);
.expect("complete identity replacement");
let CompletedCodexOAuthCredential::Reauthorize(prepared) = completed.credential else {
panic!("expected prepared reauthorization");
};
assert_eq!(
prepared.credential.account_id().as_str(),
"acct_oauth_rebind"
);
let replacement = prepared
.replacement_identity
.as_ref()
.expect("replacement identity");
assert_eq!(
replacement.upstream_account_id(),
Some("chatgpt-different-owner")
);
assert_eq!(
replacement.upstream_user_id(),
"user-chatgpt-different-owner"
);
}
@@ -55,6 +55,36 @@ fn parser_treats_any_confirmed_provider_limit_as_exhausted() {
assert!(fact.exhausted());
}
#[test]
fn parser_keeps_full_percent_as_display_without_provider_exhaustion_signal() {
let fact = parse_codex_quota_usage(&json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
}))
.expect("valid full-percent quota");
assert_eq!(fact.remaining_percent(), Some(0));
assert!(!fact.exhausted());
}
#[test]
fn parser_treats_confirmed_limit_reached_as_exhausted_even_below_full_percent() {
let fact = parse_codex_quota_usage(&json!({
"rate_limit": {
"allowed": false,
"limit_reached": true,
"primary_window": {"used_percent": 98, "reset_at": 1_900_000_000}
}
}))
.expect("valid provider-confirmed quota");
assert_eq!(fact.remaining_percent(), Some(2));
assert!(fact.exhausted());
}
#[test]
fn parser_does_not_infer_exhaustion_from_unknown_credit_fields() {
let fact = parse_codex_quota_usage(&json!({
@@ -181,6 +211,14 @@ fn blocked_network_quota_service(store: &Arc<MemoryAccountStore>) -> CodexCreden
}
async fn create_account(store: &Arc<MemoryAccountStore>, account_id: &str) {
create_account_with_enabled(store, account_id, true).await;
}
async fn create_account_with_enabled(
store: &Arc<MemoryAccountStore>,
account_id: &str,
enabled: bool,
) {
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
@@ -188,7 +226,7 @@ async fn create_account(store: &Arc<MemoryAccountStore>, account_id: &str) {
secret: secret(&format!("token-{account_id}")),
verified_account: profile(&format!("chatgpt-{account_id}")),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
enabled,
})
.await;
}
@@ -361,12 +399,66 @@ async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision
);
}
#[tokio::test]
async fn passive_full_percent_with_explicit_allowance_keeps_account_ready() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_passive_full_allowed").await;
let account = store.account("acct_passive_full_allowed").expect("account");
let service = quota_service(&store);
let headers = vec![
("x-codex-active-limit".to_owned(), "codex".to_owned()),
("x-codex-primary-used-percent".to_owned(), "100".to_owned()),
(
"x-codex-primary-reset-at".to_owned(),
"1900000000".to_owned(),
),
("x-codex-allowed".to_owned(), "true".to_owned()),
("x-codex-limit-reached".to_owned(), "false".to_owned()),
];
assert!(
service
.synchronize_passive_headers(&account, &headers)
.await
.expect("passive quota")
);
let current = store
.account("acct_passive_full_allowed")
.expect("current account");
let snapshot = service
.read_account(account.id())
.await
.expect("read quota")
.expect("quota snapshot");
assert_eq!(
(
current.availability(),
current.cooldown_until(),
snapshot.fact().remaining_percent(),
snapshot.fact().exhausted(),
),
(AccountAvailability::Ready, None, Some(0), false)
);
}
#[tokio::test]
async fn manual_quota_refresh_updates_account_state_from_provider_response() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_manual_quota_state";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Expired,
reason: Some("stale_expired_state".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("seed stale expired state");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
@@ -438,6 +530,61 @@ async fn manual_quota_refresh_updates_account_state_from_provider_response() {
);
}
#[tokio::test]
async fn manual_quota_refresh_preserves_disabled_account_state() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_disabled_quota_state";
create_account_with_enabled(&store, account_id, false).await;
let account = store.account(account_id).expect("created disabled account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::Expired,
reason: Some("disabled_expired_state".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("seed disabled account state");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header(
"authorization",
format!("Bearer token-{account_id}"),
))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 20, "reset_at": 1_900_000_000}
}
})))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
service
.refresh_account(account.id())
.await
.expect("refresh disabled account quota");
let current = store
.account(account_id)
.expect("disabled account after quota refresh");
assert!(!current.enabled());
assert_eq!(current.availability(), AccountAvailability::Expired);
assert!(store.has_quota(account_id));
}
#[tokio::test]
async fn manual_quota_auth_rejection_does_not_invalidate_refreshable_credential() {
let store = Arc::new(MemoryAccountStore::default());
@@ -697,12 +844,18 @@ async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_q
service.synchronize().await.expect("quota synchronization");
let reconciled = store
.account("acct_periodic_stuck_cooldown")
.expect("reconciled account");
assert_eq!(
store
.account("acct_periodic_stuck_cooldown")
.expect("reconciled account")
.availability(),
AccountAvailability::QuotaExhausted
(reconciled.availability(), reconciled.cooldown_until()),
(
AccountAvailability::QuotaExhausted,
Some(
SystemTime::UNIX_EPOCH
+ Duration::from_secs(u64::try_from(reset_at).expect("positive reset")),
),
)
);
}
@@ -741,12 +894,14 @@ async fn periodic_quota_synchronization_does_not_downgrade_exhaustion_on_rate_li
service.synchronize().await.expect("quota synchronization");
assert_eq!(
store
.account("acct_periodic_exhausted")
.expect("persisted account")
.availability(),
AccountAvailability::QuotaExhausted
let account = store
.account("acct_periodic_exhausted")
.expect("persisted account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
}
@@ -877,6 +1032,60 @@ async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() {
assert_eq!(summary.transient, 1);
}
#[tokio::test]
async fn periodic_quota_synchronization_waits_for_quota_cooldown_deadline_with_grace() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_exhausted_cooldown").await;
let account = store
.account("acct_periodic_exhausted_cooldown")
.expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: SystemTime::now().checked_sub(Duration::from_secs(29)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted with recent cooldown");
let service = blocked_network_quota_service(&store);
let waiting = service
.synchronize()
.await
.expect("quota synchronization before grace");
assert_eq!(waiting, CodexQuotaSyncSummary::default());
let current = store
.account("acct_periodic_exhausted_cooldown")
.expect("current account");
store
.apply_state_change(AccountStateChange {
account_id: current.id().clone(),
expected_revision: current.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: SystemTime::now().checked_sub(Duration::from_secs(31)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted past grace");
let due = service
.synchronize()
.await
.expect("quota synchronization after grace");
let repeated = service
.synchronize()
.await
.expect("quota synchronization repeated");
assert_eq!(due.transient, 1);
assert_eq!(repeated, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
let store = Arc::new(MemoryAccountStore::default());
@@ -955,8 +1164,8 @@ async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_
server.uri(),
);
let first = service.synchronize().await.expect("initial refresh");
assert_eq!(first.transient, 1);
let waiting = service.synchronize().await.expect("initial cooldown wait");
assert_eq!(waiting, CodexQuotaSyncSummary::default());
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
@@ -987,6 +1196,77 @@ async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_
assert_eq!(repeated, CodexQuotaSyncSummary::default());
}
#[tokio::test]
async fn periodic_quota_synchronization_reschedules_a_confirmed_new_reset() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_periodic_new_reset";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
let initial_reset = SystemTime::UNIX_EPOCH + Duration::from_secs(1);
let confirmed_reset = SystemTime::UNIX_EPOCH + Duration::from_secs(2);
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: Some(initial_reset),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted with due reset");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": false,
"limit_reached": true,
"primary_window": {"used_percent": 100, "reset_at": 2}
}
})))
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let first = service.synchronize().await.expect("first reset recheck");
let persisted_reset = store
.account(account_id)
.expect("account after first recheck")
.cooldown_until();
let second = service.synchronize().await.expect("new reset recheck");
let repeated = service.synchronize().await.expect("same reset throttled");
let request_count = server
.received_requests()
.await
.expect("received requests")
.len();
assert_eq!(
(
first.exhausted,
persisted_reset,
second.exhausted,
repeated,
request_count,
),
(
1,
Some(confirmed_reset),
1,
CodexQuotaSyncSummary::default(),
2,
)
);
}
#[tokio::test]
async fn invalid_quota_json_for_one_account_does_not_abort_the_batch() {
let store = Arc::new(MemoryAccountStore::default());
@@ -665,6 +665,74 @@ async fn expired_access_token_refreshes_even_when_next_refresh_is_deferred() {
);
}
#[tokio::test]
async fn expired_access_token_refreshes_during_active_cooldown() {
let store = Arc::new(MemoryAccountStore::default());
seed_due_account(&store).await;
let original = store.account("acct_refresh").expect("seeded account");
let mut expired_profile = profile("chatgpt-acct_refresh");
expired_profile.access_token_expires_at =
Some(chrono::Utc::now() - chrono::Duration::minutes(5));
store
.repository()
.rotate_oauth_secret(RotateCodexCredential {
account_id: "acct_refresh".to_owned(),
expected_credential_revision: original.revision().get(),
secret: secret("old-access"),
verified_account: expired_profile,
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::hours(12)),
})
.await
.expect("seed expired access token with deferred refresh");
let deferred = store.account("acct_refresh").expect("deferred account");
let cooldown_until = SystemTime::now() + Duration::from_secs(12 * 60 * 60);
store
.repository()
.apply_state(
&deferred,
AccountAvailability::Cooldown,
Some("rate_limited".to_owned()),
Some(cooldown_until),
SystemTime::now(),
)
.await
.expect("mark account cooling down");
let refresher = Arc::new(Refresher::new(Ok(success_tokens())));
let service = CodexCredentialRefreshService::new(
store.repository(),
refresher.clone(),
Arc::new(VerifiedIdentity),
Arc::new(RefreshLeases {
available: true,
requests: Mutex::new(Vec::new()),
}),
Arc::new(CountingCredentialState::default()),
runtime_policy(),
);
let outcomes = service.refresh_due().await.expect("refresh due");
assert!(matches!(
outcomes.as_slice(),
[CodexCredentialRefreshOutcome::Refreshed {
account_id,
credential_revision: 3,
}] if account_id == "acct_refresh"
));
assert_eq!(
refresher.seen.lock().expect("seen tokens lock").as_slice(),
["rt-old-access"]
);
let refreshed = store.account("acct_refresh").expect("refreshed account");
assert_eq!(refreshed.availability(), AccountAvailability::Cooldown);
assert_eq!(refreshed.cooldown_until(), Some(cooldown_until));
assert!(
refreshed
.access_token_expires_at()
.is_some_and(|expires_at| expires_at > SystemTime::now())
);
}
#[tokio::test]
async fn unavailable_refresh_lease_prevents_duplicate_token_exchange() {
let (store, refresher, service) = setup(Ok(success_tokens()), false).await;
@@ -7,7 +7,7 @@ use base64::{Engine as _, engine::general_purpose::STANDARD};
use chrono::Utc;
use futures::StreamExt;
use gateway_core::engine::credential::{
AccountAvailability, AccountFeedbackStats, ProviderAccountId,
AccountAvailability, AccountFeedbackStats, ProviderAccountId, ProviderAccountStore as _,
};
use gateway_core::engine::provider::{Provider as _, ProviderRequest};
use gateway_core::engine::{
@@ -29,14 +29,14 @@ use provider_openai::credential::{
use provider_openai::transport::CodexWebSocketPool;
use provider_openai::transport::profile::{CodexWireProfile, CodexWireProfileState};
use provider_openai::{CodexProvider, OFFICIAL_CODEX_BASE_URL};
use serde_json::{Map, json};
use serde_json::{Map, Value, json};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::{TcpListener, TcpStream},
sync::oneshot,
time::timeout,
};
use wiremock::matchers::{method, path};
use wiremock::matchers::{header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use crate::support::{
@@ -827,13 +827,23 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
.expect("prepare provider stream");
let mut observed_service_tier = None;
let mut upstream_service_tier = None;
while let Some(event) = stream.next().await {
let event = event.expect("provider event");
if let Some(service_tier) = event
.response_observation()
.and_then(|observation| observation.service_tier())
{
observed_service_tier = Some(service_tier.to_owned());
if let Some(observation) = event.response_observation() {
if let Some(service_tier) = observation.service_tier() {
observed_service_tier = Some(service_tier.to_owned());
}
upstream_service_tier = observation
.provider_metadata()
.and_then(|metadata| serde_json::from_str::<Value>(metadata.as_json()).ok())
.and_then(|metadata| {
metadata
.get("upstreamServiceTier")
.and_then(Value::as_str)
.map(str::to_owned)
})
.or(upstream_service_tier);
}
if event
.canonical_facts()
@@ -846,11 +856,12 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
drop(stream);
assert_eq!(affinity.binding_count(), 1);
assert_eq!(observed_service_tier.as_deref(), Some("default"));
assert_eq!(observed_service_tier.as_deref(), Some("priority"));
assert_eq!(upstream_service_tier.as_deref(), Some("default"));
}
#[tokio::test]
async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown() {
async fn response_failed_before_semantic_output_is_atomic_and_persists_quota_lock() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_atomic_failure";
create_account(&store, account_id).await;
@@ -895,7 +906,7 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown(
vec!["response.created", "response.failed"]
);
let account = store.account(account_id).expect("rate-limited account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
@@ -906,6 +917,68 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown(
let _ = server.await;
}
#[tokio::test]
async fn usage_limit_failure_marks_quota_exhausted_without_usage_probe() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_usage_limit_request_path";
create_account(&store, account_id).await;
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header("authorization", format!("Bearer at-{account_id}")))
.respond_with(ResponseTemplate::new(500))
.expect(0)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_usage_limit_confirm\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"status_code\":429,\"retry_after_seconds\":86400,\"response\":{\"id\":\"resp_usage_limit_confirm\",\"status\":\"failed\",\"error\":{\"code\":\"usage_limit_reached\",\"message\":\"usage limit reached\"}}}\n\n"
)),
)
.mount(&server)
.await;
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", http_generate_operation()),
context("req_usage_limit_confirm", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let failure = loop {
match stream.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("usage-limit failure must surface a typed failure"),
}
};
assert_eq!(failure.kind(), ProviderErrorKind::QuotaExhausted);
let account = store.account(account_id).expect("usage-limit account");
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert!(!store.has_quota(account_id));
let requests = server.received_requests().await.expect("received requests");
assert_eq!(
requests
.iter()
.filter(|request| request.url.path() == "/api/codex/usage")
.count(),
0
);
}
#[tokio::test]
async fn ordinary_request_should_forward_created_before_later_failed_chunk() {
let store = Arc::new(MemoryAccountStore::default());
@@ -1207,8 +1280,8 @@ async fn response_failed_after_semantic_output_is_exposed_and_not_replay_safe()
async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_state() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_disabled_diagnostic";
create_account_with_enabled(&store, account_id, false).await;
let account = store.account(account_id).expect("disabled test account");
create_account(&store, account_id).await;
let account = store.account(account_id).expect("test account");
store
.repository()
.apply_state(
@@ -1220,6 +1293,10 @@ async fn disabled_account_diagnostic_uses_upstream_without_persisting_account_st
)
.await
.expect("seed quota-exhausted state");
store
.set_enabled(account.id(), false)
.await
.expect("disable test account");
let affinity = Arc::new(MemorySessionAffinity::default());
let server = MockServer::start().await;
@@ -1323,6 +1400,76 @@ async fn successful_response_preserves_quota_exhaustion_from_rate_limit_headers(
assert!(store.has_quota(account_id));
}
#[tokio::test]
async fn successful_http_sse_rate_limit_event_persists_structured_exhaustion() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_http_sse_exhausted";
create_account(&store, account_id).await;
let reset_at = 1_900_000_000_u64;
let rate_limit_event = json!({
"type": "codex.rate_limits",
"rate_limits": {
"allowed": false,
"limit_reached": true,
"primary": {
"used_percent": 42,
"window_minutes": 300,
"reset_at": reset_at,
},
},
});
let completed_event = json!({
"type": "response.completed",
"response": {
"id": "resp_http_sse_exhausted",
"model": "gpt-5.4",
"status": "completed",
"output": [],
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
},
});
let body = format!(
"event: codex.rate_limits\ndata: {rate_limit_event}\n\nevent: response.completed\ndata: {completed_event}\n\n"
);
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(body),
)
.mount(&server)
.await;
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", http_generate_operation()),
context("req_http_sse_exhausted", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
while let Some(event) = stream.next().await {
event.expect("successful upstream response");
}
let account = store
.account(account_id)
.expect("account after HTTP SSE response");
assert_eq!(
(
account.availability(),
account.cooldown_until(),
store.has_quota(account_id),
),
(
AccountAvailability::QuotaExhausted,
Some(SystemTime::UNIX_EPOCH + Duration::from_secs(reset_at)),
true,
)
);
}
#[test]
fn request_observation_reads_openai_metadata_without_changing_the_operation() {
let store = Arc::new(MemoryAccountStore::default());
@@ -203,18 +203,40 @@ fn decoder_should_emit_calculated_cost_for_complete_known_model_usage() {
}
#[test]
fn decoder_should_keep_response_service_tier_when_terminal_event_omits_it() {
fn decoder_should_prefer_requested_service_tier_over_response_tier_for_billing() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"priority\"}}\n\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"default\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_fast_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
);
let mut decoder = CodexCanonicalDecoder::new("fallback");
let mut decoder =
CodexCanonicalDecoder::new("fallback").with_requested_service_tier(Some("priority"));
let events = decoder
.push(body.as_bytes())
.expect("canonical priority response");
assert_eq!(decoder.response_service_tier(), Some("default"));
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
event,
GatewayEvent::CalculatedCost(cost)
if cost.total().amount().scaled() == 6_875_000
)));
}
#[test]
fn decoder_should_use_response_service_tier_when_request_omits_it() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_response_tier_cost\",\"model\":\"gpt-5.4\",\"service_tier\":\"priority\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_response_tier_cost\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10,\"input_tokens_details\":{\"cached_tokens\":25,\"cache_write_tokens\":0},\"total_tokens\":110}}}\n\n",
);
let mut decoder = CodexCanonicalDecoder::new("fallback");
let events = decoder
.push(body.as_bytes())
.expect("canonical response-tier priority response");
assert_eq!(decoder.response_service_tier(), Some("priority"));
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
event,
@@ -223,6 +245,31 @@ fn decoder_should_keep_response_service_tier_when_terminal_event_omits_it() {
)));
}
#[test]
fn websocket_decoder_should_prefer_requested_service_tier_over_response_tier_for_billing() {
let created = websocket_event_to_sse_frame(
r#"{"type":"response.created","response":{"id":"resp_ws_fast_cost","model":"gpt-5.4","service_tier":"default"}}"#,
)
.expect("created frame");
let completed = websocket_event_to_sse_frame(
r#"{"type":"response.completed","response":{"id":"resp_ws_fast_cost","model":"gpt-5.4","status":"completed","usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":25,"cache_write_tokens":0},"total_tokens":110}}}"#,
)
.expect("completed frame");
let mut decoder = CodexCanonicalDecoder::new("fallback")
.with_requested_service_tier(Some("priority"))
.with_raw_sse_passthrough();
let events = decoder
.push(format!("{created}{completed}").as_bytes())
.expect("canonical priority WebSocket response");
assert_eq!(decoder.response_service_tier(), Some("default"));
assert!(canonical_facts(&events).into_iter().any(|event| matches!(
event,
GatewayEvent::CalculatedCost(cost)
if cost.total().amount().scaled() == 6_875_000
)));
}
#[test]
fn decoder_should_normalize_text_usage_and_completion() {
let body = concat!(
@@ -269,6 +269,54 @@ async fn codex_backend_client_should_parse_retry_after_from_rate_limit_error_bod
assert_eq!(retry_after_seconds, Some(12));
}
#[tokio::test]
async fn codex_backend_http_sse_should_capture_structured_rate_limit_event_updates() {
let server = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("POST"))
.and(wiremock::matchers::path("/codex/responses"))
.respond_with(
wiremock::ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(concat!(
"event: codex.rate_limits\n",
"data: {\"type\":\"codex.rate_limits\",\"rate_limits\":{\"allowed\":true,\"limit_reached\":false,\"primary\":{\"used_percent\":100,\"window_minutes\":300,\"reset_at\":1893456300}}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_http_rate_limits\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":1,\"total_tokens\":2}}}\n\n"
)),
)
.mount(&server)
.await;
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
test_wire_profile(),
);
let mut request = codex_request("gpt-5.5", "", Vec::new());
request.force_http_sse = true;
let response = client
.create_response(
&request,
request_context("req_http_rate_limit_event", Some("acct")),
)
.await
.expect("HTTP SSE response");
let structured = response
.rate_limit_headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-codex-allowed" | "x-codex-limit-reached"))
.cloned()
.collect::<Vec<_>>();
assert_eq!(
structured,
vec![
("x-codex-allowed".to_owned(), "true".to_owned()),
("x-codex-limit-reached".to_owned(), "false".to_owned()),
]
);
}
#[tokio::test(start_paused = true)]
async fn codex_backend_http_sse_should_fail_after_five_minutes_without_stream_data() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -942,6 +942,7 @@ fn prepared_rotation(
provider_kind,
expected_credential_revision: Revision::new(expected_revision.get())
.map_err(|_| provider_error(ProviderAdminErrorKind::Internal))?,
replacement_identity: None,
name: profile.name,
email: profile.email,
plan_type: profile.plan_type,
@@ -4,7 +4,7 @@ use std::fmt;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use chrono::Utc;
use chrono::{DateTime, Utc};
use gateway_core::engine::credential::{
AccountAvailability, AccountAvailabilityPolicy, AccountCandidate, AccountFeedbackStats,
AccountSelectionContext, AccountSelector, ProviderAccount, ProviderAccountId,
@@ -17,7 +17,7 @@ use gateway_core::provider_ports::{
use gateway_core::routing::{ProviderKind, UpstreamModelId};
use super::catalog::{GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialQuotaService};
use super::repository::GrokCredentialRepository;
use super::repository::{GrokCredentialRepository, GrokCredentialRepositoryError};
use super::types::{GrokCredentialAvailability, UpdateGrokCredentialState};
use crate::{
GrokCredentialFailure, GrokCredentialFeedbackFuture, GrokSessionBinding, GrokSessionSelection,
@@ -29,7 +29,6 @@ const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
const STREAM_INTERRUPTION_COOLDOWN: Duration = Duration::from_secs(30);
const UNAUTHORIZED_COOLDOWN: Duration = Duration::from_secs(60);
const PAYMENT_REQUIRED_COOLDOWN: Duration = Duration::from_secs(60);
const MODEL_ACCESS_DENIED_COOLDOWN: Duration = Duration::from_secs(5 * 60);
const MODEL_QUOTA_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
const MAX_MODEL_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
@@ -369,6 +368,66 @@ impl GrokAccountSessionSelector {
))
.await;
}
async fn persist_account_failure(
&self,
session: &SelectedGrokSession,
availability: GrokCredentialAvailability,
reason: &'static str,
cooldown_until: Option<DateTime<Utc>>,
observed_at: chrono::DateTime<Utc>,
) {
let update = |expected_revision| UpdateGrokCredentialState {
account_id: session.account_id().clone(),
expected_revision,
availability,
availability_reason: Some(reason.to_owned()),
cooldown_until,
observed_at,
};
match self
.repository
.update_state(&update(session.credential_revision()))
.await
{
Ok(()) => return,
Err(
GrokCredentialRepositoryError::Conflict
| GrokCredentialRepositoryError::StaleCredentialRevision,
) => {}
Err(error) => {
tracing::warn!(
account_id = %session.account_id(),
error = %error,
"xAI account failure state write failed"
);
return;
}
}
let current = match self.repository.load_current(session.account_id()).await {
Ok(current) if current.account.enabled() => current,
Ok(_) => return,
Err(error) => {
tracing::warn!(
account_id = %session.account_id(),
error = %error,
"xAI account failure state reload failed"
);
return;
}
};
if let Err(error) = self
.repository
.update_state(&update(current.account.revision()))
.await
{
tracing::warn!(
account_id = %session.account_id(),
error = %error,
"xAI account failure state retry failed"
);
}
}
}
impl fmt::Debug for GrokAccountSessionSelector {
@@ -415,15 +474,11 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
"upstream_free_quota_exhausted",
None,
),
GrokCredentialFailure::PaymentRequired { retry_after } => {
let retry_after = bounded_cooldown(
retry_after,
PAYMENT_REQUIRED_COOLDOWN,
MAX_RATE_LIMIT_COOLDOWN,
);
self.record_runtime_cooldown(session, retry_after).await;
return;
}
GrokCredentialFailure::PaymentRequired { .. } => (
GrokCredentialAvailability::QuotaExhausted,
"upstream_payment_quota_exhausted",
None,
),
GrokCredentialFailure::ModelQuotaExhausted {
upstream_model,
retry_after,
@@ -452,26 +507,31 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)
.min(MAX_RATE_LIMIT_COOLDOWN);
self.record_runtime_cooldown(session, retry_after).await;
return;
(
GrokCredentialAvailability::Cooldown,
"upstream_rate_limited",
persistent_cooldown_until(observed_at, retry_after),
)
}
GrokCredentialFailure::StreamInterrupted => {
self.record_runtime_cooldown(session, STREAM_INTERRUPTION_COOLDOWN)
.await;
return;
(
GrokCredentialAvailability::Cooldown,
"upstream_stream_interrupted",
persistent_cooldown_until(observed_at, STREAM_INTERRUPTION_COOLDOWN),
)
}
};
let (availability, reason, cooldown_until) = persistent;
let _ = self
.repository
.update_state(&UpdateGrokCredentialState {
account_id: session.account_id().clone(),
expected_revision: session.credential_revision(),
availability,
availability_reason: Some(reason.to_owned()),
cooldown_until,
observed_at,
})
.await;
self.persist_account_failure(
session,
availability,
reason,
cooldown_until,
observed_at,
)
.await;
})
}
@@ -493,7 +553,6 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
account.availability(),
AccountAvailability::Unknown
| AccountAvailability::Cooldown
| AccountAvailability::QuotaExhausted
| AccountAvailability::Expired
)
{
@@ -540,6 +599,15 @@ fn bounded_cooldown(
.min(maximum)
}
fn persistent_cooldown_until(
observed_at: DateTime<Utc>,
duration: Duration,
) -> Option<DateTime<Utc>> {
chrono::Duration::from_std(duration)
.ok()
.map(|duration| observed_at + duration)
}
fn insert_runtime_cooldown(
cooldowns: &mut std::collections::BTreeMap<ProviderAccountId, RuntimeCooldown>,
account_id: &ProviderAccountId,
+1
View File
@@ -4,6 +4,7 @@ mod admin;
pub mod config;
pub mod credential;
mod provider;
mod reasoning_replay;
pub mod transport;
use std::path::Path;
+78 -53
View File
@@ -5,7 +5,6 @@ use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime};
use async_trait::async_trait;
use base64::Engine as _;
use futures::{StreamExt, future::BoxFuture};
use gateway_core::engine::continuation::ContinuationBinding;
use gateway_core::engine::credential::{
@@ -48,9 +47,13 @@ use crate::credential::{
GrokCredentialRecovery, GrokCredentialRecoveryOutcome, GrokCredentialRefreshOutcome,
GrokCredentialRefreshService, GrokQuotaError,
};
use crate::reasoning_replay::{
GrokReasoningReplay, GrokReasoningReplayCapture, GrokReasoningReplayKey,
valid_reasoning_ciphertext,
};
use crate::transport::canonical::GrokCanonicalDecoder;
use crate::transport::config::XAI_PROVIDER_NAME;
use crate::transport::headers::{GrokClientIdentity, GrokHeader, build_grok_headers};
use crate::transport::headers::{GrokClientIdentity, build_grok_headers};
use crate::transport::profile::{GROK_CLI_RELEASE_POLL_INTERVAL, GrokCliReleaseService};
use crate::transport::{
GROK_RESPONSES_URL, GrokCompactionDecodeError, GrokCompactionRequest,
@@ -66,8 +69,6 @@ const HTTP_SSE_TRANSPORT: &str = "http_sse";
const DEFAULT_GROK_MODEL: &str = "grok-4.5";
const XAI_SESSION_STATE_MAX_BYTES: usize = 8 * 1024 * 1024;
const XAI_SESSION_OUTPUT_LIMIT: usize = 4_096;
const MIN_REASONING_CIPHERTEXT_BYTES: usize = 50;
const MIN_REASONING_CIPHERTEXT_ENTROPY: f64 = 0.85;
const REASONING_DECODE_FAILED_CODE: &str = "reasoning_decode_failed";
const RESPONSE_NOT_FOUND_CODE: &str = "not_found";
@@ -82,6 +83,7 @@ pub struct GrokBuildProvider {
credential_recovery: Arc<dyn GrokCredentialRecovery>,
account_feedback: Arc<AccountFeedbackStats>,
client_identity: GrokClientIdentity,
reasoning_replay: GrokReasoningReplay,
wire_profile: XaiWireProfileState,
responses_url: Url,
}
@@ -105,6 +107,7 @@ impl GrokBuildProvider {
credential_recovery,
account_feedback,
client_identity: GrokClientIdentity::new(),
reasoning_replay: GrokReasoningReplay::new(),
wire_profile,
responses_url,
})
@@ -217,6 +220,38 @@ impl GrokBuildProvider {
selected.account_id(),
request_input.as_slice(),
)?;
let inherited_replay_session_id = context
.continuation()
.is_none()
.then(|| {
previous_session
.as_ref()
.and_then(|state| state.session_id.as_deref())
})
.flatten();
let reasoning_replay_key = upstream_request
.reasoning_replay_session_id()
.or(inherited_replay_session_id)
.and_then(|session_id| {
self.reasoning_replay.key(
candidate.upstream_model().as_str(),
session_id,
selected.account_id().as_str(),
)
});
if !upstream_request.has_previous_response_id()
&& let (Some(key), Some(input)) = (
reasoning_replay_key.as_ref(),
upstream_request.replay_input_items(),
)
&& let Some(input) = self.reasoning_replay.apply(key, &input)
{
upstream_request
.set_replay_input(input)
.map_err(map_request_error)?;
}
let reasoning_replay_capture =
reasoning_replay_key.map(|key| self.reasoning_replay.capture(key));
let session_capture = (!matches!(
context.continuation(),
Some(ContinuationBinding::Pinned(_) | ContinuationBinding::External(_))
@@ -244,6 +279,7 @@ impl GrokBuildProvider {
context,
session: Arc::clone(&selected),
session_capture,
reasoning_replay_capture,
},
);
let stream = ProviderStream::new(metadata, events, selected);
@@ -274,7 +310,7 @@ impl GrokBuildProvider {
.map(|previous| ProviderAccountId::new(previous.account_id.clone()))
.transpose()
.map_err(|_| protocol_not_sent())?;
let upstream_session_id = previous_session
let inherited_session_id = previous_session
.as_ref()
.and_then(|previous| previous.session_id.clone());
let upstream_request = GrokCompactionRequest::encode(
@@ -283,6 +319,12 @@ impl GrokBuildProvider {
context.client_api_key_ref(),
)
.map_err(map_request_error)?;
let explicit_replay_session_id = upstream_request
.reasoning_replay_session_id()
.map(str::to_owned);
let upstream_session_id = inherited_session_id
.clone()
.or_else(|| explicit_replay_session_id.clone());
let selected = Arc::new(
select_grok_session(
self.selector.as_ref(),
@@ -293,6 +335,16 @@ impl GrokBuildProvider {
)
.await?,
);
let reasoning_replay_key = explicit_replay_session_id
.as_deref()
.or(inherited_session_id.as_deref())
.and_then(|session_id| {
self.reasoning_replay.key(
candidate.upstream_model().as_str(),
session_id,
selected.account_id().as_str(),
)
});
let allows_account_state_mutation = selected.allows_account_state_mutation();
let metadata = provider_call_metadata(candidate, &selected)?;
let events = cold_compaction_http_sse_stream(
@@ -308,6 +360,8 @@ impl GrokBuildProvider {
upstream_session_id,
context,
session: Arc::clone(&selected),
reasoning_replay: self.reasoning_replay.clone(),
reasoning_replay_key,
},
);
let stream = ProviderStream::new(metadata, events, selected);
@@ -715,46 +769,6 @@ fn portable_output_item(mut item: Value, strip_opaque: bool) -> Option<Value> {
Some(item)
}
fn valid_reasoning_ciphertext(value: &str) -> bool {
if value.is_empty()
|| value != value.trim()
|| value.len() > XAI_SESSION_STATE_MAX_BYTES
|| value.starts_with("gAAAA")
|| value.contains('=')
{
return false;
}
let Ok(decoded) = base64::engine::general_purpose::STANDARD_NO_PAD.decode(value) else {
return false;
};
decoded.len() >= MIN_REASONING_CIPHERTEXT_BYTES
&& byte_entropy_ratio(&decoded) >= MIN_REASONING_CIPHERTEXT_ENTROPY
}
fn byte_entropy_ratio(value: &[u8]) -> f64 {
if value.is_empty() {
return 0.0;
}
let mut counts = [0_u32; 256];
for byte in value {
counts[usize::from(*byte)] += 1;
}
let size = value.len() as f64;
let entropy = counts
.into_iter()
.filter(|count| *count > 0)
.map(|count| {
let probability = f64::from(count) / size;
-probability * probability.log2()
})
.sum::<f64>();
let symbols = value.len().min(256);
if symbols <= 1 {
return 0.0;
}
entropy / (symbols as f64).log2()
}
fn has_readable_reasoning(item: &Map<String, Value>) -> bool {
["summary", "content"].into_iter().any(|field| {
item.get(field)
@@ -898,6 +912,7 @@ struct GrokStreamAttempt {
context: AttemptContext,
session: Arc<SelectedGrokSession>,
session_capture: Option<GrokSessionCapture>,
reasoning_replay_capture: Option<GrokReasoningReplayCapture>,
}
struct GrokCompactionStreamAttempt {
@@ -910,6 +925,8 @@ struct GrokCompactionStreamAttempt {
upstream_session_id: Option<String>,
context: AttemptContext,
session: Arc<SelectedGrokSession>,
reasoning_replay: GrokReasoningReplay,
reasoning_replay_key: Option<GrokReasoningReplayKey>,
}
struct AcceptedGrokInference {
@@ -1044,9 +1061,11 @@ fn cold_compaction_http_sse_stream(
upstream_session_id,
context,
session,
reasoning_replay,
reasoning_replay_key,
} = attempt;
Box::pin(async_stream::try_stream! {
let mut headers = build_grok_headers(
let headers = build_grok_headers(
&wire_profile,
&session,
&client_identity,
@@ -1055,12 +1074,6 @@ fn cold_compaction_http_sse_stream(
None,
&upstream_model,
);
if let Some(upstream_session_id) = upstream_session_id.as_deref() {
headers.push(GrokHeader::sensitive(
"x-grok-session-id",
crate::credential::SecretValue::new(upstream_session_id.to_owned()),
));
}
let body = request.to_json_bytes().map_err(map_request_error)?;
let inference_request = GrokInferenceRequest::new(
responses_url,
@@ -1147,6 +1160,9 @@ fn cold_compaction_http_sse_stream(
if upstream_completed && session.allows_account_state_mutation() {
selector.record_success(&session).await;
}
if let Some(key) = reasoning_replay_key.as_ref() {
reasoning_replay.clear(key);
}
yield ProviderEvent::canonical_with_wire(vec![GatewayEvent::Started(started)], created);
yield ProviderEvent::wire(output_done);
let mut terminal_facts = facts.accounting;
@@ -1245,6 +1261,7 @@ fn cold_http_sse_stream(
context,
session,
mut session_capture,
mut reasoning_replay_capture,
} = attempt;
Box::pin(async_stream::try_stream! {
if context.cancellation().is_cancelled() {
@@ -1364,6 +1381,9 @@ fn cold_http_sse_stream(
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Completed(_)));
attach_xai_session_update(&mut events, &mut session_capture)?;
if let Some(capture) = reasoning_replay_capture.as_mut() {
capture.observe(&events);
}
if completed && session.allows_account_state_mutation() {
selector.record_success(&session).await;
}
@@ -1389,6 +1409,9 @@ fn cold_http_sse_stream(
.flat_map(ProviderEvent::canonical_facts)
.any(|event| matches!(event, GatewayEvent::Completed(_)));
attach_xai_session_update(&mut final_events, &mut session_capture)?;
if let Some(capture) = reasoning_replay_capture.as_mut() {
capture.observe(&final_events);
}
if completed && session.allows_account_state_mutation() {
selector.record_success(&session).await;
}
@@ -1871,7 +1894,9 @@ const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60);
const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60);
const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30);
const QUOTA_CATALOG_INTERVAL: Duration = Duration::from_secs(5 * 60);
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
// grok2api 对 Free/未知额度耗尽采用 24 小时恢复探测;当前公共账号状态没有
// 独立的 paid period-end 字段,因此 Build 账号统一使用同一保守下限。
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
const CLI_RELEASE_WORKER_OWNER: &str = "xai-cli-release";
pub(crate) fn worker_contributions(
@@ -0,0 +1,687 @@
//! Grok Build 的账号隔离 reasoning replay。
use std::collections::{BTreeMap, BTreeSet};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::{Duration, Instant};
use base64::Engine as _;
use gateway_core::event::ProviderEvent;
use serde_json::{Map, Value};
const REASONING_REPLAY_TTL: Duration = Duration::from_secs(60 * 60);
const REASONING_REPLAY_MAX_ENTRIES: usize = 1_024;
const REASONING_REPLAY_MAX_BYTES: usize = 8 * 1024 * 1024;
const REASONING_REPLAY_MAX_OUTPUT_ITEMS: usize = 4_096;
const MIN_REASONING_CIPHERTEXT_BYTES: usize = 50;
const MIN_REASONING_CIPHERTEXT_ENTROPY: f64 = 0.85;
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
pub(crate) struct GrokReasoningReplayKey {
model: String,
session_id: String,
account_id: String,
}
#[derive(Clone)]
pub(crate) struct GrokReasoningReplay {
state: Arc<Mutex<GrokReasoningReplayState>>,
}
#[derive(Default)]
struct GrokReasoningReplayState {
entries: BTreeMap<GrokReasoningReplayKey, GrokReasoningReplayEntry>,
}
struct GrokReasoningReplayEntry {
stored_at: Instant,
items: Vec<Value>,
}
pub(crate) struct GrokReasoningReplayCapture {
replay: GrokReasoningReplay,
key: GrokReasoningReplayKey,
output_items: BTreeMap<u32, Value>,
captured_bytes: usize,
truncated: bool,
finished: bool,
}
impl Default for GrokReasoningReplay {
fn default() -> Self {
Self::new()
}
}
impl GrokReasoningReplay {
pub(crate) fn new() -> Self {
Self {
state: Arc::new(Mutex::new(GrokReasoningReplayState::default())),
}
}
pub(crate) fn key(
&self,
model: &str,
session_id: &str,
account_id: &str,
) -> Option<GrokReasoningReplayKey> {
let model = model.trim();
let session_id = session_id.trim();
let account_id = account_id.trim();
if model.is_empty() || session_id.is_empty() || account_id.is_empty() {
return None;
}
Some(GrokReasoningReplayKey {
model: model.to_owned(),
session_id: session_id.to_owned(),
account_id: account_id.to_owned(),
})
}
pub(crate) fn apply(
&self,
key: &GrokReasoningReplayKey,
input: &[Value],
) -> Option<Vec<Value>> {
if input.is_empty() {
return None;
}
let cached = self.read(key)?;
let filtered = filter_replay_items_for_input(input, cached)?;
Some(insert_replay_items(input, &filtered))
}
pub(crate) fn capture(&self, key: GrokReasoningReplayKey) -> GrokReasoningReplayCapture {
GrokReasoningReplayCapture {
replay: self.clone(),
key,
output_items: BTreeMap::new(),
captured_bytes: 0,
truncated: false,
finished: false,
}
}
pub(crate) fn clear(&self, key: &GrokReasoningReplayKey) {
self.lock_state().entries.remove(key);
}
fn read(&self, key: &GrokReasoningReplayKey) -> Option<Vec<Value>> {
let now = Instant::now();
let mut state = self.lock_state();
prune_expired(&mut state, now);
state.entries.get(key).map(|entry| entry.items.clone())
}
fn store_output(&self, key: GrokReasoningReplayKey, output: Vec<Value>) {
let normalized = match normalize_replay_items(output) {
ReplayNormalization::Store(items) => items,
ReplayNormalization::Delete => {
self.clear(&key);
return;
}
ReplayNormalization::Oversized => return,
};
let now = Instant::now();
let mut state = self.lock_state();
prune_expired(&mut state, now);
if !state.entries.contains_key(&key)
&& state.entries.len() >= REASONING_REPLAY_MAX_ENTRIES
&& let Some(oldest) = state
.entries
.iter()
.min_by_key(|(_, entry)| entry.stored_at)
.map(|(key, _)| key.clone())
{
state.entries.remove(&oldest);
}
state.entries.insert(
key,
GrokReasoningReplayEntry {
stored_at: now,
items: normalized,
},
);
}
fn lock_state(&self) -> MutexGuard<'_, GrokReasoningReplayState> {
match self.state.lock() {
Ok(state) => state,
Err(poisoned) => poisoned.into_inner(),
}
}
}
impl GrokReasoningReplayCapture {
pub(crate) fn observe(&mut self, events: &[ProviderEvent]) {
if self.finished {
return;
}
let mut completed = false;
for event in events {
let Some(wire) = event.wire_event() else {
continue;
};
let payload = wire.data();
let event_type = wire
.event_type()
.or_else(|| payload.get("type").and_then(Value::as_str));
if event_type == Some("response.output_item.done")
&& let Some(item) = payload.get("item").cloned()
{
let index = payload
.get("output_index")
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
.unwrap_or_else(|| u32::try_from(self.output_items.len()).unwrap_or(u32::MAX));
self.capture_item(index, item);
}
if event_type == Some("response.completed") {
capture_terminal_output(payload, self);
completed = true;
}
}
if completed {
self.finished = true;
if self.truncated {
self.output_items.clear();
return;
}
let output = std::mem::take(&mut self.output_items)
.into_values()
.collect();
self.replay.store_output(self.key.clone(), output);
}
}
fn capture_item(&mut self, index: u32, item: Value) {
if self.truncated || self.output_items.contains_key(&index) {
return;
}
if self.output_items.len() >= REASONING_REPLAY_MAX_OUTPUT_ITEMS {
self.truncate();
return;
}
let Some(captured_bytes) = serde_json::to_vec(&item)
.ok()
.and_then(|encoded| self.captured_bytes.checked_add(encoded.len()))
.filter(|captured_bytes| *captured_bytes <= REASONING_REPLAY_MAX_BYTES)
else {
self.truncate();
return;
};
self.captured_bytes = captured_bytes;
self.output_items.insert(index, item);
}
fn truncate(&mut self) {
self.truncated = true;
self.output_items.clear();
self.captured_bytes = 0;
}
}
fn capture_terminal_output(payload: &Value, capture: &mut GrokReasoningReplayCapture) {
let Some(output) = payload
.get("response")
.and_then(Value::as_object)
.and_then(|response| response.get("output"))
.and_then(Value::as_array)
else {
return;
};
if output.len() > REASONING_REPLAY_MAX_OUTPUT_ITEMS {
capture.truncate();
return;
}
for (index, item) in output
.iter()
.take(REASONING_REPLAY_MAX_OUTPUT_ITEMS)
.enumerate()
{
let Some(index) = u32::try_from(index).ok() else {
break;
};
capture.capture_item(index, item.clone());
}
}
fn prune_expired(state: &mut GrokReasoningReplayState, now: Instant) {
state
.entries
.retain(|_, entry| now.saturating_duration_since(entry.stored_at) < REASONING_REPLAY_TTL);
}
enum ReplayNormalization {
Store(Vec<Value>),
Delete,
Oversized,
}
fn normalize_replay_items(items: Vec<Value>) -> ReplayNormalization {
let mut normalized = Vec::new();
let mut has_anchor = false;
for item in items {
let Some(item) = normalize_replay_item(item) else {
continue;
};
if matches!(
item.get("type").and_then(Value::as_str),
Some("reasoning" | "function_call" | "custom_tool_call")
) {
has_anchor = true;
}
normalized.push(item);
}
if !has_anchor || normalized.is_empty() {
return ReplayNormalization::Delete;
}
match serde_json::to_vec(&normalized) {
Ok(encoded) if encoded.len() <= REASONING_REPLAY_MAX_BYTES => {
ReplayNormalization::Store(normalized)
}
_ => ReplayNormalization::Oversized,
}
}
fn normalize_replay_item(item: Value) -> Option<Value> {
let object = item.as_object()?;
match object.get("type").and_then(Value::as_str)?.trim() {
"reasoning" => normalize_reasoning_item(object),
"message" => normalize_assistant_message_item(object),
"function_call" => normalize_function_call_item(object),
"custom_tool_call" => normalize_custom_tool_call_item(object),
_ => None,
}
}
fn normalize_reasoning_item(item: &Map<String, Value>) -> Option<Value> {
let encrypted = item.get("encrypted_content")?.as_str()?;
if !valid_reasoning_ciphertext(encrypted) {
return None;
}
Some(Value::Object(Map::from_iter([
("type".to_owned(), Value::String("reasoning".to_owned())),
("summary".to_owned(), Value::Array(Vec::new())),
("content".to_owned(), Value::Null),
(
"encrypted_content".to_owned(),
Value::String(encrypted.to_owned()),
),
])))
}
fn normalize_assistant_message_item(item: &Map<String, Value>) -> Option<Value> {
if !item
.get("role")
.and_then(Value::as_str)
.is_some_and(|role| role.trim().eq_ignore_ascii_case("assistant"))
{
return None;
}
let parts = item
.get("content")?
.as_array()?
.iter()
.filter_map(normalize_assistant_part)
.collect::<Vec<_>>();
if parts.is_empty() {
return None;
}
Some(Value::Object(Map::from_iter([
("type".to_owned(), Value::String("message".to_owned())),
("role".to_owned(), Value::String("assistant".to_owned())),
("content".to_owned(), Value::Array(parts)),
])))
}
fn normalize_assistant_part(part: &Value) -> Option<Value> {
let part = part.as_object()?;
match part.get("type").and_then(Value::as_str)?.trim() {
"output_text" => Some(Value::Object(Map::from_iter([
("type".to_owned(), Value::String("output_text".to_owned())),
(
"text".to_owned(),
Value::String(part.get("text")?.as_str()?.to_owned()),
),
]))),
"refusal" => Some(Value::Object(Map::from_iter([
("type".to_owned(), Value::String("refusal".to_owned())),
(
"refusal".to_owned(),
Value::String(part.get("refusal")?.as_str()?.to_owned()),
),
]))),
_ => None,
}
}
fn normalize_function_call_item(item: &Map<String, Value>) -> Option<Value> {
let call_id = nonempty_string(item, "call_id")?;
let name = nonempty_string(item, "name")?;
let arguments = item.get("arguments")?.as_str()?.to_owned();
Some(Value::Object(Map::from_iter([
("type".to_owned(), Value::String("function_call".to_owned())),
("call_id".to_owned(), Value::String(call_id)),
("name".to_owned(), Value::String(name)),
("arguments".to_owned(), Value::String(arguments)),
])))
}
fn normalize_custom_tool_call_item(item: &Map<String, Value>) -> Option<Value> {
let call_id = nonempty_string(item, "call_id")?;
let name = nonempty_string(item, "name")?;
let input = item.get("input")?.clone();
let status = item
.get("status")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("completed")
.to_owned();
Some(Value::Object(Map::from_iter([
(
"type".to_owned(),
Value::String("custom_tool_call".to_owned()),
),
("status".to_owned(), Value::String(status)),
("call_id".to_owned(), Value::String(call_id)),
("name".to_owned(), Value::String(name)),
("input".to_owned(), input),
])))
}
fn nonempty_string(item: &Map<String, Value>, field: &str) -> Option<String> {
item.get(field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn filter_replay_items_for_input(input: &[Value], cached: Vec<Value>) -> Option<Vec<Value>> {
let last_assistant = input.iter().rev().find_map(assistant_message);
let cached_assistant = cached.iter().find_map(assistant_message);
let assistant_matches = match (last_assistant, cached_assistant) {
(Some(current), Some(previous)) => {
if !assistant_content_equal(current, previous) {
return None;
}
true
}
_ => false,
};
let mut existing_calls = BTreeSet::new();
let mut existing_outputs = BTreeMap::new();
let mut existing_encrypted = BTreeSet::new();
for item in input {
let Some(item) = item.as_object() else {
continue;
};
match item.get("type").and_then(Value::as_str).map(str::trim) {
Some("reasoning") => {
if let Some(encrypted) = item
.get("encrypted_content")
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
{
existing_encrypted.insert(encrypted.to_owned());
}
}
Some("function_call_output" | "custom_tool_call_output") => {
if let Some(call_id) = item.get("call_id").and_then(Value::as_str) {
for candidate in comparable_call_ids(call_id) {
existing_outputs.insert(candidate, call_id.trim().to_owned());
}
}
}
Some(kind @ ("function_call" | "custom_tool_call")) => {
if let Some(call_id) = item.get("call_id").and_then(Value::as_str) {
existing_calls.extend(replay_call_keys(kind, call_id));
}
}
_ => {}
}
}
let mut filtered = Vec::new();
for mut item in cached {
let kind = item
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default()
.to_owned();
match kind.as_str() {
"reasoning" => {
if item
.get("encrypted_content")
.and_then(Value::as_str)
.is_some_and(|encrypted| existing_encrypted.contains(encrypted))
{
continue;
}
}
"message" if assistant_matches => continue,
"message" => {}
"function_call" | "custom_tool_call" => {
let Some(call_id) = item.get("call_id").and_then(Value::as_str) else {
continue;
};
let keys = replay_call_keys(&kind, call_id);
if keys.is_empty() || keys.iter().any(|key| existing_calls.contains(key)) {
continue;
}
let output_call_id = comparable_call_ids(call_id)
.into_iter()
.find_map(|candidate| existing_outputs.get(&candidate).cloned());
let Some(output_call_id) = output_call_id else {
continue;
};
existing_calls.extend(keys);
if output_call_id != call_id
&& let Some(item) = item.as_object_mut()
{
item.insert("call_id".to_owned(), Value::String(output_call_id));
}
}
_ => continue,
}
filtered.push(item);
}
(!filtered.is_empty()).then_some(filtered)
}
fn assistant_message(item: &Value) -> Option<&Map<String, Value>> {
let item = item.as_object()?;
let item_type = item
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
let role = item.get("role")?.as_str()?.trim();
((item_type.is_empty() || item_type == "message") && role.eq_ignore_ascii_case("assistant"))
.then_some(item)
}
fn assistant_content_equal(left: &Map<String, Value>, right: &Map<String, Value>) -> bool {
let Some(left) = assistant_parts(left.get("content")) else {
return false;
};
let Some(right) = assistant_parts(right.get("content")) else {
return false;
};
left == right
}
fn assistant_parts(content: Option<&Value>) -> Option<Vec<(String, String)>> {
match content? {
Value::String(text) => Some(vec![("output_text".to_owned(), text.clone())]),
Value::Array(parts) => {
let mut result = Vec::new();
for part in parts {
let part = part.as_object()?;
let kind = part.get("type")?.as_str()?.trim();
let value = match kind {
"output_text" => part.get("text")?.as_str()?,
"refusal" => part.get("refusal")?.as_str()?,
_ => return None,
};
result.push((kind.to_owned(), value.to_owned()));
}
(!result.is_empty()).then_some(result)
}
_ => None,
}
}
fn comparable_call_ids(call_id: &str) -> Vec<String> {
let call_id = call_id.trim();
if call_id.is_empty() {
return Vec::new();
}
if let Some(upstream) = call_id
.strip_prefix("toolu_")
.filter(|value| !value.is_empty())
{
return vec![call_id.to_owned(), upstream.to_owned()];
}
vec![call_id.to_owned(), format!("toolu_{call_id}")]
}
fn replay_call_keys(item_type: &str, call_id: &str) -> Vec<String> {
if !matches!(item_type, "function_call" | "custom_tool_call") {
return Vec::new();
}
comparable_call_ids(call_id)
.into_iter()
.map(|call_id| format!("{item_type}\0{call_id}"))
.collect()
}
fn insert_replay_items(input: &[Value], replay: &[Value]) -> Vec<Value> {
let insert_at = replay_insert_index(input, replay);
let mut result = Vec::new();
for (index, item) in input.iter().enumerate() {
if index == insert_at {
result.extend(replay.iter().cloned());
}
result.push(item.clone());
}
if insert_at == input.len() {
result.extend(replay.iter().cloned());
}
result
}
fn replay_insert_index(input: &[Value], replay: &[Value]) -> usize {
let replay_call_ids = replay
.iter()
.filter_map(Value::as_object)
.filter(|item| {
matches!(
item.get("type").and_then(Value::as_str),
Some("function_call" | "custom_tool_call")
)
})
.filter_map(|item| item.get("call_id").and_then(Value::as_str))
.flat_map(comparable_call_ids)
.collect::<BTreeSet<_>>();
if !replay_call_ids.is_empty() {
for (index, item) in input.iter().enumerate() {
let Some(item) = item.as_object() else {
continue;
};
if !matches!(
item.get("type").and_then(Value::as_str),
Some("function_call_output" | "custom_tool_call_output")
) {
continue;
}
let call_ids = item
.get("call_id")
.and_then(Value::as_str)
.map(comparable_call_ids)
.unwrap_or_default();
if call_ids.is_empty()
|| call_ids
.iter()
.any(|call_id| replay_call_ids.contains(call_id))
{
return index;
}
}
}
if let Some(index) = input
.iter()
.rposition(|item| assistant_message(item).is_some())
{
return index;
}
input
.iter()
.position(should_insert_replay_before)
.unwrap_or(input.len())
}
fn should_insert_replay_before(item: &Value) -> bool {
let Some(item) = item.as_object() else {
return true;
};
let item_type = item
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
let role = item
.get("role")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default()
.to_ascii_lowercase();
if role.is_empty() || (!item_type.is_empty() && item_type != "message") {
return true;
}
!matches!(role.as_str(), "system" | "developer")
}
pub(crate) fn valid_reasoning_ciphertext(value: &str) -> bool {
if value.is_empty()
|| value != value.trim()
|| value.len() > REASONING_REPLAY_MAX_BYTES
|| value.starts_with("gAAAA")
|| value.contains('=')
{
return false;
}
let Ok(decoded) = base64::engine::general_purpose::STANDARD_NO_PAD.decode(value) else {
return false;
};
decoded.len() >= MIN_REASONING_CIPHERTEXT_BYTES
&& byte_entropy_ratio(&decoded) >= MIN_REASONING_CIPHERTEXT_ENTROPY
}
fn byte_entropy_ratio(value: &[u8]) -> f64 {
if value.is_empty() {
return 0.0;
}
let mut counts = [0_u32; 256];
for byte in value {
counts[usize::from(*byte)] += 1;
}
let size = value.len() as f64;
let entropy = counts
.into_iter()
.filter(|count| *count > 0)
.map(|count| {
let probability = f64::from(count) / size;
-probability * probability.log2()
})
.sum::<f64>();
let symbols = value.len().min(256);
if symbols <= 1 {
return 0.0;
}
entropy / (symbols as f64).log2()
}
@@ -42,6 +42,7 @@ If the prior conversation contains a note about files at /tmp/compaction/segment
pub struct GrokCompactionRequest {
body: Map<String, Value>,
affinity: Option<GrokSessionAffinityKey>,
reasoning_replay_session_id: Option<String>,
}
impl GrokCompactionRequest {
@@ -60,6 +61,8 @@ impl GrokCompactionRequest {
upstream_model,
client_api_key_ref,
)?;
let reasoning_replay_session_id =
normalized.reasoning_replay_session_id().map(str::to_owned);
let mut body = normalized.body().clone();
let mut input = normalized.input_items();
input.push(summary_prompt_item());
@@ -95,6 +98,7 @@ impl GrokCompactionRequest {
Ok(Self {
body,
affinity: normalized.affinity().cloned(),
reasoning_replay_session_id,
})
}
@@ -110,6 +114,12 @@ impl GrokCompactionRequest {
self.affinity.as_ref()
}
/// 返回用于在成功 compact 后清除 reasoning replay 的显式会话身份。
#[must_use]
pub(crate) fn reasoning_replay_session_id(&self) -> Option<&str> {
self.reasoning_replay_session_id.as_deref()
}
/// 序列化上游请求正文。
///
/// # Errors
@@ -142,11 +142,11 @@ pub fn build_grok_headers(
"x-grok-conv-id",
SecretValue::new(upstream_session_id.to_owned()),
));
headers.push(GrokHeader::sensitive(
"x-grok-session-id",
SecretValue::new(upstream_session_id.to_owned()),
));
if let Some(turn_index) = turn_index {
headers.push(GrokHeader::sensitive(
"x-grok-session-id",
SecretValue::new(upstream_session_id.to_owned()),
));
headers.push(GrokHeader::public("x-grok-turn-idx", turn_index));
}
}
@@ -56,6 +56,7 @@ const SESSION_FIELDS: &[&str] = &[
"conversation_id",
"conversationId",
];
const PROTOCOL_CONTEXT_SESSION_FIELDS: &[&str] = &["conversation_id", "session_id", "thread_id"];
const FOREIGN_CLIENT_METADATA_FIELDS: &[&str] = &["x-openai-subagent"];
const MAX_SESSION_SEED_BYTES: usize = 1_024;
const GROK_CACHE_ROUTE_TOOLS: &[&str] = &["web_search", "x_search"];
@@ -68,6 +69,7 @@ const GROK_INTERNAL_HISTORY_KEYS: &[&str] =
pub struct GrokResponsesRequest {
body: Map<String, Value>,
session_id: Option<String>,
reasoning_replay_session_id: Option<String>,
affinity: Option<GrokSessionAffinityKey>,
response_transform: GrokResponseTransform,
}
@@ -85,6 +87,23 @@ impl GrokResponsesRequest {
self.session_id.as_deref()
}
/// 返回只有显式会话身份才启用的 reasoning replay scope。
pub(crate) fn reasoning_replay_session_id(&self) -> Option<&str> {
self.reasoning_replay_session_id.as_deref()
}
pub(crate) fn has_previous_response_id(&self) -> bool {
self.body
.get("previous_response_id")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
}
pub(crate) fn replay_input_items(&self) -> Option<Vec<Value>> {
self.body.get("input").and_then(Value::as_array).cloned()
}
/// 返回与会话一致、额外绑定模型的账号亲和键。
#[must_use]
pub const fn affinity(&self) -> Option<&GrokSessionAffinityKey> {
@@ -149,6 +168,9 @@ impl GrokResponsesRequest {
return;
};
self.session_id = Some(session_id.to_owned());
if self.reasoning_replay_session_id.is_some() {
self.reasoning_replay_session_id = Some(session_id.to_owned());
}
self.affinity = None;
self.body.insert(
"prompt_cache_key".to_owned(),
@@ -159,6 +181,7 @@ impl GrokResponsesRequest {
pub(crate) fn clear_session(&mut self) {
self.body.remove("prompt_cache_key");
self.session_id = None;
self.reasoning_replay_session_id = None;
self.affinity = None;
}
@@ -207,13 +230,16 @@ impl GrokResponsesRequest {
);
sanitize_account_identity(&mut body);
sanitize_foreign_client_metadata(&mut body);
let response_transform = normalize_responses_request(&mut body)?;
let mut response_transform = normalize_responses_request(&mut body)?;
response_transform.observe_client_cache_tools();
if enable_cache_route {
enable_grok_prompt_cache_route(&mut body);
enable_grok_prompt_cache_route(&mut body, upstream_model, &mut response_transform);
}
let (session_id, affinity) = identity.map_or((None, None), |(session_id, affinity)| {
(Some(session_id), Some(affinity))
});
let reasoning_replay_session_id =
session_seed.is_some().then(|| session_id.clone()).flatten();
match session_id.as_ref() {
Some(session_id) => {
body.insert(
@@ -230,6 +256,7 @@ impl GrokResponsesRequest {
Ok(Self {
body,
session_id,
reasoning_replay_session_id,
affinity,
response_transform,
})
@@ -321,23 +348,61 @@ fn sanitize_foreign_client_metadata(body: &mut Map<String, Value>) {
}
}
fn enable_grok_prompt_cache_route(body: &mut Map<String, Value>) {
fn enable_grok_prompt_cache_route(
body: &mut Map<String, Value>,
upstream_model: &str,
response: &mut GrokResponseTransform,
) {
let tools = body
.entry("tools".to_owned())
.or_insert_with(|| Value::Array(Vec::new()));
let Value::Array(tools) = tools else {
return;
};
// 缓存路由不得改变模型可调用的工具集:客户端带工具时无法既保留其
// tool_choice 语义又阻止注入的原生搜索被调用,因此直接不注入。
if !tools.is_empty() {
response.observe_upstream_cache_tools(tools);
if is_cache_media_model(upstream_model) || has_tool_type(tools, "image_generation") {
return;
}
for tool in GROK_CACHE_ROUTE_TOOLS {
tools.push(json_object([("type", Value::String((*tool).to_owned()))]));
if tools.is_empty() {
for tool in GROK_CACHE_ROUTE_TOOLS {
tools.push(json_object([("type", Value::String((*tool).to_owned()))]));
response.mark_injected_cache_tool(tool);
}
// 无客户端工具时以 none 选中缓存路由,同时不授予搜索能力。
body.insert("tool_choice".to_owned(), Value::String("none".to_owned()));
} else if !has_tool_type(tools, "x_search") {
tools.push(json_object([(
"type",
Value::String("x_search".to_owned()),
)]));
response.mark_injected_cache_tool("x_search");
}
append_x_search_to_allowed_tools(body);
}
fn is_cache_media_model(model: &str) -> bool {
let model = model.trim().to_ascii_lowercase();
["image", "imagine", "video"]
.into_iter()
.any(|marker| model.contains(marker))
}
fn append_x_search_to_allowed_tools(body: &mut Map<String, Value>) {
let Some(choice) = body.get_mut("tool_choice").and_then(Value::as_object_mut) else {
return;
};
if choice.get("type").and_then(Value::as_str) != Some("allowed_tools") {
return;
}
let Some(allowed) = choice.get_mut("tools").and_then(Value::as_array_mut) else {
return;
};
if !has_tool_type(allowed, "x_search") {
allowed.push(json_object([(
"type",
Value::String("x_search".to_owned()),
)]));
}
// 禁用工具调用,使原生搜索声明只承担缓存路由作用。
body.insert("tool_choice".to_owned(), Value::String("none".to_owned()));
}
fn explicit_session_seed(request: &GenerateRequest, body: &Map<String, Value>) -> Option<String> {
@@ -351,6 +416,17 @@ fn explicit_session_seed(request: &GenerateRequest, body: &Map<String, Value>) -
.and_then(Value::as_object)
.and_then(first_session_value)
})
.or_else(|| {
PROTOCOL_CONTEXT_SESSION_FIELDS.iter().find_map(|field| {
request
.protocol_payload()
.context()
.get(*field)
.and_then(Value::as_str)
.and_then(valid_session_seed)
.map(ToOwned::to_owned)
})
})
}
fn first_session_value(body: &Map<String, Value>) -> Option<String> {
@@ -674,6 +750,11 @@ pub(crate) struct GrokResponseTransform {
aliases: BTreeMap<String, ToolIdentity>,
visible_tools: Vec<Value>,
legacy_local_shell: bool,
filter_x_search: bool,
injected_tool_types: BTreeSet<String>,
client_declared_tools: BTreeSet<String>,
dropped_output_indexes: BTreeSet<u64>,
dropped_item_ids: BTreeSet<String>,
stream_calls: BTreeMap<String, StreamCallState>,
stream_keys: BTreeMap<String, String>,
}
@@ -698,7 +779,42 @@ impl GrokTransformedWireEvent {
impl GrokResponseTransform {
#[must_use]
pub(crate) fn is_empty(&self) -> bool {
self.aliases.is_empty() && self.visible_tools.is_empty() && !self.legacy_local_shell
self.aliases.is_empty()
&& self.visible_tools.is_empty()
&& !self.legacy_local_shell
&& !self.filter_x_search
&& self.injected_tool_types.is_empty()
}
fn observe_client_cache_tools(&mut self) {
for tool in &self.visible_tools {
let kind = tool.get("type").and_then(Value::as_str).unwrap_or_default();
if kind == "x_search" {
self.filter_x_search = true;
}
if matches!(kind, "function" | "custom")
&& let Some(name) = tool
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.filter(|name| !name.is_empty())
{
self.client_declared_tools.insert(name.to_owned());
}
}
}
fn observe_upstream_cache_tools(&mut self, tools: &[Value]) {
if has_tool_type(tools, "x_search") {
self.filter_x_search = true;
}
}
fn mark_injected_cache_tool(&mut self, kind: &str) {
self.injected_tool_types.insert(kind.to_owned());
if kind == "x_search" {
self.filter_x_search = true;
}
}
pub(crate) fn rewrite_stream_event(
@@ -706,6 +822,20 @@ impl GrokResponseTransform {
event_type: &str,
mut value: Value,
) -> Result<Vec<GrokTransformedWireEvent>, GrokRequestEncodeError> {
if self.cache_filter_enabled() {
if value
.get("item")
.is_some_and(|item| self.is_internal_cache_call(item))
{
self.record_dropped_cache_item(&value);
return Ok(Vec::new());
}
self.filter_cache_envelope(&mut value)?;
if self.references_dropped_cache_item(&value) {
return Ok(Vec::new());
}
self.compact_cache_output_index(&mut value);
}
if event_type == "response.output_item.added"
&& let Some(item) = value.get("item").and_then(Value::as_object)
&& let Some(primary) = self.remember_stream_call(item)
@@ -1031,9 +1161,145 @@ impl GrokResponseTransform {
])
}
fn cache_filter_enabled(&self) -> bool {
self.filter_x_search || !self.injected_tool_types.is_empty()
}
fn is_internal_cache_call(&self, item: &Value) -> bool {
let Some(item) = item.as_object() else {
return false;
};
let kind = string_field(item, "type").trim();
if kind == "web_search_call" {
return self.injected_tool_types.contains("web_search");
}
if !matches!(kind, "custom_tool_call" | "function_call") {
return false;
}
if string_field(item, "call_id").trim().starts_with("xs_call") {
return true;
}
let name = string_field(item, "name").trim();
matches!(
name,
"x_user_search" | "x_semantic_search" | "x_keyword_search" | "x_thread_fetch"
) && string_field(item, "namespace").trim().is_empty()
&& !self.client_declared_tools.contains(name)
}
fn record_dropped_cache_item(&mut self, payload: &Value) {
if let Some(index) = payload.get("output_index").and_then(Value::as_u64) {
self.dropped_output_indexes.insert(index);
}
let Some(item) = payload.get("item").and_then(Value::as_object) else {
return;
};
for field in ["id", "call_id"] {
if let Some(value) = item
.get(field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
self.dropped_item_ids.insert(value.to_owned());
}
}
}
fn references_dropped_cache_item(&self, payload: &Value) -> bool {
if payload
.get("output_index")
.and_then(Value::as_u64)
.is_some_and(|index| self.dropped_output_indexes.contains(&index))
{
return true;
}
["item_id", "call_id"].into_iter().any(|field| {
payload
.get(field)
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| self.dropped_item_ids.contains(value))
})
}
fn compact_cache_output_index(&self, payload: &mut Value) {
let Some(index) = payload.get("output_index").and_then(Value::as_u64) else {
return;
};
let removed_before = self
.dropped_output_indexes
.iter()
.filter(|dropped| **dropped < index)
.count() as u64;
if removed_before > 0
&& let Some(payload) = payload.as_object_mut()
{
payload.insert(
"output_index".to_owned(),
Value::from(index.saturating_sub(removed_before)),
);
}
}
fn filter_cache_envelope(&self, payload: &mut Value) -> Result<(), GrokRequestEncodeError> {
let Some(payload) = payload.as_object_mut() else {
return Ok(());
};
self.filter_cache_output(payload)?;
self.filter_injected_cache_tools(payload)?;
if let Some(response) = payload.get_mut("response").and_then(Value::as_object_mut) {
self.filter_cache_output(response)?;
self.filter_injected_cache_tools(response)?;
}
Ok(())
}
fn filter_cache_output(
&self,
envelope: &mut Map<String, Value>,
) -> Result<(), GrokRequestEncodeError> {
let Some(output) = envelope.get_mut("output") else {
return Ok(());
};
let output = output
.as_array_mut()
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
output.retain(|item| !self.is_internal_cache_call(item));
Ok(())
}
fn filter_injected_cache_tools(
&self,
envelope: &mut Map<String, Value>,
) -> Result<(), GrokRequestEncodeError> {
if self.injected_tool_types.is_empty() {
return Ok(());
}
let Some(tools) = envelope.get_mut("tools") else {
return Ok(());
};
let tools = tools
.as_array_mut()
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
tools.retain(|tool| {
tool.get("type")
.and_then(Value::as_str)
.is_none_or(|kind| !self.injected_tool_types.contains(kind.trim()))
});
if tools.is_empty() {
envelope.remove("tools");
}
Ok(())
}
fn restore_visible_tools(&self, response: &mut Map<String, Value>) {
if response.contains_key("tools") {
response.insert("tools".to_owned(), Value::Array(self.visible_tools.clone()));
if self.visible_tools.is_empty() {
response.remove("tools");
} else {
response.insert("tools".to_owned(), Value::Array(self.visible_tools.clone()));
}
}
}
}
@@ -1125,6 +1391,8 @@ impl fmt::Debug for GrokResponseTransform {
.field("alias_count", &self.aliases.len())
.field("visible_tool_count", &self.visible_tools.len())
.field("legacy_local_shell", &self.legacy_local_shell)
.field("filter_x_search", &self.filter_x_search)
.field("injected_tool_types", &self.injected_tool_types)
.finish()
}
}
@@ -1982,6 +2250,9 @@ impl ToolNormalizer {
]),
);
}
"allowed_tools" => {
self.normalize_allowed_tools_choice(payload, object, normalized_tools)?;
}
"function" => self.normalize_function_tool_choice(payload, object)?,
_ => {
if let Some(hosted_kind) = normalize_hosted_tool_choice_kind(&kind) {
@@ -2004,6 +2275,94 @@ impl ToolNormalizer {
Ok(())
}
fn normalize_allowed_tools_choice(
&self,
payload: &mut Map<String, Value>,
mut object: Map<String, Value>,
normalized_tools: &[Value],
) -> Result<(), GrokRequestEncodeError> {
if !matches!(
object.get("mode").and_then(Value::as_str),
Some("auto" | "required")
) {
return Err(GrokRequestEncodeError::InvalidRequestNormalization);
}
let allowed = object
.get_mut("tools")
.and_then(Value::as_array_mut)
.filter(|tools| !tools.is_empty())
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
for allowed_tool in allowed {
self.normalize_allowed_tool_reference(allowed_tool, normalized_tools)?;
}
payload.insert("tool_choice".to_owned(), Value::Object(object));
Ok(())
}
fn normalize_allowed_tool_reference(
&self,
allowed_tool: &mut Value,
normalized_tools: &[Value],
) -> Result<(), GrokRequestEncodeError> {
let object = allowed_tool
.as_object_mut()
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
let kind = string_field(object, "type").trim().to_owned();
match kind.as_str() {
"function" => {
let choice = match object.get_mut("function") {
Some(Value::Object(function)) => function,
Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization),
None => object,
};
rewrite_namespace_choice(choice, &self.identity_aliases)?;
let name = string_field(choice, "name").trim();
if name.is_empty() || !has_named_tool(normalized_tools, "function", name) {
return Err(GrokRequestEncodeError::InvalidRequestNormalization);
}
}
"custom" => {
let identity = ToolIdentity::new(
ToolKind::Custom,
string_field(object, "namespace").trim(),
string_field(object, "name").trim(),
);
let alias = self
.identity_aliases
.get(&identity)
.cloned()
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
object.insert("type".to_owned(), Value::String("function".to_owned()));
object.insert("name".to_owned(), Value::String(alias));
object.remove("namespace");
}
"apply_patch" | "tool_search" => {
let identity = if kind == "apply_patch" {
ToolIdentity::new(ToolKind::ApplyPatch, "", "apply_patch")
} else {
ToolIdentity::new(ToolKind::ToolSearch, "", "tool_search")
};
let alias = self
.identity_aliases
.get(&identity)
.cloned()
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
object.clear();
object.insert("type".to_owned(), Value::String("function".to_owned()));
object.insert("name".to_owned(), Value::String(alias));
}
_ => {
let hosted_kind = normalize_hosted_tool_choice_kind(&kind)
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
if !has_tool_type(normalized_tools, hosted_kind) {
return Err(GrokRequestEncodeError::InvalidRequestNormalization);
}
object.insert("type".to_owned(), Value::String(hosted_kind.to_owned()));
}
}
Ok(())
}
fn normalize_function_tool_choice(
&self,
payload: &mut Map<String, Value>,
@@ -2080,6 +2439,13 @@ fn has_tool_type(tools: &[Value], kind: &str) -> bool {
.any(|tool| tool.pointer("/type").and_then(Value::as_str) == Some(kind))
}
fn has_named_tool(tools: &[Value], kind: &str, name: &str) -> bool {
tools.iter().any(|tool| {
tool.pointer("/type").and_then(Value::as_str) == Some(kind)
&& tool.pointer("/name").and_then(Value::as_str) == Some(name)
})
}
fn tools_of_type(tools: &[Value], kind: &str) -> Vec<Value> {
tools
.iter()
@@ -324,7 +324,7 @@ async fn unauthorized_feedback_records_runtime_cooldown_without_persisting_accou
}
#[tokio::test]
async fn cooldown_for_superseded_credential_does_not_block_rotated_account() {
async fn account_scoped_cooldown_survives_credential_rotation() {
let fixture = SelectorFixture::new(&["revision-fence"]).await;
let id = account_id("revision-fence");
let session = fixture
@@ -369,17 +369,25 @@ async fn cooldown_for_superseded_credential_does_not_block_rotated_account() {
CredentialCasOutcome::Updated(revision) if revision.get() == 2
));
let next = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("replacement credential should not inherit cooldown");
assert_eq!(next.account_id(), &id);
assert!(fixture.cooldowns.cooldown(&id).is_none());
assert!(matches!(
fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await,
Err(GrokSessionSelectorError::NoEligibleSession)
));
let account = fixture.store.account(&id).expect("rotated account");
assert_eq!(account.revision().get(), 2);
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
}
#[tokio::test]
async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account_state() {
async fn rate_limit_feedback_persists_the_grok2api_cooldown_state() {
let fixture = SelectorFixture::new(&["rate-limit", "available"]).await;
let session = fixture
.selector
@@ -397,7 +405,12 @@ async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account
)
.await;
let account = fixture.store.account(&selected).expect("account");
assert_eq!(account.availability(), AccountAvailability::Ready);
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert!(
fixture
.cooldowns
@@ -413,7 +426,35 @@ async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account
}
#[tokio::test]
async fn unknown_payment_required_feedback_only_records_runtime_account_cooldown() {
async fn successful_request_clears_the_persisted_cooldown_state() {
let fixture = SelectorFixture::new(&["cooldown-success"]).await;
let session = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("session");
fixture
.selector
.record_failure(
&session,
GrokCredentialFailure::RateLimited {
retry_after: Some(Duration::from_secs(5)),
},
)
.await;
fixture.selector.record_success(&session).await;
let account = fixture
.store
.account(session.account_id())
.expect("account");
assert_eq!(account.availability(), AccountAvailability::Ready);
assert!(account.cooldown_until().is_none());
}
#[tokio::test]
async fn payment_required_feedback_persists_account_quota_exhaustion() {
let fixture = SelectorFixture::new(&["payment-required", "available"]).await;
let session = fixture
.selector
@@ -437,9 +478,9 @@ async fn unknown_payment_required_feedback_only_records_runtime_account_cooldown
.account(&selected)
.expect("selected account")
.availability(),
AccountAvailability::Ready
AccountAvailability::QuotaExhausted
);
assert!(fixture.cooldowns.cooldown(&selected).is_some());
assert!(fixture.cooldowns.cooldown(&selected).is_none());
let next = fixture
.selector
.select(fixture.request(BTreeSet::new()))
@@ -538,7 +579,7 @@ async fn model_access_feedback_reports_model_cooldown_without_blocking_the_accou
}
#[tokio::test]
async fn interrupted_stream_feedback_records_runtime_cooldown_without_persisting_account_state() {
async fn interrupted_stream_feedback_persists_the_grok2api_cooldown_state() {
let fixture = SelectorFixture::new(&["stream-interrupted"]).await;
let session = fixture
.selector
@@ -554,7 +595,12 @@ async fn interrupted_stream_feedback_records_runtime_cooldown_without_persisting
.store
.account(session.account_id())
.expect("account");
assert_eq!(account.availability(), AccountAvailability::Ready);
assert_eq!(account.availability(), AccountAvailability::Cooldown);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > SystemTime::now())
);
assert!(
fixture
.cooldowns
@@ -595,6 +641,75 @@ async fn quota_feedback_uses_common_quota_exhausted_state() {
);
}
#[tokio::test]
async fn quota_feedback_should_follow_the_account_across_a_credential_rotation() {
let fixture = SelectorFixture::new(&["quota-rotation"]).await;
let id = account_id("quota-rotation");
let session = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("session");
let current = fixture
.store
.load_credential(&id, CredentialRevision::new(1).expect("revision"))
.await
.expect("current credential");
let prepared = GrokCredentialAdmin
.prepare_rotation(&RotateManagedGrokCredential {
current,
secret: provider_xai::GrokOAuthSecret {
access_token: provider_xai::SecretValue::new("new-access"),
refresh_token: provider_xai::SecretValue::new("new-refresh"),
id_token: None,
scope: provider_xai::OFFICIAL_SCOPES.join(" "),
},
verified_account: crate::support::profile("subject-quota-rotation"),
next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30),
})
.expect("rotate");
fixture
.store
.compare_and_swap_credential(prepared.credential)
.await
.expect("persist rotation");
fixture
.selector
.record_failure(&session, GrokCredentialFailure::FreeQuotaExhausted)
.await;
assert_eq!(
fixture.store.account(&id).expect("account").availability(),
AccountAvailability::QuotaExhausted
);
}
#[tokio::test]
async fn ordinary_success_should_not_clear_a_confirmed_quota_exhaustion() {
let fixture = SelectorFixture::new(&["quota-success"]).await;
let session = fixture
.selector
.select(fixture.request(BTreeSet::new()))
.await
.expect("session");
fixture
.selector
.record_failure(&session, GrokCredentialFailure::FreeQuotaExhausted)
.await;
fixture.selector.record_success(&session).await;
assert_eq!(
fixture
.store
.account(session.account_id())
.expect("account")
.availability(),
AccountAvailability::QuotaExhausted
);
}
#[tokio::test]
async fn excluded_account_is_never_selected_again() {
let fixture = SelectorFixture::new(&["excluded"]).await;
@@ -106,6 +106,15 @@ const SUCCESS_SSE: &[u8] = concat!(
)
.as_bytes();
fn replay_ciphertext(seed: u8) -> String {
use base64::Engine as _;
let bytes = (0_u16..128)
.map(|value| (value as u8).wrapping_add(seed))
.collect::<Vec<_>>();
base64::engine::general_purpose::STANDARD_NO_PAD.encode(bytes)
}
fn stateful_sse(encrypted_content: &str) -> Vec<u8> {
format!(
concat!(
@@ -121,6 +130,45 @@ fn stateful_sse(encrypted_content: &str) -> Vec<u8> {
.into_bytes()
}
fn stateful_sse_with_assistant(encrypted_content: &str, assistant: &str) -> Vec<u8> {
let assistant = serde_json::to_string(assistant).expect("assistant text");
format!(
concat!(
"event: response.created\n",
"data: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\"}}}}\n\n",
"event: response.output_item.done\n",
"data: {{\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{{\"id\":\"reason_account_bound\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[],\"content\":null,\"encrypted_content\":\"{}\"}}}}\n\n",
"event: response.output_item.done\n",
"data: {{\"type\":\"response.output_item.done\",\"output_index\":1,\"item\":{{\"id\":\"message_account_bound\",\"type\":\"message\",\"role\":\"assistant\",\"status\":\"completed\",\"content\":[{{\"type\":\"output_text\",\"text\":{assistant}}}]}}}}\n\n",
"event: response.completed\n",
"data: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\",\"status\":\"completed\"}}}}\n\n"
),
encrypted_content,
assistant = assistant,
)
.into_bytes()
}
fn stateful_sse_with_custom_tool_call(encrypted_content: &str, input: &str) -> Vec<u8> {
let arguments = serde_json::to_string(&json!({"input": input})).expect("tool arguments");
let arguments_json = serde_json::to_string(&arguments).expect("arguments JSON string");
format!(
concat!(
"event: response.created\n",
"data: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\"}}}}\n\n",
"event: response.output_item.done\n",
"data: {{\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{{\"id\":\"reason_account_bound\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[],\"content\":null,\"encrypted_content\":\"{}\"}}}}\n\n",
"event: response.output_item.done\n",
"data: {{\"type\":\"response.output_item.done\",\"output_index\":1,\"item\":{{\"id\":\"item_exec\",\"type\":\"function_call\",\"status\":\"completed\",\"call_id\":\"call_exec\",\"name\":\"exec\",\"arguments\":{arguments_json}}}}}\n\n",
"event: response.completed\n",
"data: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\",\"status\":\"completed\"}}}}\n\n"
),
encrypted_content,
arguments_json = arguments_json,
)
.into_bytes()
}
fn custom_apply_patch_sse(patch: &str) -> Vec<u8> {
let arguments = serde_json::to_string(&json!({"patch": patch})).expect("patch arguments");
let arguments_json = serde_json::to_string(&arguments).expect("arguments JSON string");
@@ -278,6 +326,54 @@ impl GrokSessionSelector for StubSelector {
}
}
struct SequencedAccountSelector {
accounts: Mutex<VecDeque<gateway_core::engine::credential::ProviderAccountId>>,
}
impl SequencedAccountSelector {
fn new(
accounts: impl IntoIterator<Item = gateway_core::engine::credential::ProviderAccountId>,
) -> Arc<Self> {
Arc::new(Self {
accounts: Mutex::new(accounts.into_iter().collect()),
})
}
}
impl GrokSessionSelector for SequencedAccountSelector {
fn select(&self, request: GrokSessionSelection) -> GrokSessionSelectorFuture<'_> {
let selected = self.accounts.lock().expect("account sequence").pop_front();
Box::pin(async move {
let id = selected.ok_or(GrokSessionSelectorError::NoEligibleSession)?;
if request.excluded_accounts().contains(&id)
|| request
.required_account()
.is_some_and(|required| required != &id)
{
return Err(GrokSessionSelectorError::NoEligibleSession);
}
SelectedGrokSession::new(
id,
CredentialRevision::new(1).expect("revision"),
SecretValue::new("oauth-access"),
SecretValue::new("verified-user"),
Some(SecretValue::new("user@example.com")),
GrokSessionBinding::new("acct_provider").expect("binding"),
(),
)
.map_err(|_| GrokSessionSelectorError::InvalidSession)
})
}
fn record_failure<'a>(
&'a self,
_: &'a SelectedGrokSession,
_: GrokCredentialFailure,
) -> GrokCredentialFeedbackFuture<'a> {
Box::pin(async {})
}
}
enum InferenceMode {
Success,
SuccessWithMetrics(GrokInferenceTransportMetrics),
@@ -481,7 +577,7 @@ async fn provider_with_recovery(
}
async fn provider_with_catalog_transport(
selector: Arc<StubSelector>,
selector: Arc<dyn GrokSessionSelector>,
transport: Arc<StubInferenceTransport>,
recovery: Arc<StubRecovery>,
catalog_transport: Arc<dyn GrokModelCatalogTransport>,
@@ -569,6 +665,36 @@ fn operation() -> Operation {
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn reasoning_replay_operation(session: &str, input: serde_json::Value) -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("client-model")),
("prompt_cache_key".to_owned(), json!(session)),
("input".to_owned(), input),
]),
)
.expect("OpenAI payload");
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn contextual_reasoning_replay_operation(session: &str, input: serde_json::Value) -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("client-model")),
("input".to_owned(), input),
("tools".to_owned(), json!([{"type":"custom","name":"exec"}])),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([(
"conversation_id".to_owned(),
json!(session),
)]));
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn operation_with_reasoning_effort(effort: &str) -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
@@ -621,6 +747,26 @@ fn compaction_operation() -> Operation {
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn compaction_operation_with_prompt_cache(session: &str) -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("client-model")),
("prompt_cache_key".to_owned(), json!(session)),
(
"input".to_owned(),
json!([
{"type": "message", "role": "user", "content": "history"},
{"type": "compaction_trigger"}
]),
),
("stream".to_owned(), json!(true)),
]),
)
.expect("OpenAI payload");
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
}
fn compaction_operation_with_state(state: ProviderSessionState) -> Operation {
let payload = ProtocolPayload::json_object(
"openai",
@@ -737,6 +883,18 @@ fn context_with_continuation_attempt(
context(CancellationToken::new(), Some(continuation)).with_continuation_attempt(attempt)
}
async fn execute_successfully(provider: &GrokBuildProvider, operation: Operation) {
let mut stream = provider
.execute(
provider_request_with_operation("xai", operation),
context(CancellationToken::new(), None),
)
.await
.expect("provider stream");
let events = stream.by_ref().collect::<Vec<_>>().await;
assert!(events.iter().all(Result::is_ok));
}
#[tokio::test]
async fn request_observation_preserves_raw_xai_reasoning_effort() {
let provider = provider(StubSelector::success(), StubInferenceTransport::success()).await;
@@ -1300,7 +1458,7 @@ async fn explicit_quota_http_429_marks_provider_error_replay_safe() {
}
#[tokio::test]
async fn unknown_http_402_is_retryable_without_account_quota_feedback() {
async fn unknown_http_402_is_retryable_with_payment_quota_feedback() {
let selector = StubSelector::success();
let transport = StubInferenceTransport::error(
GrokInferenceTransportError::new(
@@ -1750,6 +1908,239 @@ async fn external_previous_response_is_rejected_before_selection() {
assert_eq!(selector.calls.load(Ordering::SeqCst), 0);
}
#[tokio::test]
async fn reasoning_replay_is_reused_only_for_the_same_explicit_session() {
let encrypted_content = replay_ciphertext(0);
let transport = StubInferenceTransport::sequence([
InferenceMode::SuccessBody(stateful_sse(&encrypted_content)),
InferenceMode::Success,
InferenceMode::Success,
]);
let provider = provider(StubSelector::success(), transport.clone()).await;
execute_successfully(
&provider,
reasoning_replay_operation(
"session-a",
json!([{"type":"message","role":"user","content":"first"}]),
),
)
.await;
execute_successfully(
&provider,
reasoning_replay_operation(
"session-b",
json!([{"type":"message","role":"user","content":"other"}]),
),
)
.await;
execute_successfully(
&provider,
reasoning_replay_operation(
"session-a",
json!([{"type":"message","role":"user","content":"next"}]),
),
)
.await;
let requests = transport.requests.lock().expect("requests");
let other: serde_json::Value =
serde_json::from_slice(requests[1].body()).expect("other session body");
let replayed: serde_json::Value =
serde_json::from_slice(requests[2].body()).expect("replayed body");
assert!(
other["input"]
.as_array()
.is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning"))
);
assert_eq!(
replayed.pointer("/input/0/encrypted_content"),
Some(&json!(encrypted_content))
);
}
#[tokio::test]
async fn reasoning_replay_should_pair_context_scoped_custom_call_with_its_output() {
let encrypted_content = replay_ciphertext(4);
let transport = StubInferenceTransport::sequence([
InferenceMode::SuccessBody(stateful_sse_with_custom_tool_call(
&encrypted_content,
"inspect skills",
)),
InferenceMode::Success,
]);
let provider = provider(StubSelector::success(), transport.clone()).await;
execute_successfully(
&provider,
contextual_reasoning_replay_operation(
"conversation-from-header",
json!([{"type":"message","role":"user","content":"find skills"}]),
),
)
.await;
execute_successfully(
&provider,
contextual_reasoning_replay_operation(
"conversation-from-header",
json!([{
"type":"custom_tool_call_output",
"call_id":"call_exec",
"output":[{"type":"input_text","text":"skills found"}]
}]),
),
)
.await;
let requests = transport.requests.lock().expect("requests");
let replayed: serde_json::Value =
serde_json::from_slice(requests[1].body()).expect("replayed body");
let item_types = replayed["input"]
.as_array()
.expect("replayed input")
.iter()
.filter_map(|item| item["type"].as_str())
.collect::<Vec<_>>();
assert_eq!(
item_types,
["reasoning", "function_call", "function_call_output"]
);
}
#[tokio::test]
async fn reasoning_replay_is_scoped_to_the_selected_account() {
let encrypted_content = replay_ciphertext(1);
let transport = StubInferenceTransport::sequence([
InferenceMode::SuccessBody(stateful_sse(&encrypted_content)),
InferenceMode::Success,
InferenceMode::Success,
]);
let selector = SequencedAccountSelector::new([
account_id("replay-a"),
account_id("replay-b"),
account_id("replay-a"),
]);
let provider = provider_with_catalog_transport(
selector,
transport.clone(),
StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable),
Arc::new(StaticCatalogTransport),
)
.await;
for prompt in ["first", "different account", "owner again"] {
execute_successfully(
&provider,
reasoning_replay_operation(
"shared-session",
json!([{"type":"message","role":"user","content":prompt}]),
),
)
.await;
}
let requests = transport.requests.lock().expect("requests");
let different_account: serde_json::Value =
serde_json::from_slice(requests[1].body()).expect("different account request body");
let owner_again: serde_json::Value =
serde_json::from_slice(requests[2].body()).expect("owner request body");
assert!(
different_account["input"]
.as_array()
.is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning"))
);
assert_eq!(
owner_again.pointer("/input/0/encrypted_content"),
Some(&json!(encrypted_content))
);
}
#[tokio::test]
async fn reasoning_replay_rejects_mismatched_assistant_history() {
let encrypted_content = replay_ciphertext(2);
let transport = StubInferenceTransport::sequence([
InferenceMode::SuccessBody(stateful_sse_with_assistant(
&encrypted_content,
"cached answer",
)),
InferenceMode::Success,
]);
let provider = provider(StubSelector::success(), transport.clone()).await;
execute_successfully(
&provider,
reasoning_replay_operation(
"assistant-session",
json!([{"type":"message","role":"user","content":"first"}]),
),
)
.await;
execute_successfully(
&provider,
reasoning_replay_operation(
"assistant-session",
json!([
{"type":"message","role":"user","content":"first"},
{"type":"message","role":"assistant","content":[{"type":"output_text","text":"different answer"}]},
{"type":"message","role":"user","content":"next"}
]),
),
)
.await;
let requests = transport.requests.lock().expect("requests");
let body: serde_json::Value = serde_json::from_slice(requests[1].body()).expect("second body");
assert!(
body["input"]
.as_array()
.is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning"))
);
}
#[tokio::test]
async fn successful_compaction_clears_reasoning_replay_for_the_session() {
let encrypted_content = replay_ciphertext(3);
let summary = valid_compaction_summary("reasoning replay cleared");
let transport = StubInferenceTransport::sequence([
InferenceMode::SuccessBody(stateful_sse(&encrypted_content)),
InferenceMode::SuccessBody(compaction_sse(&summary, None)),
InferenceMode::Success,
]);
let provider = provider(StubSelector::success(), transport.clone()).await;
execute_successfully(
&provider,
reasoning_replay_operation(
"compact-session",
json!([{"type":"message","role":"user","content":"first"}]),
),
)
.await;
execute_successfully(
&provider,
compaction_operation_with_prompt_cache("compact-session"),
)
.await;
execute_successfully(
&provider,
reasoning_replay_operation(
"compact-session",
json!([{"type":"message","role":"user","content":"after compact"}]),
),
)
.await;
let requests = transport.requests.lock().expect("requests");
let body: serde_json::Value =
serde_json::from_slice(requests[2].body()).expect("post-compaction body");
assert!(
body["input"]
.as_array()
.is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning"))
);
}
#[tokio::test]
async fn connection_state_inherits_session_and_recovers_reasoning_on_pinned_account() {
use base64::Engine as _;
@@ -642,6 +642,76 @@ fn decoder_should_preserve_image_and_hosted_tool_items_without_inventing_facts()
);
}
#[test]
fn decoder_should_hide_injected_search_calls_and_compact_later_output_indexes() {
let request = tool_request(serde_json::json!({
"model": "client",
"prompt_cache_key": "cache-session",
"input": "hello"
}));
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_cache\",\"tools\":[{\"type\":\"web_search\"},{\"type\":\"x_search\"}]}}\n\n",
"event: response.output_item.added\n",
"data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_xs\",\"call_id\":\"xs_call_1\",\"name\":\"x_semantic_search\",\"arguments\":\"{}\"}}\n\n",
"event: response.function_call_arguments.delta\n",
"data: {\"type\":\"response.function_call_arguments.delta\",\"output_index\":0,\"item_id\":\"item_xs\",\"call_id\":\"xs_call_1\",\"delta\":\"{}\"}\n\n",
"event: response.output_item.done\n",
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_xs\",\"call_id\":\"xs_call_1\",\"name\":\"x_semantic_search\",\"arguments\":\"{}\"}}\n\n",
"event: response.output_item.added\n",
"data: {\"type\":\"response.output_item.added\",\"output_index\":1,\"item\":{\"type\":\"image_generation_call\",\"id\":\"image_visible\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_cache\",\"status\":\"completed\",\"tools\":[{\"type\":\"web_search\"},{\"type\":\"x_search\"}],\"output\":[{\"type\":\"function_call\",\"call_id\":\"xs_call_1\",\"name\":\"x_semantic_search\",\"arguments\":\"{}\"},{\"type\":\"image_generation_call\",\"id\":\"image_visible\"}]}}\n\n",
);
let events = GrokCanonicalDecoder::for_request("grok-4.5", &request)
.push(body.as_bytes())
.expect("filtered cache route response");
let wire = wire_events(&events);
assert_eq!(wire.len(), 3);
assert_eq!(
wire[1].data().get("output_index"),
Some(&serde_json::json!(0))
);
assert_eq!(
wire[2].data().pointer("/response/output/0/id"),
Some(&serde_json::json!("image_visible"))
);
assert_eq!(wire[2].data().pointer("/response/output/1"), None);
assert_eq!(wire[2].data().pointer("/response/tools"), None);
}
#[test]
fn decoder_should_preserve_an_internal_search_name_declared_by_the_client() {
let request = tool_request(serde_json::json!({
"model": "client",
"prompt_cache_key": "cache-session",
"input": "hello",
"tools": [{
"type": "function",
"name": "x_user_search",
"parameters": {"type": "object"}
}]
}));
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_declared\"}}\n\n",
"event: response.output_item.added\n",
"data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_client\",\"call_id\":\"call_client\",\"name\":\"x_user_search\",\"arguments\":\"{}\"}}\n\n",
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_declared\",\"status\":\"completed\"}}\n\n",
);
let events = GrokCanonicalDecoder::for_request("grok-4.5", &request)
.push(body.as_bytes())
.expect("client-declared search function");
assert!(wire_events(&events).iter().any(|event| {
event.data().pointer("/item/name") == Some(&serde_json::json!("x_user_search"))
}));
}
#[test]
fn decoder_should_tolerate_mismatched_event_identity_and_preserve_the_wire_event() {
let body = concat!(
@@ -91,7 +91,7 @@ fn headers_should_bind_identity_to_the_selected_oauth_account() {
}
#[test]
fn headers_should_keep_openai_cache_requests_out_of_shell_session_mode() {
fn headers_should_send_both_stable_session_headers_without_inventing_a_turn_index() {
let headers = build_grok_headers(
&crate::support::xai_wire_profile(),
&selected_session(),
@@ -109,7 +109,7 @@ fn headers_should_keep_openai_cache_requests_out_of_shell_session_mode() {
};
assert_eq!(value("x-grok-conv-id"), Some("session-fixture"));
assert_eq!(value("x-grok-session-id"), None);
assert_eq!(value("x-grok-session-id"), Some("session-fixture"));
assert_eq!(value("x-grok-turn-idx"), None);
}
@@ -256,7 +256,7 @@ fn explicit_session_should_enable_the_noop_native_cache_route() {
}
#[test]
fn explicit_session_should_not_add_invokable_cache_tools_to_client_tool_requests() {
fn explicit_session_should_add_only_internal_x_search_to_client_tool_requests() {
let request = raw_request(json!({
"model": "client",
"prompt_cache_key": "conversation-42",
@@ -273,15 +273,15 @@ fn explicit_session_should_not_add_invokable_cache_tools_to_client_tool_requests
GrokResponsesRequest::encode(&request, "grok-4.5", &client_key()).expect("request");
let body = Value::Object(encoded.body().clone());
// 缓存路由不得扩大模型可调用的工具集:带客户端工具时不注入原生搜索。
assert_eq!(
body.pointer("/tools"),
Some(&json!([{
body.pointer("/tools/0"),
Some(&json!({
"type": "function",
"name": "read_file",
"parameters": {"type": "object"}
}]))
}))
);
assert_eq!(body.pointer("/tools/1"), Some(&json!({"type": "x_search"})));
assert_eq!(body.pointer("/tool_choice"), Some(&json!("auto")));
}
@@ -384,7 +384,7 @@ fn function_parameters_should_report_the_invalid_nullable_root_field() {
}
#[test]
fn explicit_session_should_not_inject_cache_tools_after_codex_additional_tools_normalization() {
fn explicit_session_should_add_x_search_after_codex_additional_tools_normalization() {
let request = raw_request(json!({
"model": "client",
"prompt_cache_key": "conversation-42",
@@ -417,11 +417,38 @@ fn explicit_session_should_not_inject_cache_tools_after_codex_additional_tools_n
Some(&json!(["patch"]))
);
assert_eq!(body.pointer("/tools/1/name"), Some(&json!("read_file")));
// 归一化装载的会话工具视同客户端工具:不再追加可被调用的原生搜索。
assert_eq!(body.pointer("/tools/2"), None);
assert_eq!(body.pointer("/tools/2"), Some(&json!({"type": "x_search"})));
assert_eq!(body.pointer("/tool_choice"), Some(&json!("auto")));
}
#[test]
fn explicit_session_should_add_x_search_to_allowed_tools_choice() {
let request = raw_request(json!({
"model": "client",
"prompt_cache_key": "conversation-42",
"input": "first",
"tools": [{
"type": "function",
"name": "read_file",
"parameters": {"type": "object"}
}],
"tool_choice": {
"type": "allowed_tools",
"mode": "auto",
"tools": [{"type": "function", "name": "read_file"}]
}
}));
let encoded =
GrokResponsesRequest::encode(&request, "grok-4.5", &client_key()).expect("request");
let body = Value::Object(encoded.body().clone());
assert_eq!(
body.pointer("/tool_choice/tools/1"),
Some(&json!({"type": "x_search"}))
);
}
#[test]
fn explicit_session_should_not_duplicate_existing_native_cache_tools() {
let request = raw_request(json!({
+1
View File
@@ -2,3 +2,4 @@
0ae3945c02461347ab8dcfc64821027cd2366505778e0450d71bd82772c4e99d 0002_drop_opaque_response_id_index.sql
e608ed89991f9d6317491714447a00a61ede90caca4fb030eb30695654afb664 0003_store_opaque_response_ids_as_bytes.sql
df4e02cf2335733326c12429ced4cbd0b35b8a22b3fc24981c3c0549477d7c43 0004_add_model_request_service_tier.sql
fe8f36913c6db7053148db450f05e008dab6cdb9ca8cdcfcb22d97ad9a87bf49 0005_allow_quota_exhausted_cooldown.sql
@@ -0,0 +1,12 @@
alter table provider_accounts
drop constraint provider_accounts_cooldown_ck,
add constraint provider_accounts_cooldown_ck check (
(
availability != 'cooldown'
or cooldown_until is not null
)
and (
cooldown_until is null
or availability in ('cooldown', 'quota_exhausted')
)
);
+1 -1
View File
@@ -1 +1 @@
version: "3.0.22"
version: "3.0.26"