mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-08 16:05:33 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a7ea957843 | ||
|
|
c92c95a733 | ||
|
|
ae7608440e | ||
|
|
5c9da8a39d | ||
|
|
e012f24ef2 | ||
|
|
756c9c7158 | ||
|
|
ab4bf75d0e | ||
|
|
7a9715ed41 | ||
|
|
c353a6159e | ||
|
|
f5dca4848a | ||
|
|
36321a9af4 | ||
|
|
b485b8bd4a | ||
|
|
cae2c8656a | ||
|
|
acf574129f | ||
|
|
6426e0d3de | ||
|
|
eedc39b33a | ||
|
|
9672a2e000 | ||
|
|
0fffd2c550 | ||
|
|
5671a57dc8 | ||
|
|
166e8754a2 | ||
|
|
fbaf6c1bdd | ||
|
|
7cdc2dba6f | ||
|
|
a321ccf737 | ||
|
|
4533ae8850 | ||
|
|
7542ed365e | ||
|
|
bf4da7a3e7 | ||
|
|
0fe354b083 |
@@ -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
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"
|
||||
}),
|
||||
]
|
||||
);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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(¤t.credential)
|
||||
.map_err(|_| CodexOAuthAdminError::Credential)?;
|
||||
if runtime.authentication.oauth().is_none() {
|
||||
return Err(CodexOAuthAdminError::Credential);
|
||||
}
|
||||
let account_id = current
|
||||
.account
|
||||
.upstream_account_id()
|
||||
.ok_or(CodexOAuthAdminError::Credential)?;
|
||||
CodexIdentityExpectation::reauthorization(
|
||||
account_id.to_owned(),
|
||||
current.account.upstream_user_id().to_owned(),
|
||||
runtime.installation_id,
|
||||
)
|
||||
.map_err(|_| CodexOAuthAdminError::Credential)
|
||||
}
|
||||
|
||||
fn authorization_url(
|
||||
pending: &CodexPendingAuthorization,
|
||||
oauth_client_id: &str,
|
||||
|
||||
@@ -30,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(¤t, observed_at).await,
|
||||
Ok(true)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
let snapshot = match self.read_snapshot_for(¤t).await {
|
||||
Ok(Some(snapshot))
|
||||
if quota_fact_is_active_exhaustion(snapshot.fact(), observed_at) =>
|
||||
{
|
||||
snapshot
|
||||
}
|
||||
_ => continue,
|
||||
};
|
||||
let cooldown_until =
|
||||
quota_reset_at_from_snapshot(Some(&snapshot)).or_else(|| current.cooldown_until());
|
||||
if let Err(error) = self
|
||||
.repository
|
||||
.apply_state(
|
||||
¤t,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
Some("quota_exhausted".to_owned()),
|
||||
None,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.await
|
||||
@@ -881,8 +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 {
|
||||
¤t,
|
||||
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(
|
||||
¤t,
|
||||
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(¤t, observed_at)
|
||||
.await
|
||||
.map_err(|_| CredentialSelectionError::Store)?;
|
||||
let (availability, reason, cooldown_until) = if quota_exhausted {
|
||||
(
|
||||
AccountAvailability::QuotaExhausted,
|
||||
Some("quota_exhausted".to_owned()),
|
||||
None,
|
||||
)
|
||||
} else {
|
||||
(
|
||||
AccountAvailability::Cooldown,
|
||||
Some(reason.to_owned()),
|
||||
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)),
|
||||
)
|
||||
let cooldown_until =
|
||||
observed_at.checked_add(retry_after.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN));
|
||||
let cooldown_until = match (current.availability(), current.cooldown_until()) {
|
||||
(AccountAvailability::QuotaExhausted, Some(existing)) => {
|
||||
cooldown_until.map_or(Some(existing), |until| Some(existing.max(until)))
|
||||
}
|
||||
_ => cooldown_until,
|
||||
};
|
||||
let reason = if current.availability() == AccountAvailability::QuotaExhausted
|
||||
&& current.cooldown_until().is_none()
|
||||
{
|
||||
Some("quota_exhausted".to_owned())
|
||||
} else {
|
||||
Some(reason.to_owned())
|
||||
};
|
||||
// 一次真实 429 证明上游已经接受当前凭据,可纠正旧的认证/封禁终态。
|
||||
// resettable 用量限制沿用 v2 的 quota lock 语义,并保留/延长复核时间。
|
||||
self.repository
|
||||
.apply_state(¤t, availability, reason, cooldown_until, observed_at)
|
||||
.apply_state(
|
||||
¤t,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
reason,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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, ¤t.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,
|
||||
|
||||
@@ -4,6 +4,7 @@ mod admin;
|
||||
pub mod config;
|
||||
pub mod credential;
|
||||
mod provider;
|
||||
mod reasoning_replay;
|
||||
pub mod transport;
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
@@ -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 +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
|
||||
|
||||
@@ -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:]]'
|
||||
)
|
||||
);
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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 @@
|
||||
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 是可丢失热缓存。
|
||||
|
||||
Reference in New Issue
Block a user