Compare commits

...
27 Commits
Author SHA1 Message Date
Zyy a7ea957843 test(openai): move provider mapping tests to integration suite 2026-07-31 17:37:20 +08:00
ZyyandCodex c92c95a733 chore: bump version to 3.0.32
Co-authored-by: Codex <noreply@openai.com>
2026-07-31 16:27:35 +08:00
ZyyandCodex ae7608440e fix(openai): replay client-owned conversation history
Co-authored-by: Codex <noreply@openai.com>
2026-07-31 16:07:32 +08:00
ZyyandCodex 5c9da8a39d fix(openai): preserve upstream failures during fallback
Co-authored-by: Codex <noreply@openai.com>
2026-07-31 15:33:22 +08:00
zyycnandCodex e012f24ef2 chore: bump version to 3.0.31
Co-authored-by: Codex <noreply@openai.com>
2026-07-31 02:20:58 +08:00
zyycnandCodex 756c9c7158 fix(openai): preserve exhausted quotas until reset
Co-authored-by: Codex <noreply@openai.com>
2026-07-31 02:20:37 +08:00
zyycnandCodex ab4bf75d0e chore: bump version to 3.0.30
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 21:43:11 +08:00
zyycnandCodex 7a9715ed41 chore(db): fold schema updates into initial migration
Collapse the completed response ID, service tier, and quota cooldown changes into the fresh-database baseline. Refresh the migration contract and state consistently that cross-major online upgrades are unsupported.

Co-authored-by: Codex <noreply@openai.com>
2026-07-30 21:42:57 +08:00
zyycnandCodex c353a6159e fix(openai): preserve failover before stream commit
Hold structural response events for a bounded replay grace so a split-frame 429 can rotate accounts without disconnecting the client. Release immediately on semantic output and keep post-commit failures non-replayable.

Co-authored-by: Codex <noreply@openai.com>
2026-07-30 21:42:30 +08:00
zyycnandCodex f5dca4848a fix(openai): harden exhausted quota recovery
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 21:07:53 +08:00
zyycnandCodex 36321a9af4 chore: bump version to 3.0.29
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 04:03:58 +08:00
zyycnandCodex b485b8bd4a fix(openai): reconcile exhausted account quotas
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 04:03:04 +08:00
zyycnandCodex cae2c8656a fix(xai): preserve quota exhaustion state
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 03:10:35 +08:00
zyycnandCodex acf574129f chore: bump version to 3.0.28
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 02:26:35 +08:00
zyycnandCodex 6426e0d3de fix(accounts): align availability summaries
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 02:22:23 +08:00
zyycnandCodex eedc39b33a chore: bump version to 3.0.27
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 01:55:51 +08:00
zyycnandCodex 9672a2e000 fix(xai): prioritize structured quota signals
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 01:53:43 +08:00
zyycnandCodex 0fffd2c550 fix(accounts): converge account state and imports
Co-authored-by: Codex <noreply@openai.com>
2026-07-30 01:23:42 +08:00
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
98 changed files with 6096 additions and 1397 deletions
+1 -1
View File
@@ -626,7 +626,7 @@ jobs:
make_latest: ${{ needs.release-info.outputs.prerelease == 'true' && 'false' || 'true' }}
body: |
> [!WARNING]
> **Upgrading across major versions?** Database schema and configuration are only compatible within a major version, and the one-click self-update refuses cross-major targets (releases before 3.0 lack this gate — do not use their built-in update to reach a newer major). Deploy the new major fresh (new `.runtime/` data directory) and re-import or re-authorize provider accounts and client keys. See the [upgrade notes](https://github.com/${{ github.repository }}/blob/main/deploy/README.md#镜像升级与源码构建).
> **Upgrading across major versions?** Online upgrades across major versions are not supported. Deploy the new major fresh with a new `.runtime/` data directory, then re-import or re-authorize provider accounts and client keys. See the [upgrade notes](https://github.com/${{ github.repository }}/blob/main/deploy/README.md#镜像升级与源码构建).
## 📥 Installation
+2 -4
View File
@@ -168,10 +168,8 @@ docker compose -f deploy/compose.yaml up -d
> 运行状态。
> [!WARNING]
> 跨大版本不是原地升级:数据库 schema 与配置只在同一大版本内兼容,内置一键自更新
> 会拒绝跨大版本目标(3.0 之前的旧版没有这道门禁,不要用旧版界面的一键更新跨大版本)。
> 升级到新的大版本需要全新部署(全新 `.runtime/` 数据目录),并重新导入或重新授权
> Provider 账号与客户端 Key。
> 不支持跨大版本在线升级。升级到新的大版本时,请使用全新的 `.runtime/` 数据目录重新部署,
> 并重新导入或重新授权 Provider 账号与客户端 Key。
Compose 默认只绑定 `127.0.0.1`。公网接入应使用 HTTPS 反向代理,转发 WebSocket
upgrade 与真实客户端 IP;不要暴露 PostgreSQL 或 Redis。
@@ -169,6 +169,8 @@ pub struct AccountPage {
}
/// 统一账号目录的全局状态计数,不受当前筛选和分页影响。
///
/// 配额耗尽是可恢复状态,始终单列;`unavailable` 只包含过期、禁用和封禁账号。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccountSummary {
pub total: u64,
@@ -460,6 +460,10 @@ pub struct RequestMetricPoint {
}
/// 账号池容量统计。
///
/// 首页概览的 `unavailable` 表示不可参与调度,包含配额耗尽及过期、禁用和封禁账号。
///
/// `quota_exhausted` 仍单列,以便状态详情展示和诊断。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AccountPoolMetrics {
pub total: u64,
@@ -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(),
@@ -373,6 +374,7 @@ impl ProviderAdmin for FakeProviderAdmin {
pub(super) struct FakeAccountStore {
events: EventLog,
account: AccountRecord,
account_after_probe: Mutex<Option<AccountRecord>>,
fail_commit: Mutex<bool>,
audit_requests: Mutex<Vec<String>>,
quota_window_usage: Mutex<Vec<AccountUsageWindowResult>>,
@@ -387,6 +389,7 @@ impl FakeAccountStore {
Arc::new(Self {
events,
account,
account_after_probe: Mutex::new(None),
fail_commit: Mutex::new(false),
audit_requests: Mutex::new(Vec::new()),
quota_window_usage: Mutex::new(Vec::new()),
@@ -405,6 +408,13 @@ impl FakeAccountStore {
*self.quota_window_usage.lock().expect("quota window usage") = usage;
}
fn set_account_after_probe(&self, account: AccountRecord) {
*self
.account_after_probe
.lock()
.expect("account after probe") = Some(account);
}
fn record(&self, event: &'static str) {
self.events.lock().expect("store events").push(event);
}
@@ -444,7 +454,13 @@ impl AccountStore for FakeAccountStore {
async fn load_account(&self, account_id: &str) -> AdminStoreResult<Option<AccountRecord>> {
self.record("store.load_account");
Ok((account_id == self.account.id).then(|| self.account.clone()))
let account = self
.account_after_probe
.lock()
.expect("account after probe")
.clone()
.unwrap_or_else(|| self.account.clone());
Ok((account_id == account.id.as_str()).then_some(account))
}
async fn load_account_usage(
@@ -801,6 +817,65 @@ async fn connection_test_should_probe_unavailable_account() {
));
}
#[tokio::test]
async fn connection_test_reloads_resettable_usage_limit_status_after_probe() {
let events = events();
let provider = FakeProviderAdmin::new("xai", events.clone());
let store = FakeAccountStore::new("xai", events);
let probe = Arc::new(FreeModelQuotaProbe {
store: Arc::clone(&store),
});
let services = accounts_service_with_probe(provider, store, probe).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::QuotaExhausted,
..
})
));
}
#[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());
@@ -1204,6 +1279,30 @@ impl AccountProbe for FailingAccountProbe {
}
}
struct FreeModelQuotaProbe {
store: Arc<FakeAccountStore>,
}
impl AccountProbe for FreeModelQuotaProbe {
fn probe(
&self,
_: AccountProbeRequest,
) -> BoxFuture<'_, Result<AccountProbeResult, AccountProbeError>> {
Box::pin(async move {
let mut account = account_record("xai");
account.availability = AccountAvailability::Cooldown;
account.availability_reason = Some("usage_limit_exhausted".to_owned());
account.cooldown_until = Some(Utc::now() + TimeDelta::hours(24));
self.store.set_account_after_probe(account);
Err(GatewayError::new(
GatewayErrorKind::RateLimited,
"xAI free model quota is exhausted",
)
.into())
})
}
}
fn store_unavailable() -> AdminStoreError {
AdminStoreError::new(
AdminStoreErrorKind::Unavailable,
@@ -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")
@@ -1653,12 +1653,7 @@ fn format_token_price(cost: &domain::CurrencyCost) -> String {
if cost.currency != "USD" {
return format!("{} {} / 1M Token", cost.currency, cost.amount.as_str());
}
let value = cost
.amount
.as_str()
.parse::<f64>()
.map_or(0.0, |value| value);
format!("${value:.4} / 1M Token")
format!("${} / 1M Token", cost.amount.as_str())
}
fn format_service_tier(service_tier: Option<&str>) -> String {
+12 -15
View File
@@ -7,8 +7,8 @@ use axum::{
response::{IntoResponse, Response},
};
use gateway_core::{
engine::{EngineError, UpstreamSendState},
error::{GatewayError, GatewayErrorKind, ProviderErrorKind},
engine::EngineError,
error::{GatewayError, GatewayErrorKind},
};
use serde_json::{Value, json};
@@ -72,19 +72,6 @@ pub fn model_not_found_response() -> (StatusCode, Json<Value>) {
#[must_use]
pub fn gateway_error_from_engine(error: &EngineError) -> GatewayError {
match error {
EngineError::Provider(error)
if error.kind() == ProviderErrorKind::Unavailable
&& error.send_state() == UpstreamSendState::NotSent =>
{
let gateway = GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream provider is currently available for this request",
);
match error.client_visible_upstream_error() {
Some(upstream) => gateway.with_client_visible_upstream_error(upstream.clone()),
None => gateway,
}
}
EngineError::Provider(error) => GatewayError::from_provider(error),
EngineError::Cancelled => {
GatewayError::new(GatewayErrorKind::Cancelled, "request was cancelled")
@@ -185,6 +172,16 @@ pub const fn gateway_error_contract(
"server_error",
"no_available_provider",
),
GatewayErrorKind::AccountCapacityUnavailable => (
StatusCode::SERVICE_UNAVAILABLE,
"server_error",
"account_capacity_unavailable",
),
GatewayErrorKind::ProviderInfrastructureUnavailable => (
StatusCode::SERVICE_UNAVAILABLE,
"server_error",
"provider_infrastructure_unavailable",
),
GatewayErrorKind::RateLimited => (
StatusCode::TOO_MANY_REQUESTS,
"rate_limit_error",
@@ -340,7 +340,7 @@ mod actions {
upstream_status: Some(429),
upstream_content_type: Some("application/json".to_owned()),
upstream_body: Some(r#"{"error":{"type":"usage_limit_reached"}}"#.to_owned()),
account_status: AccountStatus::Active,
account_status: AccountStatus::QuotaExhausted,
},
]
.map(|event| AccountConnectionTestEvent::from(event).data);
@@ -374,7 +374,7 @@ mod actions {
"upstreamStatus": 429,
"upstreamContentType": "application/json",
"upstreamBody": r#"{"error":{"type":"usage_limit_reached"}}"#,
"accountStatus": "active"
"accountStatus": "quota_exhausted"
}),
]
);
+10 -1
View File
@@ -79,6 +79,7 @@ pub(super) struct AdminTestFixture {
pub auth: Arc<MemoryAuthStore>,
pub settings: Arc<MemorySettingsStore>,
pub usage_records: Arc<Mutex<Vec<UsageRecord>>>,
pub dashboard_observation: Arc<Mutex<Option<DashboardObservation>>>,
}
impl AdminTestFixture {
@@ -92,8 +93,10 @@ impl AdminTestFixture {
let settings = Arc::new(MemorySettingsStore::new(api_key));
let client_keys = Arc::new(MemoryClientKeyStore);
let usage_records = Arc::new(Mutex::new(Vec::new()));
let dashboard_observation = Arc::new(Mutex::new(None));
let unused = Arc::new(UnusedStore {
usage_records: Arc::clone(&usage_records),
dashboard_observation: Arc::clone(&dashboard_observation),
});
let stores = AdminStorePorts::new(
unused.clone(),
@@ -125,6 +128,7 @@ impl AdminTestFixture {
auth,
settings,
usage_records,
dashboard_observation,
}
}
@@ -392,6 +396,7 @@ impl ClientKeyStore for MemoryClientKeyStore {
struct UnusedStore {
usage_records: Arc<Mutex<Vec<UsageRecord>>>,
dashboard_observation: Arc<Mutex<Option<DashboardObservation>>>,
}
#[async_trait]
@@ -503,7 +508,11 @@ impl AccountStore for UnusedStore {
#[async_trait]
impl ObservabilityStore for UnusedStore {
async fn dashboard_summary(&self, _: TimeRange) -> AdminStoreResult<DashboardObservation> {
Err(unavailable("dashboard"))
self.dashboard_observation
.lock()
.expect("dashboard observation")
.clone()
.ok_or_else(|| unavailable("dashboard"))
}
async fn dashboard_trend(&self, _: TimeRange) -> AdminStoreResult<Vec<RequestMetricPoint>> {
@@ -260,6 +260,80 @@ mod response {
}
}
#[tokio::test]
async fn dashboard_summary_should_include_quota_exhaustion_in_unavailable_headline_count() {
use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode, header},
};
use chrono::{Duration, Utc};
use gateway_admin::model::observability::{
AccountPoolMetrics, AttemptMetrics, DashboardObservation, RequestMetrics, TimeRange,
};
use gateway_api::admin::observability;
use tower::ServiceExt as _;
use super::{AdminTestFixture, AdminTestState};
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let now = Utc::now();
*fixture
.dashboard_observation
.lock()
.expect("dashboard observation") = Some(DashboardObservation {
range: TimeRange::new(now - Duration::hours(1), now).expect("dashboard range"),
requests: RequestMetrics::default(),
attempts: AttemptMetrics::default(),
provider_accounts: AccountPoolMetrics {
total: 807,
enabled: 804,
unavailable: 792,
active: 15,
expired: 2,
quota_exhausted: 785,
refreshing: None,
disabled: 3,
banned: 2,
},
trend: Vec::new(),
account_usage: Vec::new(),
recent_requests: Vec::new(),
});
let response = observability::router::<AdminTestState>()
.with_state(fixture.state())
.oneshot(
Request::builder()
.uri("/api/admin/dashboard/summary")
.header(header::COOKIE, "cpr_admin_session=valid-session")
.header("x-request-id", "req_dashboard_account_summary")
.body(Body::empty())
.expect("dashboard request"),
)
.await
.expect("dashboard response");
let status = response.status();
let body = to_bytes(response.into_body(), 1024 * 1024)
.await
.expect("dashboard response body");
let value: serde_json::Value = serde_json::from_slice(&body).expect("dashboard response JSON");
assert_eq!(status, StatusCode::OK, "{value}");
assert_eq!(
(
&value["data"]["cards"]["credentials"]["totalValue"],
&value["data"]["cards"]["credentials"]["availableValue"],
&value["data"]["cards"]["credentials"]["unavailableValue"],
),
(
&serde_json::json!(807),
&serde_json::json!(15),
&serde_json::json!(792)
),
);
}
#[tokio::test]
async fn usage_route_should_forward_a_bounded_unknown_outcome_filter() {
use axum::{
@@ -291,12 +365,17 @@ async fn usage_route_should_forward_a_bounded_unknown_outcome_filter() {
#[tokio::test]
async fn usage_route_should_expose_image_and_websocket_facts() {
use std::str::FromStr as _;
use axum::{
body::{Body, to_bytes},
http::{Request, StatusCode, header},
};
use chrono::Utc;
use gateway_admin::model::observability::{RequestOutcome, UsageRecord};
use gateway_admin::model::observability::{
CalculatedBillingBreakdown, CurrencyCost, DecimalAmount, RequestOutcome, UsageBilling,
UsageRecord,
};
use gateway_api::admin::observability;
use tower::ServiceExt as _;
@@ -304,6 +383,10 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
let fixture = AdminTestFixture::new().await;
fixture.auth.insert_session("valid-session");
let usd = |amount: &str| CurrencyCost {
currency: "USD".to_owned(),
amount: DecimalAmount::from_str(amount).expect("USD amount"),
};
fixture
.usage_records
.lock()
@@ -359,7 +442,22 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
cost_source: "unavailable".to_owned(),
cost_amount: None,
cost_currency: None,
billing: None,
billing: Some(UsageBilling::Calculated(Box::new(
CalculatedBillingBreakdown {
input_amount: usd("0.03"),
output_amount: usd("0.07"),
cache_read_amount: usd("0.00"),
cache_write_amount: usd("0.00"),
standard_amount: usd("0.10"),
total_amount: usd("0.10"),
input_price_per_million: usd("10.0000"),
output_price_per_million: usd("60.0000"),
cache_read_price_per_million: usd("1.0000"),
cache_write_price_per_million: usd("12.5000"),
service_tier: Some("priority".to_owned()),
multiplier_percent: 100,
},
))),
transport_decision_wait_ms: Some(7),
connect_ms: Some(11),
headers_ms: Some(13),
@@ -400,6 +498,19 @@ async fn usage_route_should_expose_image_and_websocket_facts() {
.expect("usage response body");
let value: serde_json::Value = serde_json::from_slice(&body).expect("usage response JSON");
assert_eq!(
value["data"]["items"][0]["billing"]["inputPriceDisplay"],
"$10 / 1M Token"
);
assert_eq!(
value["data"]["items"][0]["billing"]["outputPriceDisplay"],
"$60 / 1M Token"
);
assert_eq!(
value["data"]["items"][0]["billing"]["cacheWritePriceDisplay"],
"$12.5 / 1M Token"
);
assert_eq!(
serde_json::json!({
"route": value["data"]["items"][0]["route"],
@@ -58,6 +58,30 @@ fn no_available_provider_error_should_map_to_service_unavailable() {
);
}
#[test]
fn account_capacity_error_should_map_to_service_unavailable() {
assert_eq!(
gateway_error_contract(GatewayErrorKind::AccountCapacityUnavailable),
(
StatusCode::SERVICE_UNAVAILABLE,
"server_error",
"account_capacity_unavailable",
)
);
}
#[test]
fn provider_infrastructure_error_should_map_to_service_unavailable() {
assert_eq!(
gateway_error_contract(GatewayErrorKind::ProviderInfrastructureUnavailable),
(
StatusCode::SERVICE_UNAVAILABLE,
"server_error",
"provider_infrastructure_unavailable",
)
);
}
#[test]
fn rate_limited_error_should_map_to_openai_retryable_status() {
assert_eq!(
@@ -122,29 +146,32 @@ fn engine_provider_error_should_preserve_retry_classification() {
#[test]
fn local_provider_capacity_exhaustion_should_map_to_service_unavailable() {
let error = EngineError::Provider(ProviderError::new(
ProviderErrorKind::Unavailable,
ProviderErrorKind::AccountCapacityUnavailable,
UpstreamSendState::NotSent,
));
assert_eq!(
gateway_error_from_engine(&error),
GatewayError::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream provider is currently available for this request"
GatewayErrorKind::AccountCapacityUnavailable,
"all eligible upstream accounts are temporarily busy"
)
);
}
#[test]
fn local_provider_capacity_exhaustion_should_preserve_safe_provider_detail() {
fn no_eligible_provider_account_should_preserve_safe_provider_detail() {
let detail = ClientVisibleUpstreamError::new(
"no account is eligible for the requested model",
Some("no_eligible_account".to_owned()),
Some("account_unavailable_error".to_owned()),
);
let error = EngineError::Provider(
ProviderError::new(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
.with_client_visible_upstream_error(detail),
ProviderError::new(
ProviderErrorKind::NoEligibleAccount,
UpstreamSendState::NotSent,
)
.with_client_visible_upstream_error(detail),
);
let gateway = gateway_error_from_engine(&error);
@@ -164,6 +191,22 @@ fn local_provider_capacity_exhaustion_should_preserve_safe_provider_detail() {
);
}
#[test]
fn provider_infrastructure_failure_should_have_a_distinct_client_contract() {
let error = EngineError::Provider(ProviderError::new(
ProviderErrorKind::ProviderInfrastructureUnavailable,
UpstreamSendState::NotSent,
));
assert_eq!(
gateway_error_from_engine(&error),
GatewayError::new(
GatewayErrorKind::ProviderInfrastructureUnavailable,
"provider account infrastructure is temporarily unavailable"
)
);
}
#[test]
fn no_eligible_provider_account_should_map_to_service_unavailable() {
let error = EngineError::Provider(ProviderError::new(
@@ -181,16 +224,18 @@ fn no_eligible_provider_account_should_map_to_service_unavailable() {
}
#[test]
fn sent_provider_unavailability_should_remain_bad_gateway() {
let error = EngineError::Provider(ProviderError::new(
ProviderErrorKind::Unavailable,
UpstreamSendState::Sent,
));
fn provider_unavailability_should_remain_bad_gateway_regardless_of_send_state() {
for send_state in [UpstreamSendState::NotSent, UpstreamSendState::Sent] {
let error = EngineError::Provider(ProviderError::new(
ProviderErrorKind::Unavailable,
send_state,
));
assert_eq!(
gateway_error_from_engine(&error).kind(),
GatewayErrorKind::UpstreamUnavailable
);
assert_eq!(
gateway_error_from_engine(&error).kind(),
GatewayErrorKind::UpstreamUnavailable
);
}
}
#[test]
@@ -165,8 +165,8 @@ where
timings: ModelRequestTimings::default(),
client_response_id: None,
upstream_response_id: None,
last_account_exhaustion: None,
last_account_exhaustion_events: Vec::new(),
last_retryable_failure: None,
last_retryable_failure_events: Vec::new(),
provider_attempt_outcomes: Vec::new(),
pending_terminal_failure: None,
};
@@ -245,9 +245,10 @@ pub struct ResponseExecutionSession<S: ?Sized> {
timings: ModelRequestTimings,
client_response_id: Option<String>,
upstream_response_id: Option<String>,
last_account_exhaustion: Option<ProviderError>,
/// 最近一次账号耗尽在重试时被丢弃的原始失败批次;只用于后续空选路终态。
last_account_exhaustion_events: Vec<ProviderEvent>,
/// 最近一次为无感恢复而被丢弃的原始上游错误;只在后续空选路时成为终态。
last_retryable_failure: Option<ProviderError>,
/// 与 `last_retryable_failure` 同属一个 attempt 的原始失败批次。
last_retryable_failure_events: Vec<ProviderEvent>,
provider_attempt_outcomes: Vec<ProviderAttemptOutcome>,
/// 原子失败批次已交给协议层、但尚待下游提交后收敛的原 Provider 错误。
pending_terminal_failure: Option<PendingTerminalFailure>,
@@ -659,25 +660,30 @@ where
ProviderBoundary::Result(result) => match *result {
Ok(stream) => stream,
Err(error) => {
if error.kind() == ProviderErrorKind::NoEligibleAccount
&& let Some(exhaustion) = self.last_account_exhaustion.take()
if matches!(
error.kind(),
ProviderErrorKind::AccountCapacityUnavailable
| ProviderErrorKind::NoEligibleAccount
) && let Some(last_failure) = self.last_retryable_failure.take()
{
let send_state = self.current_send_state();
let mut events = std::mem::take(&mut self.last_account_exhaustion_events);
let mut events = std::mem::take(&mut self.last_retryable_failure_events);
if !events.is_empty() {
self.observe_atomic_terminal_events(&mut events);
return Ok(Some(PullOutcome::TerminalFailure {
events,
error: exhaustion,
error: last_failure,
send_state,
}));
}
self.finish_provider_error(&exhaustion).await?;
return Err(provider_engine_error(exhaustion));
self.finish_provider_error(&last_failure).await?;
return Err(provider_engine_error(last_failure));
}
if !(matches!(
error.kind(),
ProviderErrorKind::Unavailable | ProviderErrorKind::NoEligibleAccount
ProviderErrorKind::AccountCapacityUnavailable
| ProviderErrorKind::NoEligibleAccount
| ProviderErrorKind::ProviderInfrastructureUnavailable
) && error.send_state() == UpstreamSendState::NotSent)
{
self.record_provider_failure(candidate.provider().clone(), error.kind());
@@ -909,14 +915,10 @@ where
&mut self,
mut error: ProviderError,
) -> Result<StreamErrorOutcome, EngineError> {
// 原始 wire 只能活在本次决策栈内;后续 clone、attempt 记录与终态对象均只
// 接触已剥离的稳定错误字段。
// 原始 wire 只活在 request-local 决策状态;clone、attempt 记录与持久化终态
// 均只接触已剥离的稳定错误字段。
let mut atomic_client_events = error.take_atomic_client_events();
let current = self.current.take().ok_or(EngineError::NoActiveAttempt)?;
let account_exhausted = matches!(
error.kind(),
ProviderErrorKind::RateLimited | ProviderErrorKind::QuotaExhausted
);
self.record_provider_failure(current.metadata.provider().clone(), error.kind());
// attempt_send_state 是本 attempt 自身的发送事实,驱动重试门;
// 持久化与终态用请求级水位,二者不可混用(水位会把早先 attempt 的
@@ -976,8 +978,8 @@ where
continuation_retry || same_account_retry || ordinary_retry || pre_delivery_retry;
if retryable {
// 普通 clone 只保留稳定事实;原始失败响应只能随最后一次客户端可见
// 的账号耗尽错误继续,不能进入中间失败记录。
// 普通 clone 只保留稳定事实;原始 wire/HTTP response 由 request-local
// 所有权保留到下一次 attempt 成功,或最终空选路时返回客户端。
let persistence_error = self.request_persisted.then(|| error.clone());
if same_account_retry {
if let Some(account) = current.metadata.provider_account_id() {
@@ -992,10 +994,8 @@ where
{
self.excluded_accounts.insert(account.clone());
}
if account_exhausted {
self.last_account_exhaustion_events = atomic_client_events;
self.last_account_exhaustion = Some(error);
}
self.last_retryable_failure_events = atomic_client_events;
self.last_retryable_failure = Some(error);
if let Some(error) = persistence_error {
best_effort_store_write(
"record_intermediate_failure",
@@ -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);
+20
View File
@@ -129,8 +129,12 @@ pub enum ProviderErrorKind {
RateLimited,
/// Credential 配额耗尽。
QuotaExhausted,
/// 仍有符合条件的账号,但它们暂时没有可调度容量。
AccountCapacityUnavailable,
/// Provider 已确认当前请求无法选出可用账号。
NoEligibleAccount,
/// Provider 的账号存储、租约协调或本地凭据数据不可用。
ProviderInfrastructureUnavailable,
/// 请求超时。
Timeout,
/// 网络或 transport 失败。
@@ -156,7 +160,9 @@ impl ProviderErrorKind {
Self::PermissionDenied => "permission_denied",
Self::RateLimited => "rate_limited",
Self::QuotaExhausted => "quota_exhausted",
Self::AccountCapacityUnavailable => "account_capacity_unavailable",
Self::NoEligibleAccount => "no_eligible_account",
Self::ProviderInfrastructureUnavailable => "provider_infrastructure_unavailable",
Self::Timeout => "timeout",
Self::Transport => "transport",
Self::Protocol => "protocol",
@@ -711,6 +717,10 @@ pub enum GatewayErrorKind {
ModelNotFound,
/// 当前没有可用 target。
NoAvailableProvider,
/// 符合条件的上游账号暂时没有调度容量。
AccountCapacityUnavailable,
/// Provider 的本地账号基础设施不可用。
ProviderInfrastructureUnavailable,
/// 上游限流。
RateLimited,
/// 上游暂不可用。
@@ -734,6 +744,8 @@ impl GatewayErrorKind {
Self::PolicyDenied => "policy_denied",
Self::ModelNotFound => "model_not_found",
Self::NoAvailableProvider => "no_available_provider",
Self::AccountCapacityUnavailable => "account_capacity_unavailable",
Self::ProviderInfrastructureUnavailable => "provider_infrastructure_unavailable",
Self::RateLimited => "rate_limited",
Self::UpstreamUnavailable => "upstream_unavailable",
Self::Timeout => "timeout",
@@ -782,10 +794,18 @@ impl GatewayError {
GatewayErrorKind::RateLimited,
"upstream capacity is temporarily unavailable",
),
ProviderErrorKind::AccountCapacityUnavailable => Self::new(
GatewayErrorKind::AccountCapacityUnavailable,
"all eligible upstream accounts are temporarily busy",
),
ProviderErrorKind::NoEligibleAccount => Self::new(
GatewayErrorKind::NoAvailableProvider,
"no upstream provider is currently available for this request",
),
ProviderErrorKind::ProviderInfrastructureUnavailable => Self::new(
GatewayErrorKind::ProviderInfrastructureUnavailable,
"provider account infrastructure is temporarily unavailable",
),
ProviderErrorKind::Timeout => {
Self::new(GatewayErrorKind::Timeout, "upstream request timed out")
}
@@ -29,8 +29,9 @@ use gateway_core::engine::{
ProviderAttemptOutcome, RecoveryReport, UpstreamSendState,
};
use gateway_core::error::{
ClientVisibleUpstreamResponse, ContinuationFailure, GatewayErrorKind, OpaqueUpstreamValue,
ProviderError, ProviderErrorKind, StoreError, StoreErrorKind,
ClientVisibleUpstreamError, ClientVisibleUpstreamResponse, ContinuationFailure,
GatewayErrorKind, OpaqueUpstreamValue, ProviderError, ProviderErrorKind, StoreError,
StoreErrorKind,
};
use gateway_core::event::{
ContentItem, ContentKind, GatewayEvent, ProtocolWireEvent, ProviderEvent,
@@ -2554,6 +2555,128 @@ fn pre_delivery_transport_failure_rotates_account_for_non_idempotent_generation(
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded);
}
#[test]
fn final_capacity_exhaustion_returns_the_last_retryable_upstream_failure() {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation);
let raw_body = Bytes::from_static(b"{\"error\":\"last upstream body\"}\x00");
let upstream_error =
ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous)
.with_pre_delivery_retry()
.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
"message too big",
Some("1009".to_owned()),
Some("websocket_close_error".to_owned()),
))
.with_client_visible_upstream_response(ClientVisibleUpstreamResponse::new(
502,
Some(b"application/json".to_vec()),
raw_body.clone(),
));
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
items: vec![Err(upstream_error)],
},
Script::Error(ProviderError::new(
ProviderErrorKind::AccountCapacityUnavailable,
UpstreamSendState::NotSent,
)),
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let error = block_on(session.collect_uncommitted())
.expect_err("empty selection must return the retained upstream failure");
let EngineError::Provider(error) = error else {
panic!("expected retained provider error");
};
assert_eq!(error.kind(), ProviderErrorKind::Transport);
let detail = error
.client_visible_upstream_error()
.expect("structured upstream detail");
assert_eq!(detail.message(), "message too big");
assert_eq!(detail.code(), Some("1009"));
assert_eq!(detail.error_type(), Some("websocket_close_error"));
let response = error
.client_visible_upstream_response()
.expect("original upstream response");
assert_eq!(response.status(), 502);
assert_eq!(response.body(), &raw_body);
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
let state = store.state.lock().expect("store lock");
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
assert_eq!(
state.finalizations[0].send_state,
UpstreamSendState::Ambiguous
);
}
#[test]
fn provider_infrastructure_failure_does_not_restore_the_last_upstream_failure() {
let operation = operation(RetrySafety::NonIdempotent);
let route_plan = plan(&operation);
let upstream_error =
ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous)
.with_pre_delivery_retry()
.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
"upstream marker",
Some("upstream_marker".to_owned()),
Some("upstream_error".to_owned()),
))
.with_client_visible_upstream_response(ClientVisibleUpstreamResponse::new(
502,
Some(b"application/json".to_vec()),
Bytes::from_static(b"{\"error\":\"upstream marker\"}"),
));
let (coordinator, store, provider) = coordinator(vec![
Script::Stream {
account_id: "acct_first",
items: vec![Err(upstream_error)],
},
Script::Error(ProviderError::new(
ProviderErrorKind::ProviderInfrastructureUnavailable,
UpstreamSendState::NotSent,
)),
]);
let mut session = block_on(coordinator.start(
model_request(&operation, SystemTime::now() + Duration::from_secs(30)),
operation,
route_plan,
None,
None,
CancellationToken::new(),
))
.expect("start execution");
let error = block_on(session.collect_uncommitted())
.expect_err("provider infrastructure failure must remain terminal");
let EngineError::Provider(error) = error else {
panic!("expected provider error");
};
assert_eq!(
error.kind(),
ProviderErrorKind::ProviderInfrastructureUnavailable
);
assert!(error.client_visible_upstream_error().is_none());
assert!(error.client_visible_upstream_response().is_none());
assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 2);
let state = store.state.lock().expect("store lock");
assert_eq!(state.intermediate_failures, 1);
assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed);
}
#[test]
fn pre_delivery_retry_marker_is_ignored_after_downstream_commit() {
let operation = operation(RetrySafety::NonIdempotent);
@@ -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,34 @@ 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 <> 'active')::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 +2042,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,9 +622,16 @@ 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)
where id = $1 and credential_revision = $2",
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
and (availability_observed_at is null or availability_observed_at <= $6)",
)
.bind(update.account_id)
.bind(to_i64(update.expected_revision.get())?)
@@ -659,7 +676,8 @@ impl ProviderAccountRepository for PgProviderAccountRepository {
"update provider_accounts
set provider_quota_json = $3, quota_observed_at = $4,
updated_at = greatest(now(), coalesce($4, now()))
where id = $1 and credential_revision = $2",
where id = $1 and credential_revision = $2
and ($4 is null or quota_observed_at is null or quota_observed_at <= $4)",
)
.bind(account_id)
.bind(to_i64(expected_revision.get())?)
@@ -755,6 +773,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 +794,7 @@ impl ProviderAccountAdminRepository for PgProviderAccountRepository {
&mut transaction,
&command.scope,
&command.profile,
command.replacement_identity.as_ref(),
&command.credential,
)
.await?;
@@ -1040,6 +1065,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)?,
@@ -1586,43 +1612,26 @@ fn admin_account_summary(
accounts: &[ProviderAccountSummary],
now: DateTime<Utc>,
) -> AccountSummary {
let total = u64::try_from(accounts.len()).unwrap_or(u64::MAX);
let active = u64::try_from(
accounts
.iter()
.filter(|account| admin_account_status(account, now) == AdminAccountStatus::Active)
.count(),
)
.unwrap_or(u64::MAX);
let quota_exhausted = u64::try_from(
accounts
.iter()
.filter(|account| {
admin_account_status(account, now) == AdminAccountStatus::QuotaExhausted
})
.count(),
)
.unwrap_or(u64::MAX);
let unavailable = u64::try_from(
accounts
.iter()
.filter(|account| {
matches!(
admin_account_status(account, now),
AdminAccountStatus::Expired
| AdminAccountStatus::Disabled
| AdminAccountStatus::Banned
)
})
.count(),
)
.unwrap_or(u64::MAX);
AccountSummary {
total,
active,
quota_exhausted,
unavailable,
let mut summary = AccountSummary {
total: u64::try_from(accounts.len()).unwrap_or(u64::MAX),
active: 0,
quota_exhausted: 0,
unavailable: 0,
};
for account in accounts {
match admin_account_status(account, now) {
AdminAccountStatus::Active => summary.active = summary.active.saturating_add(1),
AdminAccountStatus::QuotaExhausted => {
summary.quota_exhausted = summary.quota_exhausted.saturating_add(1);
}
AdminAccountStatus::Expired
| AdminAccountStatus::Disabled
| AdminAccountStatus::Banned => {
summary.unavailable = summary.unavailable.saturating_add(1);
}
}
}
summary
}
fn sort_admin_account_items(items: &mut [AdminAccountListItem], sort: Option<AdminAccountSort>) {
@@ -1975,8 +1984,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 +2001,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 +2038,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,
@@ -58,24 +58,6 @@ impl TestDatabase {
.execute(&mut *migration)
.await
.expect("apply terminal migration");
sqlx::raw_sql(include_str!(
"../../../../migrations/0002_drop_opaque_response_id_index.sql"
))
.execute(&mut *migration)
.await
.expect("apply opaque response ID migration");
sqlx::raw_sql(include_str!(
"../../../../migrations/0003_store_opaque_response_ids_as_bytes.sql"
))
.execute(&mut *migration)
.await
.expect("apply opaque response storage migration");
sqlx::raw_sql(include_str!(
"../../../../migrations/0004_add_model_request_service_tier.sql"
))
.execute(&mut *migration)
.await
.expect("apply model request service tier migration");
migration.commit().await.expect("commit terminal migration");
Some(Self {
admin,
@@ -171,24 +153,19 @@ 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, 1));
assert_eq!(response_id_types, ["bytea", "bytea"]);
assert!(!raw_response_id_index_exists);
}
#[test]
fn migrations_should_leave_transaction_ownership_to_sqlx() {
let transaction_statements = [
include_str!("../../../../migrations/0001_initial.sql"),
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"),
]
.into_iter()
.flat_map(str::lines)
.map(str::trim)
.filter(|line| matches!(*line, "begin;" | "commit;"))
.count();
let transaction_statements = [include_str!("../../../../migrations/0001_initial.sql")]
.into_iter()
.flat_map(str::lines)
.map(str::trim)
.filter(|line| matches!(*line, "begin;" | "commit;"))
.count();
assert_eq!(transaction_statements, 0);
}
@@ -112,7 +112,7 @@ async fn observability_preserves_and_filters_opaque_response_ids() {
}
#[tokio::test]
async fn dashboard_account_metrics_should_partition_available_and_unavailable_accounts() {
async fn dashboard_account_metrics_should_partition_account_statuses() {
let Some(database) = TestDatabase::create("dashboard_account_metrics").await else {
return;
};
@@ -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', 'xai', '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,11 +165,18 @@ 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, 4, 1, 1, 1, 1),
);
// 首页按调度可用性聚合:配额耗尽也属于不可用。
assert_eq!(
metrics.unavailable,
metrics.quota_exhausted + metrics.expired + metrics.disabled + metrics.banned
);
assert_eq!(metrics.total, metrics.active + metrics.unavailable);
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::{
@@ -110,6 +111,97 @@ async fn core_quota_batch_reads_only_observed_accounts_in_one_contract_call() {
database.close().await;
}
#[tokio::test]
async fn delayed_provider_observations_cannot_overwrite_newer_state_or_quota() {
let Some(database) = TestDatabase::create("provider_account_observation_fence").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let account_id = ProviderAccountId::new("acct_observation_fence").expect("account id");
repository
.insert_provider_account(account(account_id.as_str(), "user-observation-fence"))
.await
.expect("insert observation fence fixture");
let revision = CredentialRevision::new(1).expect("revision");
let newer = SystemTime::now()
.checked_add(Duration::from_secs(60))
.expect("newer observation time");
let older = newer
.checked_sub(Duration::from_secs(30))
.expect("older observation time");
let quota = |marker: &str, observed_at| QuotaObservation {
account_id: account_id.clone(),
expected_revision: revision,
quota: Some(OpaqueProviderData::new(
json!({"marker": marker})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(observed_at),
};
assert_eq!(
repository
.compare_and_swap_quota(quota("newer", newer))
.await
.expect("persist newer quota"),
QuotaWriteOutcome::Updated
);
assert_eq!(
repository
.compare_and_swap_quota(quota("older", older))
.await
.expect("reject older quota"),
QuotaWriteOutcome::Conflict
);
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: revision,
availability: AccountAvailability::QuotaExhausted,
reason: Some("quota_exhausted".to_owned()),
cooldown_until: Some(newer),
observed_at: newer,
})
.await
.expect("persist newer account state");
assert!(
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: revision,
availability: AccountAvailability::Ready,
reason: None,
cooldown_until: None,
observed_at: older,
})
.await
.is_err(),
"a delayed state observation must conflict"
);
let current = repository
.get_account(&account_id)
.await
.expect("load fenced account")
.expect("fenced account");
assert_eq!(current.availability(), AccountAvailability::QuotaExhausted);
let observed = repository
.get_quotas(std::slice::from_ref(&account_id))
.await
.expect("load fenced quota")
.pop()
.expect("fenced quota");
assert_eq!(
observed.quota.expect("quota document").expose_to_provider()["marker"],
"newer"
);
database.close().await;
}
#[tokio::test]
async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_usage() {
let Some(database) = TestDatabase::create("provider_account_terminal_list").await else {
@@ -129,10 +221,13 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
let mut invalid = account("acct_invalid", "user-invalid");
invalid.email = Some("invalid@example.invalid".to_owned());
invalid.availability = AccountAvailability::Invalid;
let mut disabled = account("acct_disabled", "user-disabled");
disabled.email = Some("disabled@example.invalid".to_owned());
disabled.enabled = false;
let mut quota_exhausted = account("acct_quota_exhausted", "user-quota-exhausted");
quota_exhausted.email = Some("quota-exhausted@example.invalid".to_owned());
quota_exhausted.availability = AccountAvailability::QuotaExhausted;
for account in [alpha, beta, charlie, invalid, quota_exhausted] {
for account in [alpha, beta, charlie, invalid, disabled, quota_exhausted] {
repository
.insert_provider_account(account)
.await
@@ -212,10 +307,17 @@ async fn terminal_admin_list_filters_and_sorts_before_pagination_with_retained_u
.await
.expect("sort accounts by retained usage");
assert_eq!(usage_page.config_revision.get(), 1);
assert_eq!(usage_page.total, 5);
assert_eq!(usage_page.summary.total, 5);
assert_eq!(usage_page.total, 6);
assert_eq!(usage_page.summary.total, 6);
assert_eq!(usage_page.summary.active, 2);
assert_eq!(usage_page.summary.quota_exhausted, 1);
assert_eq!(usage_page.summary.unavailable, 3);
assert_eq!(
usage_page.summary.total,
usage_page.summary.active
+ usage_page.summary.quota_exhausted
+ usage_page.summary.unavailable
);
assert_eq!(
usage_page
.items
@@ -898,6 +1000,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 +1029,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 +1040,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 +1176,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 +1205,174 @@ async fn verified_credential_rotation_preserves_quota_exhaustion() {
database.close().await;
}
#[tokio::test]
async fn xai_free_quota_exhaustion_state_is_persisted() {
let Some(database) = TestDatabase::create("provider_account_xai_free_quota").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let account_id = ProviderAccountId::new("acct_xai_free_quota").expect("account ID");
let mut imported = account(account_id.as_str(), "user-xai-free-quota");
imported.provider_kind = "xai".to_owned();
repository
.insert_provider_account(imported)
.await
.expect("insert xAI account");
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
availability: AccountAvailability::QuotaExhausted,
reason: Some("upstream_free_quota_exhausted".to_owned()),
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("persist xAI free quota exhaustion");
let current: (String, String, Option<String>) = sqlx::query_as(
"select provider_kind, availability, availability_reason
from provider_accounts where id = $1",
)
.bind(account_id.as_str())
.fetch_one(&database.pool)
.await
.expect("load persisted xAI quota state");
assert_eq!(
current,
(
"xai".to_owned(),
"quota_exhausted".to_owned(),
Some("upstream_free_quota_exhausted".to_owned()),
),
);
database.close().await;
}
#[tokio::test]
async fn xai_resettable_usage_limit_state_is_persisted() {
let Some(database) = TestDatabase::create("provider_account_xai_usage_limit").await else {
return;
};
let repository = PgProviderAccountRepository::new(database.pool.clone());
let account_id = ProviderAccountId::new("acct_xai_usage_limit").expect("account ID");
let mut imported = account(account_id.as_str(), "user-xai-usage-limit");
imported.provider_kind = "xai".to_owned();
repository
.insert_provider_account(imported)
.await
.expect("insert xAI account");
let observed_at = SystemTime::now();
let cooldown_until = observed_at
.checked_add(Duration::from_secs(24 * 60 * 60))
.expect("cooldown deadline");
repository
.apply_state_change(AccountStateChange {
account_id: account_id.clone(),
expected_revision: CredentialRevision::new(1).expect("credential revision"),
availability: AccountAvailability::Cooldown,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: Some(cooldown_until),
observed_at,
})
.await
.expect("persist xAI resettable usage limit");
let current: (String, String, Option<String>, bool) = sqlx::query_as(
"select provider_kind, availability, availability_reason, cooldown_until is not null
from provider_accounts where id = $1",
)
.bind(account_id.as_str())
.fetch_one(&database.pool)
.await
.expect("load persisted xAI usage limit state");
assert_eq!(
current,
(
"xai".to_owned(),
"cooldown".to_owned(),
Some("usage_limit_exhausted".to_owned()),
true,
),
);
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 +1532,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);
+85 -8
View File
@@ -25,8 +25,9 @@ use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::{ProviderAdmin, ProviderAdminError, ProviderAdminErrorKind};
use gateway_core::accounting::Money;
use gateway_core::engine::credential::{
CredentialRevision, LoadedCredential, NewProviderAccount, OpaqueProviderData,
PlaintextCredential, ProviderAccount, ProviderAccountId, ProviderAccountStore,
AccountAvailability, CredentialRevision, LoadedCredential, NewProviderAccount,
OpaqueProviderData, PlaintextCredential, ProviderAccount, ProviderAccountId,
ProviderAccountStore,
};
use gateway_core::error::StoreErrorKind;
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
@@ -458,21 +459,23 @@ impl ProviderAdmin for OpenAiAdminProvider {
refresh,
rolling_usage: _,
} = request;
self.account(&account_id).await?;
let mut account = self.account(&account_id).await?;
let snapshot = if refresh {
Some(
let snapshot = Some(
self.quota
.refresh_account(&account_id)
.await
.map_err(map_quota_error)?,
)
);
account = self.account(&account_id).await?;
snapshot
} else {
self.quota
.read_account(&account_id)
.await
.map_err(map_quota_error)?
};
Ok(snapshot.map_or_else(empty_quota, project_quota))
Ok(project_quota(snapshot, &account))
}
async fn models(
@@ -586,7 +589,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 +609,7 @@ fn prepared_rotation(
account_id,
provider_kind,
expected_credential_revision,
replacement_identity,
name: profile.name,
email: profile.email,
plan_type: profile.plan_type,
@@ -745,7 +749,18 @@ fn empty_quota() -> ProviderQuota {
}
}
fn project_quota(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota {
fn project_quota(
snapshot: Option<CodexAccountQuotaSnapshot>,
account: &ProviderAccount,
) -> ProviderQuota {
let mut quota = snapshot.map_or_else(empty_quota, project_quota_snapshot);
if account.availability() == AccountAvailability::QuotaExhausted {
force_confirmed_exhaustion_projection(&mut quota, account.cooldown_until());
}
quota
}
fn project_quota_snapshot(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota {
let mut provider_data = Map::new();
provider_data.insert(
"remaining_percent".to_owned(),
@@ -798,6 +813,68 @@ fn project_quota(snapshot: CodexAccountQuotaSnapshot) -> ProviderQuota {
}
}
fn force_confirmed_exhaustion_projection(
quota: &mut ProviderQuota,
cooldown_until: Option<SystemTime>,
) {
// 真实 429 对“当前已耗尽”的证明强于可能延迟结算的 usage 文档。这里只覆盖
// Admin 展示投影,原始 Provider JSON 仍作为后续判断窗口是否真正恢复的基线。
quota.provider_data = Some(ProviderDocument::new(OpaqueProviderData::new(
Map::from_iter([
(
"remaining_percent".to_owned(),
Value::Number(Number::from(0)),
),
("exhausted".to_owned(), Value::Bool(true)),
]),
)));
let reset_at = cooldown_until.map(DateTime::<Utc>::from);
if quota.windows.is_empty() {
quota.windows.push(ProviderQuotaWindow {
key: "confirmed-quota-exhaustion".to_owned(),
group: "other".to_owned(),
label: "额度".to_owned(),
source: None,
window_seconds: None,
used_percent: Some(100.0),
reset_at,
local_usage: None,
provider_data: None,
});
return;
}
let matched_reset = reset_at.and_then(|reset_at| {
quota
.windows
.iter()
.enumerate()
.filter_map(|(index, window)| {
let distance = window.reset_at.map(|window_reset| {
(window_reset.timestamp() - reset_at.timestamp()).unsigned_abs()
})?;
Some((index, distance))
})
.min_by_key(|(_, distance)| *distance)
.map(|(index, _)| index)
});
let index = matched_reset
.or_else(|| {
quota
.windows
.iter()
.enumerate()
.max_by(|(_, left), (_, right)| {
left.used_percent
.unwrap_or(-1.0)
.total_cmp(&right.used_percent.unwrap_or(-1.0))
})
.map(|(index, _)| index)
})
.unwrap_or(0);
quota.windows[index].used_percent = Some(100.0);
}
const fn quota_group(kind: CodexQuotaWindowKind) -> &'static str {
match kind {
CodexQuotaWindowKind::Monthly => "monthly",
@@ -217,7 +217,7 @@ pub struct CodexQuotaSettings {
impl Default for CodexQuotaSettings {
fn default() -> Self {
Self {
refresh_interval_minutes: 5,
refresh_interval_minutes: 15,
skip_exhausted: true,
}
}
@@ -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,
})
}
@@ -16,9 +16,16 @@ pub(crate) fn derive_codex_session_affinity_key(
.map(str::trim)
.filter(|value| !value.is_empty())
{
return ProviderSessionAffinityKey::try_new(conversation_id.to_owned()).ok();
if let Ok(key) = ProviderSessionAffinityKey::try_new(conversation_id.to_owned()) {
return Some(key);
}
return opaque_affinity_key("local-conversation", conversation_id);
}
let (domain, value) = derive_conversation_anchor(request)?;
opaque_affinity_key(domain, &value)
}
fn opaque_affinity_key(domain: &str, value: &str) -> Option<ProviderSessionAffinityKey> {
let mut hasher = Sha256::new();
hasher.update(b"codex-session-affinity-v1\0");
hasher.update(domain.as_bytes());
@@ -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,14 +30,12 @@ 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);
// ready 账号主要依赖真实请求的被动观察;这里只提供低频、定额的状态兜底。
const QUOTA_READY_REFRESH_MAX_AGE: Duration = Duration::from_secs(2 * 60 * 60);
const QUOTA_READY_REFRESH_BATCH_SIZE: usize = 8;
/// OpenAI Provider 主动额度刷新的调度策略。
///
/// 账号筛选仍由 core 的可调度状态统一处理;该策略只控制 Provider 自己的
/// quota worker 频率,避免把 OpenAI 额度语义泄漏到公共调度层。
/// 正常账号依赖请求响应的被动额度同步;周期 worker 仅复核已耗尽账号。
/// 该策略只控制 Provider 自己的 quota worker 频率,避免把 OpenAI
/// 额度语义泄漏到公共调度层。
#[derive(Debug, Clone, Copy)]
pub struct CodexQuotaRefreshPolicy {
interval: Duration,
@@ -196,6 +194,12 @@ pub struct CodexQuotaSyncSummary {
pub stale: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum QuotaRecoveryEvidence {
SuccessfulResponse,
RefreshedSnapshot,
}
impl CodexQuotaSyncSummary {
#[must_use]
pub const fn has_operational_failures(self) -> bool {
@@ -248,8 +252,6 @@ struct CodexQuotaProjectionState {
next_version: u64,
entries: BTreeMap<ProviderAccountId, CodexQuotaSchedulingEntry>,
last_periodic_refresh_at: BTreeMap<ProviderAccountId, Instant>,
last_exhausted_reset_recheck: BTreeMap<ProviderAccountId, CodexQuotaResetMarker>,
ready_refresh_cursor: Option<ProviderAccountId>,
}
#[derive(Debug, Clone, Copy)]
@@ -260,24 +262,6 @@ struct CodexQuotaSchedulingEntry {
signals: Option<AccountQuotaSignals>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct CodexQuotaResetMarker {
revision: CredentialRevision,
reset_at: SystemTime,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CodexQuotaRefreshKind {
Exhausted { reset_at: Option<SystemTime> },
Ready,
}
#[derive(Clone)]
struct CodexQuotaRefreshCandidate {
account: ProviderAccount,
kind: CodexQuotaRefreshKind,
}
#[derive(Clone)]
struct CodexQuotaHydrationTarget {
account: ProviderAccount,
@@ -307,7 +291,7 @@ enum CodexQuotaFetchError {
enum CodexQuotaStateTransition {
Banned { reason: &'static str },
Exhausted,
Cooldown { until: Option<SystemTime> },
ResettableLimit { until: Option<SystemTime> },
}
impl CodexQuotaStateTransition {
@@ -315,7 +299,7 @@ impl CodexQuotaStateTransition {
match self {
Self::Banned { .. } => AccountAvailability::Banned,
Self::Exhausted => AccountAvailability::QuotaExhausted,
Self::Cooldown { .. } => AccountAvailability::Cooldown,
Self::ResettableLimit { .. } => AccountAvailability::QuotaExhausted,
}
}
@@ -323,13 +307,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,
}
}
@@ -460,16 +444,15 @@ impl CodexQuotaSchedulingProjection {
fn reserve_periodic_refreshes(
&self,
accounts: Vec<ProviderAccount>,
observations: Option<&BTreeMap<ProviderAccountId, QuotaObservation>>,
now: SystemTime,
) -> Vec<ProviderAccount> {
let candidates = accounts
.into_iter()
.filter_map(|account| quota_refresh_candidate(account, observations, now))
.filter_map(|account| quota_refresh_candidate(account, now))
.collect::<Vec<_>>();
let candidate_ids = candidates
.iter()
.map(|candidate| candidate.account.id().clone())
.map(|account| account.id().clone())
.collect::<BTreeSet<_>>();
let refreshed_at = Instant::now();
let mut state = self
@@ -479,93 +462,27 @@ impl CodexQuotaSchedulingProjection {
state
.last_periodic_refresh_at
.retain(|account_id, _| candidate_ids.contains(account_id));
state
.last_exhausted_reset_recheck
.retain(|account_id, _| candidate_ids.contains(account_id));
// 正常账号只由真实请求的响应头和 `codex.rate_limits` 被动同步。定时器只兜底
// 已耗尽账号,且不依据 reset 时间跳过复核;每个账号仍至少间隔 30 分钟。
let mut reserved = Vec::new();
let mut ready = Vec::new();
for candidate in candidates {
match candidate.kind {
CodexQuotaRefreshKind::Exhausted { reset_at } => {
if reserve_exhausted_quota_refresh(
&mut state,
&candidate.account,
reset_at,
now,
refreshed_at,
) {
reserved.push(candidate.account);
}
}
CodexQuotaRefreshKind::Ready => ready.push(candidate.account),
for account in candidates {
if !periodic_quota_refresh_due(&state, account.id(), refreshed_at) {
continue;
}
}
ready.sort_by(|left, right| left.id().cmp(right.id()));
let ready = ready
.into_iter()
.filter(|account| periodic_quota_refresh_due(&state, account.id(), refreshed_at))
.collect::<Vec<_>>();
let ready_len = ready.len();
if ready_len == 0 {
return reserved;
}
let start = state
.ready_refresh_cursor
.as_ref()
.and_then(|cursor| ready.iter().position(|account| account.id() > cursor))
.unwrap_or(0);
let ready = ready
.into_iter()
.cycle()
.skip(start)
.take(ready_len.min(QUOTA_READY_REFRESH_BATCH_SIZE))
.collect::<Vec<_>>();
if let Some(last) = ready.last() {
state.ready_refresh_cursor = Some(last.id().clone());
}
for account in &ready {
state
.last_periodic_refresh_at
.insert(account.id().clone(), refreshed_at);
reserved.push(account);
}
reserved.extend(ready);
reserved
}
}
fn quota_refresh_candidate(
account: ProviderAccount,
observations: Option<&BTreeMap<ProviderAccountId, QuotaObservation>>,
now: SystemTime,
) -> Option<CodexQuotaRefreshCandidate> {
if !eligible_periodic_quota_refresh(&account, now) {
return None;
}
let observation = observations.and_then(|observations| observations.get(account.id()));
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()),
},
AccountAvailability::Cooldown
if snapshot
.as_ref()
.is_some_and(|snapshot| snapshot.fact().exhausted()) =>
{
CodexQuotaRefreshKind::Exhausted {
reset_at: quota_reset_at_from_snapshot(snapshot.as_ref()),
}
}
AccountAvailability::Ready
if observations.is_some() && quota_observation_is_stale(&account, observation, now) =>
{
CodexQuotaRefreshKind::Ready
}
_ => return None,
};
Some(CodexQuotaRefreshCandidate { account, kind })
fn quota_refresh_candidate(account: ProviderAccount, now: SystemTime) -> Option<ProviderAccount> {
(eligible_periodic_quota_refresh(&account, now)
&& account.availability() == AccountAvailability::QuotaExhausted)
.then_some(account)
}
fn quota_snapshot_for_account(
@@ -594,57 +511,6 @@ fn quota_observation_confirms_active_exhaustion(
.is_some_and(|snapshot| quota_fact_is_active_exhaustion(snapshot.fact(), now))
}
fn quota_observation_is_stale(
account: &ProviderAccount,
observation: Option<&QuotaObservation>,
now: SystemTime,
) -> bool {
let Some(observation) = observation.filter(|observation| {
observation.expected_revision == account.revision()
&& quota_snapshot_from_observation(observation).is_some()
}) else {
return true;
};
observation
.observed_at
.and_then(|observed_at| now.duration_since(observed_at).ok())
.is_none_or(|age| age >= QUOTA_READY_REFRESH_MAX_AGE)
}
fn reserve_exhausted_quota_refresh(
state: &mut CodexQuotaProjectionState,
account: &ProviderAccount,
reset_at: Option<SystemTime>,
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) {
state.last_exhausted_reset_recheck.insert(
account.id().clone(),
CodexQuotaResetMarker {
revision: account.revision(),
reset_at,
},
);
}
true
}
fn periodic_quota_refresh_due(
state: &CodexQuotaProjectionState,
account_id: &ProviderAccountId,
@@ -756,16 +622,14 @@ impl CodexCredentialQuotaService {
.collect::<BTreeMap<_, _>>(),
),
Err(error) => {
// 持久观察不可用时不盲扫 ready;保留已耗尽账号的既有恢复轮询。
tracing::warn!(error = %error, "OpenAI quota refresh skipped ready-account scan");
// 持久观察不可用时跳过旧状态修正;仍复核当前已标记耗尽的账号。
tracing::warn!(error = %error, "OpenAI quota refresh skipped persisted exhaustion reconciliation");
None
}
};
self.reconcile_persisted_exhaustion(&mut accounts, observations.as_ref(), now)
.await;
let accounts =
self.scheduling
.reserve_periodic_refreshes(accounts, observations.as_ref(), now);
let accounts = self.scheduling.reserve_periodic_refreshes(accounts, now);
if accounts.is_empty() {
return Ok(summary);
}
@@ -802,16 +666,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 +716,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 +745,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,
);
}
}
@@ -894,7 +761,13 @@ impl CodexCredentialQuotaService {
observed_at: SystemTime,
summary: &mut CodexQuotaSyncSummary,
) -> Result<(), CodexCredentialQuotaError> {
let fact = parse_codex_quota_usage(value)?;
let snapshot = parse_account_quota_snapshot(
account.id().clone(),
account.revision(),
observed_at,
value,
)?;
let fact = snapshot.fact();
let object = value
.as_object()
.cloned()
@@ -927,11 +800,12 @@ 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(),
QuotaRecoveryEvidence::RefreshedSnapshot,
) {
let _ = self
.repository
@@ -939,7 +813,7 @@ impl CodexCredentialQuotaService {
&current,
availability,
fact.exhausted().then_some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await;
@@ -993,18 +867,19 @@ 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,
QuotaRecoveryEvidence::SuccessfulResponse,
) {
self.repository
.apply_state(
&current,
availability,
fact.exhausted().then_some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await?;
@@ -1026,19 +901,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,
@@ -1086,6 +948,21 @@ impl CodexCredentialQuotaService {
pub async fn refresh_account(
&self,
account_id: &ProviderAccountId,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
self.refresh_account_snapshot(account_id).await
}
/// 真实限额失败后的异步刷新只补齐展示快照,不允许 usage 快照撤销已确认的失败状态。
pub(crate) async fn refresh_account_after_failure(
&self,
account_id: &ProviderAccountId,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
self.refresh_account_snapshot(account_id).await
}
async fn refresh_account_snapshot(
&self,
account_id: &ProviderAccountId,
) -> Result<CodexAccountQuotaSnapshot, CodexCredentialQuotaError> {
let account = self
.store
@@ -1164,11 +1041,12 @@ 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(),
QuotaRecoveryEvidence::RefreshedSnapshot,
) {
self.repository
.apply_state(
@@ -1178,7 +1056,7 @@ impl CodexCredentialQuotaService {
.fact()
.exhausted()
.then_some("quota_exhausted".to_owned()),
None,
cooldown_until,
observed_at,
)
.await?;
@@ -1386,30 +1264,58 @@ 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 {
recovery_evidence: QuotaRecoveryEvidence,
) -> Option<(AccountAvailability, Option<SystemTime>)> {
if fact.exhausted() {
let observed_cooldown = match (fact.resets_at().map(SystemTime::from), cooldown_until) {
(Some(observed), Some(confirmed)) => Some(observed.max(confirmed)),
(observed, confirmed) => observed.or(confirmed),
};
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 => Some(AccountAvailability::Ready),
AccountAvailability::QuotaExhausted
if !quota_recovery_confirmed(recovery_evidence, cooldown_until, now) =>
{
None
}
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_recovery_confirmed(
evidence: QuotaRecoveryEvidence,
cooldown_until: Option<SystemTime>,
now: SystemTime,
) -> bool {
if cooldown_until.is_some_and(|cooldown_until| cooldown_until <= now) {
return true;
}
match evidence {
QuotaRecoveryEvidence::SuccessfulResponse => true,
// used_percent 是可滞后、会取整的观测值;确认 429 后的 100% 也是
// 管理端投影。它们的回落不能覆盖仍未到期的上游 reset 时间。
QuotaRecoveryEvidence::RefreshedSnapshot => false,
}
}
@@ -1438,7 +1344,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 +1782,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(())
}
+172 -173
View File
@@ -64,10 +64,10 @@ use crate::transport::protocol::responses::{
CodexResponsesRequest, PreviousResponseScope, ResponseEventSignals,
};
use crate::transport::request::{
CodexRequestEncodeError, encode_generate_request, sanitize_cross_account_item,
scope_request_to_account,
CodexRequestEncodeError, encode_generate_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,
@@ -83,6 +83,12 @@ const MAX_COOKIE_HEADER_BYTES: usize = 16 * 1024;
/// 提交边界前最多保留 64 KiB 原始上游 chunk;达到阈值后结束无感换号窗口,
/// 但不会把上游数据改写成协议失败。
const MAX_STREAM_PREFETCH_BYTES: usize = 64 * 1024;
/// 短暂保留 response.created 等结构事件,让随后到达的明确拒绝可以无感换号;
/// 到期即放行,避免模型长时间思考时让客户端一直收不到首事件。
const STREAM_REPLAY_GRACE: Duration = Duration::from_millis(1_200);
// 额度拒绝后先给上游额度结算留出时间,再以受限时长同步 usage 快照。
const QUOTA_FAILURE_REFRESH_DELAY: Duration = Duration::from_secs(2);
const QUOTA_FAILURE_REFRESH_TIMEOUT: Duration = Duration::from_secs(5);
pub const OFFICIAL_CODEX_BASE_PATH: &str = "/backend-api";
pub const OFFICIAL_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api";
@@ -170,17 +176,7 @@ impl Provider for CodexProvider {
return ProviderRequestObservation::default();
};
let (semantics, reasoning_effort) = encode_generate_request(request, "observability")
.map(|mut encoded| {
if let Some(previous) = decode_openai_session_state(request) {
let mut input = previous
.transcript
.iter()
.map(OpenAiReplayItem::value)
.cloned()
.collect::<Vec<_>>();
input.extend(encoded.input().iter().cloned());
encoded.set_input(input);
}
.map(|encoded| {
let semantics = encoded.semantics();
let reasoning_effort = semantics.reasoning_effort.clone();
(semantics, reasoning_effort)
@@ -260,7 +256,6 @@ impl Provider for CodexProvider {
let session_affinity_key = derive_codex_session_affinity_key(&upstream_request);
let cyber_policy_session_key =
derive_codex_cyber_policy_session_key(&upstream_request, context.client_api_key_ref());
let request_input = upstream_request.input().to_vec();
let transport = selected_transport(&upstream_request);
apply_transport(&mut upstream_request, transport);
@@ -292,11 +287,8 @@ impl Provider for CodexProvider {
) || previous_session.as_ref().is_some_and(|state| {
state.continuation_scope == OpenAiContinuationScope::ReplayRequired
});
if replay_previous_response && let Some(state) = previous_session.as_ref() {
let mut input = replay_input_for_account(state, lease.account_id().as_str());
input.reserve(request_input.len());
input.extend(request_input.iter().cloned());
upstream_request.set_input(input);
if replay_previous_response && previous_session.is_some() {
// Codex 客户端 input 是完整历史;恢复时只解除上游续链绑定,避免重复叠加代理副本。
upstream_request.set_previous_response_id(None);
upstream_request.previous_response_scope = None;
upstream_request.turn_state = None;
@@ -362,8 +354,6 @@ impl Provider for CodexProvider {
let response_store = upstream_request.store();
let session_capture =
(!continuation_requested || previous_session.is_some()).then(|| OpenAiSessionCapture {
previous: previous_session,
request_input,
account_id: lease.account_id().as_str().to_owned(),
conversation_id: upstream_request.local_conversation_id.clone(),
turn_state: upstream_request.turn_state.clone(),
@@ -418,7 +408,6 @@ struct OpenAiSessionState {
#[serde(default)]
turn_state: Option<String>,
continuation_scope: OpenAiContinuationScope,
transcript: Vec<OpenAiReplayItem>,
}
#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -429,27 +418,7 @@ enum OpenAiContinuationScope {
ReplayRequired,
}
#[derive(Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
enum OpenAiReplayItem {
ClientInput(Value),
SanitizedOutput(Value),
AccountOutput { account_id: String, item: Value },
}
impl OpenAiReplayItem {
fn value(&self) -> &Value {
match self {
Self::ClientInput(value)
| Self::SanitizedOutput(value)
| Self::AccountOutput { item: value, .. } => value,
}
}
}
struct OpenAiSessionCapture {
previous: Option<OpenAiSessionState>,
request_input: Vec<Value>,
account_id: String,
conversation_id: Option<String>,
turn_state: Option<String>,
@@ -482,33 +451,15 @@ fn attach_openai_session_update(
events: &mut [ProviderEvent],
capture: &mut Option<OpenAiSessionCapture>,
) {
let Some((terminal_index, output)) = events.iter().enumerate().find_map(|(index, event)| {
terminal_response_output(event).map(|output| (index, output.to_vec()))
}) else {
let Some(terminal_index) = events
.iter()
.position(|event| terminal_response_output(event).is_some())
else {
return;
};
let Some(capture) = capture.take() else {
return;
};
let mut transcript = capture
.previous
.map(|state| state.transcript)
.unwrap_or_default();
project_transcript_to_account(&mut transcript, &capture.account_id);
transcript.extend(
capture
.request_input
.into_iter()
.map(OpenAiReplayItem::ClientInput),
);
transcript.extend(
output
.into_iter()
.map(|item| OpenAiReplayItem::AccountOutput {
account_id: capture.account_id.clone(),
item,
}),
);
let Some(continuation_scope) = capture.continuation_scope else {
return;
};
@@ -517,7 +468,6 @@ fn attach_openai_session_update(
conversation_id: capture.conversation_id,
turn_state: capture.turn_state,
continuation_scope,
transcript,
}) else {
return;
};
@@ -545,47 +495,6 @@ fn terminal_response_output(event: &ProviderEvent) -> Option<&[Value]> {
.flatten()
}
fn project_transcript_to_account(transcript: &mut Vec<OpenAiReplayItem>, account_id: &str) {
*transcript = transcript
.drain(..)
.filter_map(|item| match item {
OpenAiReplayItem::AccountOutput {
account_id: owner,
item,
} if owner != account_id => {
sanitize_cross_account_item(item).map(OpenAiReplayItem::SanitizedOutput)
}
item => Some(item),
})
.collect();
}
fn replay_input_for_account(state: &OpenAiSessionState, account_id: &str) -> Vec<Value> {
state
.transcript
.iter()
.filter_map(|item| match item {
OpenAiReplayItem::ClientInput(value) | OpenAiReplayItem::SanitizedOutput(value) => {
Some(value.clone())
}
OpenAiReplayItem::AccountOutput {
account_id: owner,
item,
} if owner == account_id => Some(without_output_id(item.clone())),
OpenAiReplayItem::AccountOutput { item, .. } => {
sanitize_cross_account_item(item.clone())
}
})
.collect()
}
fn without_output_id(mut item: Value) -> Value {
if let Value::Object(object) = &mut item {
object.remove("id");
}
item
}
enum CodexHandshakeAttemptError {
Client(CodexClientError),
Cancelled,
@@ -678,8 +587,6 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
lease.escape_reason(),
lease.account_switch(),
);
let stream_commit_policy =
StreamCommitPolicy::for_continuation(context.continuation_attempt());
let mut response = create_response_attempt(
&client,
&request,
@@ -822,14 +729,16 @@ 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 pre_commit_events = PreCommitClientEvents::new(stream_commit_policy);
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();
loop {
let Some(stream_deadline) = remaining(context.deadline()) else {
Err(provider_error(ProviderErrorKind::Timeout, UpstreamSendState::Sent))?;
return;
};
let replay_grace_deadline = pre_commit_events.replay_grace_deadline();
let next = tokio::select! {
biased;
_ = cancellation.cancelled() => Err(MappedProviderFailure::plain(provider_error(
@@ -840,14 +749,21 @@ fn cold_response_stream(response: ColdResponse) -> EventStream {
ProviderErrorKind::Timeout,
UpstreamSendState::Sent,
))),
_ = wait_for_replay_grace(replay_grace_deadline) => Ok(PreCommitPoll::GraceElapsed),
chunk = body.next() => match chunk {
Some(Ok(chunk)) => Ok(Some(chunk)),
Some(Ok(chunk)) => Ok(PreCommitPoll::Upstream(Some(chunk))),
Some(Err(error)) => Err(map_stream_error(error)),
None => Ok(None),
None => Ok(PreCommitPoll::Upstream(None)),
},
};
let next = match next {
Ok(next) => next,
Ok(PreCommitPoll::Upstream(next)) => next,
Ok(PreCommitPoll::GraceElapsed) => {
for event in pre_commit_events.commit_pending() {
yield event;
}
continue;
}
Err(failure) => {
let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await;
if !updates.is_empty() {
@@ -895,7 +811,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 +906,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 +986,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 +1019,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 +1044,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 +1075,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 +1168,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()),
);
}
@@ -1734,7 +1669,10 @@ fn map_request_error(error: CodexRequestEncodeError) -> ProviderError {
fn map_selection_error(error: CredentialSelectionError) -> ProviderError {
match error {
CredentialSelectionError::CapacityUnavailable { retry_after } => {
let error = provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent);
let error = provider_error(
ProviderErrorKind::AccountCapacityUnavailable,
UpstreamSendState::NotSent,
);
match retry_after {
Some(retry) => error.with_retry_after(retry),
None => error,
@@ -1747,9 +1685,10 @@ fn map_selection_error(error: CredentialSelectionError) -> ProviderError {
CredentialSelectionError::InvalidCredential
| CredentialSelectionError::Store
| CredentialSelectionError::Coordinator
| CredentialSelectionError::CookiePolicy => {
provider_error(ProviderErrorKind::Unavailable, UpstreamSendState::NotSent)
}
| CredentialSelectionError::CookiePolicy => provider_error(
ProviderErrorKind::ProviderInfrastructureUnavailable,
UpstreamSendState::NotSent,
),
}
}
@@ -1777,54 +1716,37 @@ impl MappedProviderFailure {
}
}
/// 与 2.x 一致的流提交策略:普通请求优先首事件,续写/重放优先可恢复性。
#[derive(Clone, Copy)]
enum StreamCommitPolicy {
FirstForwardableEvent,
UntilOutputOrTerminal,
enum PreCommitPoll<T> {
Upstream(T),
GraceElapsed,
}
impl StreamCommitPolicy {
const fn for_continuation(attempt: ContinuationAttempt) -> Self {
match attempt {
ContinuationAttempt::None => Self::FirstForwardableEvent,
ContinuationAttempt::Native
| ContinuationAttempt::ReplayOwner
| ContinuationAttempt::ReplayAny => Self::UntilOutputOrTerminal,
}
}
fn boundary_reached(
self,
events: &[ProviderEvent],
timing_signals: ResponseEventSignals,
completed: bool,
) -> bool {
match self {
Self::FirstForwardableEvent => events.iter().any(ProviderEvent::has_client_event),
Self::UntilOutputOrTerminal => timing_signals.semantic_output || completed,
}
async fn wait_for_replay_grace(deadline: Option<Instant>) {
if let Some(deadline) = deadline {
tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await;
} else {
std::future::pending::<()>().await;
}
}
/// 提交边界前的上游事件预取。
///
/// 原始 chunk 计数而不是重编码后的 event 大小。64 KiB 只限定无感换号窗口;达到
/// 阈值会提交已缓存 wire,绝不能因为网关私有资源规则伪造上游协议失败。一旦提交,
/// 后续事件不再具备无痕重放资格。
/// 原始 chunk 计数而不是重编码后的 event 大小。时间与 64 KiB 共同限定无感换号
/// 窗口;任一边界到达都会提交已缓存 wire,不能因网关私有资源规则伪造上游协议
/// 失败。一旦提交,后续事件不再具备无痕重放资格。
struct PreCommitClientEvents {
policy: StreamCommitPolicy,
pending: Vec<ProviderEvent>,
prefetched_bytes: usize,
replay_grace_deadline: Option<Instant>,
committed: bool,
}
impl PreCommitClientEvents {
const fn new(policy: StreamCommitPolicy) -> Self {
const fn new() -> Self {
Self {
policy,
pending: Vec::new(),
prefetched_bytes: 0,
replay_grace_deadline: None,
committed: false,
}
}
@@ -1844,13 +1766,17 @@ impl PreCommitClientEvents {
if self.committed {
return incoming;
}
let boundary_reached = self
.policy
.boundary_reached(&incoming, timing_signals, completed);
let starts_replay_grace = incoming.iter().any(ProviderEvent::has_client_event);
self.pending.extend(incoming);
if boundary_reached || self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES {
if timing_signals.semantic_output
|| completed
|| self.prefetched_bytes > MAX_STREAM_PREFETCH_BYTES
{
return self.commit_pending();
}
if starts_replay_grace && self.replay_grace_deadline.is_none() {
self.replay_grace_deadline = Instant::now().checked_add(STREAM_REPLAY_GRACE);
}
Vec::new()
}
@@ -1879,9 +1805,14 @@ impl PreCommitClientEvents {
fn take_for_failure(&mut self, incoming: Vec<ProviderEvent>) -> Vec<ProviderEvent> {
self.pending.extend(incoming);
self.prefetched_bytes = 0;
self.replay_grace_deadline = None;
std::mem::take(&mut self.pending)
}
const fn replay_grace_deadline(&self) -> Option<Instant> {
self.replay_grace_deadline
}
const fn is_committed(&self) -> bool {
self.committed
}
@@ -1889,6 +1820,7 @@ impl PreCommitClientEvents {
fn commit_pending(&mut self) -> Vec<ProviderEvent> {
self.committed = true;
self.prefetched_bytes = 0;
self.replay_grace_deadline = None;
std::mem::take(&mut self.pending)
}
}
@@ -1916,7 +1848,7 @@ impl ReplayBoundary {
struct OpenAiFailureContext<'a> {
client: &'a CodexBackendClient,
selector: &'a CodexCredentialSelector,
quota: &'a CodexCredentialQuotaService,
quota: &'a Arc<CodexCredentialQuotaService>,
response_origin: &'a Url,
cyber_policy_scope: Option<&'a CodexCyberPolicyScope>,
allows_account_state_mutation: bool,
@@ -1931,6 +1863,10 @@ async fn apply_failure(
return;
}
synchronize_passive_quota(context.quota, account, &failure.rate_limit_headers).await;
let needs_authoritative_quota_refresh = matches!(
failure.account_failure,
Some(CodexAccountFailure::QuotaExhausted | CodexAccountFailure::UsageLimitExhausted { .. })
);
if failure.cyber_policy_failure {
context
.selector
@@ -1954,6 +1890,9 @@ async fn apply_failure(
);
}
}
if needs_authoritative_quota_refresh {
schedule_authoritative_quota_refresh_after_failure(context.quota, account);
}
if failure.capture_response_cookies
&& !failure.set_cookie_headers.is_empty()
&& let Err(error) = context
@@ -1973,6 +1912,42 @@ async fn apply_failure(
}
}
fn schedule_authoritative_quota_refresh_after_failure(
quota: &Arc<CodexCredentialQuotaService>,
account: &ProviderAccount,
) {
// 限额错误可以确认账号状态,但 Responses 事件中的 used_percent 可能仍停在上一结算点。
// usage 快照在后台补齐展示基线,不得撤销同一轮真实失败,也不能阻塞原始响应。
// 先等待 2 秒让上游结算,再查询;5 秒仅限制实际查询自身。
let quota = Arc::clone(quota);
let account_id = account.id().clone();
drop(tokio::spawn(async move {
tokio::time::sleep(QUOTA_FAILURE_REFRESH_DELAY).await;
match tokio::time::timeout(
QUOTA_FAILURE_REFRESH_TIMEOUT,
quota.refresh_account_after_failure(&account_id),
)
.await
{
Ok(Ok(_)) => {}
Ok(Err(error)) => {
tracing::warn!(
account_id = %account_id,
error = %error,
"OpenAI authoritative quota synchronization after quota rejection failed"
);
}
Err(_) => {
tracing::warn!(
account_id = %account_id,
timeout_ms = QUOTA_FAILURE_REFRESH_TIMEOUT.as_millis(),
"OpenAI authoritative quota synchronization after quota rejection timed out"
);
}
}
}));
}
fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure {
map_client_error(error, UpstreamSendState::Ambiguous, true)
}
@@ -1994,7 +1969,7 @@ fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool
CodexWebSocketExchangeError::Transport(_)
| CodexWebSocketExchangeError::PostSendAmbiguous { .. }
| CodexWebSocketExchangeError::SendTimeout { .. }
| CodexWebSocketExchangeError::ClosedBeforeTerminal
| CodexWebSocketExchangeError::ClosedBeforeTerminal(_)
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
| CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. }
| CodexWebSocketExchangeError::InitialEventTimeout { .. }
@@ -2082,10 +2057,19 @@ fn map_client_error(
send_state,
))
}
CodexClientError::WebSocket(error) => MappedProviderFailure::plain(provider_error(
websocket_error_kind(&error),
websocket_send_state(&error),
)),
CodexClientError::WebSocket(error) => {
let client_visible_error = websocket_client_visible_error(&error);
let mut failure = MappedProviderFailure::plain(provider_error(
websocket_error_kind(&error),
websocket_send_state(&error),
));
if let Some(client_visible_error) = client_visible_error {
failure.error = failure
.error
.with_client_visible_upstream_error(client_visible_error);
}
failure
}
};
if let Some(continuation_failure) = continuation_failure {
failure.error = failure
@@ -2096,6 +2080,21 @@ fn map_client_error(
failure
}
fn websocket_client_visible_error(
error: &CodexWebSocketExchangeError,
) -> Option<ClientVisibleUpstreamError> {
let close = error.close_before_terminal()?;
let message = close
.reason()
.filter(|reason| !reason.is_empty())
.map_or_else(|| close.to_string(), str::to_owned);
Some(ClientVisibleUpstreamError::new(
message,
close.code().map(|code| code.to_string()),
Some("websocket_close_error".to_owned()),
))
}
fn map_upstream_failure(
mut failure: CodexUpstreamFailure,
observation: Option<ProviderResponseObservation>,
@@ -2264,7 +2263,7 @@ const fn websocket_send_state(error: &CodexWebSocketExchangeError) -> UpstreamSe
CodexWebSocketExchangeError::Transport(_)
| CodexWebSocketExchangeError::PostSendAmbiguous { .. }
| CodexWebSocketExchangeError::SendTimeout { .. }
| CodexWebSocketExchangeError::ClosedBeforeTerminal
| CodexWebSocketExchangeError::ClosedBeforeTerminal(_)
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
| CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. }
| CodexWebSocketExchangeError::InitialEventTimeout { .. } => UpstreamSendState::Ambiguous,
@@ -2291,7 +2290,7 @@ const fn websocket_error_kind(error: &CodexWebSocketExchangeError) -> ProviderEr
CodexWebSocketExchangeError::Transport(_)
| CodexWebSocketExchangeError::Connect(_)
| CodexWebSocketExchangeError::PostSendAmbiguous { .. }
| CodexWebSocketExchangeError::ClosedBeforeTerminal
| CodexWebSocketExchangeError::ClosedBeforeTerminal(_)
| CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => {
ProviderErrorKind::Transport
}
@@ -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))
@@ -477,6 +477,7 @@ pub enum CodexTransportDecision {
Http2BreakerOpen,
Http2PoolUnavailable,
Http2PreSendFailure,
Http2PreDeliveryFailure,
}
impl CodexTransportDecision {
@@ -491,6 +492,7 @@ impl CodexTransportDecision {
Self::Http2BreakerOpen => "http2_breaker_open",
Self::Http2PoolUnavailable => "http2_pool_unavailable",
Self::Http2PreSendFailure => "http2_ws_pre_send_failure",
Self::Http2PreDeliveryFailure => "http2_ws_pre_delivery_failure",
}
}
}
@@ -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,
@@ -31,10 +32,11 @@ use crate::transport::{
},
response_meta,
websocket::{
CodexWebSocketConnection, CodexWebSocketPool, CodexWebSocketPoolKey,
DEFAULT_INITIAL_EVENT_TIMEOUT, WEBSOCKET_FAST_PATH_BUDGET, WebSocketOriginBreaker,
WebSocketPoolDecision, execute_prepared_response_create_request_stream,
post_send_ambiguous, prepare_response_create_request_with_pool, websocket_audit_dir,
CodexWebSocketConnection, CodexWebSocketExchangeError, CodexWebSocketPool,
CodexWebSocketPoolKey, CodexWebSocketStreamingExchange, DEFAULT_INITIAL_EVENT_TIMEOUT,
WEBSOCKET_FAST_PATH_BUDGET, WebSocketOriginBreaker, WebSocketPoolDecision,
execute_prepared_response_create_request_stream, post_send_ambiguous,
prepare_response_create_request_with_pool, websocket_audit_dir,
write_websocket_audit_artifact_from_env,
},
};
@@ -130,13 +132,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,
@@ -306,7 +309,7 @@ impl CodexBackendClient {
route,
metrics,
} = prepared;
let result = match route {
match route {
PreparedResponseRoute::Http => self
.create_response_stream_http_sse(request, context)
.await
@@ -315,50 +318,124 @@ impl CodexBackendClient {
response
}),
PreparedResponseRoute::WebSocket(route) => {
let PreparedWebSocketRoute { request, prepared } = *route;
execute_prepared_response_create_request_stream(&request, prepared)
.await
.map_err(websocket_exchange_error_to_client_error)
.map(|exchange| {
tracing::info!(
request_id = %context.request_id,
websocket_connection_id = %exchange.websocket_connection_id,
ws_pool = exchange.pool_decision.map_or("unpooled", WebSocketPoolDecision::kind),
"WebSocket response stream established"
);
CodexBackendStreamingResponse {
body: Box::pin(
exchange
.body
.map_err(post_send_ambiguous)
.map_err(websocket_exchange_error_to_client_error),
),
transport: CodexBackendTransport::WebSocket,
turn_state: exchange.turn_state,
set_cookie_headers: exchange.set_cookie_headers,
rate_limit_headers: exchange.rate_limit_headers,
rate_limit_header_updates: Some(exchange.rate_limit_header_updates),
turn_state_update: Some(exchange.turn_state_update),
websocket_pool_decision: exchange.pool_decision,
diagnostics: exchange.diagnostics,
response_metadata: exchange.response_metadata,
transport_metrics: metrics,
connection_local_continuation: exchange.connection_local_continuation,
}
})
let PreparedWebSocketRoute {
request: websocket_request,
prepared,
} = *route;
let delivery_wait_started_at = Instant::now();
let mut exchange = match execute_prepared_response_create_request_stream(
&websocket_request,
prepared,
)
.await
{
Ok(exchange) => exchange,
Err(error)
if requirement.allows_pre_delivery_http_fallback()
&& error.allows_pre_delivery_http_fallback() =>
{
return self
.fallback_to_http_before_websocket_delivery(
request,
context,
requirement,
metrics,
delivery_wait_started_at,
error,
)
.await;
}
Err(error) => return Err(websocket_exchange_error_to_client_error(error)),
};
if requirement.allows_pre_delivery_http_fallback()
&& let Err(error) = await_websocket_delivery_boundary(&mut exchange).await
{
if error.allows_pre_delivery_http_fallback() {
return self
.fallback_to_http_before_websocket_delivery(
request,
context,
requirement,
metrics,
delivery_wait_started_at,
error,
)
.await;
}
return Err(websocket_exchange_error_to_client_error(
post_send_ambiguous(error),
));
}
tracing::info!(
request_id = %context.request_id,
websocket_connection_id = %exchange.websocket_connection_id,
ws_pool = exchange.pool_decision.map_or("unpooled", WebSocketPoolDecision::kind),
"WebSocket response stream established"
);
Ok(CodexBackendStreamingResponse {
body: Box::pin(
exchange
.body
.map_err(post_send_ambiguous)
.map_err(websocket_exchange_error_to_client_error),
),
transport: CodexBackendTransport::WebSocket,
turn_state: exchange.turn_state,
set_cookie_headers: exchange.set_cookie_headers,
rate_limit_headers: exchange.rate_limit_headers,
rate_limit_header_updates: Some(exchange.rate_limit_header_updates),
turn_state_update: Some(exchange.turn_state_update),
websocket_pool_decision: exchange.pool_decision,
diagnostics: exchange.diagnostics,
response_metadata: exchange.response_metadata,
transport_metrics: metrics,
connection_local_continuation: exchange.connection_local_continuation,
})
}
};
result.inspect_err(|error| {
}
.inspect_err(|error| {
tracing::warn!(
request_id = %context.request_id,
transport_requirement = requirement.as_str(),
failure_phase = "post_send_or_explicit_response",
error = %error,
"Responses stream transport failed after preparation; automatic fallback is disabled"
"Responses stream transport failed after preparation"
);
})
}
async fn fallback_to_http_before_websocket_delivery(
&self,
request: &CodexResponsesRequest,
context: CodexRequestContext<'_>,
requirement: TransportRequirement,
mut metrics: CodexTransportMetrics,
delivery_wait_started_at: Instant,
error: CodexWebSocketExchangeError,
) -> CodexClientResult<CodexBackendStreamingResponse> {
let decision_wait_ms = metrics
.transport_decision_wait_ms
.unwrap_or_default()
.saturating_add(elapsed_duration_millis(delivery_wait_started_at.elapsed()));
metrics.decision = Some(CodexTransportDecision::Http2PreDeliveryFailure);
metrics.transport_decision_wait_ms = Some(decision_wait_ms);
tracing::warn!(
request_id = %context.request_id,
account_id = context.account_id.unwrap_or_default(),
transport_requirement = requirement.as_str(),
transport_decision = CodexTransportDecision::Http2PreDeliveryFailure.as_str(),
transport_decision_wait_ms = decision_wait_ms,
error = %error,
"WebSocket failed before first deliverable event; using same-account HTTP"
);
self.create_response_stream_http_sse(request, context)
.await
.map(|mut response| {
merge_preparation_metrics(&mut response.transport_metrics, metrics);
response
})
}
fn websocket_pool_key(
&self,
request: &CodexResponsesRequest,
@@ -561,6 +638,36 @@ impl CodexBackendClient {
}
}
async fn await_websocket_delivery_boundary(
exchange: &mut CodexWebSocketStreamingExchange,
) -> Result<(), CodexWebSocketExchangeError> {
let mut prelude = Vec::new();
loop {
match exchange.body.next().await {
Some(Ok(frame)) if is_websocket_lifecycle_prelude(&frame) => prelude.push(frame),
Some(Ok(frame)) => {
prelude.push(frame);
let remaining =
std::mem::replace(&mut exchange.body, Box::pin(futures::stream::empty()));
exchange.body =
Box::pin(futures::stream::iter(prelude.into_iter().map(Ok)).chain(remaining));
return Ok(());
}
Some(Err(error)) => return Err(error),
None => {
return Err(CodexWebSocketExchangeError::closed_before_terminal(
None, None,
));
}
}
}
}
fn is_websocket_lifecycle_prelude(frame: &[u8]) -> bool {
frame.starts_with(b"event: response.created\n")
|| frame.starts_with(b"event: response.in_progress\n")
}
async fn read_model_catalog_body(response: ReqwestResponse) -> CodexClientResult<Vec<u8>> {
if response
.content_length()
@@ -594,20 +701,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);
}
}
@@ -139,6 +139,11 @@ impl TransportRequirement {
)
}
/// WebSocket 已发送 payload、但尚未交付业务事件时,是否允许同账号 HTTP 重放。
pub fn allows_pre_delivery_http_fallback(self) -> bool {
matches!(self, Self::NewChain)
}
/// 用于审计与遥测的稳定名称。
pub fn as_str(self) -> &'static str {
match self {
@@ -161,13 +161,14 @@ pub(crate) fn derive_conversation_anchor(
request: &CodexResponsesRequest,
) -> Option<(&'static str, String)> {
request
.prompt_cache_key()
.map(|value| ("prompt-cache", value.to_owned()))
.client_session_id
.as_deref()
.map(|value| ("session", value.to_owned()))
.or_else(|| {
request
.client_session_id
.client_conversation_id
.as_deref()
.map(|value| ("session", value.to_owned()))
.map(|value| ("conversation", value.to_owned()))
})
.or_else(|| {
request
@@ -177,9 +178,8 @@ pub(crate) fn derive_conversation_anchor(
})
.or_else(|| {
request
.client_conversation_id
.as_deref()
.map(|value| ("conversation", value.to_owned()))
.prompt_cache_key()
.map(|value| ("prompt-cache", value.to_owned()))
})
.or_else(|| derive_stable_conversation_key(request).map(|value| ("request", value)))
}
@@ -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)
}
@@ -71,8 +71,8 @@ pub enum CodexWebSocketExchangeError {
reason: PreviousResponseUnavailableReason,
},
/// 上游在 terminal 事件前关闭。
#[error("websocket closed before terminal event")]
ClosedBeforeTerminal,
#[error("{0}")]
ClosedBeforeTerminal(CodexWebSocketCloseError),
/// 上游在指定时间内没有发送任何事件。
#[error("websocket receive idle timeout after {timeout:?}")]
ReceiveIdleTimeout {
@@ -87,6 +87,9 @@ pub enum CodexWebSocketExchangeError {
ReusedConnectionDiedBeforeFirstEvent {
/// 底层失效原因。
message: String,
/// 原始 typed transport failure。
#[source]
source: Option<Box<CodexWebSocketExchangeError>>,
},
/// 建连并发送后,上游在配置时间内没有产生任何事件。
#[error("websocket first upstream event not received within {timeout:?}")]
@@ -96,6 +99,55 @@ pub enum CodexWebSocketExchangeError {
},
}
/// 上游在 terminal 事件前发送的 WebSocket close 信息。
///
/// close reason 只供原请求的客户端错误响应使用;`Debug` 与 `Display` 均不会输出它。
#[derive(Clone, PartialEq, Eq)]
pub struct CodexWebSocketCloseError {
code: Option<u16>,
reason: Option<String>,
}
impl CodexWebSocketCloseError {
pub(crate) fn new(code: Option<u16>, reason: Option<String>) -> Self {
Self { code, reason }
}
/// 返回上游 close code;没有 close frame 时为 `None`。
#[must_use]
pub const fn code(&self) -> Option<u16> {
self.code
}
/// 返回上游 close reason;仅可用于当前请求的客户端协议响应。
#[must_use]
pub fn reason(&self) -> Option<&str> {
self.reason.as_deref()
}
}
impl fmt::Debug for CodexWebSocketCloseError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("CodexWebSocketCloseError")
.field("code", &self.code)
.field("reason", &self.reason.as_ref().map(|_| "<redacted>"))
.finish()
}
}
impl fmt::Display for CodexWebSocketCloseError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.code {
Some(code) => write!(
formatter,
"websocket closed before terminal event (code {code})"
),
None => formatter.write_str("websocket closed before terminal event"),
}
}
}
/// WebSocket 上游错误帧载荷。
#[derive(Clone, PartialEq, Eq)]
pub struct CodexWebSocketUpstreamError {
@@ -140,6 +192,25 @@ impl fmt::Display for CodexWebSocketUpstreamError {
}
impl CodexWebSocketExchangeError {
pub(crate) fn closed_before_terminal(code: Option<u16>, reason: Option<String>) -> Self {
Self::ClosedBeforeTerminal(CodexWebSocketCloseError::new(code, reason))
}
pub(crate) fn close_before_terminal(&self) -> Option<&CodexWebSocketCloseError> {
match self {
Self::ClosedBeforeTerminal(close) => Some(close),
Self::PostSendAmbiguous {
source: Some(source),
..
}
| Self::ReusedConnectionDiedBeforeFirstEvent {
source: Some(source),
..
} => source.close_before_terminal(),
_ => None,
}
}
pub(super) fn upstream(
status_code: u16,
retry_after_seconds: Option<u64>,
@@ -173,4 +244,16 @@ impl CodexWebSocketExchangeError {
| Self::ContinuationUnavailable { .. }
)
}
/// 首个业务事件交付前,普通 WebSocket 失败可切到同账号 HTTP。
pub(in crate::transport) fn allows_pre_delivery_http_fallback(&self) -> bool {
match self {
Self::Upstream(_) => false,
Self::PostSendAmbiguous {
source: Some(source),
..
} => source.allows_pre_delivery_http_fallback(),
_ => true,
}
}
}
@@ -40,19 +40,20 @@ pub(super) fn reused_stream_receive_error(
error: CodexWebSocketExchangeError,
) -> CodexWebSocketExchangeError {
match error {
CodexWebSocketExchangeError::ClosedBeforeTerminal
CodexWebSocketExchangeError::ClosedBeforeTerminal(_)
| CodexWebSocketExchangeError::ReceiveIdleTimeout { .. }
| CodexWebSocketExchangeError::Transport(_) => {
reused_connection_died_before_first_event(&error)
reused_connection_died_before_first_event(error)
}
error => error,
}
}
pub(super) fn reused_connection_died_before_first_event(
error: &CodexWebSocketExchangeError,
error: CodexWebSocketExchangeError,
) -> CodexWebSocketExchangeError {
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent {
message: error.to_string(),
source: Some(Box::new(error)),
}
}
@@ -208,14 +208,17 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
.await;
return;
}
tungstenite::Message::Close(_) => {
tungstenite::Message::Close(frame) => {
discard_stream_websocket(
websocket,
pool_return,
StreamWebSocketDiscardReason::UpstreamClosed,
)
.await;
let error = CodexWebSocketExchangeError::ClosedBeforeTerminal;
let error = CodexWebSocketExchangeError::closed_before_terminal(
frame.as_ref().map(|frame| u16::from(frame.code)),
frame.map(|frame| frame.reason.to_string()),
);
let error = if reused_connection {
reused_stream_receive_error(error)
} else {
@@ -285,13 +288,19 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) {
}
}
let upstream_close = websocket
.exit_reason()
.and_then(|reason| reason.upstream_close().cloned());
discard_stream_websocket(
websocket,
pool_return,
StreamWebSocketDiscardReason::UpstreamClosed,
)
.await;
let error = CodexWebSocketExchangeError::ClosedBeforeTerminal;
let error = upstream_close.map_or_else(
|| CodexWebSocketExchangeError::closed_before_terminal(None, None),
CodexWebSocketExchangeError::ClosedBeforeTerminal,
);
let error = if reused_connection {
reused_stream_receive_error(error)
} else {
@@ -19,7 +19,7 @@ pub use self::{
WebSocketOriginBreaker, WebSocketOriginBreakerConfig, WebSocketOriginBreakerDecision,
WebSocketOriginBreakerPermit,
},
error::{CodexWebSocketExchangeError, CodexWebSocketUpstreamError},
error::{CodexWebSocketCloseError, CodexWebSocketExchangeError, CodexWebSocketUpstreamError},
exchange::{
CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketSseStream,
CodexWebSocketStreamingExchange, CodexWebSocketTurnStateUpdate,
@@ -29,6 +29,8 @@ use tokio::{
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, tungstenite::Message};
use uuid::Uuid;
use super::CodexWebSocketCloseError;
/// 底层 tungstenite WebSocket 流。
pub(crate) type RawWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
@@ -69,13 +71,25 @@ const PUMP_MESSAGE_BUFFER: usize = 64;
pub(crate) enum PumpExitReason {
CommandChannelClosed,
LocalClose,
OutboundTransportError { message: String },
UpstreamCloseFrame { frame: Option<String> },
OutboundTransportError {
message: String,
},
UpstreamCloseFrame {
close: Option<CodexWebSocketCloseError>,
},
UpstreamEof,
InboundTransportError { message: String },
PongTimeout { timeout: Duration },
LivenessTimeout { timeout: Duration },
KeepaliveTransportError { message: String },
InboundTransportError {
message: String,
},
PongTimeout {
timeout: Duration,
},
LivenessTimeout {
timeout: Duration,
},
KeepaliveTransportError {
message: String,
},
MessageReceiverClosed,
}
@@ -100,7 +114,7 @@ impl PumpExitReason {
Self::OutboundTransportError { message }
| Self::InboundTransportError { message }
| Self::KeepaliveTransportError { message } => Some(message.clone()),
Self::UpstreamCloseFrame { frame } => frame.clone(),
Self::UpstreamCloseFrame { close } => close.as_ref().map(ToString::to_string),
Self::PongTimeout { timeout } | Self::LivenessTimeout { timeout } => {
Some(format!("{timeout:?}"))
}
@@ -111,6 +125,13 @@ impl PumpExitReason {
}
}
pub(crate) fn upstream_close(&self) -> Option<&CodexWebSocketCloseError> {
match self {
Self::UpstreamCloseFrame { close: Some(close) } => Some(close),
_ => None,
}
}
fn is_unexpected(&self) -> bool {
matches!(
self,
@@ -336,7 +357,12 @@ async fn pump_loop(
message => {
let terminal_reason = match &message {
Message::Close(frame) => Some(PumpExitReason::UpstreamCloseFrame {
frame: frame.as_ref().map(|frame| format!("{frame:?}")),
close: frame.as_ref().map(|frame| {
CodexWebSocketCloseError::new(
Some(u16::from(frame.code)),
Some(frame.reason.to_string()),
)
}),
}),
_ => None,
};
+104 -2
View File
@@ -20,8 +20,8 @@ use gateway_admin::model::provider_credentials::{
use gateway_admin::model::{MutationActor, MutationContext, Revision};
use gateway_admin::ports::provider::ProviderAdminErrorKind;
use gateway_core::engine::credential::{
CredentialRevision, OpaqueProviderData, ProviderAccount, ProviderAccountId,
ProviderAccountStore, QuotaObservation,
AccountAvailability as CoreAccountAvailability, AccountStateChange, CredentialRevision,
OpaqueProviderData, ProviderAccount, ProviderAccountId, ProviderAccountStore, QuotaObservation,
};
use gateway_core::operation::{GenerateRequest, Operation, ProtocolPayload};
use gateway_core::provider_ports::{
@@ -503,6 +503,108 @@ async fn openai_admin_provider_projects_codex_additional_limit_as_the_primary_qu
assert_eq!(monthly[0].used_percent, Some(2.0));
}
#[tokio::test]
async fn openai_admin_projects_confirmed_quota_exhaustion_as_full_without_mutating_raw_usage() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_admin_confirmed_exhaustion";
store
.seed_oauth_credential(ImportCodexOAuthCredential {
account_id: account_id.to_owned(),
name: "admin confirmed exhaustion".to_owned(),
secret: secret("admin-confirmed-exhaustion-access"),
verified_account: profile("chatgpt-admin-confirmed-exhaustion"),
next_refresh_at: Some(Utc::now() + chrono::Duration::minutes(30)),
enabled: true,
})
.await;
let account = store.account(account_id).expect("stored account");
let reset_at = 1_900_000_000_u64;
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 86, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist raw quota");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: CoreAccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: Some(SystemTime::UNIX_EPOCH + Duration::from_secs(reset_at)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted");
let config = valid_config();
let bundle = provider_openai::initialize(
config.config.clone(),
provider_ports_with(Arc::clone(&store), Arc::new(TestOAuthPending::default())),
)
.await
.expect("OpenAI bundle");
let exhausted = bundle
.admin_provider()
.quota(ProviderQuotaRequest {
account_id: account.id().clone(),
refresh: false,
rolling_usage: None,
})
.await
.expect("project exhausted quota");
assert_eq!(exhausted.windows.len(), 1);
assert_eq!(exhausted.windows[0].used_percent, Some(100.0));
let raw = store
.get_quotas(std::slice::from_ref(account.id()))
.await
.expect("read raw quota")
.pop()
.expect("raw quota");
assert_eq!(
raw.quota.expect("raw document").expose_to_provider()["rate_limit"]["primary_window"]["used_percent"],
86
);
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: CoreAccountAvailability::Ready,
reason: None,
cooldown_until: None,
observed_at: SystemTime::now(),
})
.await
.expect("recover account");
let recovered = bundle
.admin_provider()
.quota(ProviderQuotaRequest {
account_id: account.id().clone(),
refresh: false,
rolling_usage: None,
})
.await
.expect("project recovered quota");
assert_eq!(recovered.windows[0].used_percent, Some(86.0));
}
#[tokio::test]
async fn openai_admin_provider_rejects_unprepared_mutations_before_store_commit() {
let store = Arc::new(MemoryAccountStore::default());
@@ -117,7 +117,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() {
64,
8,
20_000,
5,
15,
true,
true,
"app_EMoamEEZ73f0CkXaXp7hrann",
@@ -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;
}
@@ -362,11 +400,65 @@ async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision
}
#[tokio::test]
async fn manual_quota_refresh_updates_account_state_from_provider_response() {
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_keeps_confirmed_exhaustion_until_the_recorded_reset() {
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"))
@@ -417,25 +509,113 @@ async fn manual_quota_refresh_updates_account_state_from_provider_response() {
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 20, "reset_at": 1_900_000_000}
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
})))
.mount(&server)
.await;
let recovered = service
let stale = service
.refresh_account(account.id())
.await
.expect("refresh recovered quota");
.expect("refresh contradictory full quota");
assert!(!recovered.fact().exhausted());
assert!(!stale.fact().exhausted());
assert_eq!(
store
.account(account_id)
.expect("recovered account")
.expect("account after contradictory quota")
.availability(),
AccountAvailability::Ready
AccountAvailability::QuotaExhausted,
"an unchanged usage window cannot override a confirmed quota failure"
);
server.reset().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": 99, "reset_at": 1_900_000_000}
}
})))
.mount(&server)
.await;
let unchanged_reset = service
.refresh_account(account.id())
.await
.expect("refresh lower usage quota");
assert!(!unchanged_reset.fact().exhausted());
assert_eq!(
store
.account(account_id)
.expect("account after lower usage")
.availability(),
AccountAvailability::QuotaExhausted,
"a lower used_percent before the recorded reset is not quota recovery"
);
}
#[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]
@@ -608,7 +788,7 @@ async fn synchronize_without_accounts_is_a_noop_before_network_io() {
}
#[tokio::test]
async fn periodic_quota_synchronization_skips_fresh_ready_accounts() {
async fn periodic_quota_synchronization_does_not_scan_stale_ready_accounts() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_ready").await;
let account = store
@@ -629,15 +809,16 @@ async fn periodic_quota_synchronization_skips_fresh_ready_accounts() {
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
// 正常账号即使缓存已旧,也只等待真实请求被动同步,不能被定时任务主动拉取。
observed_at: Some(SystemTime::UNIX_EPOCH),
})
.await
.expect("persist fresh quota");
.expect("persist stale quota");
let summary = blocked_network_quota_service(&store)
.synchronize()
.await
.expect("fresh ready account must not query upstream");
.expect("ready account must not query upstream");
assert_eq!(summary, CodexQuotaSyncSummary::default());
}
@@ -683,7 +864,13 @@ async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_q
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(429))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": false,
"limit_reached": true,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})))
.mount(&server)
.await;
let service = quota_service_with_base_url(
@@ -697,12 +884,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 +934,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())
);
}
@@ -784,72 +979,6 @@ async fn periodic_quota_synchronization_skips_ordinary_cooldown_without_exhauste
);
}
#[tokio::test]
async fn periodic_quota_synchronization_limits_stale_ready_account_scans() {
let store = Arc::new(MemoryAccountStore::default());
for index in 0..9 {
let account_id = format!("acct_periodic_stale_ready_{index}");
create_account(&store, &account_id).await;
let account = store.account(&account_id).expect("created account");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"primary_window": {"used_percent": 10, "reset_at": 1_900_000_000}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::UNIX_EPOCH),
})
.await
.expect("persist stale quota");
}
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,
"primary_window": {"used_percent": 100, "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(),
);
let summary = service.synchronize().await.expect("quota synchronization");
assert_eq!(summary.exhausted, 8);
assert_eq!(
store
.account("acct_periodic_stale_ready_0")
.expect("scanned account")
.availability(),
AccountAvailability::QuotaExhausted
);
assert_eq!(
store
.account("acct_periodic_stale_ready_8")
.expect("deferred account")
.availability(),
AccountAvailability::Ready
);
}
#[tokio::test]
async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() {
let store = Arc::new(MemoryAccountStore::default());
@@ -877,6 +1006,215 @@ async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() {
assert_eq!(summary.transient, 1);
}
#[tokio::test]
async fn periodic_quota_synchronization_rechecks_exhausted_accounts_without_waiting_for_reset() {
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_add(Duration::from_secs(24 * 60 * 60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted with a future reset");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(500))
.expect(1)
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
// reset 仅用于展示/调度状态,不能推迟额度耗尽账号的定时复核。
let due = service.synchronize().await.expect("quota synchronization");
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_does_not_recover_from_percent_drop_before_reset() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_periodic_recovery_evidence";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
let reset_at = 1_900_000_000;
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("seed post-failure quota snapshot");
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_add(Duration::from_secs(24 * 60 * 60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 100, "reset_at": reset_at}
}
})))
.expect(1)
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let unchanged = service
.synchronize()
.await
.expect("unchanged quota recheck");
assert_eq!(unchanged.updated, 1);
assert_eq!(
store
.account(account_id)
.expect("account after unchanged snapshot")
.availability(),
AccountAvailability::QuotaExhausted
);
server.reset().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 99, "reset_at": reset_at}
}
})))
.expect(1)
.mount(&server)
.await;
// 新 worker 实例模拟下一次 30 分钟复核,避免测试等待真实节流周期。
let next_cycle = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
let lower_percent = next_cycle
.synchronize()
.await
.expect("lower percent quota recheck");
assert_eq!(lower_percent.updated, 1);
assert_eq!(
store
.account(account_id)
.expect("account after lower percent snapshot")
.availability(),
AccountAvailability::QuotaExhausted,
"a provider percentage change is not sufficient to release a confirmed exhaustion"
);
}
#[tokio::test]
async fn periodic_quota_synchronization_recovers_after_recorded_reset() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_periodic_reset_elapsed";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
store
.apply_state_change(AccountStateChange {
account_id: account.id().clone(),
expected_revision: account.revision(),
availability: AccountAvailability::QuotaExhausted,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: SystemTime::now().checked_sub(Duration::from_secs(60)),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted with an elapsed reset");
let server = MockServer::start().await;
let next_reset = Utc::now().timestamp() + 24 * 60 * 60;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 1, "reset_at": next_reset}
}
})))
.expect(1)
.mount(&server)
.await;
let service = quota_service_with_base_url(
&store,
reqwest::Client::builder()
.no_proxy()
.build()
.expect("client"),
server.uri(),
);
// 定时复核可提前执行,但只有已记录的 reset 到期后才释放额度耗尽状态。
let summary = service.synchronize().await.expect("quota synchronization");
assert_eq!(summary.updated, 1);
let current = store
.account(account_id)
.expect("account after elapsed reset");
assert_eq!(current.availability(), AccountAvailability::Ready);
assert_eq!(current.cooldown_until(), None);
}
#[tokio::test]
async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
let store = Arc::new(MemoryAccountStore::default());
@@ -904,46 +1242,34 @@ async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
}
#[tokio::test]
async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_its_reset_is_due() {
async fn periodic_quota_synchronization_does_not_bypass_throttle_for_a_new_reset() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_periodic_reset_due").await;
let account = store
.account("acct_periodic_reset_due")
.expect("created account");
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("quota_exhausted".to_owned()),
cooldown_until: None,
reason: Some("usage_limit_exhausted".to_owned()),
cooldown_until: Some(initial_reset),
observed_at: SystemTime::now(),
})
.await
.expect("mark account exhausted");
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": 1_900_000_000}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist future reset");
.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(500))
.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(
@@ -955,36 +1281,38 @@ 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);
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": 1}
}
})
.as_object()
.expect("quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("persist due reset");
let reset_recheck = service.synchronize().await.expect("reset recheck");
let repeated = service
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("throttled reset recheck");
let repeated = service.synchronize().await.expect("same reset throttled");
let request_count = server
.received_requests()
.await
.expect("received requests")
.len();
assert_eq!(reset_recheck.transient, 1);
assert_eq!(repeated, CodexQuotaSyncSummary::default());
assert_eq!(
(
first.exhausted,
persisted_reset,
second.exhausted,
repeated,
request_count,
),
(
1,
Some(confirmed_reset),
0,
CodexQuotaSyncSummary::default(),
1,
)
);
}
#[tokio::test]
@@ -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;
@@ -5,9 +5,10 @@ use std::time::{Duration, SystemTime};
use base64::{Engine as _, engine::general_purpose::STANDARD};
use chrono::Utc;
use futures::StreamExt;
use futures::{SinkExt, StreamExt};
use gateway_core::engine::credential::{
AccountAvailability, AccountFeedbackStats, ProviderAccountId,
AccountAvailability, AccountFeedbackStats, OpaqueProviderData, ProviderAccountId,
ProviderAccountStore as _, QuotaObservation,
};
use gateway_core::engine::provider::{Provider as _, ProviderRequest};
use gateway_core::engine::{
@@ -29,20 +30,23 @@ 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 tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::protocol::{CloseFrame, frame::coding::CloseCode};
use wiremock::matchers::{header, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use crate::support::{
MemoryAccountStore, MemorySessionAffinity, MemorySessionExclusions, TestLeaseCoordinator,
account_policy, agent_identity_service_with_pool, catalog_cache, profile, secret,
};
use crate::transport::accept_codex_test_websocket;
const OFFICIAL_FIXTURE: &[u8] =
include_bytes!("../transport/fixtures/official_models_snapshot.json");
@@ -80,10 +84,36 @@ fn provider_with_base_url(store: &Arc<MemoryAccountStore>, base_url: String) ->
provider_with_affinity_and_base_url(store, Arc::new(MemorySessionAffinity::default()), base_url)
}
fn provider_with_leases(
store: &Arc<MemoryAccountStore>,
leases: Arc<TestLeaseCoordinator>,
) -> CodexProvider {
provider_with_affinity_and_base_url_and_leases(
store,
Arc::new(MemorySessionAffinity::default()),
OFFICIAL_CODEX_BASE_URL.to_owned(),
leases,
)
}
fn provider_with_affinity_and_base_url(
store: &Arc<MemoryAccountStore>,
session_affinity: Arc<MemorySessionAffinity>,
base_url: String,
) -> CodexProvider {
provider_with_affinity_and_base_url_and_leases(
store,
session_affinity,
base_url,
Arc::new(TestLeaseCoordinator::default()),
)
}
fn provider_with_affinity_and_base_url_and_leases(
store: &Arc<MemoryAccountStore>,
session_affinity: Arc<MemorySessionAffinity>,
base_url: String,
leases: Arc<TestLeaseCoordinator>,
) -> CodexProvider {
let profile = wire_profile();
let http = reqwest::Client::builder()
@@ -111,7 +141,7 @@ fn provider_with_affinity_and_base_url(
let selector = Arc::new(CodexCredentialSelector::new(
ProviderKind::new("openai").expect("provider"),
store.repository(),
Arc::new(TestLeaseCoordinator::default()),
leases,
session_affinity,
Arc::new(MemorySessionExclusions::default()),
Arc::clone(&catalog),
@@ -516,6 +546,130 @@ async fn generate_without_an_eligible_openai_account_fails_before_network_io() {
assert_eq!(error.send_state(), UpstreamSendState::NotSent);
}
#[tokio::test]
async fn capacity_selection_error_preserves_classification_and_retry_after() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_capacity_busy").await;
let leases = Arc::new(TestLeaseCoordinator::default());
*leases.busy.lock().expect("lease busy lock") = true;
let error = match provider_with_leases(&store, leases)
.execute(
planned_request("openai", generate_operation()),
context("req_capacity_busy", CancellationToken::new()),
)
.await
{
Ok(_) => panic!("busy account selection must fail"),
Err(error) => error,
};
assert_eq!(
(error.kind(), error.send_state(), error.retry_after()),
(
ProviderErrorKind::AccountCapacityUnavailable,
UpstreamSendState::NotSent,
Some(Duration::from_millis(25)),
)
);
}
#[tokio::test]
async fn selection_infrastructure_errors_have_a_distinct_classification() {
let store = Arc::new(MemoryAccountStore::default());
store.fail_provider_listing();
let error = match provider(&store)
.execute(
planned_request("openai", generate_operation()),
context("req_selection_store_failure", CancellationToken::new()),
)
.await
{
Ok(_) => panic!("account store failure must fail selection"),
Err(error) => error,
};
assert_eq!(
(error.kind(), error.send_state()),
(
ProviderErrorKind::ProviderInfrastructureUnavailable,
UpstreamSendState::NotSent,
)
);
}
#[tokio::test]
async fn websocket_close_details_are_only_exposed_through_the_client_error() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_websocket_close").await;
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind WebSocket listener");
let base_url = format!(
"http://{}",
listener.local_addr().expect("listener address")
);
let server = tokio::spawn(async move {
let (stream, _) = listener
.accept()
.await
.expect("accept WebSocket connection");
let mut websocket = accept_codex_test_websocket(stream).await;
let _request = websocket
.next()
.await
.expect("WebSocket request")
.expect("valid WebSocket request");
websocket
.send(Message::Text(
json!({
"type": "response.output_text.delta",
"delta": "partial output"
})
.to_string()
.into(),
))
.await
.expect("send semantic WebSocket event");
websocket
.close(Some(CloseFrame {
code: CloseCode::Size,
reason: "message too big".into(),
}))
.await
.expect("close WebSocket");
});
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", generate_operation()),
context("req_websocket_close", CancellationToken::new()),
)
.await
.expect("prepare WebSocket provider stream");
let error = loop {
match stream.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("WebSocket close must surface a provider error"),
}
};
server.await.expect("WebSocket server");
assert!(!format!("{error:?}").contains("message too big"));
assert!(!error.to_string().contains("message too big"));
assert_eq!(error.kind(), ProviderErrorKind::Transport);
assert_eq!(error.send_state(), UpstreamSendState::Ambiguous);
assert!(error.allows_pre_delivery_retry());
let detail = error
.client_visible_upstream_error()
.expect("WebSocket close detail");
assert_eq!(detail.message(), "message too big");
assert_eq!(detail.code(), Some("1009"));
assert_eq!(detail.error_type(), Some("websocket_close_error"));
}
#[tokio::test]
async fn disabled_account_is_excluded_from_normal_scheduling() {
let store = Arc::new(MemoryAccountStore::default());
@@ -700,6 +854,109 @@ async fn cross_account_scope_sanitizes_only_known_turn_metadata_fields() {
);
}
#[tokio::test]
async fn continuation_replay_should_use_the_client_full_input_without_legacy_transcript() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_client_history").await;
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(CAPTURE_COMPLETED_SSE),
)
.expect(1)
.mount(&server)
.await;
let client_input = json!([
{"role": "user", "content": "complete client history"},
{"role": "assistant", "content": "previous answer"},
{"role": "user", "content": "current turn"}
]);
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), client_input.clone()),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([("use_websocket".to_owned(), json!(false))])),
)
.with_provider_session_state(
ProviderSessionState::new(
"openai",
Map::from_iter([
("account_id".to_owned(), json!("acct_client_history")),
("conversation_id".to_owned(), json!("conversation")),
("continuation_scope".to_owned(), json!("replay_required")),
(
"transcript".to_owned(),
json!([{"client_input": {"role": "user", "content": "stale proxy copy"}}]),
),
]),
)
.expect("provider session state"),
);
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", Operation::Generate(generation)),
context("req_client_history", CancellationToken::new())
.with_continuation_attempt(ContinuationAttempt::ReplayAny),
)
.await
.expect("prepare replay stream");
while let Some(event) = stream.next().await {
event.expect("replay response");
}
let requests = server
.received_requests()
.await
.expect("captured replay request");
let body: Value = serde_json::from_slice(&requests[0].body).expect("captured JSON body");
assert_eq!(body.get("input"), Some(&client_input));
}
#[tokio::test]
async fn completed_response_session_state_should_not_copy_the_conversation_transcript() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_bounded_session_state").await;
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(CAPTURE_COMPLETED_SSE),
)
.expect(1)
.mount(&server)
.await;
let mut stream = provider_with_base_url(&store, server.uri())
.execute(
planned_request("openai", http_generate_operation()),
context("req_bounded_session_state", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let mut session_payload = None;
while let Some(event) = stream.next().await {
let event = event.expect("provider response");
if let Some(update) = event.session_update() {
session_payload = Some(update.payload().clone());
}
}
assert!(
!session_payload
.expect("completed response session update")
.contains_key("transcript")
);
}
#[tokio::test]
async fn account_change_drops_only_account_bound_opaque_headers() {
let protocol_context = Map::from_iter([(
@@ -784,6 +1041,94 @@ async fn prompt_cache_key_should_become_an_opaque_session_affinity_lookup_key()
assert!(keys[0].bytes().all(|byte| byte.is_ascii_hexdigit()));
}
#[tokio::test]
async fn explicit_session_id_should_override_turn_specific_prompt_cache_keys_for_affinity() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_session_affinity").await;
let affinity = Arc::new(MemorySessionAffinity::default());
let provider = provider_with_affinity(&store, Arc::clone(&affinity));
for (request_id, prompt_cache_key) in [
("req_session_affinity_first", "turn-cache-first"),
("req_session_affinity_second", "turn-cache-second"),
] {
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("hello")),
("prompt_cache_key".to_owned(), json!(prompt_cache_key)),
]),
)
.expect("OpenAI payload")
.with_context(Map::from_iter([(
"session_id".to_owned(),
json!("stable-client-session"),
)])),
);
let stream = provider
.execute(
planned_request("openai", Operation::Generate(generation)),
context(request_id, CancellationToken::new()),
)
.await
.expect("prepare provider stream");
drop(stream);
}
let keys = affinity.lookup_keys();
assert_eq!(keys.len(), 2);
assert_eq!(keys[0], keys[1]);
}
#[tokio::test]
async fn invalid_local_conversation_id_should_fall_back_to_an_opaque_affinity_key() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_local_affinity").await;
let affinity = Arc::new(MemorySessionAffinity::default());
let generation = GenerateRequest::from_protocol_payload(
ProtocolPayload::json_object(
"openai",
Map::from_iter([
("model".to_owned(), json!("gpt-5.4")),
("input".to_owned(), json!("hello")),
]),
)
.expect("OpenAI payload"),
)
.with_provider_session_state(
ProviderSessionState::new(
"openai",
Map::from_iter([
("account_id".to_owned(), json!("acct_local_affinity")),
(
"conversation_id".to_owned(),
json!("lc_UppercaseBase64ConversationId"),
),
("continuation_scope".to_owned(), json!("replay_required")),
("transcript".to_owned(), json!([])),
]),
)
.expect("provider session state"),
);
let stream = provider_with_affinity(&store, Arc::clone(&affinity))
.execute(
planned_request("openai", Operation::Generate(generation)),
context("req_local_affinity", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
drop(stream);
let keys = affinity.lookup_keys();
assert_eq!(keys.len(), 1);
assert_eq!(keys[0].len(), 64);
assert!(keys[0].bytes().all(|byte| byte.is_ascii_hexdigit()));
}
#[tokio::test]
async fn completed_response_persists_session_affinity_before_stream_consumer_stops_polling() {
let store = Arc::new(MemoryAccountStore::default());
@@ -827,13 +1172,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 +1201,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 +1251,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()
@@ -907,10 +1263,157 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown(
}
#[tokio::test]
async fn ordinary_request_should_forward_created_before_later_failed_chunk() {
async fn usage_limit_failure_returns_promptly_and_refreshes_the_authoritative_quota_snapshot() {
let store = Arc::new(MemoryAccountStore::default());
let account_id = "acct_usage_limit_request_path";
create_account(&store, account_id).await;
let account = store.account(account_id).expect("created account");
store
.compare_and_swap_quota(QuotaObservation {
account_id: account.id().clone(),
expected_revision: account.revision(),
quota: Some(OpaqueProviderData::new(
json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 99, "reset_at": 1_900_000_000}
}
})
.as_object()
.expect("stale quota object")
.clone(),
)),
observed_at: Some(SystemTime::now()),
})
.await
.expect("seed stale passive quota");
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/api/codex/usage"))
.and(header("authorization", format!("Bearer at-{account_id}")))
// 生产环境会短暂返回互相矛盾的快照:数字已到 100%,但布尔位仍声称可用。
// 后台同步应覆盖展示数据,不能据此撤销真实请求已经确认的额度耗尽。
.respond_with(
ResponseTemplate::new(200)
// 验证后台 usage 同步不能把原始的额度错误响应拖到查询完成之后。
.set_delay(Duration::from_millis(750))
.set_body_json(json!({
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {"used_percent": 100, "reset_at": 1_900_000_000}
}
})),
)
.expect(1)
.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 = timeout(Duration::from_millis(300), async {
loop {
match stream.next().await {
Some(Ok(_)) => {}
Some(Err(error)) => break error,
None => panic!("usage-limit failure must surface a typed failure"),
}
}
})
.await
.expect("usage refresh must not block the original quota 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())
);
tokio::time::sleep(Duration::from_millis(200)).await;
let requests = server.received_requests().await.expect("received requests");
assert_eq!(
requests
.iter()
.filter(|request| request.url.path() == "/api/codex/usage")
.count(),
0,
"usage refresh must wait for the upstream quota settlement delay"
);
let quota = timeout(Duration::from_secs(5), async {
loop {
let quota = store
.get_quotas(&[account.id().clone()])
.await
.expect("read refreshed quota")
.into_iter()
.next()
.and_then(|observation| observation.quota)
.expect("authoritative quota snapshot");
let quota = Value::Object(quota.into_inner());
if quota
.pointer("/rate_limit/primary_window/used_percent")
.and_then(Value::as_u64)
== Some(100)
{
break quota;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("background usage refresh must replace the stale quota snapshot");
assert_eq!(
quota
.pointer("/rate_limit/primary_window/used_percent")
.and_then(Value::as_u64),
Some(100)
);
let requests = server.received_requests().await.expect("received requests");
assert_eq!(
requests
.iter()
.filter(|request| request.url.path() == "/api/codex/usage")
.count(),
1
);
assert_eq!(
store
.account(account_id)
.expect("account after stale full usage refresh")
.availability(),
AccountAvailability::QuotaExhausted,
"a contradictory full usage snapshot must not unlock a confirmed exhausted account"
);
}
#[tokio::test]
async fn ordinary_request_should_hold_created_until_later_failure_can_rotate() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_first_event_latency").await;
let (base_url, release, _first_chunk_sent, server) = paused_chunked_sse_server(
let (base_url, release, first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_first_event\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n"
@@ -931,41 +1434,111 @@ async fn ordinary_request_should_forward_created_before_later_failed_chunk() {
.await
.expect("prepare provider stream");
let first_event = loop {
let next = timeout(Duration::from_secs(1), stream.next())
.await
.expect("first chunk must not wait for the withheld failure")
.expect("provider stream must stay open")
.expect("provider event");
if next.has_client_event() {
break next;
let mut first_chunk_sent = Box::pin(first_chunk_sent);
loop {
tokio::select! {
sent = &mut first_chunk_sent => {
sent.expect("first upstream chunk");
break;
}
next = stream.next() => {
let event = next
.expect("provider stream must stay open")
.expect("provider event");
assert!(!event.has_client_event(), "response.created must remain replayable");
}
}
};
assert_eq!(
first_event.wire_event().and_then(|wire| wire.event_type()),
Some("response.created")
}
let exposed = timeout(Duration::from_millis(100), async {
loop {
let next = stream
.next()
.await
.expect("provider stream must stay open")
.expect("provider event");
if next.has_client_event() {
return next;
}
}
})
.await;
assert!(
exposed.is_err(),
"a structural event must not commit the downstream before a later 429"
);
release.send(()).expect("release second upstream chunk");
let mut later_wire_types = Vec::new();
let failure = loop {
let mut failure = loop {
let next = timeout(Duration::from_secs(1), stream.next())
.await
.expect("later failure must arrive after release")
.expect("provider stream must return typed failure");
match next {
Ok(event) => {
if let Some(event_type) = event.wire_event().and_then(|wire| wire.event_type()) {
later_wire_types.push(event_type.to_owned());
}
}
Ok(event) => assert!(
!event.has_client_event(),
"failure attempt leaked downstream"
),
Err(error) => break error,
}
};
assert_eq!(later_wire_types, vec!["response.failed"]);
assert!(!failure.replay_is_safe());
assert!(!failure.has_atomic_client_events());
assert!(failure.replay_is_safe());
assert_eq!(
failure
.take_atomic_client_events()
.iter()
.filter_map(|event| event.wire_event()?.event_type())
.collect::<Vec<_>>(),
vec!["response.created", "response.failed"]
);
server.await.expect("chunked SSE server");
}
#[tokio::test]
async fn ordinary_request_should_bound_structural_event_replay_grace() {
let store = Arc::new(MemoryAccountStore::default());
create_account(&store, "acct_bounded_replay_grace").await;
let (base_url, release, _first_chunk_sent, server) = paused_chunked_sse_server(
concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_bounded_grace\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n"
)
.to_owned(),
String::new(),
)
.await;
let mut stream = provider_with_base_url(&store, base_url)
.execute(
planned_request("openai", http_generate_operation()),
context("req_bounded_replay_grace", CancellationToken::new()),
)
.await
.expect("prepare provider stream");
let first_event = timeout(Duration::from_secs(2), async {
loop {
let event = stream
.next()
.await
.expect("provider stream must stay open")
.expect("provider event");
if event.has_client_event() {
return event;
}
}
})
.await
.expect("response.created must be released after the bounded grace period");
assert_eq!(
first_event.wire_event().and_then(|wire| wire.event_type()),
Some("response.created")
);
release.send(()).expect("finish upstream response");
while let Some(event) = stream.next().await {
event.expect("clean upstream EOF");
}
server.await.expect("chunked SSE server");
}
@@ -1207,8 +1780,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 +1793,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 +1900,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());
@@ -2,7 +2,7 @@
use std::collections::{BTreeMap, BTreeSet};
use std::num::NonZeroU32;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime};
@@ -36,12 +36,14 @@ struct StoredAccount {
account: ProviderAccount,
credential: gateway_core::engine::credential::PlaintextCredential,
quota: Option<QuotaObservation>,
state_observed_at: Option<SystemTime>,
}
#[derive(Default)]
pub(crate) struct MemoryAccountStore {
accounts: Mutex<BTreeMap<ProviderAccountId, StoredAccount>>,
quota_reads: AtomicUsize,
fail_provider_listing: AtomicBool,
}
impl MemoryAccountStore {
@@ -114,6 +116,10 @@ impl MemoryAccountStore {
.get(&id)
.is_some_and(|stored| stored.quota.is_some())
}
pub(crate) fn fail_provider_listing(&self) {
self.fail_provider_listing.store(true, Ordering::SeqCst);
}
}
#[async_trait]
@@ -129,6 +135,7 @@ impl ProviderAccountStore for MemoryAccountStore {
account: input.account,
credential: input.credential,
quota: None,
state_observed_at: None,
},
);
Ok(())
@@ -160,6 +167,9 @@ impl ProviderAccountStore for MemoryAccountStore {
&self,
provider: &ProviderKind,
) -> Result<Vec<ProviderAccount>, StoreError> {
if self.fail_provider_listing.load(Ordering::SeqCst) {
return Err(store_error(StoreErrorKind::Unavailable));
}
Ok(self
.accounts
.lock()
@@ -260,6 +270,15 @@ impl ProviderAccountStore for MemoryAccountStore {
if stored.account.revision() != observation.expected_revision {
return Ok(QuotaWriteOutcome::Conflict);
}
if stored
.quota
.as_ref()
.and_then(|quota| quota.observed_at)
.zip(observation.observed_at)
.is_some_and(|(current, incoming)| current > incoming)
{
return Ok(QuotaWriteOutcome::Conflict);
}
stored.quota = Some(observation);
Ok(QuotaWriteOutcome::Updated)
}
@@ -272,6 +291,12 @@ impl ProviderAccountStore for MemoryAccountStore {
if stored.account.revision() != change.expected_revision {
return Err(store_error(StoreErrorKind::Conflict));
}
if stored
.state_observed_at
.is_some_and(|current| current > change.observed_at)
{
return Err(store_error(StoreErrorKind::Conflict));
}
stored.account = rebuild_account(
&stored.account,
AccountRebuild {
@@ -285,6 +310,7 @@ impl ProviderAccountStore for MemoryAccountStore {
profile: None,
},
);
stored.state_observed_at = Some(change.observed_at);
Ok(())
}
@@ -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();
@@ -5,6 +5,7 @@ use provider_openai::transport::{
WebSocketOriginBreaker, WebSocketOriginBreakerConfig, WebSocketOriginBreakerDecision,
},
};
use tokio_tungstenite::tungstenite::protocol::{CloseFrame, frame::coding::CloseCode};
use super::*;
@@ -913,14 +914,133 @@ async fn concurrent_same_key_should_singleflight_websocket_opening() {
}
#[tokio::test]
async fn payload_send_failure_should_not_open_http_fallback() {
async fn websocket_failure_before_first_delivery_should_use_same_account_http() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket_with(stream, |request, response| {
assert_eq!(
request
.headers()
.get("authorization")
.and_then(|value| value.to_str().ok()),
Some("Bearer access-token")
);
assert_eq!(
request
.headers()
.get("chatgpt-account-id")
.and_then(|value| value.to_str().ok()),
Some("chatgpt-account")
);
response.headers_mut().insert(
"sec-websocket-extensions",
"permessage-deflate".parse().unwrap(),
);
})
.await;
let _payload = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
json!({
"type": "response.created",
"response": {
"id": "resp_abandoned_websocket",
"model": "gpt-test"
}
})
.to_string()
.into(),
))
.await
.unwrap();
websocket
.send(Message::Text(
json!({
"type": "response.in_progress",
"response": {
"id": "resp_abandoned_websocket",
"model": "gpt-test"
}
})
.to_string()
.into(),
))
.await
.unwrap();
websocket
.close(Some(CloseFrame {
code: CloseCode::Size,
reason: "message too big".into(),
}))
.await
.unwrap();
let (mut http, _) = listener.accept().await.unwrap();
let request = read_http_request(&mut http).await;
assert!(request.starts_with("POST /codex/responses HTTP/1.1"));
assert_eq!(
read_header_value(&request, "authorization"),
Some("Bearer access-token")
);
assert_eq!(
read_header_value(&request, "chatgpt-account-id"),
Some("chatgpt-account")
);
write_completed_sse_response(&mut http).await;
});
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
test_wire_profile(),
)
.with_websocket_pool(Arc::new(CodexWebSocketPool::default()));
let response = backend
.create_response(
&new_chain_request("conversation-pre-delivery-fallback"),
request_context("req_pre_delivery_fallback", Some("chatgpt-account")),
)
.await
.expect("pre-delivery WebSocket failure should use same-account HTTP");
server.await.unwrap();
assert_eq!(response.transport, CodexBackendTransport::HttpSse);
assert_eq!(
response.transport_metrics.decision,
Some(CodexTransportDecision::Http2PreDeliveryFailure)
);
assert!(response.body.contains("resp_http_fallback"));
assert!(!response.body.contains("resp_abandoned_websocket"));
}
#[tokio::test]
async fn websocket_failure_after_first_delivery_should_not_open_http_fallback() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _payload = websocket.next().await.unwrap().unwrap();
websocket.close(None).await.unwrap();
websocket
.send(Message::Text(
json!({
"type": "response.output_text.delta",
"delta": "delivered"
})
.to_string()
.into(),
))
.await
.unwrap();
websocket
.close(Some(CloseFrame {
code: CloseCode::Size,
reason: "message too big after output".into(),
}))
.await
.unwrap();
assert!(
timeout(Duration::from_millis(100), listener.accept())
.await
@@ -936,17 +1056,25 @@ async fn payload_send_failure_should_not_open_http_fallback() {
let error = backend
.create_response(
&new_chain_request("conversation-post-send"),
request_context("req_post_send", Some("chatgpt-account")),
&new_chain_request("conversation-post-delivery"),
request_context("req_post_delivery", Some("chatgpt-account")),
)
.await
.expect_err("post-send close must not be replayed");
.expect_err("post-delivery close must not be replayed");
server.await.unwrap();
std::assert_matches!(
error,
CodexClientError::WebSocket(CodexWebSocketExchangeError::PostSendAmbiguous { .. })
);
let CodexClientError::WebSocket(CodexWebSocketExchangeError::PostSendAmbiguous {
source: Some(source),
..
}) = error
else {
panic!("expected typed post-delivery WebSocket failure");
};
let CodexWebSocketExchangeError::ClosedBeforeTerminal(close) = *source else {
panic!("expected upstream close frame source");
};
assert_eq!(close.code(), Some(1009));
assert_eq!(close.reason(), Some("message too big after output"));
}
#[tokio::test]
@@ -218,6 +218,17 @@ fn transport_should_prefer_websocket_when_requested_without_history() {
assert!(transport_requirement(&request).allows_pre_send_http_fallback());
}
#[test]
fn new_chain_should_allow_pre_delivery_http_fallback() {
let mut request = codex_request("gpt-test", "be brief", Vec::new());
request.use_websocket = true;
assert!(
transport_requirement(&request).allows_pre_delivery_http_fallback(),
"new chains may follow the TS same-account fallback behavior"
);
}
#[test]
fn transport_should_mark_unknown_previous_response_as_external() {
let mut request = codex_request("gpt-test", "be brief", Vec::new());
@@ -230,6 +241,17 @@ fn transport_should_mark_unknown_previous_response_as_external() {
assert!(transport_requirement(&request).allows_pre_send_http_fallback());
}
#[test]
fn previous_response_should_forbid_pre_delivery_http_fallback() {
let mut request = codex_request("gpt-test", "be brief", Vec::new());
request.set_previous_response_id(Some("resp_previous".to_owned()));
assert!(
!transport_requirement(&request).allows_pre_delivery_http_fallback(),
"a sent continuation must not be replayed on another transport"
);
}
#[test]
fn transport_should_allow_forced_http_sse() {
let mut request = codex_request("gpt-test", "be brief", Vec::new());
@@ -649,7 +649,15 @@ async fn websocket_execute_response_create_request_should_reject_binary_event()
.unwrap();
websocket.close(None).await.unwrap();
});
let prepared = prepared_websocket_request(&format!("http://{addr}"));
let mut request = codex_request("gpt-test", "be brief", Vec::new());
request.set_previous_response_id(Some("resp_binary_previous".to_owned()));
let prepared = CodexWebSocketConnection::responses_create_request(
&format!("http://{addr}"),
"dGhlIHNhbXBsZSBub25jZQ==",
vec![("authorization".to_owned(), "Bearer access-token".to_owned())],
&request,
)
.expect("prepare continuation WebSocket request");
let error = execute_response_create_request(&prepared)
.await
@@ -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,
@@ -469,7 +469,7 @@ impl GrokCredentialQuotaService {
if quota_refresh_may_update_state(&loaded.account, &current.account)
&& let Some(availability) = quota_refresh_availability(
current.account.availability(),
quota_is_exhausted(snapshot.billing()),
authoritative_quota_exhaustion(snapshot.billing()),
)
{
self.repository
@@ -1329,6 +1329,13 @@ fn quota_is_exhausted(billing: &GrokBillingPresentation) -> bool {
)
}
/// `None` means the billing document cannot prove either quota state.
fn authoritative_quota_exhaustion(billing: &GrokBillingPresentation) -> Option<bool> {
billing
.has_authoritative_quota()
.then(|| quota_is_exhausted(billing))
}
fn quota_refresh_may_update_state(before: &ProviderAccount, current: &ProviderAccount) -> bool {
current.enabled()
&& current
@@ -1341,22 +1348,26 @@ fn quota_refresh_may_update_state(before: &ProviderAccount, current: &ProviderAc
fn quota_refresh_availability(
current: AccountAvailability,
exhausted: bool,
exhausted: Option<bool>,
) -> Option<GrokCredentialAvailability> {
match current {
AccountAvailability::Invalid
| AccountAvailability::Expired
| AccountAvailability::Banned => None,
AccountAvailability::QuotaExhausted => {
(!exhausted).then_some(GrokCredentialAvailability::Ready)
matches!(exhausted, Some(false)).then_some(GrokCredentialAvailability::Ready)
}
AccountAvailability::Cooldown | AccountAvailability::Unknown => {
exhausted.map(|exhausted| {
if exhausted {
GrokCredentialAvailability::QuotaExhausted
} else {
GrokCredentialAvailability::Ready
}
})
}
AccountAvailability::Cooldown | AccountAvailability::Unknown => Some(if exhausted {
GrokCredentialAvailability::QuotaExhausted
} else {
GrokCredentialAvailability::Ready
}),
AccountAvailability::Ready => {
exhausted.then_some(GrokCredentialAvailability::QuotaExhausted)
matches!(exhausted, Some(true)).then_some(GrokCredentialAvailability::QuotaExhausted)
}
}
}
@@ -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,24 +474,21 @@ 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::ModelQuotaExhausted {
upstream_model,
retry_after,
} => {
GrokCredentialFailure::PaymentRequired { .. } => (
GrokCredentialAvailability::QuotaExhausted,
"upstream_payment_quota_exhausted",
None,
),
GrokCredentialFailure::ModelQuotaExhausted { retry_after, .. } => {
let retry_after =
bounded_cooldown(retry_after, MODEL_QUOTA_COOLDOWN, MAX_MODEL_COOLDOWN);
self.record_model_runtime_cooldown(session, upstream_model, retry_after)
.await;
return;
// xAI 免费模型额度按账号滚动窗口恢复;持久化使页面和各调度节点收敛。
self.record_runtime_cooldown(session, retry_after).await;
(
GrokCredentialAvailability::Cooldown,
"usage_limit_exhausted",
persistent_cooldown_until(observed_at, retry_after),
)
}
GrokCredentialFailure::ModelAccessDenied {
upstream_model,
@@ -452,26 +508,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 +554,6 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
account.availability(),
AccountAvailability::Unknown
| AccountAvailability::Cooldown
| AccountAvailability::QuotaExhausted
| AccountAvailability::Expired
)
{
@@ -540,6 +600,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;
+165 -93
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,17 +47,22 @@ 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,
GrokCompactionSummaryDecoder, GrokCredentialFailure, GrokInferenceChunkStream,
GrokInferenceRequest, GrokInferenceResponse, GrokInferenceTransport,
GrokInferenceTransportError, GrokInferenceTransportErrorKind, GrokInferenceTransportMetrics,
GrokProviderConfigError, GrokRequestEncodeError, GrokResponsesRequest, GrokSessionAffinityKey,
GrokSessionSelection, GrokSessionSelector, GrokSessionSelectorError, SelectedGrokSession,
GrokProviderConfigError, GrokQuotaFailureKind, GrokRequestEncodeError, GrokResponsesRequest,
GrokSessionAffinityKey, GrokSessionSelection, GrokSessionSelector, GrokSessionSelectorError,
SelectedGrokSession, classify_grok_quota_failure,
};
use crate::{GrokCatalogCapabilityEvidence, GrokCatalogModel};
@@ -66,8 +70,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 +84,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 +108,7 @@ impl GrokBuildProvider {
credential_recovery,
account_feedback,
client_identity: GrokClientIdentity::new(),
reasoning_replay: GrokReasoningReplay::new(),
wire_profile,
responses_url,
})
@@ -217,6 +221,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 +280,7 @@ impl GrokBuildProvider {
context,
session: Arc::clone(&selected),
session_capture,
reasoning_replay_capture,
},
);
let stream = ProviderStream::new(metadata, events, selected);
@@ -274,7 +311,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 +320,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 +336,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 +361,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 +770,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 +913,7 @@ struct GrokStreamAttempt {
context: AttemptContext,
session: Arc<SelectedGrokSession>,
session_capture: Option<GrokSessionCapture>,
reasoning_replay_capture: Option<GrokReasoningReplayCapture>,
}
struct GrokCompactionStreamAttempt {
@@ -910,6 +926,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 +1062,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 +1075,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,
@@ -1105,9 +1119,21 @@ fn cold_compaction_http_sse_stream(
.await
.map_err(mark_transient_compaction_failure)?
{
let events = canonical.push(&chunk).map_err(|error| {
mark_transient_compaction_failure(map_continuation_failure(&context, error))
})?;
let events = match canonical.push(&chunk) {
Ok(events) => events,
Err(error) => {
let error = map_continuation_failure(&context, error);
let error = record_stream_failure(
selector.as_ref(),
&session,
error,
&upstream_model,
)
.await;
Err(mark_transient_compaction_failure(error))?;
return;
}
};
for event in events {
summary.observe(&event).map_err(map_compaction_decode_error)?;
facts.observe(&event);
@@ -1118,9 +1144,22 @@ fn cold_compaction_http_sse_stream(
}
if facts.completed.is_none() {
for event in canonical.finish_without_terminal().map_err(|error| {
mark_transient_compaction_failure(map_continuation_failure(&context, error))
})? {
let events = match canonical.finish_without_terminal() {
Ok(events) => events,
Err(error) => {
let error = map_continuation_failure(&context, error);
let error = record_stream_failure(
selector.as_ref(),
&session,
error,
&upstream_model,
)
.await;
Err(mark_transient_compaction_failure(error))?;
return;
}
};
for event in events {
summary.observe(&event).map_err(map_compaction_decode_error)?;
facts.observe(&event);
if facts.completed.is_some() {
@@ -1147,6 +1186,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 +1287,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() {
@@ -1354,7 +1397,13 @@ fn cold_http_sse_stream(
Ok(events) => events,
Err(error) => {
let error = map_continuation_failure(&context, error);
let error = record_stream_failure(selector.as_ref(), &session, error).await;
let error = record_stream_failure(
selector.as_ref(),
&session,
error,
&upstream_model,
)
.await;
Err(error)?;
return;
}
@@ -1364,6 +1413,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;
}
@@ -1379,7 +1431,13 @@ fn cold_http_sse_stream(
Ok(events) => events,
Err(error) => {
let error = map_continuation_failure(&context, error);
let error = record_stream_failure(selector.as_ref(), &session, error).await;
let error = record_stream_failure(
selector.as_ref(),
&session,
error,
&upstream_model,
)
.await;
Err(error)?;
return;
}
@@ -1389,6 +1447,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;
}
@@ -1485,28 +1546,6 @@ fn xai_transport_metadata(
ProviderResponseMetadata::new(serde_json::to_string(&Value::Object(metadata)).ok()?)
}
async fn record_failure(
selector: &dyn GrokSessionSelector,
session: &SelectedGrokSession,
error: ProviderError,
) -> ProviderError {
if !session.allows_account_state_mutation() {
return error;
}
let failure = match error.kind() {
ProviderErrorKind::Unauthorized => Some(GrokCredentialFailure::Unauthorized),
ProviderErrorKind::RateLimited => Some(GrokCredentialFailure::RateLimited {
retry_after: error.retry_after(),
}),
ProviderErrorKind::QuotaExhausted => Some(GrokCredentialFailure::QuotaExhausted),
_ => None,
};
if let Some(failure) = failure {
selector.record_failure(session, failure).await;
}
error
}
async fn record_credential_failure(
selector: &dyn GrokSessionSelector,
session: &SelectedGrokSession,
@@ -1523,17 +1562,13 @@ async fn record_stream_failure(
selector: &dyn GrokSessionSelector,
session: &SelectedGrokSession,
error: ProviderError,
upstream_model: &UpstreamModelId,
) -> ProviderError {
if !session.allows_account_state_mutation() {
return error;
}
if matches!(
error.kind(),
ProviderErrorKind::Unauthorized
| ProviderErrorKind::RateLimited
| ProviderErrorKind::QuotaExhausted
) {
return record_failure(selector, session, error).await;
if let Some(failure) = stream_credential_failure(&error, upstream_model) {
return record_credential_failure(selector, session, error, failure).await;
}
if error.kind() != ProviderErrorKind::Cancelled {
selector
@@ -1543,6 +1578,41 @@ async fn record_stream_failure(
error
}
fn stream_credential_failure(
error: &ProviderError,
upstream_model: &UpstreamModelId,
) -> Option<GrokCredentialFailure> {
match error.kind() {
ProviderErrorKind::Unauthorized => Some(GrokCredentialFailure::Unauthorized),
ProviderErrorKind::RateLimited => Some(GrokCredentialFailure::RateLimited {
retry_after: error.retry_after(),
}),
ProviderErrorKind::QuotaExhausted => {
let visible = error.client_visible_upstream_error();
let quota_kind = classify_grok_quota_failure(
visible
.and_then(ClientVisibleUpstreamError::code)
.or_else(|| error.upstream_code().map(|code| code.as_str())),
visible.and_then(ClientVisibleUpstreamError::error_type),
visible.map(ClientVisibleUpstreamError::message),
);
Some(match quota_kind {
Some(GrokQuotaFailureKind::FreeAccount) => {
GrokCredentialFailure::FreeQuotaExhausted
}
Some(GrokQuotaFailureKind::FreeModelUsage) => {
GrokCredentialFailure::ModelQuotaExhausted {
upstream_model: upstream_model.clone(),
retry_after: error.retry_after(),
}
}
Some(GrokQuotaFailureKind::Account) | None => GrokCredentialFailure::QuotaExhausted,
})
}
_ => None,
}
}
async fn recover_or_record_failure(
selector: &dyn GrokSessionSelector,
recovery: &dyn GrokCredentialRecovery,
@@ -1621,7 +1691,7 @@ async fn map_and_record_stream_transport_failure(
match credential_failure {
Some(failure) => record_credential_failure(selector, session, error, failure).await,
None if request_scoped => error,
None => record_stream_failure(selector, session, error).await,
None => record_stream_failure(selector, session, error, upstream_model).await,
}
}
@@ -1871,7 +1941,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()
}
@@ -19,7 +19,7 @@ use gateway_protocol::openai::sse::{SseEvent, SseEventDecoder};
use serde_json::Value;
use super::request::{GrokResponseTransform, GrokResponsesRequest};
use super::scrub_account_fingerprints;
use super::{classify_grok_quota_failure, scrub_account_fingerprints};
const CONTENTS_PER_OUTPUT: u32 = 1_024;
const LONG_CONTEXT_THRESHOLD: u64 = 200_000;
@@ -785,14 +785,19 @@ fn incomplete_finish_reason(response: &Value) -> FinishReason {
fn upstream_event_error(value: &Value) -> ProviderError {
let code = upstream_error_field(value, "code");
let kind = match code {
Some("invalid_request" | "invalid_prompt") => ProviderErrorKind::InvalidRequest,
Some("unsupported" | "unsupported_feature") => ProviderErrorKind::Unsupported,
Some("unauthorized" | "invalid_token") => ProviderErrorKind::Unauthorized,
Some("permission_denied") => ProviderErrorKind::PermissionDenied,
Some("rate_limit_exceeded") => ProviderErrorKind::RateLimited,
Some("quota_exceeded" | "insufficient_quota") => ProviderErrorKind::QuotaExhausted,
_ => ProviderErrorKind::Unavailable,
let error_type = upstream_error_field(value, "type");
let message = upstream_error_field(value, "message");
let kind = if classify_grok_quota_failure(code, error_type, message).is_some() {
ProviderErrorKind::QuotaExhausted
} else {
match code {
Some("invalid_request" | "invalid_prompt") => ProviderErrorKind::InvalidRequest,
Some("unsupported" | "unsupported_feature") => ProviderErrorKind::Unsupported,
Some("unauthorized" | "invalid_token") => ProviderErrorKind::Unauthorized,
Some("permission_denied") => ProviderErrorKind::PermissionDenied,
Some("rate_limit_exceeded") => ProviderErrorKind::RateLimited,
_ => ProviderErrorKind::Unavailable,
}
};
let mut error = ProviderError::new(kind, UpstreamSendState::Sent)
.redact_sensitive_context("upstream event");
@@ -801,11 +806,11 @@ fn upstream_event_error(value: &Value) -> ProviderError {
}
// 结构化 message/code/type 供原客户端展示与重试分类;message 先脱去
// 账号指纹(上游限流文案内嵌 team UUID),非结构化正文不透出。
if let Some(message) = upstream_error_field(value, "message") {
if let Some(message) = message {
error = error.with_client_visible_upstream_error(ClientVisibleUpstreamError::new(
scrub_account_fingerprints(message),
code.map(str::to_owned),
upstream_error_field(value, "type").map(str::to_owned),
error_type.map(str::to_owned),
));
}
error
@@ -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,85 @@ pub use transport::{
};
const UUID_TEXT_LEN: usize = 36;
const FREE_USAGE_SIGNALS: &[&str] = &[
"subscription:free-usage-exhausted",
"subscription_free_usage_exhausted",
"free-usage-exhausted",
"free_usage_exhausted",
"used all the included free usage",
"used all your free usage",
];
const ACCOUNT_QUOTA_SIGNALS: &[&str] = &[
"personal-team-blocked:spending-limit",
"personal_team_blocked_spending_limit",
"quota_exceeded",
"insufficient_quota",
];
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum GrokQuotaFailureKind {
Account,
FreeAccount,
FreeModelUsage,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum GrokQuotaSignal {
Account,
FreeUsage,
}
pub(crate) fn classify_grok_quota_failure(
code: Option<&str>,
error_type: Option<&str>,
message: Option<&str>,
) -> Option<GrokQuotaFailureKind> {
let code = code.map(str::to_ascii_lowercase);
let error_type = error_type.map(str::to_ascii_lowercase);
let message = message.map(str::to_ascii_lowercase);
// `code` and `type` are stable upstream fields. Do not let a conflicting
// human-readable message override them; use the message only as a
// fallback or to refine a confirmed free-usage signal's rolling window.
let structured_signal = code
.as_deref()
.and_then(classify_quota_signal)
.or_else(|| error_type.as_deref().and_then(classify_quota_signal));
match structured_signal {
Some(GrokQuotaSignal::Account) => Some(GrokQuotaFailureKind::Account),
Some(GrokQuotaSignal::FreeUsage) => Some(free_usage_failure_kind(message.as_deref())),
None => classify_message_quota_failure(message.as_deref()),
}
}
fn classify_quota_signal(value: &str) -> Option<GrokQuotaSignal> {
if contains_any(value, ACCOUNT_QUOTA_SIGNALS) {
Some(GrokQuotaSignal::Account)
} else if contains_any(value, FREE_USAGE_SIGNALS) {
Some(GrokQuotaSignal::FreeUsage)
} else {
None
}
}
fn classify_message_quota_failure(message: Option<&str>) -> Option<GrokQuotaFailureKind> {
match message.and_then(classify_quota_signal) {
Some(GrokQuotaSignal::Account) => Some(GrokQuotaFailureKind::Account),
Some(GrokQuotaSignal::FreeUsage) => Some(free_usage_failure_kind(message)),
None => None,
}
}
fn free_usage_failure_kind(message: Option<&str>) -> GrokQuotaFailureKind {
if message.is_some_and(|value| value.contains("for model")) {
GrokQuotaFailureKind::FreeModelUsage
} else {
GrokQuotaFailureKind::FreeAccount
}
}
fn contains_any(value: &str, needles: &[&str]) -> bool {
needles.iter().any(|needle| value.contains(needle))
}
/// 把 message 中的 UUID 替换为占位符,控制字符归一为空格,
/// 使文本满足客户端可见错误的约束且不携带可定位账号的标识。
@@ -33,7 +33,7 @@ use gateway_core::engine::UpstreamSendState;
use gateway_core::error::{ClientVisibleUpstreamError, OpaqueUpstreamValue};
use gateway_core::event::UpstreamHttpVersion;
use super::scrub_account_fingerprints;
use super::{GrokQuotaFailureKind, classify_grok_quota_failure, scrub_account_fingerprints};
pub(crate) const OFFICIAL_OAUTH_HOST: &str = "auth.x.ai";
const OFFICIAL_INFERENCE_HOST: &str = "cli-chat-proxy.grok.com";
@@ -1082,15 +1082,10 @@ fn first_string(object: &serde_json::Map<String, Value>, fields: &[&str]) -> Opt
}
fn classify_payment_required(metadata: &InferenceErrorMetadata) -> GrokInferenceTransportErrorKind {
if model_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::ModelQuotaExhausted
} else if paid_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::QuotaExhausted
} else if free_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::FreeQuotaExhausted
} else {
GrokInferenceTransportErrorKind::PaymentRequired
}
quota_failure(metadata).map_or(
GrokInferenceTransportErrorKind::PaymentRequired,
quota_transport_error_kind,
)
}
fn classify_forbidden(
@@ -1103,12 +1098,8 @@ fn classify_forbidden(
GrokInferenceTransportErrorKind::SafetyRejected
} else if credential_rejected(&text) {
GrokInferenceTransportErrorKind::Unauthorized
} else if model_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::ModelQuotaExhausted
} else if paid_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::QuotaExhausted
} else if free_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::FreeQuotaExhausted
} else if let Some(failure) = quota_failure(metadata) {
quota_transport_error_kind(failure)
} else if model_access_denied(metadata) {
GrokInferenceTransportErrorKind::ModelAccessDenied
} else {
@@ -1119,15 +1110,10 @@ fn classify_forbidden(
fn classify_too_many_requests(
metadata: &InferenceErrorMetadata,
) -> GrokInferenceTransportErrorKind {
if model_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::ModelQuotaExhausted
} else if paid_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::QuotaExhausted
} else if free_quota_exhausted(metadata) {
GrokInferenceTransportErrorKind::FreeQuotaExhausted
} else {
GrokInferenceTransportErrorKind::RateLimited
}
quota_failure(metadata).map_or(
GrokInferenceTransportErrorKind::RateLimited,
quota_transport_error_kind,
)
}
fn reasoning_decode_failed(metadata: &InferenceErrorMetadata) -> bool {
@@ -1137,32 +1123,22 @@ fn reasoning_decode_failed(metadata: &InferenceErrorMetadata) -> bool {
})
}
fn paid_quota_exhausted(metadata: &InferenceErrorMetadata) -> bool {
contains_any(
&inference_metadata_text(metadata),
&[
"personal-team-blocked:spending-limit",
"personal_team_blocked_spending_limit",
],
fn quota_failure(metadata: &InferenceErrorMetadata) -> Option<GrokQuotaFailureKind> {
classify_grok_quota_failure(
metadata.code.as_deref(),
metadata.error_type.as_deref(),
metadata.message.as_deref(),
)
}
fn free_quota_exhausted(metadata: &InferenceErrorMetadata) -> bool {
contains_any(
&inference_metadata_text(metadata),
&[
"subscription:free-usage-exhausted",
"subscription_free_usage_exhausted",
"free-usage-exhausted",
"free_usage_exhausted",
"used all the included free usage",
"used all your free usage",
],
)
}
fn model_quota_exhausted(metadata: &InferenceErrorMetadata) -> bool {
inference_metadata_text(metadata).contains("used all the included free usage for model")
fn quota_transport_error_kind(failure: GrokQuotaFailureKind) -> GrokInferenceTransportErrorKind {
match failure {
GrokQuotaFailureKind::Account => GrokInferenceTransportErrorKind::QuotaExhausted,
GrokQuotaFailureKind::FreeAccount => GrokInferenceTransportErrorKind::FreeQuotaExhausted,
GrokQuotaFailureKind::FreeModelUsage => {
GrokInferenceTransportErrorKind::ModelQuotaExhausted
}
}
}
fn model_access_denied(metadata: &InferenceErrorMetadata) -> bool {
@@ -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()
@@ -319,7 +319,7 @@ pub enum GrokCredentialFailure {
/// 已解析并脱敏的延迟;缺失时由 selector 使用短暂默认值。
retry_after: Option<Duration>,
},
/// 选中账号只耗尽了指定模型的额度。
/// 上游声明指定模型的免费用量耗尽;selector 会收敛为账号级可恢复额度状态。
ModelQuotaExhausted {
upstream_model: UpstreamModelId,
retry_after: Option<Duration>,
@@ -282,7 +282,7 @@ pub enum GrokInferenceTransportErrorKind {
QuotaExhausted,
/// 账号订阅级免费额度耗尽。
FreeQuotaExhausted,
/// 当前账号只耗尽了指定模型的免费额度。
/// 上游声明指定模型的免费用量耗尽;selector 会收敛为账号级可恢复额度状态。
ModelQuotaExhausted,
/// 当前账号只缺少指定模型的访问权限。
ModelAccessDenied,
@@ -874,17 +874,24 @@ async fn quota_projection_preserves_unknown_period_for_dynamic_duration_fallback
}
#[tokio::test]
async fn weekly_period_without_reported_allowance_is_not_authoritative_quota() {
let (_, repository) = repository_with_accounts(&[("free-quota", "subject-free-quota")]).await;
async fn non_authoritative_quota_refresh_does_not_clear_existing_quota_exhaustion() {
let (store, repository) =
repository_with_accounts(&[("free-quota", "subject-free-quota")]).await;
let id = account_id("free-quota");
set_account_state(&store, &id, AccountAvailability::QuotaExhausted, None).await;
let transport = QueueBillingTransport::success(
br#"{"config":{"currentPeriod":{"type":"USAGE_PERIOD_TYPE_WEEKLY","start":"2026-07-15T00:00:00Z","end":"2026-07-22T00:00:00Z"},"onDemandCap":{"val":0},"onDemandUsed":{"val":0},"prepaidBalance":{"val":0}}}"#,
);
let snapshot = crate::support::grok_quota_service(repository, transport)
.refresh_account(&account_id("free-quota"))
.refresh_account(&id)
.await
.expect("refresh Free quota");
assert!(!snapshot.billing().has_authoritative_quota());
assert_eq!(
store.account(&id).expect("account").availability(),
AccountAvailability::QuotaExhausted
);
}
#[tokio::test]
@@ -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()))
@@ -449,7 +490,7 @@ async fn unknown_payment_required_feedback_only_records_runtime_account_cooldown
}
#[tokio::test]
async fn model_quota_feedback_blocks_only_the_failed_account_and_model() {
async fn free_model_quota_feedback_persists_resettable_account_exhaustion() {
let fixture = SelectorFixture::new(&["model-quota", "available"]).await;
let failed_model = UpstreamModelId::new("grok-4.5").expect("failed model");
let session = fixture
@@ -464,38 +505,47 @@ async fn model_quota_feedback_blocks_only_the_failed_account_and_model() {
&session,
GrokCredentialFailure::ModelQuotaExhausted {
upstream_model: failed_model.clone(),
retry_after: Some(Duration::from_secs(5)),
retry_after: None,
},
)
.await;
assert_eq!(
fixture
.store
.account(&selected)
.expect("selected account")
.availability(),
AccountAvailability::Ready
let account = fixture.store.account(&selected).expect("selected account");
assert_eq!(account.availability(), AccountAvailability::Cooldown);
let minimum_until = SystemTime::now() + Duration::from_secs(23 * 60 * 60);
assert!(
account
.cooldown_until()
.is_some_and(|until| until > minimum_until)
);
let scope = ProviderCooldownScope::upstream_model(failed_model);
assert!(
fixture
.cooldowns
.scoped_cooldown(&selected, &scope)
.is_some()
.is_none()
);
let unrelated_model = fixture
assert!(
fixture
.cooldowns
.cooldown(&selected)
.is_some_and(|cooldown| cooldown.until() > minimum_until)
);
assert!(matches!(
fixture
.selector
.select(fixture.request_for_model("grok-4.6", Some(selected.clone())))
.await,
Err(GrokSessionSelectorError::AccountCoolingDown {
retry_after: Some(_)
})
));
let next = fixture
.selector
.select(fixture.request_for_model("grok-4.6", Some(selected.clone())))
.select(fixture.request_for_model("grok-4.6", None))
.await
.expect("same account remains valid for another model");
assert_eq!(unrelated_model.account_id(), &selected);
let failed_model_retry = fixture
.selector
.select(fixture.request_for_model("grok-4.5", None))
.await
.expect("another account serves the failed model");
assert_ne!(failed_model_retry.account_id(), &selected);
.expect("another account remains schedulable");
assert_ne!(next.account_id(), &selected);
}
#[tokio::test]
@@ -538,7 +588,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 +604,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 +650,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(
@@ -1360,6 +1518,73 @@ async fn model_quota_transport_failure_preserves_model_scoped_feedback() {
);
}
#[tokio::test]
async fn free_quota_sse_failure_marks_the_account_exhausted() {
let selector = StubSelector::success();
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_free_quota\",\"model\":\"grok-4.5\"}}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"error\":{\"code\":\"subscription:free-usage-exhausted\",\"message\":\"You have used all your free usage\"}}\n\n",
)
.as_bytes()
.to_vec();
let transport = StubInferenceTransport::sequence([InferenceMode::SuccessBody(body)]);
let provider = provider(selector.clone(), transport).await;
let mut stream = provider
.execute(
provider_request("xai"),
context(CancellationToken::new(), None),
)
.await
.expect("stream");
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::QuotaExhausted);
assert_eq!(
selector.feedback.lock().expect("feedback").as_slice(),
&[GrokCredentialFailure::FreeQuotaExhausted]
);
}
#[tokio::test]
async fn free_model_quota_sse_failure_is_classified_for_resettable_account_cooldown() {
let selector = StubSelector::success();
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_model_quota\",\"model\":\"grok-4.5\"}}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"error\":{\"status\":null,\"contentType\":null,\"body\":null,\"code\":\"subscription_free_usage_exhausted\",\"type\":null,\"message\":\"You've used all the included free usage for model grok-4.5-0722 for now. Usage resets over a rolling 24-hour window \u{2014} tokens (actual/limit): 500505/500000. Upgrade to a Grok subscription for higher limits: https://grok.com/supergrok\"}}\n\n",
)
.as_bytes()
.to_vec();
let transport = StubInferenceTransport::sequence([InferenceMode::SuccessBody(body)]);
let provider = provider(selector.clone(), transport).await;
let mut stream = provider
.execute(
provider_request("xai"),
context(CancellationToken::new(), None),
)
.await
.expect("stream");
let error = next_provider_error(&mut stream).await;
assert_eq!(error.kind(), ProviderErrorKind::QuotaExhausted);
assert_eq!(
error.upstream_code().map(|code| code.as_str()),
Some("subscription_free_usage_exhausted")
);
assert_eq!(
selector.feedback.lock().expect("feedback").as_slice(),
&[GrokCredentialFailure::ModelQuotaExhausted {
upstream_model: UpstreamModelId::new("grok-4.5").expect("model"),
retry_after: None,
}]
);
}
#[tokio::test]
async fn model_access_denial_is_retryable_without_credential_recovery() {
let selector = StubSelector::success();
@@ -1750,6 +1975,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!(
@@ -732,6 +802,28 @@ fn decoder_should_classify_failed_event_without_retaining_body() {
assert!(!format!("{error:?}").contains("secret"));
}
#[test]
fn decoder_should_classify_free_quota_failed_event() {
let body = concat!(
"event: response.created\n",
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_quota\"}}\n\n",
"event: response.failed\n",
"data: {\"type\":\"response.failed\",\"error\":{\"code\":\"subscription:free-usage-exhausted\",\"type\":\"billing_error\",\"message\":\"You have used all your free usage\"}}\n\n",
);
let error = GrokCanonicalDecoder::new("fallback")
.push(body.as_bytes())
.expect_err("free quota failure must surface");
assert_eq!(
error.kind(),
gateway_core::error::ProviderErrorKind::QuotaExhausted
);
assert_eq!(
error.upstream_code().map(|code| code.as_str()),
Some("subscription:free-usage-exhausted")
);
}
#[test]
fn decoder_should_scrub_account_fingerprints_from_failure_messages() {
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);
}
@@ -471,8 +471,12 @@ async fn inference_transport_should_classify_model_quota_for_403_and_429() {
.and(path("/v1/responses"))
.respond_with(ResponseTemplate::new(status).set_body_json(json!({
"error": {
"code": "subscription:free-usage-exhausted",
"message": "You have used all the included free usage for model grok-4.5"
"status": null,
"contentType": null,
"body": null,
"code": "subscription_free_usage_exhausted",
"type": null,
"message": "You've used all the included free usage for model grok-4.5-0722 for now. Usage resets over a rolling 24-hour window - tokens (actual/limit): 500505/500000. Upgrade to a Grok subscription for higher limits: https://grok.com/supergrok"
}
})))
.expect(1)
@@ -491,6 +495,38 @@ async fn inference_transport_should_classify_model_quota_for_403_and_429() {
}
}
#[tokio::test]
async fn inference_transport_should_prefer_structured_quota_fields_over_conflicting_message() {
for (code, error_type) in [
("quota_exceeded", Some("subscription_free_usage_exhausted")),
("rate_limit_exceeded", Some("quota_exceeded")),
] {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/v1/responses"))
.respond_with(ResponseTemplate::new(429).set_body_json(json!({
"error": {
"code": code,
"type": error_type,
"message": "You've used all the included free usage for model grok-4.5-0722 for now."
}
})))
.expect(1)
.mount(&server)
.await;
let origin = Url::parse(&server.uri()).expect("wiremock origin");
let error = inference_transport(&origin)
.execute(inference_request(&origin))
.await
.expect_err("quota response must fail");
assert_eq!(
error.kind(),
GrokInferenceTransportErrorKind::QuotaExhausted
);
}
}
#[tokio::test]
async fn inference_transport_should_classify_free_usage_429_as_account_free_quota_exhausted() {
let server = MockServer::start().await;
@@ -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 -4
View File
@@ -1,4 +1 @@
0554fed953bd1440b50bc4033dc8bad96b51e132c57c2ccf1553728a991d5cb2 0001_initial.sql
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
a4681a5c040e80c5234e879293fc34972080e5251dfbc7c40e1374f8e4170a4c 0001_initial.sql
+20 -7
View File
@@ -166,7 +166,14 @@ create table provider_accounts (
)
),
constraint provider_accounts_cooldown_ck check (
(availability = 'cooldown') = (cooldown_until is not null)
(
availability != 'cooldown'
or cooldown_until is not null
)
and (
cooldown_until is null
or availability in ('cooldown', 'quota_exhausted')
)
),
constraint provider_accounts_time_ck check (
created_at <= updated_at
@@ -208,6 +215,7 @@ create table model_requests (
requested_model_id text not null,
provider_kind text,
upstream_model_id text,
service_tier text,
provider_account_id text,
provider_account_ref text,
upstream_transport text,
@@ -220,9 +228,11 @@ create table model_requests (
outcome text not null default 'running',
client_status_code integer,
upstream_status_code integer,
client_response_id text,
-- Provider response IDs are opaque JSON strings. Bytes preserve embedded NUL
-- and avoid imposing PostgreSQL text-index limits on protocol handles.
client_response_id bytea,
upstream_request_id text,
upstream_response_id text,
upstream_response_id bytea,
error_kind text,
provider_error_code text,
error_message text,
@@ -321,6 +331,13 @@ create table model_requests (
provider_observation_json is null
or jsonb_typeof(provider_observation_json) = 'object'
),
constraint model_requests_service_tier_ck check (
service_tier is null
or (
octet_length(service_tier) between 1 and 64
and service_tier !~ '[[:cntrl:]]'
)
),
constraint model_requests_cost_ck check (
cost_source in ('provider_reported', 'calculated', 'unavailable')
and (
@@ -414,10 +431,6 @@ create index model_requests_running_deadline_idx
create index model_requests_retention_idx
on model_requests (completed_at)
where outcome <> 'running';
create unique index model_requests_client_response_uq
on model_requests (client_response_id)
where client_response_id is not null;
create table ops_events (
id text primary key,
model_request_id text,
@@ -1,3 +0,0 @@
-- Provider response IDs are opaque protocol values, not relational identifiers.
-- Indexing the raw text makes PostgreSQL reject otherwise valid long handles.
drop index model_requests_client_response_uq;
@@ -1,7 +0,0 @@
-- Response IDs are opaque JSON strings. PostgreSQL text rejects NUL even though JSON accepts it.
-- Store their UTF-8 bytes without an index; Rust decodes them only at the audit/API boundary.
alter table model_requests
alter column client_response_id type bytea
using convert_to(client_response_id, 'UTF8'),
alter column upstream_response_id type bytea
using convert_to(upstream_response_id, 'UTF8');
@@ -1,9 +0,0 @@
alter table model_requests
add column service_tier text,
add constraint model_requests_service_tier_ck check (
service_tier is null
or (
octet_length(service_tier) between 1 and 64
and service_tier !~ '[[:cntrl:]]'
)
);
+7 -5
View File
@@ -6,14 +6,16 @@ sqlx 在服务启动时(`serve` 监听之前)按编号顺序执行本目录
## 冻结规则
- 已合入 main 的迁移文件**永久冻结**,一个字节都不能改——包括注释和空白。
任何改动都会让所有已部署实例(含自更新链路上的存量用户)在下次重启时
因 checksum 失配而无法启动,且没有自动恢复路径。
- schema 变更一律新增编号迁移(`0003_...`、`0004_...`),永远不回改旧文件。
- 同一大版本内,已合入 main 的迁移文件**永久冻结**,一个字节都不能改——包括注释和
空白。任何改动都会让该大版本的已部署实例在下次重启时因 checksum 失配而无法启动,
且没有自动恢复路径。
- 同一大版本内的 schema 变更一律新增编号迁移(`0002_...`、`0003_...`),永远不回改
旧文件。切换大版本时允许把已完成升级的 schema 折叠成新的 `0001` 基线,因为跨大版本
不支持在线升级,目标版本必须使用全新数据目录部署。
- `.frozen-sha256` 是冻结清单(`sha256sum` 格式)。CI 会校验:清单内文件
字节未变,且目录里每个 `*.sql` 都已入册。新增迁移时在同一提交里执行
`sha256sum 000N_xxx.sql >> .frozen-sha256` 入册。
- 若确需修正已冻结迁移里的错误,用新迁移做补偿性变更(`alter` / 回填),
- 若同一大版本内确需修正已冻结迁移里的错误,用新迁移做补偿性变更(`alter` / 回填),
不要动原文件。
## 本地测试库
+2 -4
View File
@@ -130,10 +130,8 @@ Provider schema 以明文 JSON 保存于 PostgreSQL。Redis 只保存可重建
## 镜像升级与源码构建
> [!WARNING]
> 以下命令只适用于同一大版本内的升级。数据库 schema 与配置只在同一大版本内兼容,
> 跨大版本直接换镜像会导致服务无法启动;内置一键自更新会拒绝跨大版本目标(3.0 之前
> 的旧版没有这道门禁,不要用旧版界面的一键更新跨大版本)。跨大版本请全新部署(全新
> `.runtime/` 数据目录),并重新导入或重新授权 Provider 账号与客户端 Key。
> 以下命令只适用于同一大版本内的升级,不支持跨大版本在线升级。跨大版本请使用全新的
> `.runtime/` 数据目录重新部署,并重新导入或重新授权 Provider 账号与客户端 Key。
```bash
docker compose -f deploy/compose.yaml build codex-proxy-rs
+5 -5
View File
@@ -170,11 +170,11 @@ Admin API 不要求客户端提交配置 revision,也不向客户端暴露配
客户端可见协议结果。队列不是第二份业务权威,也不做无幂等键的盲目重试。
- retention 只删除已满足保留规则的历史事实,不改变运行中请求。
业务 schema 由 `backend/migrations/` 下的编号迁移按序定义:`0001_initial.sql` 建表,`0002` 删除
opaque response ID 唯一索引,`0003` 将 client/upstream response ID 改为无格式假设的 UTF-8 bytes,
`0004` 增加响应观测得到的 `model_requests.service_tier`。已合入的迁移永久冻结——
`.frozen-sha256` 是冻结清单,CI 校验清单相对 PR base 只增不改,启动时 sqlx 重校验已应用迁移的
checksum,不一致直接拒绝启动。规则见 `backend/migrations/README.md`。
业务 schema 由 `backend/migrations/0001_initial.sql` 直接定义当前大版本的完整新库基线,其中
response ID 使用无格式假设的 UTF-8 bytes,`model_requests.service_tier` 保存响应观测到的服务档位,
且额度耗尽账号可以携带冷却截止时间。同一大版本内已合入的迁移永久冻结——
`.frozen-sha256` 是冻结清单,CI 校验清单相对 PR base 只增不改,启动时 sqlx 重校验已应用
迁移的 checksum,不一致直接拒绝启动。规则见 `backend/migrations/README.md`。
## 9. Redis 终态
@@ -193,19 +193,6 @@ function connectionLogClass(tone: string) {
</div>
</BaseScrollbar>
</div>
<div v-if="error" class="mt-3 rounded-lg bg-(--cp-danger-bg) px-3 py-2.5">
<p class="m-0 text-[11px] font-heavy text-(--cp-danger-text)">
错误信息
</p>
<BaseScrollbar max-height="118px" view-class="pt-1.5 pr-2">
<p
class="m-0 wrap-break-word text-[12px] leading-[1.55] font-emphasis text-(--cp-danger-text)"
>
{{ error }}
</p>
</BaseScrollbar>
</div>
</section>
</div>
@@ -33,7 +33,7 @@ const overviewItems = computed(() => [
icon: Gauge,
},
{
label: '不可用账号',
label: '异常账号',
value: formatCount(props.summary.unavailable),
caption: '过期 / 禁用 / 封禁',
tone: 'danger',
@@ -280,17 +280,7 @@ export function useAccountConnectionTest(options: AccountConnectionTestOptions)
if (event.type === 'error') {
applyAccountStatus(event.accountStatus)
connectionTestError.value = connectionTestErrorText(event)
const detail = event.upstreamStatus || event.upstreamContentType || event.upstreamBody
|| event.providerErrorCode || event.providerErrorType
? {
status: event.upstreamStatus,
contentType: event.upstreamContentType,
body: event.upstreamBody,
code: event.providerErrorCode,
type: event.providerErrorType,
}
: undefined
appendConnectionTestLog(connectionTestError.value, 'danger', detail)
appendConnectionTestLog('上游返回错误', 'danger', connectionTestError.value)
finishConnectionTest('error')
clearConnectionTestRun()
}
@@ -148,15 +148,20 @@ export function useAccountOnboarding(options: {
}
async function importAccountDocument() {
const data = parseImportJson(createForm.value.importText)
if (Array.isArray(data) || typeof data !== 'object' || data === null)
throw new Error('导入文件必须是 JSON object')
const result = await importAccounts({
provider: createForm.value.provider,
data,
})
const provider = createForm.value.provider === 'xai' ? 'xAI' : 'OpenAI'
return `${provider} 账号已导入 ${result.importedCount} 个`
const provider = requireImportProvider(createForm.value.provider)
const documents = providerImportDocuments(
parseImportJson(createForm.value.importText),
provider,
)
let importedCount = 0
for (const entry of documents) {
const result = await importAccounts({
provider,
data: entry.document,
})
importedCount += result.importedCount
}
return `${providerDisplayName(provider) ?? provider} 账号已导入 ${importedCount} 个`
}
async function importMixedAccountDocument() {
@@ -259,6 +264,27 @@ function parseImportJson(value: string) {
}
}
function requireImportProvider(value: string): ImportProvider {
if (value === 'openai' || value === 'xai')
return value
throw new Error('请选择要导入的账号平台')
}
function providerImportDocuments(value: unknown, provider: ImportProvider): MixedImportDocument[] {
if (isRecord(value) && Array.isArray(value.documents)) {
const documents = parseMixedImportDocuments(value)
.filter(entry => entry.provider === provider)
if (documents.length === 0) {
const label = providerDisplayName(provider) ?? provider
throw new Error(`批量导入文件不包含 ${label} 账号文档`)
}
return documents
}
if (!isRecord(value))
throw new Error('导入文件必须是 JSON object')
return [{ provider, document: value }]
}
function parseMixedImportDocuments(value: unknown): MixedImportDocument[] {
if (!isRecord(value) || !Array.isArray(value.documents))
throw new Error('批量导入文件必须是 CPR 多平台导出文件')
@@ -56,17 +56,17 @@ function itemValueClass(tone?: string, accent?: boolean) {
<UsageDetailPopover v-if="billing" title="计费明细" width="248px" trigger-label="查看费用明细">
<div class="grid gap-1.5 text-(--cp-text-secondary)">
<div v-for="item in amountItems" :key="item.label" class="flex justify-between gap-4">
<span>{{ item.label }}</span>
<span class="font-mono font-heavy" :class="itemValueClass(undefined, item.accent)">
<div v-for="item in amountItems" :key="item.label" class="grid grid-cols-[auto_minmax(0,1fr)] items-center gap-4">
<span class="whitespace-nowrap">{{ item.label }}</span>
<span class="justify-self-end whitespace-nowrap font-mono font-heavy" :class="itemValueClass(undefined, item.accent)">
{{ item.value }}
</span>
</div>
</div>
<div class="mt-1 grid gap-1.5 rounded-(--cp-input-radius-base) bg-(--cp-bg-subtle) p-2 text-(--cp-text-secondary)">
<div v-for="item in billingItems" :key="item.label" class="flex justify-between gap-4">
<span>{{ item.label }}</span>
<span class="font-mono font-heavy" :class="itemValueClass(item.tone)">
<div v-for="item in billingItems" :key="item.label" class="grid grid-cols-[auto_minmax(0,1fr)] items-center gap-4">
<span class="whitespace-nowrap">{{ item.label }}</span>
<span class="justify-self-end whitespace-nowrap font-mono font-heavy" :class="itemValueClass(item.tone)">
{{ item.value }}
</span>
</div>
+1 -1
View File
@@ -1 +1 @@
version: "3.0.22"
version: "3.0.32"
@@ -40,9 +40,9 @@
## 存储
- PostgreSQL 业务表只有 `0001_initial.sql` 定义的七张;后续编号迁移只做增量。当前 `0002` 删除
opaque response ID 索引,`0003` 把 response ID 改为 bytes,`0004` 增加响应观测的 service tier。
- 已应用迁移按字节冻结:`backend/migrations/.frozen-sha256` 是冻结清单,CI 相对 PR base 做 append-only 校验;schema 变更一律新增编号迁移,规则见 `backend/migrations/README.md`。
- PostgreSQL 业务表只有 `0001_initial.sql` 定义的七张;该文件直接表达当前大版本的完整新库基线,
包括 bytes response ID、响应观测 service tier 与额度耗尽冷却约束。
- 同一大版本内已应用迁移按字节冻结:`backend/migrations/.frozen-sha256` 是冻结清单,CI 相对 PR base 做 append-only 校验;schema 变更一律新增编号迁移,规则见 `backend/migrations/README.md`。
- `config_revision` 只用于会改变调度快照或安全配置的管理 mutation。
- quota、cooldown、catalog generation、自动 refresh 不推进全局 revision。
- refresh 只推进账号 `credential_revision`;Redis cooldown 是可丢失热缓存。