mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-08 16:05:33 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5671a57dc8 | ||
|
|
166e8754a2 | ||
|
|
fbaf6c1bdd | ||
|
|
7cdc2dba6f | ||
|
|
a321ccf737 | ||
|
|
4533ae8850 |
@@ -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)) = (
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -392,6 +392,12 @@ pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(Strin
|
||||
&format!("{prefix}-secondary"),
|
||||
details.secondary,
|
||||
);
|
||||
if let Some(allowed) = details.allowed {
|
||||
headers.push((format!("{prefix}-allowed"), allowed.to_string()));
|
||||
}
|
||||
if let Some(limit_reached) = details.limit_reached {
|
||||
headers.push((format!("{prefix}-limit-reached"), limit_reached.to_string()));
|
||||
}
|
||||
if let Some(limit_name) = &details.limit_name {
|
||||
headers.push((format!("{prefix}-limit-name"), limit_name.clone()));
|
||||
}
|
||||
|
||||
@@ -440,6 +440,8 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
|
||||
"metered_limit_name": "codex-other",
|
||||
"limit_name": "Other requests",
|
||||
"rate_limits": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary": {
|
||||
"used_percent": 33,
|
||||
"window_minutes": 60,
|
||||
@@ -450,7 +452,8 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
|
||||
"has_credits": true,
|
||||
"unlimited": false,
|
||||
"balance": "9.25"
|
||||
}
|
||||
},
|
||||
"rate_limit_reached_type": "usage_limit_reached"
|
||||
});
|
||||
|
||||
let parsed = parse_rate_limits_event(&event).expect("rate-limit event");
|
||||
@@ -460,16 +463,22 @@ fn parse_rate_limits_event_should_extract_websocket_metered_limit() {
|
||||
(
|
||||
parsed.active_limit.as_deref(),
|
||||
details.limit_name.as_deref(),
|
||||
details.allowed,
|
||||
details.limit_reached,
|
||||
details.primary.expect("primary window").window_minutes,
|
||||
parsed.plan_type.as_deref(),
|
||||
parsed.credits.expect("credits").balance,
|
||||
parsed.rate_limit_reached_type.as_deref(),
|
||||
),
|
||||
(
|
||||
Some("codex_other"),
|
||||
Some("Other requests"),
|
||||
Some(false),
|
||||
Some(true),
|
||||
Some(60),
|
||||
Some("team"),
|
||||
Some("9.25".to_owned()),
|
||||
Some("usage_limit_reached"),
|
||||
)
|
||||
);
|
||||
}
|
||||
@@ -498,9 +507,15 @@ fn rate_limit_header_round_trip_should_preserve_wire_facts() {
|
||||
"300".to_owned(),
|
||||
),
|
||||
("x-codex-active-limit".to_owned(), "codex".to_owned()),
|
||||
("x-codex-allowed".to_owned(), "false".to_owned()),
|
||||
("x-codex-limit-reached".to_owned(), "true".to_owned()),
|
||||
("x-codex-credits-has-credits".to_owned(), "true".to_owned()),
|
||||
("x-codex-credits-unlimited".to_owned(), "false".to_owned()),
|
||||
("x-codex-plan-type".to_owned(), "plus".to_owned()),
|
||||
(
|
||||
"x-codex-rate-limit-reached-type".to_owned(),
|
||||
"usage_limit_reached".to_owned(),
|
||||
),
|
||||
];
|
||||
let parsed = parse_rate_limit_headers(&source).expect("initial parse");
|
||||
let encoded = rate_limits_to_header_pairs(&parsed);
|
||||
|
||||
@@ -1992,27 +1992,35 @@ async fn provider_account_metrics(
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<ProviderAccountMetrics> {
|
||||
let row = sqlx::query(
|
||||
"select count(*)::bigint as total,
|
||||
"with normalized as (
|
||||
select enabled,
|
||||
case
|
||||
when not enabled then 'disabled'
|
||||
when availability = 'banned' then 'banned'
|
||||
when availability = 'quota_exhausted'
|
||||
or (
|
||||
availability = 'cooldown'
|
||||
and availability_reason = 'usage_limit_exhausted'
|
||||
and (cooldown_until is null or cooldown_until > $1)
|
||||
)
|
||||
then 'quota_exhausted'
|
||||
when availability in ('expired', 'invalid')
|
||||
or (access_token_expires_at is not null and access_token_expires_at <= $1)
|
||||
then 'expired'
|
||||
else 'active'
|
||||
end as status
|
||||
from provider_accounts
|
||||
)
|
||||
select count(*)::bigint as total,
|
||||
count(*) filter (where enabled)::bigint as enabled,
|
||||
count(*) filter (
|
||||
where not (
|
||||
enabled and availability = 'ready'
|
||||
and (access_token_expires_at is null or access_token_expires_at > $1)
|
||||
)
|
||||
)::bigint as unavailable,
|
||||
count(*) filter (
|
||||
where enabled and availability = 'ready'
|
||||
and (access_token_expires_at is null or access_token_expires_at > $1)
|
||||
)::bigint as active,
|
||||
count(*) filter (
|
||||
where availability = 'expired'
|
||||
or (access_token_expires_at is not null and access_token_expires_at <= $1)
|
||||
)::bigint as expired,
|
||||
count(*) filter (where availability = 'quota_exhausted')::bigint
|
||||
as quota_exhausted,
|
||||
count(*) filter (where not enabled)::bigint as disabled,
|
||||
count(*) filter (where availability = 'banned')::bigint as banned
|
||||
from provider_accounts",
|
||||
count(*) filter (where status in ('expired', 'disabled', 'banned'))::bigint
|
||||
as unavailable,
|
||||
count(*) filter (where status = 'active')::bigint as active,
|
||||
count(*) filter (where status = 'expired')::bigint as expired,
|
||||
count(*) filter (where status = 'quota_exhausted')::bigint as quota_exhausted,
|
||||
count(*) filter (where status = 'disabled')::bigint as disabled,
|
||||
count(*) filter (where status = 'banned')::bigint as banned
|
||||
from normalized",
|
||||
)
|
||||
.bind(observed_at)
|
||||
.fetch_one(pool)
|
||||
@@ -2035,11 +2043,29 @@ async fn active_provider_account_ids(
|
||||
observed_at: DateTime<Utc>,
|
||||
) -> StoreResult<Vec<String>> {
|
||||
sqlx::query(
|
||||
"select id
|
||||
from provider_accounts
|
||||
where enabled and availability = 'ready'
|
||||
and (access_token_expires_at is null or access_token_expires_at > $1)
|
||||
order by id",
|
||||
"with normalized as (
|
||||
select id,
|
||||
case
|
||||
when not enabled then 'disabled'
|
||||
when availability = 'banned' then 'banned'
|
||||
when availability = 'quota_exhausted'
|
||||
or (
|
||||
availability = 'cooldown'
|
||||
and availability_reason = 'usage_limit_exhausted'
|
||||
and (cooldown_until is null or cooldown_until > $1)
|
||||
)
|
||||
then 'quota_exhausted'
|
||||
when availability in ('expired', 'invalid')
|
||||
or (access_token_expires_at is not null and access_token_expires_at <= $1)
|
||||
then 'expired'
|
||||
else 'active'
|
||||
end as status
|
||||
from provider_accounts
|
||||
)
|
||||
select id
|
||||
from normalized
|
||||
where status = 'active'
|
||||
order by id",
|
||||
)
|
||||
.bind(observed_at)
|
||||
.fetch_all(pool)
|
||||
|
||||
@@ -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,
|
||||
@@ -380,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(())
|
||||
}
|
||||
}
|
||||
@@ -392,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"));
|
||||
}
|
||||
|
||||
@@ -76,6 +76,12 @@ impl TestDatabase {
|
||||
.execute(&mut *migration)
|
||||
.await
|
||||
.expect("apply model request service tier migration");
|
||||
sqlx::raw_sql(include_str!(
|
||||
"../../../../migrations/0005_allow_quota_exhausted_cooldown.sql"
|
||||
))
|
||||
.execute(&mut *migration)
|
||||
.await
|
||||
.expect("apply quota exhausted cooldown migration");
|
||||
migration.commit().await.expect("commit terminal migration");
|
||||
Some(Self {
|
||||
admin,
|
||||
@@ -171,7 +177,7 @@ async fn connect_and_migrate_should_apply_all_migrations_once_and_reopen_cleanly
|
||||
.expect("drop migration test database");
|
||||
admin.close().await;
|
||||
|
||||
assert_eq!((first_table_count, migration_count), (8, 4));
|
||||
assert_eq!((first_table_count, migration_count), (8, 5));
|
||||
assert_eq!(response_id_types, ["bytea", "bytea"]);
|
||||
assert!(!raw_response_id_index_exists);
|
||||
}
|
||||
@@ -183,6 +189,7 @@ fn migrations_should_leave_transaction_ownership_to_sqlx() {
|
||||
include_str!("../../../../migrations/0002_drop_opaque_response_id_index.sql"),
|
||||
include_str!("../../../../migrations/0003_store_opaque_response_ids_as_bytes.sql"),
|
||||
include_str!("../../../../migrations/0004_add_model_request_service_tier.sql"),
|
||||
include_str!("../../../../migrations/0005_allow_quota_exhausted_cooldown.sql"),
|
||||
]
|
||||
.into_iter()
|
||||
.flat_map(str::lines)
|
||||
|
||||
@@ -122,16 +122,27 @@ async fn dashboard_account_metrics_should_partition_available_and_unavailable_ac
|
||||
id, provider_kind, name, upstream_user_id, authentication_kind,
|
||||
provider_credentials_json, credential_revision, has_refresh_token,
|
||||
access_token_expires_at, enabled, availability,
|
||||
availability_reason, cooldown_until,
|
||||
availability_observed_at, created_at, updated_at
|
||||
) values
|
||||
('acct_available', 'openai', 'available', 'user-available', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready', $1, $1, $1),
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'ready',
|
||||
null, null, $1, $1, $1),
|
||||
('acct_expired', 'openai', 'expired', 'user-expired', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready', $1, $1, $1),
|
||||
'{}'::jsonb, 1, false, $1 - interval '1 day', true, 'ready',
|
||||
null, null, $1, $1, $1),
|
||||
('acct_rate_cooldown', 'openai', 'rate-cooldown', 'user-rate-cooldown', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'cooldown',
|
||||
'rate_limited', $1 + interval '2 hours', $1, $1, $1),
|
||||
('acct_usage_limit', 'openai', 'usage-limit', 'user-usage-limit', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'cooldown',
|
||||
'usage_limit_exhausted', $1 + interval '2 hours', $1, $1, $1),
|
||||
('acct_banned', 'xai', 'banned', 'user-banned', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned', $1, $1, $1),
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', true, 'banned',
|
||||
null, null, $1, $1, $1),
|
||||
('acct_disabled', 'xai', 'disabled', 'user-disabled', 'oauth',
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready', $1, $1, $1)",
|
||||
'{}'::jsonb, 1, false, $1 + interval '1 day', false, 'ready',
|
||||
null, null, $1, $1, $1)",
|
||||
)
|
||||
.bind(now)
|
||||
.execute(&database.pool)
|
||||
@@ -154,10 +165,11 @@ async fn dashboard_account_metrics_should_partition_available_and_unavailable_ac
|
||||
metrics.active,
|
||||
metrics.unavailable,
|
||||
metrics.expired,
|
||||
metrics.quota_exhausted,
|
||||
metrics.disabled,
|
||||
metrics.banned,
|
||||
),
|
||||
(4, 3, 1, 3, 1, 1, 1),
|
||||
(6, 5, 2, 3, 1, 1, 1, 1),
|
||||
);
|
||||
database.close().await;
|
||||
}
|
||||
|
||||
@@ -1334,7 +1334,15 @@ fn cooldown_observation_requires_expiry() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn imported_account_rejects_mismatched_cooldown_runtime_facts() {
|
||||
fn quota_exhausted_account_accepts_cooldown_runtime_facts() {
|
||||
let mut imported = account("acct_quota_cooldown", "user-quota-cooldown");
|
||||
imported.availability = AccountAvailability::QuotaExhausted;
|
||||
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
|
||||
assert!(imported.validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn imported_account_rejects_cooldown_until_for_ready_runtime_facts() {
|
||||
let mut imported = account("acct_invalid_cooldown", "user-invalid-cooldown");
|
||||
imported.cooldown_until = Some(Utc::now() + TimeDelta::minutes(5));
|
||||
assert!(imported.validate().is_err());
|
||||
|
||||
@@ -30,6 +30,7 @@ const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60);
|
||||
const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60);
|
||||
const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5);
|
||||
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
|
||||
const QUOTA_RESET_RECHECK_GRACE: Duration = Duration::from_secs(30);
|
||||
// ready 账号主要依赖真实请求的被动观察;这里只提供低频、定额的状态兜底。
|
||||
const QUOTA_READY_REFRESH_MAX_AGE: Duration = Duration::from_secs(2 * 60 * 60);
|
||||
const QUOTA_READY_REFRESH_BATCH_SIZE: usize = 8;
|
||||
@@ -307,7 +308,7 @@ enum CodexQuotaFetchError {
|
||||
enum CodexQuotaStateTransition {
|
||||
Banned { reason: &'static str },
|
||||
Exhausted,
|
||||
Cooldown { until: Option<SystemTime> },
|
||||
ResettableLimit { until: Option<SystemTime> },
|
||||
}
|
||||
|
||||
impl CodexQuotaStateTransition {
|
||||
@@ -315,7 +316,7 @@ impl CodexQuotaStateTransition {
|
||||
match self {
|
||||
Self::Banned { .. } => AccountAvailability::Banned,
|
||||
Self::Exhausted => AccountAvailability::QuotaExhausted,
|
||||
Self::Cooldown { .. } => AccountAvailability::Cooldown,
|
||||
Self::ResettableLimit { .. } => AccountAvailability::QuotaExhausted,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -323,13 +324,13 @@ impl CodexQuotaStateTransition {
|
||||
match self {
|
||||
Self::Banned { reason } => reason,
|
||||
Self::Exhausted => "quota_exhausted",
|
||||
Self::Cooldown { .. } => "quota_rate_limited",
|
||||
Self::ResettableLimit { .. } => "quota_rate_limited",
|
||||
}
|
||||
}
|
||||
|
||||
const fn cooldown_until(self) -> Option<SystemTime> {
|
||||
match self {
|
||||
Self::Cooldown { until } => until,
|
||||
Self::ResettableLimit { until } => until,
|
||||
Self::Banned { .. } | Self::Exhausted => None,
|
||||
}
|
||||
}
|
||||
@@ -547,7 +548,9 @@ fn quota_refresh_candidate(
|
||||
let snapshot = quota_snapshot_for_account(&account, observation);
|
||||
let kind = match account.availability() {
|
||||
AccountAvailability::QuotaExhausted => CodexQuotaRefreshKind::Exhausted {
|
||||
reset_at: quota_reset_at_from_snapshot(snapshot.as_ref()),
|
||||
reset_at: account
|
||||
.cooldown_until()
|
||||
.or_else(|| quota_reset_at_from_snapshot(snapshot.as_ref())),
|
||||
},
|
||||
AccountAvailability::Cooldown
|
||||
if snapshot
|
||||
@@ -618,22 +621,24 @@ fn reserve_exhausted_quota_refresh(
|
||||
now: SystemTime,
|
||||
refreshed_at: Instant,
|
||||
) -> bool {
|
||||
let reset_recheck_due = reset_at.is_some_and(|reset_at| {
|
||||
reset_at <= now
|
||||
&& state
|
||||
.last_exhausted_reset_recheck
|
||||
.get(account.id())
|
||||
.is_none_or(|marker| {
|
||||
marker.revision != account.revision() || marker.reset_at != reset_at
|
||||
})
|
||||
});
|
||||
if !reset_recheck_due && !periodic_quota_refresh_due(state, account.id(), refreshed_at) {
|
||||
return false;
|
||||
}
|
||||
state
|
||||
.last_periodic_refresh_at
|
||||
.insert(account.id().clone(), refreshed_at);
|
||||
if let Some(reset_at) = reset_at.filter(|reset_at| *reset_at <= now) {
|
||||
if let Some(reset_at) = reset_at {
|
||||
if !quota_reset_recheck_ready(reset_at, now) {
|
||||
return false;
|
||||
}
|
||||
let same_reset_already_checked = state
|
||||
.last_exhausted_reset_recheck
|
||||
.get(account.id())
|
||||
.is_some_and(|marker| {
|
||||
marker.revision == account.revision() && marker.reset_at == reset_at
|
||||
});
|
||||
if same_reset_already_checked
|
||||
&& !periodic_quota_refresh_due(state, account.id(), refreshed_at)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
state
|
||||
.last_periodic_refresh_at
|
||||
.insert(account.id().clone(), refreshed_at);
|
||||
state.last_exhausted_reset_recheck.insert(
|
||||
account.id().clone(),
|
||||
CodexQuotaResetMarker {
|
||||
@@ -641,10 +646,23 @@ fn reserve_exhausted_quota_refresh(
|
||||
reset_at,
|
||||
},
|
||||
);
|
||||
return true;
|
||||
}
|
||||
if !periodic_quota_refresh_due(state, account.id(), refreshed_at) {
|
||||
return false;
|
||||
}
|
||||
state
|
||||
.last_periodic_refresh_at
|
||||
.insert(account.id().clone(), refreshed_at);
|
||||
true
|
||||
}
|
||||
|
||||
fn quota_reset_recheck_ready(reset_at: SystemTime, now: SystemTime) -> bool {
|
||||
reset_at
|
||||
.checked_add(QUOTA_RESET_RECHECK_GRACE)
|
||||
.is_some_and(|ready_at| ready_at <= now)
|
||||
}
|
||||
|
||||
fn periodic_quota_refresh_due(
|
||||
state: &CodexQuotaProjectionState,
|
||||
account_id: &ProviderAccountId,
|
||||
@@ -802,16 +820,12 @@ impl CodexCredentialQuotaService {
|
||||
summary.transient += 1;
|
||||
continue;
|
||||
};
|
||||
let preserve_exhaustion = account.availability()
|
||||
== AccountAvailability::QuotaExhausted
|
||||
&& transition.availability() == AccountAvailability::Cooldown;
|
||||
match transition {
|
||||
CodexQuotaStateTransition::Banned { .. } => summary.banned += 1,
|
||||
CodexQuotaStateTransition::Exhausted => summary.exhausted += 1,
|
||||
CodexQuotaStateTransition::Cooldown { .. } => summary.cooldown += 1,
|
||||
}
|
||||
if preserve_exhaustion {
|
||||
continue;
|
||||
CodexQuotaStateTransition::Exhausted
|
||||
| CodexQuotaStateTransition::ResettableLimit { .. } => {
|
||||
summary.exhausted += 1;
|
||||
}
|
||||
}
|
||||
let _ = self
|
||||
.repository
|
||||
@@ -856,19 +870,23 @@ impl CodexCredentialQuotaService {
|
||||
{
|
||||
continue;
|
||||
}
|
||||
if !matches!(
|
||||
self.has_active_exhaustion(¤t, observed_at).await,
|
||||
Ok(true)
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
let snapshot = match self.read_snapshot_for(¤t).await {
|
||||
Ok(Some(snapshot))
|
||||
if quota_fact_is_active_exhaustion(snapshot.fact(), observed_at) =>
|
||||
{
|
||||
snapshot
|
||||
}
|
||||
_ => continue,
|
||||
};
|
||||
let cooldown_until =
|
||||
quota_reset_at_from_snapshot(Some(&snapshot)).or_else(|| current.cooldown_until());
|
||||
if let Err(error) = self
|
||||
.repository
|
||||
.apply_state(
|
||||
¤t,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
Some("quota_exhausted".to_owned()),
|
||||
None,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.await
|
||||
@@ -881,8 +899,11 @@ impl CodexCredentialQuotaService {
|
||||
continue;
|
||||
}
|
||||
let enabled = current.enabled();
|
||||
*account =
|
||||
current.with_runtime_state(enabled, AccountAvailability::QuotaExhausted, None);
|
||||
*account = current.with_runtime_state(
|
||||
enabled,
|
||||
AccountAvailability::QuotaExhausted,
|
||||
cooldown_until,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -927,10 +948,10 @@ impl CodexCredentialQuotaService {
|
||||
} else {
|
||||
summary.updated += 1;
|
||||
}
|
||||
if let Some(availability) = quota_success_availability(
|
||||
if let Some((availability, cooldown_until)) = quota_success_state(
|
||||
current.availability(),
|
||||
current.cooldown_until(),
|
||||
fact.exhausted(),
|
||||
fact,
|
||||
SystemTime::now(),
|
||||
) {
|
||||
let _ = self
|
||||
@@ -939,7 +960,7 @@ impl CodexCredentialQuotaService {
|
||||
¤t,
|
||||
availability,
|
||||
fact.exhausted().then_some("quota_exhausted".to_owned()),
|
||||
None,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.await;
|
||||
@@ -993,10 +1014,10 @@ impl CodexCredentialQuotaService {
|
||||
}
|
||||
self.scheduling
|
||||
.observe(current.id().clone(), current.revision(), observed_at, fact);
|
||||
if let Some(availability) = quota_success_availability(
|
||||
if let Some((availability, cooldown_until)) = quota_success_state(
|
||||
current.availability(),
|
||||
current.cooldown_until(),
|
||||
fact.exhausted(),
|
||||
fact,
|
||||
observed_at,
|
||||
) {
|
||||
self.repository
|
||||
@@ -1004,7 +1025,7 @@ impl CodexCredentialQuotaService {
|
||||
¤t,
|
||||
availability,
|
||||
fact.exhausted().then_some("quota_exhausted".to_owned()),
|
||||
None,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.await?;
|
||||
@@ -1026,19 +1047,6 @@ impl CodexCredentialQuotaService {
|
||||
self.read_snapshot_for(&account).await
|
||||
}
|
||||
|
||||
pub(super) async fn has_active_exhaustion(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
now: SystemTime,
|
||||
) -> Result<bool, CodexCredentialQuotaError> {
|
||||
match self.read_snapshot_for(account).await {
|
||||
Ok(snapshot) => Ok(snapshot
|
||||
.is_some_and(|snapshot| quota_fact_is_active_exhaustion(snapshot.fact(), now))),
|
||||
Err(CodexCredentialQuotaError::RevisionConflict) => Ok(false),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_snapshot_for(
|
||||
&self,
|
||||
account: &ProviderAccount,
|
||||
@@ -1164,10 +1172,10 @@ impl CodexCredentialQuotaService {
|
||||
snapshot.observed_at(),
|
||||
snapshot.fact(),
|
||||
);
|
||||
if let Some(availability) = quota_success_availability(
|
||||
if let Some((availability, cooldown_until)) = quota_success_state(
|
||||
current.availability(),
|
||||
current.cooldown_until(),
|
||||
snapshot.fact().exhausted(),
|
||||
snapshot.fact(),
|
||||
SystemTime::now(),
|
||||
) {
|
||||
self.repository
|
||||
@@ -1178,7 +1186,7 @@ impl CodexCredentialQuotaService {
|
||||
.fact()
|
||||
.exhausted()
|
||||
.then_some("quota_exhausted".to_owned()),
|
||||
None,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.await?;
|
||||
@@ -1386,32 +1394,34 @@ fn passive_rate_limit_window(window: RateLimitWindow) -> Map<String, Value> {
|
||||
snapshot
|
||||
}
|
||||
|
||||
fn quota_success_availability(
|
||||
fn quota_success_state(
|
||||
current: AccountAvailability,
|
||||
cooldown_until: Option<SystemTime>,
|
||||
exhausted: bool,
|
||||
fact: CodexQuotaFact,
|
||||
now: SystemTime,
|
||||
) -> Option<AccountAvailability> {
|
||||
let observed = if exhausted {
|
||||
AccountAvailability::QuotaExhausted
|
||||
} else {
|
||||
AccountAvailability::Ready
|
||||
};
|
||||
match current {
|
||||
) -> Option<(AccountAvailability, Option<SystemTime>)> {
|
||||
if fact.exhausted() {
|
||||
let observed_cooldown = fact.resets_at().map(SystemTime::from).or(cooldown_until);
|
||||
return (current != AccountAvailability::QuotaExhausted
|
||||
|| cooldown_until != observed_cooldown)
|
||||
.then_some((AccountAvailability::QuotaExhausted, observed_cooldown));
|
||||
}
|
||||
|
||||
let availability = match current {
|
||||
AccountAvailability::Invalid
|
||||
| AccountAvailability::Expired
|
||||
| AccountAvailability::Banned
|
||||
| AccountAvailability::Unknown => Some(observed),
|
||||
AccountAvailability::QuotaExhausted => (!exhausted).then_some(AccountAvailability::Ready),
|
||||
AccountAvailability::Ready => exhausted.then_some(AccountAvailability::QuotaExhausted),
|
||||
AccountAvailability::Cooldown if exhausted => Some(AccountAvailability::QuotaExhausted),
|
||||
| AccountAvailability::Unknown
|
||||
| AccountAvailability::QuotaExhausted => Some(AccountAvailability::Ready),
|
||||
AccountAvailability::Ready => None,
|
||||
AccountAvailability::Cooldown
|
||||
if cooldown_until.is_some_and(|cooldown_until| cooldown_until <= now) =>
|
||||
{
|
||||
Some(AccountAvailability::Ready)
|
||||
}
|
||||
AccountAvailability::Cooldown => None,
|
||||
}
|
||||
};
|
||||
availability.map(|availability| (availability, None))
|
||||
}
|
||||
|
||||
fn quota_state_transition(
|
||||
@@ -1439,7 +1449,7 @@ fn quota_state_transition(
|
||||
retry_after_seconds.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN.as_secs()),
|
||||
)
|
||||
.min(MAX_RATE_LIMIT_COOLDOWN);
|
||||
Some(CodexQuotaStateTransition::Cooldown {
|
||||
Some(CodexQuotaStateTransition::ResettableLimit {
|
||||
until: observed_at.checked_add(duration),
|
||||
})
|
||||
}
|
||||
@@ -1877,7 +1887,6 @@ impl QuotaAggregate {
|
||||
self.remaining_percent
|
||||
.map_or(remaining, |current| current.min(remaining)),
|
||||
);
|
||||
self.exhausted |= used >= 100.0;
|
||||
}
|
||||
if let Some(reset) = object.get("reset_at") {
|
||||
let seconds = reset
|
||||
|
||||
@@ -256,8 +256,6 @@ impl CodexCredentialRefreshService {
|
||||
| AccountAvailability::Banned
|
||||
| AccountAvailability::Invalid
|
||||
)
|
||||
&& (account.availability() != AccountAvailability::Cooldown
|
||||
|| account.cooldown_until().is_some_and(|until| until <= now))
|
||||
&& refresh_due_at(account, now).is_some()
|
||||
});
|
||||
accounts.sort_by_key(|account| (refresh_due_at(account, now), account.id().clone()));
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
@@ -3,7 +3,8 @@ use std::{sync::Arc, time::Instant};
|
||||
use futures::{StreamExt, TryStreamExt};
|
||||
use gateway_protocol::openai::{
|
||||
X_OPENAI_INTERNAL_CODEX_RESPONSES_LITE_HEADER, X_OPENAI_MEMGEN_REQUEST_HEADER,
|
||||
events::retry_after_seconds_from_body,
|
||||
events::{self, retry_after_seconds_from_body},
|
||||
sse::{SseEventDecoder, SseFrame},
|
||||
};
|
||||
use reqwest::{
|
||||
Client, Response as ReqwestResponse,
|
||||
@@ -130,13 +131,14 @@ impl CodexBackendClient {
|
||||
});
|
||||
}
|
||||
|
||||
let rate_limit_header_updates = Arc::new(tokio::sync::Mutex::new(Vec::new()));
|
||||
Ok(CodexBackendStreamingResponse {
|
||||
body: http_sse_stream(response),
|
||||
body: http_sse_stream(response, Arc::clone(&rate_limit_header_updates)),
|
||||
transport: CodexBackendTransport::HttpSse,
|
||||
turn_state,
|
||||
set_cookie_headers,
|
||||
rate_limit_headers,
|
||||
rate_limit_header_updates: None,
|
||||
rate_limit_header_updates: Some(rate_limit_header_updates),
|
||||
turn_state_update: None,
|
||||
websocket_pool_decision: None,
|
||||
diagnostics,
|
||||
@@ -594,20 +596,74 @@ fn websocket_connection_profile(headers: &HeaderMap) -> String {
|
||||
.join("\0")
|
||||
}
|
||||
|
||||
fn http_sse_stream(response: ReqwestResponse) -> CodexBackendSseStream {
|
||||
fn http_sse_stream(
|
||||
response: ReqwestResponse,
|
||||
rate_limit_updates: CodexRateLimitHeaderUpdates,
|
||||
) -> CodexBackendSseStream {
|
||||
let stream: CodexBackendSseStream =
|
||||
Box::pin(response.bytes_stream().map_err(CodexClientError::Http));
|
||||
Box::pin(futures::stream::unfold(Some(stream), |stream| async move {
|
||||
let mut stream = stream?;
|
||||
match tokio::time::timeout(UPSTREAM_STREAM_IDLE_TIMEOUT, stream.next()).await {
|
||||
Ok(Some(chunk)) => Some((chunk, Some(stream))),
|
||||
Ok(None) => None,
|
||||
Err(_) => Some((
|
||||
Err(CodexClientError::StreamIdleTimeout {
|
||||
timeout: UPSTREAM_STREAM_IDLE_TIMEOUT,
|
||||
}),
|
||||
None,
|
||||
)),
|
||||
let stream: CodexBackendSseStream =
|
||||
Box::pin(futures::stream::unfold(Some(stream), |stream| async move {
|
||||
let mut stream = stream?;
|
||||
match tokio::time::timeout(UPSTREAM_STREAM_IDLE_TIMEOUT, stream.next()).await {
|
||||
Ok(Some(chunk)) => Some((chunk, Some(stream))),
|
||||
Ok(None) => None,
|
||||
Err(_) => Some((
|
||||
Err(CodexClientError::StreamIdleTimeout {
|
||||
timeout: UPSTREAM_STREAM_IDLE_TIMEOUT,
|
||||
}),
|
||||
None,
|
||||
)),
|
||||
}
|
||||
}));
|
||||
observe_http_sse_rate_limits(stream, rate_limit_updates)
|
||||
}
|
||||
|
||||
fn observe_http_sse_rate_limits(
|
||||
stream: CodexBackendSseStream,
|
||||
updates: CodexRateLimitHeaderUpdates,
|
||||
) -> CodexBackendSseStream {
|
||||
Box::pin(futures::stream::unfold(
|
||||
(stream, SseEventDecoder::default(), updates),
|
||||
|(mut stream, mut decoder, updates)| async move {
|
||||
match stream.next().await {
|
||||
Some(chunk) => {
|
||||
if let Ok(bytes) = &chunk {
|
||||
append_http_sse_rate_limit_updates(decoder.push_frames(bytes), &updates)
|
||||
.await;
|
||||
}
|
||||
Some((chunk, (stream, decoder, updates)))
|
||||
}
|
||||
None => {
|
||||
append_http_sse_rate_limit_updates(decoder.finish_frames(), &updates).await;
|
||||
None
|
||||
}
|
||||
}
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
async fn append_http_sse_rate_limit_updates(
|
||||
frames: Vec<SseFrame>,
|
||||
updates: &CodexRateLimitHeaderUpdates,
|
||||
) {
|
||||
let mut headers = Vec::new();
|
||||
for frame in frames {
|
||||
for event in frame.events() {
|
||||
if event
|
||||
.event
|
||||
.as_deref()
|
||||
.is_some_and(|event| event != "codex.rate_limits")
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let Some(rate_limits) = events::parse_rate_limits_event_raw(&event.data) else {
|
||||
continue;
|
||||
};
|
||||
headers.extend(events::rate_limits_to_header_pairs(&rate_limits));
|
||||
}
|
||||
}))
|
||||
}
|
||||
if !headers.is_empty() {
|
||||
updates.lock().await.extend(headers);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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]
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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!({
|
||||
@@ -369,6 +399,49 @@ async fn passive_rate_limit_headers_update_quota_and_account_state_with_revision
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn passive_full_percent_with_explicit_allowance_keeps_account_ready() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_passive_full_allowed").await;
|
||||
let account = store.account("acct_passive_full_allowed").expect("account");
|
||||
let service = quota_service(&store);
|
||||
let headers = vec![
|
||||
("x-codex-active-limit".to_owned(), "codex".to_owned()),
|
||||
("x-codex-primary-used-percent".to_owned(), "100".to_owned()),
|
||||
(
|
||||
"x-codex-primary-reset-at".to_owned(),
|
||||
"1900000000".to_owned(),
|
||||
),
|
||||
("x-codex-allowed".to_owned(), "true".to_owned()),
|
||||
("x-codex-limit-reached".to_owned(), "false".to_owned()),
|
||||
];
|
||||
|
||||
assert!(
|
||||
service
|
||||
.synchronize_passive_headers(&account, &headers)
|
||||
.await
|
||||
.expect("passive quota")
|
||||
);
|
||||
let current = store
|
||||
.account("acct_passive_full_allowed")
|
||||
.expect("current account");
|
||||
let snapshot = service
|
||||
.read_account(account.id())
|
||||
.await
|
||||
.expect("read quota")
|
||||
.expect("quota snapshot");
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
current.availability(),
|
||||
current.cooldown_until(),
|
||||
snapshot.fact().remaining_percent(),
|
||||
snapshot.fact().exhausted(),
|
||||
),
|
||||
(AccountAvailability::Ready, None, Some(0), false)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manual_quota_refresh_updates_account_state_from_provider_response() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -771,12 +844,18 @@ async fn periodic_quota_synchronization_repairs_cooldown_with_active_exhausted_q
|
||||
|
||||
service.synchronize().await.expect("quota synchronization");
|
||||
|
||||
let reconciled = store
|
||||
.account("acct_periodic_stuck_cooldown")
|
||||
.expect("reconciled account");
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_periodic_stuck_cooldown")
|
||||
.expect("reconciled account")
|
||||
.availability(),
|
||||
AccountAvailability::QuotaExhausted
|
||||
(reconciled.availability(), reconciled.cooldown_until()),
|
||||
(
|
||||
AccountAvailability::QuotaExhausted,
|
||||
Some(
|
||||
SystemTime::UNIX_EPOCH
|
||||
+ Duration::from_secs(u64::try_from(reset_at).expect("positive reset")),
|
||||
),
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
@@ -815,12 +894,14 @@ async fn periodic_quota_synchronization_does_not_downgrade_exhaustion_on_rate_li
|
||||
|
||||
service.synchronize().await.expect("quota synchronization");
|
||||
|
||||
assert_eq!(
|
||||
store
|
||||
.account("acct_periodic_exhausted")
|
||||
.expect("persisted account")
|
||||
.availability(),
|
||||
AccountAvailability::QuotaExhausted
|
||||
let account = store
|
||||
.account("acct_periodic_exhausted")
|
||||
.expect("persisted account");
|
||||
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
);
|
||||
}
|
||||
|
||||
@@ -951,6 +1032,60 @@ async fn periodic_quota_synchronization_attempts_quota_exhausted_accounts() {
|
||||
assert_eq!(summary.transient, 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn periodic_quota_synchronization_waits_for_quota_cooldown_deadline_with_grace() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
create_account(&store, "acct_periodic_exhausted_cooldown").await;
|
||||
let account = store
|
||||
.account("acct_periodic_exhausted_cooldown")
|
||||
.expect("created account");
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
reason: Some("usage_limit_exhausted".to_owned()),
|
||||
cooldown_until: SystemTime::now().checked_sub(Duration::from_secs(29)),
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
.expect("mark account exhausted with recent cooldown");
|
||||
let service = blocked_network_quota_service(&store);
|
||||
|
||||
let waiting = service
|
||||
.synchronize()
|
||||
.await
|
||||
.expect("quota synchronization before grace");
|
||||
|
||||
assert_eq!(waiting, CodexQuotaSyncSummary::default());
|
||||
let current = store
|
||||
.account("acct_periodic_exhausted_cooldown")
|
||||
.expect("current account");
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
account_id: current.id().clone(),
|
||||
expected_revision: current.revision(),
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
reason: Some("usage_limit_exhausted".to_owned()),
|
||||
cooldown_until: SystemTime::now().checked_sub(Duration::from_secs(31)),
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
.expect("mark account exhausted past grace");
|
||||
|
||||
let due = service
|
||||
.synchronize()
|
||||
.await
|
||||
.expect("quota synchronization after grace");
|
||||
let repeated = service
|
||||
.synchronize()
|
||||
.await
|
||||
.expect("quota synchronization repeated");
|
||||
|
||||
assert_eq!(due.transient, 1);
|
||||
assert_eq!(repeated, CodexQuotaSyncSummary::default());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn periodic_quota_synchronization_throttles_the_same_exhausted_account() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -1029,8 +1164,8 @@ async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_
|
||||
server.uri(),
|
||||
);
|
||||
|
||||
let first = service.synchronize().await.expect("initial refresh");
|
||||
assert_eq!(first.transient, 1);
|
||||
let waiting = service.synchronize().await.expect("initial cooldown wait");
|
||||
assert_eq!(waiting, CodexQuotaSyncSummary::default());
|
||||
store
|
||||
.compare_and_swap_quota(QuotaObservation {
|
||||
account_id: account.id().clone(),
|
||||
@@ -1061,6 +1196,77 @@ async fn periodic_quota_synchronization_rechecks_an_exhausted_account_once_when_
|
||||
assert_eq!(repeated, CodexQuotaSyncSummary::default());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn periodic_quota_synchronization_reschedules_a_confirmed_new_reset() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_periodic_new_reset";
|
||||
create_account(&store, account_id).await;
|
||||
let account = store.account(account_id).expect("created account");
|
||||
let initial_reset = SystemTime::UNIX_EPOCH + Duration::from_secs(1);
|
||||
let confirmed_reset = SystemTime::UNIX_EPOCH + Duration::from_secs(2);
|
||||
store
|
||||
.apply_state_change(AccountStateChange {
|
||||
account_id: account.id().clone(),
|
||||
expected_revision: account.revision(),
|
||||
availability: AccountAvailability::QuotaExhausted,
|
||||
reason: Some("usage_limit_exhausted".to_owned()),
|
||||
cooldown_until: Some(initial_reset),
|
||||
observed_at: SystemTime::now(),
|
||||
})
|
||||
.await
|
||||
.expect("mark account exhausted with due reset");
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"rate_limit": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary_window": {"used_percent": 100, "reset_at": 2}
|
||||
}
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let service = quota_service_with_base_url(
|
||||
&store,
|
||||
reqwest::Client::builder()
|
||||
.no_proxy()
|
||||
.build()
|
||||
.expect("client"),
|
||||
server.uri(),
|
||||
);
|
||||
|
||||
let first = service.synchronize().await.expect("first reset recheck");
|
||||
let persisted_reset = store
|
||||
.account(account_id)
|
||||
.expect("account after first recheck")
|
||||
.cooldown_until();
|
||||
let second = service.synchronize().await.expect("new reset recheck");
|
||||
let repeated = service.synchronize().await.expect("same reset throttled");
|
||||
let request_count = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("received requests")
|
||||
.len();
|
||||
|
||||
assert_eq!(
|
||||
(
|
||||
first.exhausted,
|
||||
persisted_reset,
|
||||
second.exhausted,
|
||||
repeated,
|
||||
request_count,
|
||||
),
|
||||
(
|
||||
1,
|
||||
Some(confirmed_reset),
|
||||
1,
|
||||
CodexQuotaSyncSummary::default(),
|
||||
2,
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_quota_json_for_one_account_does_not_abort_the_batch() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
|
||||
@@ -665,6 +665,74 @@ async fn expired_access_token_refreshes_even_when_next_refresh_is_deferred() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn expired_access_token_refreshes_during_active_cooldown() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
seed_due_account(&store).await;
|
||||
let original = store.account("acct_refresh").expect("seeded account");
|
||||
let mut expired_profile = profile("chatgpt-acct_refresh");
|
||||
expired_profile.access_token_expires_at =
|
||||
Some(chrono::Utc::now() - chrono::Duration::minutes(5));
|
||||
store
|
||||
.repository()
|
||||
.rotate_oauth_secret(RotateCodexCredential {
|
||||
account_id: "acct_refresh".to_owned(),
|
||||
expected_credential_revision: original.revision().get(),
|
||||
secret: secret("old-access"),
|
||||
verified_account: expired_profile,
|
||||
next_refresh_at: Some(chrono::Utc::now() + chrono::Duration::hours(12)),
|
||||
})
|
||||
.await
|
||||
.expect("seed expired access token with deferred refresh");
|
||||
let deferred = store.account("acct_refresh").expect("deferred account");
|
||||
let cooldown_until = SystemTime::now() + Duration::from_secs(12 * 60 * 60);
|
||||
store
|
||||
.repository()
|
||||
.apply_state(
|
||||
&deferred,
|
||||
AccountAvailability::Cooldown,
|
||||
Some("rate_limited".to_owned()),
|
||||
Some(cooldown_until),
|
||||
SystemTime::now(),
|
||||
)
|
||||
.await
|
||||
.expect("mark account cooling down");
|
||||
let refresher = Arc::new(Refresher::new(Ok(success_tokens())));
|
||||
let service = CodexCredentialRefreshService::new(
|
||||
store.repository(),
|
||||
refresher.clone(),
|
||||
Arc::new(VerifiedIdentity),
|
||||
Arc::new(RefreshLeases {
|
||||
available: true,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
}),
|
||||
Arc::new(CountingCredentialState::default()),
|
||||
runtime_policy(),
|
||||
);
|
||||
|
||||
let outcomes = service.refresh_due().await.expect("refresh due");
|
||||
|
||||
assert!(matches!(
|
||||
outcomes.as_slice(),
|
||||
[CodexCredentialRefreshOutcome::Refreshed {
|
||||
account_id,
|
||||
credential_revision: 3,
|
||||
}] if account_id == "acct_refresh"
|
||||
));
|
||||
assert_eq!(
|
||||
refresher.seen.lock().expect("seen tokens lock").as_slice(),
|
||||
["rt-old-access"]
|
||||
);
|
||||
let refreshed = store.account("acct_refresh").expect("refreshed account");
|
||||
assert_eq!(refreshed.availability(), AccountAvailability::Cooldown);
|
||||
assert_eq!(refreshed.cooldown_until(), Some(cooldown_until));
|
||||
assert!(
|
||||
refreshed
|
||||
.access_token_expires_at()
|
||||
.is_some_and(|expires_at| expires_at > SystemTime::now())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unavailable_refresh_lease_prevents_duplicate_token_exchange() {
|
||||
let (store, refresher, service) = setup(Ok(success_tokens()), false).await;
|
||||
|
||||
@@ -36,7 +36,7 @@ use tokio::{
|
||||
sync::oneshot,
|
||||
time::timeout,
|
||||
};
|
||||
use wiremock::matchers::{method, path};
|
||||
use wiremock::matchers::{header, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
use crate::support::{
|
||||
@@ -861,7 +861,7 @@ async fn completed_response_persists_session_affinity_before_stream_consumer_sto
|
||||
}
|
||||
|
||||
#[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;
|
||||
@@ -906,7 +906,7 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown(
|
||||
vec!["response.created", "response.failed"]
|
||||
);
|
||||
let account = store.account(account_id).expect("rate-limited account");
|
||||
assert_eq!(account.availability(), AccountAvailability::Cooldown);
|
||||
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
@@ -917,6 +917,68 @@ async fn response_failed_before_semantic_output_is_atomic_and_persists_cooldown(
|
||||
let _ = server.await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_limit_failure_marks_quota_exhausted_without_usage_probe() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_usage_limit_request_path";
|
||||
create_account(&store, account_id).await;
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/api/codex/usage"))
|
||||
.and(header("authorization", format!("Bearer at-{account_id}")))
|
||||
.respond_with(ResponseTemplate::new(500))
|
||||
.expect(0)
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_usage_limit_confirm\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\"}}\n\n",
|
||||
"event: response.failed\n",
|
||||
"data: {\"type\":\"response.failed\",\"status_code\":429,\"retry_after_seconds\":86400,\"response\":{\"id\":\"resp_usage_limit_confirm\",\"status\":\"failed\",\"error\":{\"code\":\"usage_limit_reached\",\"message\":\"usage limit reached\"}}}\n\n"
|
||||
)),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let mut stream = provider_with_base_url(&store, server.uri())
|
||||
.execute(
|
||||
planned_request("openai", http_generate_operation()),
|
||||
context("req_usage_limit_confirm", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("prepare provider stream");
|
||||
let failure = loop {
|
||||
match stream.next().await {
|
||||
Some(Ok(_)) => {}
|
||||
Some(Err(error)) => break error,
|
||||
None => panic!("usage-limit failure must surface a typed failure"),
|
||||
}
|
||||
};
|
||||
|
||||
assert_eq!(failure.kind(), ProviderErrorKind::QuotaExhausted);
|
||||
let account = store.account(account_id).expect("usage-limit account");
|
||||
assert_eq!(account.availability(), AccountAvailability::QuotaExhausted);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
);
|
||||
assert!(!store.has_quota(account_id));
|
||||
let requests = server.received_requests().await.expect("received requests");
|
||||
assert_eq!(
|
||||
requests
|
||||
.iter()
|
||||
.filter(|request| request.url.path() == "/api/codex/usage")
|
||||
.count(),
|
||||
0
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ordinary_request_should_forward_created_before_later_failed_chunk() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
@@ -1338,6 +1400,76 @@ async fn successful_response_preserves_quota_exhaustion_from_rate_limit_headers(
|
||||
assert!(store.has_quota(account_id));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_http_sse_rate_limit_event_persists_structured_exhaustion() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
let account_id = "acct_http_sse_exhausted";
|
||||
create_account(&store, account_id).await;
|
||||
let reset_at = 1_900_000_000_u64;
|
||||
let rate_limit_event = json!({
|
||||
"type": "codex.rate_limits",
|
||||
"rate_limits": {
|
||||
"allowed": false,
|
||||
"limit_reached": true,
|
||||
"primary": {
|
||||
"used_percent": 42,
|
||||
"window_minutes": 300,
|
||||
"reset_at": reset_at,
|
||||
},
|
||||
},
|
||||
});
|
||||
let completed_event = json!({
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_http_sse_exhausted",
|
||||
"model": "gpt-5.4",
|
||||
"status": "completed",
|
||||
"output": [],
|
||||
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
});
|
||||
let body = format!(
|
||||
"event: codex.rate_limits\ndata: {rate_limit_event}\n\nevent: response.completed\ndata: {completed_event}\n\n"
|
||||
);
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/codex/responses"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.insert_header("content-type", "text/event-stream")
|
||||
.set_body_string(body),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let mut stream = provider_with_base_url(&store, server.uri())
|
||||
.execute(
|
||||
planned_request("openai", http_generate_operation()),
|
||||
context("req_http_sse_exhausted", CancellationToken::new()),
|
||||
)
|
||||
.await
|
||||
.expect("prepare provider stream");
|
||||
while let Some(event) = stream.next().await {
|
||||
event.expect("successful upstream response");
|
||||
}
|
||||
|
||||
let account = store
|
||||
.account(account_id)
|
||||
.expect("account after HTTP SSE response");
|
||||
assert_eq!(
|
||||
(
|
||||
account.availability(),
|
||||
account.cooldown_until(),
|
||||
store.has_quota(account_id),
|
||||
),
|
||||
(
|
||||
AccountAvailability::QuotaExhausted,
|
||||
Some(SystemTime::UNIX_EPOCH + Duration::from_secs(reset_at)),
|
||||
true,
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_observation_reads_openai_metadata_without_changing_the_operation() {
|
||||
let store = Arc::new(MemoryAccountStore::default());
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -4,7 +4,7 @@ use std::fmt;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime};
|
||||
|
||||
use chrono::Utc;
|
||||
use chrono::{DateTime, Utc};
|
||||
use gateway_core::engine::credential::{
|
||||
AccountAvailability, AccountAvailabilityPolicy, AccountCandidate, AccountFeedbackStats,
|
||||
AccountSelectionContext, AccountSelector, ProviderAccount, ProviderAccountId,
|
||||
@@ -17,7 +17,7 @@ use gateway_core::provider_ports::{
|
||||
use gateway_core::routing::{ProviderKind, UpstreamModelId};
|
||||
|
||||
use super::catalog::{GrokCatalogScope, GrokCredentialCatalogCache, GrokCredentialQuotaService};
|
||||
use super::repository::GrokCredentialRepository;
|
||||
use super::repository::{GrokCredentialRepository, GrokCredentialRepositoryError};
|
||||
use super::types::{GrokCredentialAvailability, UpdateGrokCredentialState};
|
||||
use crate::{
|
||||
GrokCredentialFailure, GrokCredentialFeedbackFuture, GrokSessionBinding, GrokSessionSelection,
|
||||
@@ -29,7 +29,6 @@ const DEFAULT_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60);
|
||||
const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
const STREAM_INTERRUPTION_COOLDOWN: Duration = Duration::from_secs(30);
|
||||
const UNAUTHORIZED_COOLDOWN: Duration = Duration::from_secs(60);
|
||||
const PAYMENT_REQUIRED_COOLDOWN: Duration = Duration::from_secs(60);
|
||||
const MODEL_ACCESS_DENIED_COOLDOWN: Duration = Duration::from_secs(5 * 60);
|
||||
const MODEL_QUOTA_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
const MAX_MODEL_COOLDOWN: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
@@ -369,6 +368,66 @@ impl GrokAccountSessionSelector {
|
||||
))
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn persist_account_failure(
|
||||
&self,
|
||||
session: &SelectedGrokSession,
|
||||
availability: GrokCredentialAvailability,
|
||||
reason: &'static str,
|
||||
cooldown_until: Option<DateTime<Utc>>,
|
||||
observed_at: chrono::DateTime<Utc>,
|
||||
) {
|
||||
let update = |expected_revision| UpdateGrokCredentialState {
|
||||
account_id: session.account_id().clone(),
|
||||
expected_revision,
|
||||
availability,
|
||||
availability_reason: Some(reason.to_owned()),
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
};
|
||||
match self
|
||||
.repository
|
||||
.update_state(&update(session.credential_revision()))
|
||||
.await
|
||||
{
|
||||
Ok(()) => return,
|
||||
Err(
|
||||
GrokCredentialRepositoryError::Conflict
|
||||
| GrokCredentialRepositoryError::StaleCredentialRevision,
|
||||
) => {}
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id = %session.account_id(),
|
||||
error = %error,
|
||||
"xAI account failure state write failed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
}
|
||||
let current = match self.repository.load_current(session.account_id()).await {
|
||||
Ok(current) if current.account.enabled() => current,
|
||||
Ok(_) => return,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id = %session.account_id(),
|
||||
error = %error,
|
||||
"xAI account failure state reload failed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
if let Err(error) = self
|
||||
.repository
|
||||
.update_state(&update(current.account.revision()))
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id = %session.account_id(),
|
||||
error = %error,
|
||||
"xAI account failure state retry failed"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for GrokAccountSessionSelector {
|
||||
@@ -415,15 +474,11 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
|
||||
"upstream_free_quota_exhausted",
|
||||
None,
|
||||
),
|
||||
GrokCredentialFailure::PaymentRequired { retry_after } => {
|
||||
let retry_after = bounded_cooldown(
|
||||
retry_after,
|
||||
PAYMENT_REQUIRED_COOLDOWN,
|
||||
MAX_RATE_LIMIT_COOLDOWN,
|
||||
);
|
||||
self.record_runtime_cooldown(session, retry_after).await;
|
||||
return;
|
||||
}
|
||||
GrokCredentialFailure::PaymentRequired { .. } => (
|
||||
GrokCredentialAvailability::QuotaExhausted,
|
||||
"upstream_payment_quota_exhausted",
|
||||
None,
|
||||
),
|
||||
GrokCredentialFailure::ModelQuotaExhausted {
|
||||
upstream_model,
|
||||
retry_after,
|
||||
@@ -452,26 +507,31 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
|
||||
.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN)
|
||||
.min(MAX_RATE_LIMIT_COOLDOWN);
|
||||
self.record_runtime_cooldown(session, retry_after).await;
|
||||
return;
|
||||
(
|
||||
GrokCredentialAvailability::Cooldown,
|
||||
"upstream_rate_limited",
|
||||
persistent_cooldown_until(observed_at, retry_after),
|
||||
)
|
||||
}
|
||||
GrokCredentialFailure::StreamInterrupted => {
|
||||
self.record_runtime_cooldown(session, STREAM_INTERRUPTION_COOLDOWN)
|
||||
.await;
|
||||
return;
|
||||
(
|
||||
GrokCredentialAvailability::Cooldown,
|
||||
"upstream_stream_interrupted",
|
||||
persistent_cooldown_until(observed_at, STREAM_INTERRUPTION_COOLDOWN),
|
||||
)
|
||||
}
|
||||
};
|
||||
let (availability, reason, cooldown_until) = persistent;
|
||||
let _ = self
|
||||
.repository
|
||||
.update_state(&UpdateGrokCredentialState {
|
||||
account_id: session.account_id().clone(),
|
||||
expected_revision: session.credential_revision(),
|
||||
availability,
|
||||
availability_reason: Some(reason.to_owned()),
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
})
|
||||
.await;
|
||||
self.persist_account_failure(
|
||||
session,
|
||||
availability,
|
||||
reason,
|
||||
cooldown_until,
|
||||
observed_at,
|
||||
)
|
||||
.await;
|
||||
})
|
||||
}
|
||||
|
||||
@@ -493,7 +553,6 @@ impl GrokSessionSelector for GrokAccountSessionSelector {
|
||||
account.availability(),
|
||||
AccountAvailability::Unknown
|
||||
| AccountAvailability::Cooldown
|
||||
| AccountAvailability::QuotaExhausted
|
||||
| AccountAvailability::Expired
|
||||
)
|
||||
{
|
||||
@@ -540,6 +599,15 @@ fn bounded_cooldown(
|
||||
.min(maximum)
|
||||
}
|
||||
|
||||
fn persistent_cooldown_until(
|
||||
observed_at: DateTime<Utc>,
|
||||
duration: Duration,
|
||||
) -> Option<DateTime<Utc>> {
|
||||
chrono::Duration::from_std(duration)
|
||||
.ok()
|
||||
.map(|duration| observed_at + duration)
|
||||
}
|
||||
|
||||
fn insert_runtime_cooldown(
|
||||
cooldowns: &mut std::collections::BTreeMap<ProviderAccountId, RuntimeCooldown>,
|
||||
account_id: &ProviderAccountId,
|
||||
|
||||
@@ -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,9 +47,13 @@ use crate::credential::{
|
||||
GrokCredentialRecovery, GrokCredentialRecoveryOutcome, GrokCredentialRefreshOutcome,
|
||||
GrokCredentialRefreshService, GrokQuotaError,
|
||||
};
|
||||
use crate::reasoning_replay::{
|
||||
GrokReasoningReplay, GrokReasoningReplayCapture, GrokReasoningReplayKey,
|
||||
valid_reasoning_ciphertext,
|
||||
};
|
||||
use crate::transport::canonical::GrokCanonicalDecoder;
|
||||
use crate::transport::config::XAI_PROVIDER_NAME;
|
||||
use crate::transport::headers::{GrokClientIdentity, GrokHeader, build_grok_headers};
|
||||
use crate::transport::headers::{GrokClientIdentity, build_grok_headers};
|
||||
use crate::transport::profile::{GROK_CLI_RELEASE_POLL_INTERVAL, GrokCliReleaseService};
|
||||
use crate::transport::{
|
||||
GROK_RESPONSES_URL, GrokCompactionDecodeError, GrokCompactionRequest,
|
||||
@@ -66,8 +69,6 @@ const HTTP_SSE_TRANSPORT: &str = "http_sse";
|
||||
const DEFAULT_GROK_MODEL: &str = "grok-4.5";
|
||||
const XAI_SESSION_STATE_MAX_BYTES: usize = 8 * 1024 * 1024;
|
||||
const XAI_SESSION_OUTPUT_LIMIT: usize = 4_096;
|
||||
const MIN_REASONING_CIPHERTEXT_BYTES: usize = 50;
|
||||
const MIN_REASONING_CIPHERTEXT_ENTROPY: f64 = 0.85;
|
||||
const REASONING_DECODE_FAILED_CODE: &str = "reasoning_decode_failed";
|
||||
const RESPONSE_NOT_FOUND_CODE: &str = "not_found";
|
||||
|
||||
@@ -82,6 +83,7 @@ pub struct GrokBuildProvider {
|
||||
credential_recovery: Arc<dyn GrokCredentialRecovery>,
|
||||
account_feedback: Arc<AccountFeedbackStats>,
|
||||
client_identity: GrokClientIdentity,
|
||||
reasoning_replay: GrokReasoningReplay,
|
||||
wire_profile: XaiWireProfileState,
|
||||
responses_url: Url,
|
||||
}
|
||||
@@ -105,6 +107,7 @@ impl GrokBuildProvider {
|
||||
credential_recovery,
|
||||
account_feedback,
|
||||
client_identity: GrokClientIdentity::new(),
|
||||
reasoning_replay: GrokReasoningReplay::new(),
|
||||
wire_profile,
|
||||
responses_url,
|
||||
})
|
||||
@@ -217,6 +220,38 @@ impl GrokBuildProvider {
|
||||
selected.account_id(),
|
||||
request_input.as_slice(),
|
||||
)?;
|
||||
let inherited_replay_session_id = context
|
||||
.continuation()
|
||||
.is_none()
|
||||
.then(|| {
|
||||
previous_session
|
||||
.as_ref()
|
||||
.and_then(|state| state.session_id.as_deref())
|
||||
})
|
||||
.flatten();
|
||||
let reasoning_replay_key = upstream_request
|
||||
.reasoning_replay_session_id()
|
||||
.or(inherited_replay_session_id)
|
||||
.and_then(|session_id| {
|
||||
self.reasoning_replay.key(
|
||||
candidate.upstream_model().as_str(),
|
||||
session_id,
|
||||
selected.account_id().as_str(),
|
||||
)
|
||||
});
|
||||
if !upstream_request.has_previous_response_id()
|
||||
&& let (Some(key), Some(input)) = (
|
||||
reasoning_replay_key.as_ref(),
|
||||
upstream_request.replay_input_items(),
|
||||
)
|
||||
&& let Some(input) = self.reasoning_replay.apply(key, &input)
|
||||
{
|
||||
upstream_request
|
||||
.set_replay_input(input)
|
||||
.map_err(map_request_error)?;
|
||||
}
|
||||
let reasoning_replay_capture =
|
||||
reasoning_replay_key.map(|key| self.reasoning_replay.capture(key));
|
||||
let session_capture = (!matches!(
|
||||
context.continuation(),
|
||||
Some(ContinuationBinding::Pinned(_) | ContinuationBinding::External(_))
|
||||
@@ -244,6 +279,7 @@ impl GrokBuildProvider {
|
||||
context,
|
||||
session: Arc::clone(&selected),
|
||||
session_capture,
|
||||
reasoning_replay_capture,
|
||||
},
|
||||
);
|
||||
let stream = ProviderStream::new(metadata, events, selected);
|
||||
@@ -274,7 +310,7 @@ impl GrokBuildProvider {
|
||||
.map(|previous| ProviderAccountId::new(previous.account_id.clone()))
|
||||
.transpose()
|
||||
.map_err(|_| protocol_not_sent())?;
|
||||
let upstream_session_id = previous_session
|
||||
let inherited_session_id = previous_session
|
||||
.as_ref()
|
||||
.and_then(|previous| previous.session_id.clone());
|
||||
let upstream_request = GrokCompactionRequest::encode(
|
||||
@@ -283,6 +319,12 @@ impl GrokBuildProvider {
|
||||
context.client_api_key_ref(),
|
||||
)
|
||||
.map_err(map_request_error)?;
|
||||
let explicit_replay_session_id = upstream_request
|
||||
.reasoning_replay_session_id()
|
||||
.map(str::to_owned);
|
||||
let upstream_session_id = inherited_session_id
|
||||
.clone()
|
||||
.or_else(|| explicit_replay_session_id.clone());
|
||||
let selected = Arc::new(
|
||||
select_grok_session(
|
||||
self.selector.as_ref(),
|
||||
@@ -293,6 +335,16 @@ impl GrokBuildProvider {
|
||||
)
|
||||
.await?,
|
||||
);
|
||||
let reasoning_replay_key = explicit_replay_session_id
|
||||
.as_deref()
|
||||
.or(inherited_session_id.as_deref())
|
||||
.and_then(|session_id| {
|
||||
self.reasoning_replay.key(
|
||||
candidate.upstream_model().as_str(),
|
||||
session_id,
|
||||
selected.account_id().as_str(),
|
||||
)
|
||||
});
|
||||
let allows_account_state_mutation = selected.allows_account_state_mutation();
|
||||
let metadata = provider_call_metadata(candidate, &selected)?;
|
||||
let events = cold_compaction_http_sse_stream(
|
||||
@@ -308,6 +360,8 @@ impl GrokBuildProvider {
|
||||
upstream_session_id,
|
||||
context,
|
||||
session: Arc::clone(&selected),
|
||||
reasoning_replay: self.reasoning_replay.clone(),
|
||||
reasoning_replay_key,
|
||||
},
|
||||
);
|
||||
let stream = ProviderStream::new(metadata, events, selected);
|
||||
@@ -715,46 +769,6 @@ fn portable_output_item(mut item: Value, strip_opaque: bool) -> Option<Value> {
|
||||
Some(item)
|
||||
}
|
||||
|
||||
fn valid_reasoning_ciphertext(value: &str) -> bool {
|
||||
if value.is_empty()
|
||||
|| value != value.trim()
|
||||
|| value.len() > XAI_SESSION_STATE_MAX_BYTES
|
||||
|| value.starts_with("gAAAA")
|
||||
|| value.contains('=')
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let Ok(decoded) = base64::engine::general_purpose::STANDARD_NO_PAD.decode(value) else {
|
||||
return false;
|
||||
};
|
||||
decoded.len() >= MIN_REASONING_CIPHERTEXT_BYTES
|
||||
&& byte_entropy_ratio(&decoded) >= MIN_REASONING_CIPHERTEXT_ENTROPY
|
||||
}
|
||||
|
||||
fn byte_entropy_ratio(value: &[u8]) -> f64 {
|
||||
if value.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
let mut counts = [0_u32; 256];
|
||||
for byte in value {
|
||||
counts[usize::from(*byte)] += 1;
|
||||
}
|
||||
let size = value.len() as f64;
|
||||
let entropy = counts
|
||||
.into_iter()
|
||||
.filter(|count| *count > 0)
|
||||
.map(|count| {
|
||||
let probability = f64::from(count) / size;
|
||||
-probability * probability.log2()
|
||||
})
|
||||
.sum::<f64>();
|
||||
let symbols = value.len().min(256);
|
||||
if symbols <= 1 {
|
||||
return 0.0;
|
||||
}
|
||||
entropy / (symbols as f64).log2()
|
||||
}
|
||||
|
||||
fn has_readable_reasoning(item: &Map<String, Value>) -> bool {
|
||||
["summary", "content"].into_iter().any(|field| {
|
||||
item.get(field)
|
||||
@@ -898,6 +912,7 @@ struct GrokStreamAttempt {
|
||||
context: AttemptContext,
|
||||
session: Arc<SelectedGrokSession>,
|
||||
session_capture: Option<GrokSessionCapture>,
|
||||
reasoning_replay_capture: Option<GrokReasoningReplayCapture>,
|
||||
}
|
||||
|
||||
struct GrokCompactionStreamAttempt {
|
||||
@@ -910,6 +925,8 @@ struct GrokCompactionStreamAttempt {
|
||||
upstream_session_id: Option<String>,
|
||||
context: AttemptContext,
|
||||
session: Arc<SelectedGrokSession>,
|
||||
reasoning_replay: GrokReasoningReplay,
|
||||
reasoning_replay_key: Option<GrokReasoningReplayKey>,
|
||||
}
|
||||
|
||||
struct AcceptedGrokInference {
|
||||
@@ -1044,9 +1061,11 @@ fn cold_compaction_http_sse_stream(
|
||||
upstream_session_id,
|
||||
context,
|
||||
session,
|
||||
reasoning_replay,
|
||||
reasoning_replay_key,
|
||||
} = attempt;
|
||||
Box::pin(async_stream::try_stream! {
|
||||
let mut headers = build_grok_headers(
|
||||
let headers = build_grok_headers(
|
||||
&wire_profile,
|
||||
&session,
|
||||
&client_identity,
|
||||
@@ -1055,12 +1074,6 @@ fn cold_compaction_http_sse_stream(
|
||||
None,
|
||||
&upstream_model,
|
||||
);
|
||||
if let Some(upstream_session_id) = upstream_session_id.as_deref() {
|
||||
headers.push(GrokHeader::sensitive(
|
||||
"x-grok-session-id",
|
||||
crate::credential::SecretValue::new(upstream_session_id.to_owned()),
|
||||
));
|
||||
}
|
||||
let body = request.to_json_bytes().map_err(map_request_error)?;
|
||||
let inference_request = GrokInferenceRequest::new(
|
||||
responses_url,
|
||||
@@ -1147,6 +1160,9 @@ fn cold_compaction_http_sse_stream(
|
||||
if upstream_completed && session.allows_account_state_mutation() {
|
||||
selector.record_success(&session).await;
|
||||
}
|
||||
if let Some(key) = reasoning_replay_key.as_ref() {
|
||||
reasoning_replay.clear(key);
|
||||
}
|
||||
yield ProviderEvent::canonical_with_wire(vec![GatewayEvent::Started(started)], created);
|
||||
yield ProviderEvent::wire(output_done);
|
||||
let mut terminal_facts = facts.accounting;
|
||||
@@ -1245,6 +1261,7 @@ fn cold_http_sse_stream(
|
||||
context,
|
||||
session,
|
||||
mut session_capture,
|
||||
mut reasoning_replay_capture,
|
||||
} = attempt;
|
||||
Box::pin(async_stream::try_stream! {
|
||||
if context.cancellation().is_cancelled() {
|
||||
@@ -1364,6 +1381,9 @@ fn cold_http_sse_stream(
|
||||
.flat_map(ProviderEvent::canonical_facts)
|
||||
.any(|event| matches!(event, GatewayEvent::Completed(_)));
|
||||
attach_xai_session_update(&mut events, &mut session_capture)?;
|
||||
if let Some(capture) = reasoning_replay_capture.as_mut() {
|
||||
capture.observe(&events);
|
||||
}
|
||||
if completed && session.allows_account_state_mutation() {
|
||||
selector.record_success(&session).await;
|
||||
}
|
||||
@@ -1389,6 +1409,9 @@ fn cold_http_sse_stream(
|
||||
.flat_map(ProviderEvent::canonical_facts)
|
||||
.any(|event| matches!(event, GatewayEvent::Completed(_)));
|
||||
attach_xai_session_update(&mut final_events, &mut session_capture)?;
|
||||
if let Some(capture) = reasoning_replay_capture.as_mut() {
|
||||
capture.observe(&final_events);
|
||||
}
|
||||
if completed && session.allows_account_state_mutation() {
|
||||
selector.record_success(&session).await;
|
||||
}
|
||||
@@ -1871,7 +1894,9 @@ const WORKER_LEASE_TTL: Duration = Duration::from_secs(15 * 60);
|
||||
const WORKER_LEASE_RENEWAL: Duration = Duration::from_secs(5 * 60);
|
||||
const OAUTH_REFRESH_INTERVAL: Duration = Duration::from_secs(30);
|
||||
const QUOTA_CATALOG_INTERVAL: Duration = Duration::from_secs(5 * 60);
|
||||
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60);
|
||||
// grok2api 对 Free/未知额度耗尽采用 24 小时恢复探测;当前公共账号状态没有
|
||||
// 独立的 paid period-end 字段,因此 Build 账号统一使用同一保守下限。
|
||||
const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
|
||||
const CLI_RELEASE_WORKER_OWNER: &str = "xai-cli-release";
|
||||
|
||||
pub(crate) fn worker_contributions(
|
||||
|
||||
@@ -0,0 +1,687 @@
|
||||
//! Grok Build 的账号隔离 reasoning replay。
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::sync::{Arc, Mutex, MutexGuard};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use base64::Engine as _;
|
||||
use gateway_core::event::ProviderEvent;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
const REASONING_REPLAY_TTL: Duration = Duration::from_secs(60 * 60);
|
||||
const REASONING_REPLAY_MAX_ENTRIES: usize = 1_024;
|
||||
const REASONING_REPLAY_MAX_BYTES: usize = 8 * 1024 * 1024;
|
||||
const REASONING_REPLAY_MAX_OUTPUT_ITEMS: usize = 4_096;
|
||||
const MIN_REASONING_CIPHERTEXT_BYTES: usize = 50;
|
||||
const MIN_REASONING_CIPHERTEXT_ENTROPY: f64 = 0.85;
|
||||
|
||||
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub(crate) struct GrokReasoningReplayKey {
|
||||
model: String,
|
||||
session_id: String,
|
||||
account_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct GrokReasoningReplay {
|
||||
state: Arc<Mutex<GrokReasoningReplayState>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct GrokReasoningReplayState {
|
||||
entries: BTreeMap<GrokReasoningReplayKey, GrokReasoningReplayEntry>,
|
||||
}
|
||||
|
||||
struct GrokReasoningReplayEntry {
|
||||
stored_at: Instant,
|
||||
items: Vec<Value>,
|
||||
}
|
||||
|
||||
pub(crate) struct GrokReasoningReplayCapture {
|
||||
replay: GrokReasoningReplay,
|
||||
key: GrokReasoningReplayKey,
|
||||
output_items: BTreeMap<u32, Value>,
|
||||
captured_bytes: usize,
|
||||
truncated: bool,
|
||||
finished: bool,
|
||||
}
|
||||
|
||||
impl Default for GrokReasoningReplay {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl GrokReasoningReplay {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
state: Arc::new(Mutex::new(GrokReasoningReplayState::default())),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn key(
|
||||
&self,
|
||||
model: &str,
|
||||
session_id: &str,
|
||||
account_id: &str,
|
||||
) -> Option<GrokReasoningReplayKey> {
|
||||
let model = model.trim();
|
||||
let session_id = session_id.trim();
|
||||
let account_id = account_id.trim();
|
||||
if model.is_empty() || session_id.is_empty() || account_id.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(GrokReasoningReplayKey {
|
||||
model: model.to_owned(),
|
||||
session_id: session_id.to_owned(),
|
||||
account_id: account_id.to_owned(),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn apply(
|
||||
&self,
|
||||
key: &GrokReasoningReplayKey,
|
||||
input: &[Value],
|
||||
) -> Option<Vec<Value>> {
|
||||
if input.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let cached = self.read(key)?;
|
||||
let filtered = filter_replay_items_for_input(input, cached)?;
|
||||
Some(insert_replay_items(input, &filtered))
|
||||
}
|
||||
|
||||
pub(crate) fn capture(&self, key: GrokReasoningReplayKey) -> GrokReasoningReplayCapture {
|
||||
GrokReasoningReplayCapture {
|
||||
replay: self.clone(),
|
||||
key,
|
||||
output_items: BTreeMap::new(),
|
||||
captured_bytes: 0,
|
||||
truncated: false,
|
||||
finished: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn clear(&self, key: &GrokReasoningReplayKey) {
|
||||
self.lock_state().entries.remove(key);
|
||||
}
|
||||
|
||||
fn read(&self, key: &GrokReasoningReplayKey) -> Option<Vec<Value>> {
|
||||
let now = Instant::now();
|
||||
let mut state = self.lock_state();
|
||||
prune_expired(&mut state, now);
|
||||
state.entries.get(key).map(|entry| entry.items.clone())
|
||||
}
|
||||
|
||||
fn store_output(&self, key: GrokReasoningReplayKey, output: Vec<Value>) {
|
||||
let normalized = match normalize_replay_items(output) {
|
||||
ReplayNormalization::Store(items) => items,
|
||||
ReplayNormalization::Delete => {
|
||||
self.clear(&key);
|
||||
return;
|
||||
}
|
||||
ReplayNormalization::Oversized => return,
|
||||
};
|
||||
let now = Instant::now();
|
||||
let mut state = self.lock_state();
|
||||
prune_expired(&mut state, now);
|
||||
if !state.entries.contains_key(&key)
|
||||
&& state.entries.len() >= REASONING_REPLAY_MAX_ENTRIES
|
||||
&& let Some(oldest) = state
|
||||
.entries
|
||||
.iter()
|
||||
.min_by_key(|(_, entry)| entry.stored_at)
|
||||
.map(|(key, _)| key.clone())
|
||||
{
|
||||
state.entries.remove(&oldest);
|
||||
}
|
||||
state.entries.insert(
|
||||
key,
|
||||
GrokReasoningReplayEntry {
|
||||
stored_at: now,
|
||||
items: normalized,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
fn lock_state(&self) -> MutexGuard<'_, GrokReasoningReplayState> {
|
||||
match self.state.lock() {
|
||||
Ok(state) => state,
|
||||
Err(poisoned) => poisoned.into_inner(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl GrokReasoningReplayCapture {
|
||||
pub(crate) fn observe(&mut self, events: &[ProviderEvent]) {
|
||||
if self.finished {
|
||||
return;
|
||||
}
|
||||
let mut completed = false;
|
||||
for event in events {
|
||||
let Some(wire) = event.wire_event() else {
|
||||
continue;
|
||||
};
|
||||
let payload = wire.data();
|
||||
let event_type = wire
|
||||
.event_type()
|
||||
.or_else(|| payload.get("type").and_then(Value::as_str));
|
||||
if event_type == Some("response.output_item.done")
|
||||
&& let Some(item) = payload.get("item").cloned()
|
||||
{
|
||||
let index = payload
|
||||
.get("output_index")
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|value| u32::try_from(value).ok())
|
||||
.unwrap_or_else(|| u32::try_from(self.output_items.len()).unwrap_or(u32::MAX));
|
||||
self.capture_item(index, item);
|
||||
}
|
||||
if event_type == Some("response.completed") {
|
||||
capture_terminal_output(payload, self);
|
||||
completed = true;
|
||||
}
|
||||
}
|
||||
if completed {
|
||||
self.finished = true;
|
||||
if self.truncated {
|
||||
self.output_items.clear();
|
||||
return;
|
||||
}
|
||||
let output = std::mem::take(&mut self.output_items)
|
||||
.into_values()
|
||||
.collect();
|
||||
self.replay.store_output(self.key.clone(), output);
|
||||
}
|
||||
}
|
||||
|
||||
fn capture_item(&mut self, index: u32, item: Value) {
|
||||
if self.truncated || self.output_items.contains_key(&index) {
|
||||
return;
|
||||
}
|
||||
if self.output_items.len() >= REASONING_REPLAY_MAX_OUTPUT_ITEMS {
|
||||
self.truncate();
|
||||
return;
|
||||
}
|
||||
let Some(captured_bytes) = serde_json::to_vec(&item)
|
||||
.ok()
|
||||
.and_then(|encoded| self.captured_bytes.checked_add(encoded.len()))
|
||||
.filter(|captured_bytes| *captured_bytes <= REASONING_REPLAY_MAX_BYTES)
|
||||
else {
|
||||
self.truncate();
|
||||
return;
|
||||
};
|
||||
self.captured_bytes = captured_bytes;
|
||||
self.output_items.insert(index, item);
|
||||
}
|
||||
|
||||
fn truncate(&mut self) {
|
||||
self.truncated = true;
|
||||
self.output_items.clear();
|
||||
self.captured_bytes = 0;
|
||||
}
|
||||
}
|
||||
|
||||
fn capture_terminal_output(payload: &Value, capture: &mut GrokReasoningReplayCapture) {
|
||||
let Some(output) = payload
|
||||
.get("response")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|response| response.get("output"))
|
||||
.and_then(Value::as_array)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if output.len() > REASONING_REPLAY_MAX_OUTPUT_ITEMS {
|
||||
capture.truncate();
|
||||
return;
|
||||
}
|
||||
for (index, item) in output
|
||||
.iter()
|
||||
.take(REASONING_REPLAY_MAX_OUTPUT_ITEMS)
|
||||
.enumerate()
|
||||
{
|
||||
let Some(index) = u32::try_from(index).ok() else {
|
||||
break;
|
||||
};
|
||||
capture.capture_item(index, item.clone());
|
||||
}
|
||||
}
|
||||
|
||||
fn prune_expired(state: &mut GrokReasoningReplayState, now: Instant) {
|
||||
state
|
||||
.entries
|
||||
.retain(|_, entry| now.saturating_duration_since(entry.stored_at) < REASONING_REPLAY_TTL);
|
||||
}
|
||||
|
||||
enum ReplayNormalization {
|
||||
Store(Vec<Value>),
|
||||
Delete,
|
||||
Oversized,
|
||||
}
|
||||
|
||||
fn normalize_replay_items(items: Vec<Value>) -> ReplayNormalization {
|
||||
let mut normalized = Vec::new();
|
||||
let mut has_anchor = false;
|
||||
for item in items {
|
||||
let Some(item) = normalize_replay_item(item) else {
|
||||
continue;
|
||||
};
|
||||
if matches!(
|
||||
item.get("type").and_then(Value::as_str),
|
||||
Some("reasoning" | "function_call" | "custom_tool_call")
|
||||
) {
|
||||
has_anchor = true;
|
||||
}
|
||||
normalized.push(item);
|
||||
}
|
||||
if !has_anchor || normalized.is_empty() {
|
||||
return ReplayNormalization::Delete;
|
||||
}
|
||||
match serde_json::to_vec(&normalized) {
|
||||
Ok(encoded) if encoded.len() <= REASONING_REPLAY_MAX_BYTES => {
|
||||
ReplayNormalization::Store(normalized)
|
||||
}
|
||||
_ => ReplayNormalization::Oversized,
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_replay_item(item: Value) -> Option<Value> {
|
||||
let object = item.as_object()?;
|
||||
match object.get("type").and_then(Value::as_str)?.trim() {
|
||||
"reasoning" => normalize_reasoning_item(object),
|
||||
"message" => normalize_assistant_message_item(object),
|
||||
"function_call" => normalize_function_call_item(object),
|
||||
"custom_tool_call" => normalize_custom_tool_call_item(object),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_reasoning_item(item: &Map<String, Value>) -> Option<Value> {
|
||||
let encrypted = item.get("encrypted_content")?.as_str()?;
|
||||
if !valid_reasoning_ciphertext(encrypted) {
|
||||
return None;
|
||||
}
|
||||
Some(Value::Object(Map::from_iter([
|
||||
("type".to_owned(), Value::String("reasoning".to_owned())),
|
||||
("summary".to_owned(), Value::Array(Vec::new())),
|
||||
("content".to_owned(), Value::Null),
|
||||
(
|
||||
"encrypted_content".to_owned(),
|
||||
Value::String(encrypted.to_owned()),
|
||||
),
|
||||
])))
|
||||
}
|
||||
|
||||
fn normalize_assistant_message_item(item: &Map<String, Value>) -> Option<Value> {
|
||||
if !item
|
||||
.get("role")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|role| role.trim().eq_ignore_ascii_case("assistant"))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let parts = item
|
||||
.get("content")?
|
||||
.as_array()?
|
||||
.iter()
|
||||
.filter_map(normalize_assistant_part)
|
||||
.collect::<Vec<_>>();
|
||||
if parts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(Value::Object(Map::from_iter([
|
||||
("type".to_owned(), Value::String("message".to_owned())),
|
||||
("role".to_owned(), Value::String("assistant".to_owned())),
|
||||
("content".to_owned(), Value::Array(parts)),
|
||||
])))
|
||||
}
|
||||
|
||||
fn normalize_assistant_part(part: &Value) -> Option<Value> {
|
||||
let part = part.as_object()?;
|
||||
match part.get("type").and_then(Value::as_str)?.trim() {
|
||||
"output_text" => Some(Value::Object(Map::from_iter([
|
||||
("type".to_owned(), Value::String("output_text".to_owned())),
|
||||
(
|
||||
"text".to_owned(),
|
||||
Value::String(part.get("text")?.as_str()?.to_owned()),
|
||||
),
|
||||
]))),
|
||||
"refusal" => Some(Value::Object(Map::from_iter([
|
||||
("type".to_owned(), Value::String("refusal".to_owned())),
|
||||
(
|
||||
"refusal".to_owned(),
|
||||
Value::String(part.get("refusal")?.as_str()?.to_owned()),
|
||||
),
|
||||
]))),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_function_call_item(item: &Map<String, Value>) -> Option<Value> {
|
||||
let call_id = nonempty_string(item, "call_id")?;
|
||||
let name = nonempty_string(item, "name")?;
|
||||
let arguments = item.get("arguments")?.as_str()?.to_owned();
|
||||
Some(Value::Object(Map::from_iter([
|
||||
("type".to_owned(), Value::String("function_call".to_owned())),
|
||||
("call_id".to_owned(), Value::String(call_id)),
|
||||
("name".to_owned(), Value::String(name)),
|
||||
("arguments".to_owned(), Value::String(arguments)),
|
||||
])))
|
||||
}
|
||||
|
||||
fn normalize_custom_tool_call_item(item: &Map<String, Value>) -> Option<Value> {
|
||||
let call_id = nonempty_string(item, "call_id")?;
|
||||
let name = nonempty_string(item, "name")?;
|
||||
let input = item.get("input")?.clone();
|
||||
let status = item
|
||||
.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("completed")
|
||||
.to_owned();
|
||||
Some(Value::Object(Map::from_iter([
|
||||
(
|
||||
"type".to_owned(),
|
||||
Value::String("custom_tool_call".to_owned()),
|
||||
),
|
||||
("status".to_owned(), Value::String(status)),
|
||||
("call_id".to_owned(), Value::String(call_id)),
|
||||
("name".to_owned(), Value::String(name)),
|
||||
("input".to_owned(), input),
|
||||
])))
|
||||
}
|
||||
|
||||
fn nonempty_string(item: &Map<String, Value>, field: &str) -> Option<String> {
|
||||
item.get(field)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn filter_replay_items_for_input(input: &[Value], cached: Vec<Value>) -> Option<Vec<Value>> {
|
||||
let last_assistant = input.iter().rev().find_map(assistant_message);
|
||||
let cached_assistant = cached.iter().find_map(assistant_message);
|
||||
let assistant_matches = match (last_assistant, cached_assistant) {
|
||||
(Some(current), Some(previous)) => {
|
||||
if !assistant_content_equal(current, previous) {
|
||||
return None;
|
||||
}
|
||||
true
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
|
||||
let mut existing_calls = BTreeSet::new();
|
||||
let mut existing_outputs = BTreeMap::new();
|
||||
let mut existing_encrypted = BTreeSet::new();
|
||||
for item in input {
|
||||
let Some(item) = item.as_object() else {
|
||||
continue;
|
||||
};
|
||||
match item.get("type").and_then(Value::as_str).map(str::trim) {
|
||||
Some("reasoning") => {
|
||||
if let Some(encrypted) = item
|
||||
.get("encrypted_content")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
existing_encrypted.insert(encrypted.to_owned());
|
||||
}
|
||||
}
|
||||
Some("function_call_output" | "custom_tool_call_output") => {
|
||||
if let Some(call_id) = item.get("call_id").and_then(Value::as_str) {
|
||||
for candidate in comparable_call_ids(call_id) {
|
||||
existing_outputs.insert(candidate, call_id.trim().to_owned());
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(kind @ ("function_call" | "custom_tool_call")) => {
|
||||
if let Some(call_id) = item.get("call_id").and_then(Value::as_str) {
|
||||
existing_calls.extend(replay_call_keys(kind, call_id));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
let mut filtered = Vec::new();
|
||||
for mut item in cached {
|
||||
let kind = item
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default()
|
||||
.to_owned();
|
||||
match kind.as_str() {
|
||||
"reasoning" => {
|
||||
if item
|
||||
.get("encrypted_content")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|encrypted| existing_encrypted.contains(encrypted))
|
||||
{
|
||||
continue;
|
||||
}
|
||||
}
|
||||
"message" if assistant_matches => continue,
|
||||
"message" => {}
|
||||
"function_call" | "custom_tool_call" => {
|
||||
let Some(call_id) = item.get("call_id").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
let keys = replay_call_keys(&kind, call_id);
|
||||
if keys.is_empty() || keys.iter().any(|key| existing_calls.contains(key)) {
|
||||
continue;
|
||||
}
|
||||
let output_call_id = comparable_call_ids(call_id)
|
||||
.into_iter()
|
||||
.find_map(|candidate| existing_outputs.get(&candidate).cloned());
|
||||
let Some(output_call_id) = output_call_id else {
|
||||
continue;
|
||||
};
|
||||
existing_calls.extend(keys);
|
||||
if output_call_id != call_id
|
||||
&& let Some(item) = item.as_object_mut()
|
||||
{
|
||||
item.insert("call_id".to_owned(), Value::String(output_call_id));
|
||||
}
|
||||
}
|
||||
_ => continue,
|
||||
}
|
||||
filtered.push(item);
|
||||
}
|
||||
(!filtered.is_empty()).then_some(filtered)
|
||||
}
|
||||
|
||||
fn assistant_message(item: &Value) -> Option<&Map<String, Value>> {
|
||||
let item = item.as_object()?;
|
||||
let item_type = item
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
let role = item.get("role")?.as_str()?.trim();
|
||||
((item_type.is_empty() || item_type == "message") && role.eq_ignore_ascii_case("assistant"))
|
||||
.then_some(item)
|
||||
}
|
||||
|
||||
fn assistant_content_equal(left: &Map<String, Value>, right: &Map<String, Value>) -> bool {
|
||||
let Some(left) = assistant_parts(left.get("content")) else {
|
||||
return false;
|
||||
};
|
||||
let Some(right) = assistant_parts(right.get("content")) else {
|
||||
return false;
|
||||
};
|
||||
left == right
|
||||
}
|
||||
|
||||
fn assistant_parts(content: Option<&Value>) -> Option<Vec<(String, String)>> {
|
||||
match content? {
|
||||
Value::String(text) => Some(vec![("output_text".to_owned(), text.clone())]),
|
||||
Value::Array(parts) => {
|
||||
let mut result = Vec::new();
|
||||
for part in parts {
|
||||
let part = part.as_object()?;
|
||||
let kind = part.get("type")?.as_str()?.trim();
|
||||
let value = match kind {
|
||||
"output_text" => part.get("text")?.as_str()?,
|
||||
"refusal" => part.get("refusal")?.as_str()?,
|
||||
_ => return None,
|
||||
};
|
||||
result.push((kind.to_owned(), value.to_owned()));
|
||||
}
|
||||
(!result.is_empty()).then_some(result)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn comparable_call_ids(call_id: &str) -> Vec<String> {
|
||||
let call_id = call_id.trim();
|
||||
if call_id.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
if let Some(upstream) = call_id
|
||||
.strip_prefix("toolu_")
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return vec![call_id.to_owned(), upstream.to_owned()];
|
||||
}
|
||||
vec![call_id.to_owned(), format!("toolu_{call_id}")]
|
||||
}
|
||||
|
||||
fn replay_call_keys(item_type: &str, call_id: &str) -> Vec<String> {
|
||||
if !matches!(item_type, "function_call" | "custom_tool_call") {
|
||||
return Vec::new();
|
||||
}
|
||||
comparable_call_ids(call_id)
|
||||
.into_iter()
|
||||
.map(|call_id| format!("{item_type}\0{call_id}"))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn insert_replay_items(input: &[Value], replay: &[Value]) -> Vec<Value> {
|
||||
let insert_at = replay_insert_index(input, replay);
|
||||
let mut result = Vec::new();
|
||||
for (index, item) in input.iter().enumerate() {
|
||||
if index == insert_at {
|
||||
result.extend(replay.iter().cloned());
|
||||
}
|
||||
result.push(item.clone());
|
||||
}
|
||||
if insert_at == input.len() {
|
||||
result.extend(replay.iter().cloned());
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn replay_insert_index(input: &[Value], replay: &[Value]) -> usize {
|
||||
let replay_call_ids = replay
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.filter(|item| {
|
||||
matches!(
|
||||
item.get("type").and_then(Value::as_str),
|
||||
Some("function_call" | "custom_tool_call")
|
||||
)
|
||||
})
|
||||
.filter_map(|item| item.get("call_id").and_then(Value::as_str))
|
||||
.flat_map(comparable_call_ids)
|
||||
.collect::<BTreeSet<_>>();
|
||||
if !replay_call_ids.is_empty() {
|
||||
for (index, item) in input.iter().enumerate() {
|
||||
let Some(item) = item.as_object() else {
|
||||
continue;
|
||||
};
|
||||
if !matches!(
|
||||
item.get("type").and_then(Value::as_str),
|
||||
Some("function_call_output" | "custom_tool_call_output")
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
let call_ids = item
|
||||
.get("call_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(comparable_call_ids)
|
||||
.unwrap_or_default();
|
||||
if call_ids.is_empty()
|
||||
|| call_ids
|
||||
.iter()
|
||||
.any(|call_id| replay_call_ids.contains(call_id))
|
||||
{
|
||||
return index;
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(index) = input
|
||||
.iter()
|
||||
.rposition(|item| assistant_message(item).is_some())
|
||||
{
|
||||
return index;
|
||||
}
|
||||
input
|
||||
.iter()
|
||||
.position(should_insert_replay_before)
|
||||
.unwrap_or(input.len())
|
||||
}
|
||||
|
||||
fn should_insert_replay_before(item: &Value) -> bool {
|
||||
let Some(item) = item.as_object() else {
|
||||
return true;
|
||||
};
|
||||
let item_type = item
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
let role = item
|
||||
.get("role")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
if role.is_empty() || (!item_type.is_empty() && item_type != "message") {
|
||||
return true;
|
||||
}
|
||||
!matches!(role.as_str(), "system" | "developer")
|
||||
}
|
||||
|
||||
pub(crate) fn valid_reasoning_ciphertext(value: &str) -> bool {
|
||||
if value.is_empty()
|
||||
|| value != value.trim()
|
||||
|| value.len() > REASONING_REPLAY_MAX_BYTES
|
||||
|| value.starts_with("gAAAA")
|
||||
|| value.contains('=')
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let Ok(decoded) = base64::engine::general_purpose::STANDARD_NO_PAD.decode(value) else {
|
||||
return false;
|
||||
};
|
||||
decoded.len() >= MIN_REASONING_CIPHERTEXT_BYTES
|
||||
&& byte_entropy_ratio(&decoded) >= MIN_REASONING_CIPHERTEXT_ENTROPY
|
||||
}
|
||||
|
||||
fn byte_entropy_ratio(value: &[u8]) -> f64 {
|
||||
if value.is_empty() {
|
||||
return 0.0;
|
||||
}
|
||||
let mut counts = [0_u32; 256];
|
||||
for byte in value {
|
||||
counts[usize::from(*byte)] += 1;
|
||||
}
|
||||
let size = value.len() as f64;
|
||||
let entropy = counts
|
||||
.into_iter()
|
||||
.filter(|count| *count > 0)
|
||||
.map(|count| {
|
||||
let probability = f64::from(count) / size;
|
||||
-probability * probability.log2()
|
||||
})
|
||||
.sum::<f64>();
|
||||
let symbols = value.len().min(256);
|
||||
if symbols <= 1 {
|
||||
return 0.0;
|
||||
}
|
||||
entropy / (symbols as f64).log2()
|
||||
}
|
||||
@@ -42,6 +42,7 @@ If the prior conversation contains a note about files at /tmp/compaction/segment
|
||||
pub struct GrokCompactionRequest {
|
||||
body: Map<String, Value>,
|
||||
affinity: Option<GrokSessionAffinityKey>,
|
||||
reasoning_replay_session_id: Option<String>,
|
||||
}
|
||||
|
||||
impl GrokCompactionRequest {
|
||||
@@ -60,6 +61,8 @@ impl GrokCompactionRequest {
|
||||
upstream_model,
|
||||
client_api_key_ref,
|
||||
)?;
|
||||
let reasoning_replay_session_id =
|
||||
normalized.reasoning_replay_session_id().map(str::to_owned);
|
||||
let mut body = normalized.body().clone();
|
||||
let mut input = normalized.input_items();
|
||||
input.push(summary_prompt_item());
|
||||
@@ -95,6 +98,7 @@ impl GrokCompactionRequest {
|
||||
Ok(Self {
|
||||
body,
|
||||
affinity: normalized.affinity().cloned(),
|
||||
reasoning_replay_session_id,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -110,6 +114,12 @@ impl GrokCompactionRequest {
|
||||
self.affinity.as_ref()
|
||||
}
|
||||
|
||||
/// 返回用于在成功 compact 后清除 reasoning replay 的显式会话身份。
|
||||
#[must_use]
|
||||
pub(crate) fn reasoning_replay_session_id(&self) -> Option<&str> {
|
||||
self.reasoning_replay_session_id.as_deref()
|
||||
}
|
||||
|
||||
/// 序列化上游请求正文。
|
||||
///
|
||||
/// # Errors
|
||||
|
||||
@@ -142,11 +142,11 @@ pub fn build_grok_headers(
|
||||
"x-grok-conv-id",
|
||||
SecretValue::new(upstream_session_id.to_owned()),
|
||||
));
|
||||
headers.push(GrokHeader::sensitive(
|
||||
"x-grok-session-id",
|
||||
SecretValue::new(upstream_session_id.to_owned()),
|
||||
));
|
||||
if let Some(turn_index) = turn_index {
|
||||
headers.push(GrokHeader::sensitive(
|
||||
"x-grok-session-id",
|
||||
SecretValue::new(upstream_session_id.to_owned()),
|
||||
));
|
||||
headers.push(GrokHeader::public("x-grok-turn-idx", turn_index));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -56,6 +56,7 @@ const SESSION_FIELDS: &[&str] = &[
|
||||
"conversation_id",
|
||||
"conversationId",
|
||||
];
|
||||
const PROTOCOL_CONTEXT_SESSION_FIELDS: &[&str] = &["conversation_id", "session_id", "thread_id"];
|
||||
const FOREIGN_CLIENT_METADATA_FIELDS: &[&str] = &["x-openai-subagent"];
|
||||
const MAX_SESSION_SEED_BYTES: usize = 1_024;
|
||||
const GROK_CACHE_ROUTE_TOOLS: &[&str] = &["web_search", "x_search"];
|
||||
@@ -68,6 +69,7 @@ const GROK_INTERNAL_HISTORY_KEYS: &[&str] =
|
||||
pub struct GrokResponsesRequest {
|
||||
body: Map<String, Value>,
|
||||
session_id: Option<String>,
|
||||
reasoning_replay_session_id: Option<String>,
|
||||
affinity: Option<GrokSessionAffinityKey>,
|
||||
response_transform: GrokResponseTransform,
|
||||
}
|
||||
@@ -85,6 +87,23 @@ impl GrokResponsesRequest {
|
||||
self.session_id.as_deref()
|
||||
}
|
||||
|
||||
/// 返回只有显式会话身份才启用的 reasoning replay scope。
|
||||
pub(crate) fn reasoning_replay_session_id(&self) -> Option<&str> {
|
||||
self.reasoning_replay_session_id.as_deref()
|
||||
}
|
||||
|
||||
pub(crate) fn has_previous_response_id(&self) -> bool {
|
||||
self.body
|
||||
.get("previous_response_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
pub(crate) fn replay_input_items(&self) -> Option<Vec<Value>> {
|
||||
self.body.get("input").and_then(Value::as_array).cloned()
|
||||
}
|
||||
|
||||
/// 返回与会话一致、额外绑定模型的账号亲和键。
|
||||
#[must_use]
|
||||
pub const fn affinity(&self) -> Option<&GrokSessionAffinityKey> {
|
||||
@@ -149,6 +168,9 @@ impl GrokResponsesRequest {
|
||||
return;
|
||||
};
|
||||
self.session_id = Some(session_id.to_owned());
|
||||
if self.reasoning_replay_session_id.is_some() {
|
||||
self.reasoning_replay_session_id = Some(session_id.to_owned());
|
||||
}
|
||||
self.affinity = None;
|
||||
self.body.insert(
|
||||
"prompt_cache_key".to_owned(),
|
||||
@@ -159,6 +181,7 @@ impl GrokResponsesRequest {
|
||||
pub(crate) fn clear_session(&mut self) {
|
||||
self.body.remove("prompt_cache_key");
|
||||
self.session_id = None;
|
||||
self.reasoning_replay_session_id = None;
|
||||
self.affinity = None;
|
||||
}
|
||||
|
||||
@@ -207,13 +230,16 @@ impl GrokResponsesRequest {
|
||||
);
|
||||
sanitize_account_identity(&mut body);
|
||||
sanitize_foreign_client_metadata(&mut body);
|
||||
let response_transform = normalize_responses_request(&mut body)?;
|
||||
let mut response_transform = normalize_responses_request(&mut body)?;
|
||||
response_transform.observe_client_cache_tools();
|
||||
if enable_cache_route {
|
||||
enable_grok_prompt_cache_route(&mut body);
|
||||
enable_grok_prompt_cache_route(&mut body, upstream_model, &mut response_transform);
|
||||
}
|
||||
let (session_id, affinity) = identity.map_or((None, None), |(session_id, affinity)| {
|
||||
(Some(session_id), Some(affinity))
|
||||
});
|
||||
let reasoning_replay_session_id =
|
||||
session_seed.is_some().then(|| session_id.clone()).flatten();
|
||||
match session_id.as_ref() {
|
||||
Some(session_id) => {
|
||||
body.insert(
|
||||
@@ -230,6 +256,7 @@ impl GrokResponsesRequest {
|
||||
Ok(Self {
|
||||
body,
|
||||
session_id,
|
||||
reasoning_replay_session_id,
|
||||
affinity,
|
||||
response_transform,
|
||||
})
|
||||
@@ -321,23 +348,61 @@ fn sanitize_foreign_client_metadata(body: &mut Map<String, Value>) {
|
||||
}
|
||||
}
|
||||
|
||||
fn enable_grok_prompt_cache_route(body: &mut Map<String, Value>) {
|
||||
fn enable_grok_prompt_cache_route(
|
||||
body: &mut Map<String, Value>,
|
||||
upstream_model: &str,
|
||||
response: &mut GrokResponseTransform,
|
||||
) {
|
||||
let tools = body
|
||||
.entry("tools".to_owned())
|
||||
.or_insert_with(|| Value::Array(Vec::new()));
|
||||
let Value::Array(tools) = tools else {
|
||||
return;
|
||||
};
|
||||
// 缓存路由不得改变模型可调用的工具集:客户端带工具时无法既保留其
|
||||
// tool_choice 语义又阻止注入的原生搜索被调用,因此直接不注入。
|
||||
if !tools.is_empty() {
|
||||
response.observe_upstream_cache_tools(tools);
|
||||
if is_cache_media_model(upstream_model) || has_tool_type(tools, "image_generation") {
|
||||
return;
|
||||
}
|
||||
for tool in GROK_CACHE_ROUTE_TOOLS {
|
||||
tools.push(json_object([("type", Value::String((*tool).to_owned()))]));
|
||||
if tools.is_empty() {
|
||||
for tool in GROK_CACHE_ROUTE_TOOLS {
|
||||
tools.push(json_object([("type", Value::String((*tool).to_owned()))]));
|
||||
response.mark_injected_cache_tool(tool);
|
||||
}
|
||||
// 无客户端工具时以 none 选中缓存路由,同时不授予搜索能力。
|
||||
body.insert("tool_choice".to_owned(), Value::String("none".to_owned()));
|
||||
} else if !has_tool_type(tools, "x_search") {
|
||||
tools.push(json_object([(
|
||||
"type",
|
||||
Value::String("x_search".to_owned()),
|
||||
)]));
|
||||
response.mark_injected_cache_tool("x_search");
|
||||
}
|
||||
append_x_search_to_allowed_tools(body);
|
||||
}
|
||||
|
||||
fn is_cache_media_model(model: &str) -> bool {
|
||||
let model = model.trim().to_ascii_lowercase();
|
||||
["image", "imagine", "video"]
|
||||
.into_iter()
|
||||
.any(|marker| model.contains(marker))
|
||||
}
|
||||
|
||||
fn append_x_search_to_allowed_tools(body: &mut Map<String, Value>) {
|
||||
let Some(choice) = body.get_mut("tool_choice").and_then(Value::as_object_mut) else {
|
||||
return;
|
||||
};
|
||||
if choice.get("type").and_then(Value::as_str) != Some("allowed_tools") {
|
||||
return;
|
||||
}
|
||||
let Some(allowed) = choice.get_mut("tools").and_then(Value::as_array_mut) else {
|
||||
return;
|
||||
};
|
||||
if !has_tool_type(allowed, "x_search") {
|
||||
allowed.push(json_object([(
|
||||
"type",
|
||||
Value::String("x_search".to_owned()),
|
||||
)]));
|
||||
}
|
||||
// 禁用工具调用,使原生搜索声明只承担缓存路由作用。
|
||||
body.insert("tool_choice".to_owned(), Value::String("none".to_owned()));
|
||||
}
|
||||
|
||||
fn explicit_session_seed(request: &GenerateRequest, body: &Map<String, Value>) -> Option<String> {
|
||||
@@ -351,6 +416,17 @@ fn explicit_session_seed(request: &GenerateRequest, body: &Map<String, Value>) -
|
||||
.and_then(Value::as_object)
|
||||
.and_then(first_session_value)
|
||||
})
|
||||
.or_else(|| {
|
||||
PROTOCOL_CONTEXT_SESSION_FIELDS.iter().find_map(|field| {
|
||||
request
|
||||
.protocol_payload()
|
||||
.context()
|
||||
.get(*field)
|
||||
.and_then(Value::as_str)
|
||||
.and_then(valid_session_seed)
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn first_session_value(body: &Map<String, Value>) -> Option<String> {
|
||||
@@ -674,6 +750,11 @@ pub(crate) struct GrokResponseTransform {
|
||||
aliases: BTreeMap<String, ToolIdentity>,
|
||||
visible_tools: Vec<Value>,
|
||||
legacy_local_shell: bool,
|
||||
filter_x_search: bool,
|
||||
injected_tool_types: BTreeSet<String>,
|
||||
client_declared_tools: BTreeSet<String>,
|
||||
dropped_output_indexes: BTreeSet<u64>,
|
||||
dropped_item_ids: BTreeSet<String>,
|
||||
stream_calls: BTreeMap<String, StreamCallState>,
|
||||
stream_keys: BTreeMap<String, String>,
|
||||
}
|
||||
@@ -698,7 +779,42 @@ impl GrokTransformedWireEvent {
|
||||
impl GrokResponseTransform {
|
||||
#[must_use]
|
||||
pub(crate) fn is_empty(&self) -> bool {
|
||||
self.aliases.is_empty() && self.visible_tools.is_empty() && !self.legacy_local_shell
|
||||
self.aliases.is_empty()
|
||||
&& self.visible_tools.is_empty()
|
||||
&& !self.legacy_local_shell
|
||||
&& !self.filter_x_search
|
||||
&& self.injected_tool_types.is_empty()
|
||||
}
|
||||
|
||||
fn observe_client_cache_tools(&mut self) {
|
||||
for tool in &self.visible_tools {
|
||||
let kind = tool.get("type").and_then(Value::as_str).unwrap_or_default();
|
||||
if kind == "x_search" {
|
||||
self.filter_x_search = true;
|
||||
}
|
||||
if matches!(kind, "function" | "custom")
|
||||
&& let Some(name) = tool
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|name| !name.is_empty())
|
||||
{
|
||||
self.client_declared_tools.insert(name.to_owned());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn observe_upstream_cache_tools(&mut self, tools: &[Value]) {
|
||||
if has_tool_type(tools, "x_search") {
|
||||
self.filter_x_search = true;
|
||||
}
|
||||
}
|
||||
|
||||
fn mark_injected_cache_tool(&mut self, kind: &str) {
|
||||
self.injected_tool_types.insert(kind.to_owned());
|
||||
if kind == "x_search" {
|
||||
self.filter_x_search = true;
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn rewrite_stream_event(
|
||||
@@ -706,6 +822,20 @@ impl GrokResponseTransform {
|
||||
event_type: &str,
|
||||
mut value: Value,
|
||||
) -> Result<Vec<GrokTransformedWireEvent>, GrokRequestEncodeError> {
|
||||
if self.cache_filter_enabled() {
|
||||
if value
|
||||
.get("item")
|
||||
.is_some_and(|item| self.is_internal_cache_call(item))
|
||||
{
|
||||
self.record_dropped_cache_item(&value);
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.filter_cache_envelope(&mut value)?;
|
||||
if self.references_dropped_cache_item(&value) {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
self.compact_cache_output_index(&mut value);
|
||||
}
|
||||
if event_type == "response.output_item.added"
|
||||
&& let Some(item) = value.get("item").and_then(Value::as_object)
|
||||
&& let Some(primary) = self.remember_stream_call(item)
|
||||
@@ -1031,9 +1161,145 @@ impl GrokResponseTransform {
|
||||
])
|
||||
}
|
||||
|
||||
fn cache_filter_enabled(&self) -> bool {
|
||||
self.filter_x_search || !self.injected_tool_types.is_empty()
|
||||
}
|
||||
|
||||
fn is_internal_cache_call(&self, item: &Value) -> bool {
|
||||
let Some(item) = item.as_object() else {
|
||||
return false;
|
||||
};
|
||||
let kind = string_field(item, "type").trim();
|
||||
if kind == "web_search_call" {
|
||||
return self.injected_tool_types.contains("web_search");
|
||||
}
|
||||
if !matches!(kind, "custom_tool_call" | "function_call") {
|
||||
return false;
|
||||
}
|
||||
if string_field(item, "call_id").trim().starts_with("xs_call") {
|
||||
return true;
|
||||
}
|
||||
let name = string_field(item, "name").trim();
|
||||
matches!(
|
||||
name,
|
||||
"x_user_search" | "x_semantic_search" | "x_keyword_search" | "x_thread_fetch"
|
||||
) && string_field(item, "namespace").trim().is_empty()
|
||||
&& !self.client_declared_tools.contains(name)
|
||||
}
|
||||
|
||||
fn record_dropped_cache_item(&mut self, payload: &Value) {
|
||||
if let Some(index) = payload.get("output_index").and_then(Value::as_u64) {
|
||||
self.dropped_output_indexes.insert(index);
|
||||
}
|
||||
let Some(item) = payload.get("item").and_then(Value::as_object) else {
|
||||
return;
|
||||
};
|
||||
for field in ["id", "call_id"] {
|
||||
if let Some(value) = item
|
||||
.get(field)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
self.dropped_item_ids.insert(value.to_owned());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn references_dropped_cache_item(&self, payload: &Value) -> bool {
|
||||
if payload
|
||||
.get("output_index")
|
||||
.and_then(Value::as_u64)
|
||||
.is_some_and(|index| self.dropped_output_indexes.contains(&index))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
["item_id", "call_id"].into_iter().any(|field| {
|
||||
payload
|
||||
.get(field)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| self.dropped_item_ids.contains(value))
|
||||
})
|
||||
}
|
||||
|
||||
fn compact_cache_output_index(&self, payload: &mut Value) {
|
||||
let Some(index) = payload.get("output_index").and_then(Value::as_u64) else {
|
||||
return;
|
||||
};
|
||||
let removed_before = self
|
||||
.dropped_output_indexes
|
||||
.iter()
|
||||
.filter(|dropped| **dropped < index)
|
||||
.count() as u64;
|
||||
if removed_before > 0
|
||||
&& let Some(payload) = payload.as_object_mut()
|
||||
{
|
||||
payload.insert(
|
||||
"output_index".to_owned(),
|
||||
Value::from(index.saturating_sub(removed_before)),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn filter_cache_envelope(&self, payload: &mut Value) -> Result<(), GrokRequestEncodeError> {
|
||||
let Some(payload) = payload.as_object_mut() else {
|
||||
return Ok(());
|
||||
};
|
||||
self.filter_cache_output(payload)?;
|
||||
self.filter_injected_cache_tools(payload)?;
|
||||
if let Some(response) = payload.get_mut("response").and_then(Value::as_object_mut) {
|
||||
self.filter_cache_output(response)?;
|
||||
self.filter_injected_cache_tools(response)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn filter_cache_output(
|
||||
&self,
|
||||
envelope: &mut Map<String, Value>,
|
||||
) -> Result<(), GrokRequestEncodeError> {
|
||||
let Some(output) = envelope.get_mut("output") else {
|
||||
return Ok(());
|
||||
};
|
||||
let output = output
|
||||
.as_array_mut()
|
||||
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
|
||||
output.retain(|item| !self.is_internal_cache_call(item));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn filter_injected_cache_tools(
|
||||
&self,
|
||||
envelope: &mut Map<String, Value>,
|
||||
) -> Result<(), GrokRequestEncodeError> {
|
||||
if self.injected_tool_types.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let Some(tools) = envelope.get_mut("tools") else {
|
||||
return Ok(());
|
||||
};
|
||||
let tools = tools
|
||||
.as_array_mut()
|
||||
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
|
||||
tools.retain(|tool| {
|
||||
tool.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_none_or(|kind| !self.injected_tool_types.contains(kind.trim()))
|
||||
});
|
||||
if tools.is_empty() {
|
||||
envelope.remove("tools");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn restore_visible_tools(&self, response: &mut Map<String, Value>) {
|
||||
if response.contains_key("tools") {
|
||||
response.insert("tools".to_owned(), Value::Array(self.visible_tools.clone()));
|
||||
if self.visible_tools.is_empty() {
|
||||
response.remove("tools");
|
||||
} else {
|
||||
response.insert("tools".to_owned(), Value::Array(self.visible_tools.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1125,6 +1391,8 @@ impl fmt::Debug for GrokResponseTransform {
|
||||
.field("alias_count", &self.aliases.len())
|
||||
.field("visible_tool_count", &self.visible_tools.len())
|
||||
.field("legacy_local_shell", &self.legacy_local_shell)
|
||||
.field("filter_x_search", &self.filter_x_search)
|
||||
.field("injected_tool_types", &self.injected_tool_types)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -1982,6 +2250,9 @@ impl ToolNormalizer {
|
||||
]),
|
||||
);
|
||||
}
|
||||
"allowed_tools" => {
|
||||
self.normalize_allowed_tools_choice(payload, object, normalized_tools)?;
|
||||
}
|
||||
"function" => self.normalize_function_tool_choice(payload, object)?,
|
||||
_ => {
|
||||
if let Some(hosted_kind) = normalize_hosted_tool_choice_kind(&kind) {
|
||||
@@ -2004,6 +2275,94 @@ impl ToolNormalizer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn normalize_allowed_tools_choice(
|
||||
&self,
|
||||
payload: &mut Map<String, Value>,
|
||||
mut object: Map<String, Value>,
|
||||
normalized_tools: &[Value],
|
||||
) -> Result<(), GrokRequestEncodeError> {
|
||||
if !matches!(
|
||||
object.get("mode").and_then(Value::as_str),
|
||||
Some("auto" | "required")
|
||||
) {
|
||||
return Err(GrokRequestEncodeError::InvalidRequestNormalization);
|
||||
}
|
||||
let allowed = object
|
||||
.get_mut("tools")
|
||||
.and_then(Value::as_array_mut)
|
||||
.filter(|tools| !tools.is_empty())
|
||||
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
|
||||
for allowed_tool in allowed {
|
||||
self.normalize_allowed_tool_reference(allowed_tool, normalized_tools)?;
|
||||
}
|
||||
payload.insert("tool_choice".to_owned(), Value::Object(object));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn normalize_allowed_tool_reference(
|
||||
&self,
|
||||
allowed_tool: &mut Value,
|
||||
normalized_tools: &[Value],
|
||||
) -> Result<(), GrokRequestEncodeError> {
|
||||
let object = allowed_tool
|
||||
.as_object_mut()
|
||||
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
|
||||
let kind = string_field(object, "type").trim().to_owned();
|
||||
match kind.as_str() {
|
||||
"function" => {
|
||||
let choice = match object.get_mut("function") {
|
||||
Some(Value::Object(function)) => function,
|
||||
Some(_) => return Err(GrokRequestEncodeError::InvalidRequestNormalization),
|
||||
None => object,
|
||||
};
|
||||
rewrite_namespace_choice(choice, &self.identity_aliases)?;
|
||||
let name = string_field(choice, "name").trim();
|
||||
if name.is_empty() || !has_named_tool(normalized_tools, "function", name) {
|
||||
return Err(GrokRequestEncodeError::InvalidRequestNormalization);
|
||||
}
|
||||
}
|
||||
"custom" => {
|
||||
let identity = ToolIdentity::new(
|
||||
ToolKind::Custom,
|
||||
string_field(object, "namespace").trim(),
|
||||
string_field(object, "name").trim(),
|
||||
);
|
||||
let alias = self
|
||||
.identity_aliases
|
||||
.get(&identity)
|
||||
.cloned()
|
||||
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
|
||||
object.insert("type".to_owned(), Value::String("function".to_owned()));
|
||||
object.insert("name".to_owned(), Value::String(alias));
|
||||
object.remove("namespace");
|
||||
}
|
||||
"apply_patch" | "tool_search" => {
|
||||
let identity = if kind == "apply_patch" {
|
||||
ToolIdentity::new(ToolKind::ApplyPatch, "", "apply_patch")
|
||||
} else {
|
||||
ToolIdentity::new(ToolKind::ToolSearch, "", "tool_search")
|
||||
};
|
||||
let alias = self
|
||||
.identity_aliases
|
||||
.get(&identity)
|
||||
.cloned()
|
||||
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
|
||||
object.clear();
|
||||
object.insert("type".to_owned(), Value::String("function".to_owned()));
|
||||
object.insert("name".to_owned(), Value::String(alias));
|
||||
}
|
||||
_ => {
|
||||
let hosted_kind = normalize_hosted_tool_choice_kind(&kind)
|
||||
.ok_or(GrokRequestEncodeError::InvalidRequestNormalization)?;
|
||||
if !has_tool_type(normalized_tools, hosted_kind) {
|
||||
return Err(GrokRequestEncodeError::InvalidRequestNormalization);
|
||||
}
|
||||
object.insert("type".to_owned(), Value::String(hosted_kind.to_owned()));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn normalize_function_tool_choice(
|
||||
&self,
|
||||
payload: &mut Map<String, Value>,
|
||||
@@ -2080,6 +2439,13 @@ fn has_tool_type(tools: &[Value], kind: &str) -> bool {
|
||||
.any(|tool| tool.pointer("/type").and_then(Value::as_str) == Some(kind))
|
||||
}
|
||||
|
||||
fn has_named_tool(tools: &[Value], kind: &str, name: &str) -> bool {
|
||||
tools.iter().any(|tool| {
|
||||
tool.pointer("/type").and_then(Value::as_str) == Some(kind)
|
||||
&& tool.pointer("/name").and_then(Value::as_str) == Some(name)
|
||||
})
|
||||
}
|
||||
|
||||
fn tools_of_type(tools: &[Value], kind: &str) -> Vec<Value> {
|
||||
tools
|
||||
.iter()
|
||||
|
||||
@@ -324,7 +324,7 @@ async fn unauthorized_feedback_records_runtime_cooldown_without_persisting_accou
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cooldown_for_superseded_credential_does_not_block_rotated_account() {
|
||||
async fn account_scoped_cooldown_survives_credential_rotation() {
|
||||
let fixture = SelectorFixture::new(&["revision-fence"]).await;
|
||||
let id = account_id("revision-fence");
|
||||
let session = fixture
|
||||
@@ -369,17 +369,25 @@ async fn cooldown_for_superseded_credential_does_not_block_rotated_account() {
|
||||
CredentialCasOutcome::Updated(revision) if revision.get() == 2
|
||||
));
|
||||
|
||||
let next = fixture
|
||||
.selector
|
||||
.select(fixture.request(BTreeSet::new()))
|
||||
.await
|
||||
.expect("replacement credential should not inherit cooldown");
|
||||
assert_eq!(next.account_id(), &id);
|
||||
assert!(fixture.cooldowns.cooldown(&id).is_none());
|
||||
assert!(matches!(
|
||||
fixture
|
||||
.selector
|
||||
.select(fixture.request(BTreeSet::new()))
|
||||
.await,
|
||||
Err(GrokSessionSelectorError::NoEligibleSession)
|
||||
));
|
||||
let account = fixture.store.account(&id).expect("rotated account");
|
||||
assert_eq!(account.revision().get(), 2);
|
||||
assert_eq!(account.availability(), AccountAvailability::Cooldown);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account_state() {
|
||||
async fn rate_limit_feedback_persists_the_grok2api_cooldown_state() {
|
||||
let fixture = SelectorFixture::new(&["rate-limit", "available"]).await;
|
||||
let session = fixture
|
||||
.selector
|
||||
@@ -397,7 +405,12 @@ async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account
|
||||
)
|
||||
.await;
|
||||
let account = fixture.store.account(&selected).expect("account");
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
assert_eq!(account.availability(), AccountAvailability::Cooldown);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
);
|
||||
assert!(
|
||||
fixture
|
||||
.cooldowns
|
||||
@@ -413,7 +426,35 @@ async fn rate_limit_feedback_records_runtime_cooldown_without_persisting_account
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_payment_required_feedback_only_records_runtime_account_cooldown() {
|
||||
async fn successful_request_clears_the_persisted_cooldown_state() {
|
||||
let fixture = SelectorFixture::new(&["cooldown-success"]).await;
|
||||
let session = fixture
|
||||
.selector
|
||||
.select(fixture.request(BTreeSet::new()))
|
||||
.await
|
||||
.expect("session");
|
||||
fixture
|
||||
.selector
|
||||
.record_failure(
|
||||
&session,
|
||||
GrokCredentialFailure::RateLimited {
|
||||
retry_after: Some(Duration::from_secs(5)),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
fixture.selector.record_success(&session).await;
|
||||
|
||||
let account = fixture
|
||||
.store
|
||||
.account(session.account_id())
|
||||
.expect("account");
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
assert!(account.cooldown_until().is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn payment_required_feedback_persists_account_quota_exhaustion() {
|
||||
let fixture = SelectorFixture::new(&["payment-required", "available"]).await;
|
||||
let session = fixture
|
||||
.selector
|
||||
@@ -437,9 +478,9 @@ async fn unknown_payment_required_feedback_only_records_runtime_account_cooldown
|
||||
.account(&selected)
|
||||
.expect("selected account")
|
||||
.availability(),
|
||||
AccountAvailability::Ready
|
||||
AccountAvailability::QuotaExhausted
|
||||
);
|
||||
assert!(fixture.cooldowns.cooldown(&selected).is_some());
|
||||
assert!(fixture.cooldowns.cooldown(&selected).is_none());
|
||||
let next = fixture
|
||||
.selector
|
||||
.select(fixture.request(BTreeSet::new()))
|
||||
@@ -538,7 +579,7 @@ async fn model_access_feedback_reports_model_cooldown_without_blocking_the_accou
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn interrupted_stream_feedback_records_runtime_cooldown_without_persisting_account_state() {
|
||||
async fn interrupted_stream_feedback_persists_the_grok2api_cooldown_state() {
|
||||
let fixture = SelectorFixture::new(&["stream-interrupted"]).await;
|
||||
let session = fixture
|
||||
.selector
|
||||
@@ -554,7 +595,12 @@ async fn interrupted_stream_feedback_records_runtime_cooldown_without_persisting
|
||||
.store
|
||||
.account(session.account_id())
|
||||
.expect("account");
|
||||
assert_eq!(account.availability(), AccountAvailability::Ready);
|
||||
assert_eq!(account.availability(), AccountAvailability::Cooldown);
|
||||
assert!(
|
||||
account
|
||||
.cooldown_until()
|
||||
.is_some_and(|until| until > SystemTime::now())
|
||||
);
|
||||
assert!(
|
||||
fixture
|
||||
.cooldowns
|
||||
@@ -595,6 +641,75 @@ async fn quota_feedback_uses_common_quota_exhausted_state() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quota_feedback_should_follow_the_account_across_a_credential_rotation() {
|
||||
let fixture = SelectorFixture::new(&["quota-rotation"]).await;
|
||||
let id = account_id("quota-rotation");
|
||||
let session = fixture
|
||||
.selector
|
||||
.select(fixture.request(BTreeSet::new()))
|
||||
.await
|
||||
.expect("session");
|
||||
let current = fixture
|
||||
.store
|
||||
.load_credential(&id, CredentialRevision::new(1).expect("revision"))
|
||||
.await
|
||||
.expect("current credential");
|
||||
let prepared = GrokCredentialAdmin
|
||||
.prepare_rotation(&RotateManagedGrokCredential {
|
||||
current,
|
||||
secret: provider_xai::GrokOAuthSecret {
|
||||
access_token: provider_xai::SecretValue::new("new-access"),
|
||||
refresh_token: provider_xai::SecretValue::new("new-refresh"),
|
||||
id_token: None,
|
||||
scope: provider_xai::OFFICIAL_SCOPES.join(" "),
|
||||
},
|
||||
verified_account: crate::support::profile("subject-quota-rotation"),
|
||||
next_refresh_at: chrono::Utc::now() + chrono::Duration::minutes(30),
|
||||
})
|
||||
.expect("rotate");
|
||||
fixture
|
||||
.store
|
||||
.compare_and_swap_credential(prepared.credential)
|
||||
.await
|
||||
.expect("persist rotation");
|
||||
|
||||
fixture
|
||||
.selector
|
||||
.record_failure(&session, GrokCredentialFailure::FreeQuotaExhausted)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
fixture.store.account(&id).expect("account").availability(),
|
||||
AccountAvailability::QuotaExhausted
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ordinary_success_should_not_clear_a_confirmed_quota_exhaustion() {
|
||||
let fixture = SelectorFixture::new(&["quota-success"]).await;
|
||||
let session = fixture
|
||||
.selector
|
||||
.select(fixture.request(BTreeSet::new()))
|
||||
.await
|
||||
.expect("session");
|
||||
fixture
|
||||
.selector
|
||||
.record_failure(&session, GrokCredentialFailure::FreeQuotaExhausted)
|
||||
.await;
|
||||
|
||||
fixture.selector.record_success(&session).await;
|
||||
|
||||
assert_eq!(
|
||||
fixture
|
||||
.store
|
||||
.account(session.account_id())
|
||||
.expect("account")
|
||||
.availability(),
|
||||
AccountAvailability::QuotaExhausted
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn excluded_account_is_never_selected_again() {
|
||||
let fixture = SelectorFixture::new(&["excluded"]).await;
|
||||
|
||||
@@ -106,6 +106,15 @@ const SUCCESS_SSE: &[u8] = concat!(
|
||||
)
|
||||
.as_bytes();
|
||||
|
||||
fn replay_ciphertext(seed: u8) -> String {
|
||||
use base64::Engine as _;
|
||||
|
||||
let bytes = (0_u16..128)
|
||||
.map(|value| (value as u8).wrapping_add(seed))
|
||||
.collect::<Vec<_>>();
|
||||
base64::engine::general_purpose::STANDARD_NO_PAD.encode(bytes)
|
||||
}
|
||||
|
||||
fn stateful_sse(encrypted_content: &str) -> Vec<u8> {
|
||||
format!(
|
||||
concat!(
|
||||
@@ -121,6 +130,45 @@ fn stateful_sse(encrypted_content: &str) -> Vec<u8> {
|
||||
.into_bytes()
|
||||
}
|
||||
|
||||
fn stateful_sse_with_assistant(encrypted_content: &str, assistant: &str) -> Vec<u8> {
|
||||
let assistant = serde_json::to_string(assistant).expect("assistant text");
|
||||
format!(
|
||||
concat!(
|
||||
"event: response.created\n",
|
||||
"data: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\"}}}}\n\n",
|
||||
"event: response.output_item.done\n",
|
||||
"data: {{\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{{\"id\":\"reason_account_bound\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[],\"content\":null,\"encrypted_content\":\"{}\"}}}}\n\n",
|
||||
"event: response.output_item.done\n",
|
||||
"data: {{\"type\":\"response.output_item.done\",\"output_index\":1,\"item\":{{\"id\":\"message_account_bound\",\"type\":\"message\",\"role\":\"assistant\",\"status\":\"completed\",\"content\":[{{\"type\":\"output_text\",\"text\":{assistant}}}]}}}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\",\"status\":\"completed\"}}}}\n\n"
|
||||
),
|
||||
encrypted_content,
|
||||
assistant = assistant,
|
||||
)
|
||||
.into_bytes()
|
||||
}
|
||||
|
||||
fn stateful_sse_with_custom_tool_call(encrypted_content: &str, input: &str) -> Vec<u8> {
|
||||
let arguments = serde_json::to_string(&json!({"input": input})).expect("tool arguments");
|
||||
let arguments_json = serde_json::to_string(&arguments).expect("arguments JSON string");
|
||||
format!(
|
||||
concat!(
|
||||
"event: response.created\n",
|
||||
"data: {{\"type\":\"response.created\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\"}}}}\n\n",
|
||||
"event: response.output_item.done\n",
|
||||
"data: {{\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{{\"id\":\"reason_account_bound\",\"type\":\"reasoning\",\"status\":\"completed\",\"summary\":[],\"content\":null,\"encrypted_content\":\"{}\"}}}}\n\n",
|
||||
"event: response.output_item.done\n",
|
||||
"data: {{\"type\":\"response.output_item.done\",\"output_index\":1,\"item\":{{\"id\":\"item_exec\",\"type\":\"function_call\",\"status\":\"completed\",\"call_id\":\"call_exec\",\"name\":\"exec\",\"arguments\":{arguments_json}}}}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {{\"type\":\"response.completed\",\"response\":{{\"id\":\"resp_state\",\"model\":\"grok-4.5\",\"status\":\"completed\"}}}}\n\n"
|
||||
),
|
||||
encrypted_content,
|
||||
arguments_json = arguments_json,
|
||||
)
|
||||
.into_bytes()
|
||||
}
|
||||
|
||||
fn custom_apply_patch_sse(patch: &str) -> Vec<u8> {
|
||||
let arguments = serde_json::to_string(&json!({"patch": patch})).expect("patch arguments");
|
||||
let arguments_json = serde_json::to_string(&arguments).expect("arguments JSON string");
|
||||
@@ -278,6 +326,54 @@ impl GrokSessionSelector for StubSelector {
|
||||
}
|
||||
}
|
||||
|
||||
struct SequencedAccountSelector {
|
||||
accounts: Mutex<VecDeque<gateway_core::engine::credential::ProviderAccountId>>,
|
||||
}
|
||||
|
||||
impl SequencedAccountSelector {
|
||||
fn new(
|
||||
accounts: impl IntoIterator<Item = gateway_core::engine::credential::ProviderAccountId>,
|
||||
) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
accounts: Mutex::new(accounts.into_iter().collect()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl GrokSessionSelector for SequencedAccountSelector {
|
||||
fn select(&self, request: GrokSessionSelection) -> GrokSessionSelectorFuture<'_> {
|
||||
let selected = self.accounts.lock().expect("account sequence").pop_front();
|
||||
Box::pin(async move {
|
||||
let id = selected.ok_or(GrokSessionSelectorError::NoEligibleSession)?;
|
||||
if request.excluded_accounts().contains(&id)
|
||||
|| request
|
||||
.required_account()
|
||||
.is_some_and(|required| required != &id)
|
||||
{
|
||||
return Err(GrokSessionSelectorError::NoEligibleSession);
|
||||
}
|
||||
SelectedGrokSession::new(
|
||||
id,
|
||||
CredentialRevision::new(1).expect("revision"),
|
||||
SecretValue::new("oauth-access"),
|
||||
SecretValue::new("verified-user"),
|
||||
Some(SecretValue::new("user@example.com")),
|
||||
GrokSessionBinding::new("acct_provider").expect("binding"),
|
||||
(),
|
||||
)
|
||||
.map_err(|_| GrokSessionSelectorError::InvalidSession)
|
||||
})
|
||||
}
|
||||
|
||||
fn record_failure<'a>(
|
||||
&'a self,
|
||||
_: &'a SelectedGrokSession,
|
||||
_: GrokCredentialFailure,
|
||||
) -> GrokCredentialFeedbackFuture<'a> {
|
||||
Box::pin(async {})
|
||||
}
|
||||
}
|
||||
|
||||
enum InferenceMode {
|
||||
Success,
|
||||
SuccessWithMetrics(GrokInferenceTransportMetrics),
|
||||
@@ -481,7 +577,7 @@ async fn provider_with_recovery(
|
||||
}
|
||||
|
||||
async fn provider_with_catalog_transport(
|
||||
selector: Arc<StubSelector>,
|
||||
selector: Arc<dyn GrokSessionSelector>,
|
||||
transport: Arc<StubInferenceTransport>,
|
||||
recovery: Arc<StubRecovery>,
|
||||
catalog_transport: Arc<dyn GrokModelCatalogTransport>,
|
||||
@@ -569,6 +665,36 @@ fn operation() -> Operation {
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
|
||||
}
|
||||
|
||||
fn reasoning_replay_operation(session: &str, input: serde_json::Value) -> Operation {
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("client-model")),
|
||||
("prompt_cache_key".to_owned(), json!(session)),
|
||||
("input".to_owned(), input),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload");
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
|
||||
}
|
||||
|
||||
fn contextual_reasoning_replay_operation(session: &str, input: serde_json::Value) -> Operation {
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("client-model")),
|
||||
("input".to_owned(), input),
|
||||
("tools".to_owned(), json!([{"type":"custom","name":"exec"}])),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload")
|
||||
.with_context(Map::from_iter([(
|
||||
"conversation_id".to_owned(),
|
||||
json!(session),
|
||||
)]));
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
|
||||
}
|
||||
|
||||
fn operation_with_reasoning_effort(effort: &str) -> Operation {
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
@@ -621,6 +747,26 @@ fn compaction_operation() -> Operation {
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
|
||||
}
|
||||
|
||||
fn compaction_operation_with_prompt_cache(session: &str) -> Operation {
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
Map::from_iter([
|
||||
("model".to_owned(), json!("client-model")),
|
||||
("prompt_cache_key".to_owned(), json!(session)),
|
||||
(
|
||||
"input".to_owned(),
|
||||
json!([
|
||||
{"type": "message", "role": "user", "content": "history"},
|
||||
{"type": "compaction_trigger"}
|
||||
]),
|
||||
),
|
||||
("stream".to_owned(), json!(true)),
|
||||
]),
|
||||
)
|
||||
.expect("OpenAI payload");
|
||||
Operation::Generate(GenerateRequest::from_protocol_payload(payload))
|
||||
}
|
||||
|
||||
fn compaction_operation_with_state(state: ProviderSessionState) -> Operation {
|
||||
let payload = ProtocolPayload::json_object(
|
||||
"openai",
|
||||
@@ -737,6 +883,18 @@ fn context_with_continuation_attempt(
|
||||
context(CancellationToken::new(), Some(continuation)).with_continuation_attempt(attempt)
|
||||
}
|
||||
|
||||
async fn execute_successfully(provider: &GrokBuildProvider, operation: Operation) {
|
||||
let mut stream = provider
|
||||
.execute(
|
||||
provider_request_with_operation("xai", operation),
|
||||
context(CancellationToken::new(), None),
|
||||
)
|
||||
.await
|
||||
.expect("provider stream");
|
||||
let events = stream.by_ref().collect::<Vec<_>>().await;
|
||||
assert!(events.iter().all(Result::is_ok));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn request_observation_preserves_raw_xai_reasoning_effort() {
|
||||
let provider = provider(StubSelector::success(), StubInferenceTransport::success()).await;
|
||||
@@ -1300,7 +1458,7 @@ async fn explicit_quota_http_429_marks_provider_error_replay_safe() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_http_402_is_retryable_without_account_quota_feedback() {
|
||||
async fn unknown_http_402_is_retryable_with_payment_quota_feedback() {
|
||||
let selector = StubSelector::success();
|
||||
let transport = StubInferenceTransport::error(
|
||||
GrokInferenceTransportError::new(
|
||||
@@ -1750,6 +1908,239 @@ async fn external_previous_response_is_rejected_before_selection() {
|
||||
assert_eq!(selector.calls.load(Ordering::SeqCst), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_replay_is_reused_only_for_the_same_explicit_session() {
|
||||
let encrypted_content = replay_ciphertext(0);
|
||||
let transport = StubInferenceTransport::sequence([
|
||||
InferenceMode::SuccessBody(stateful_sse(&encrypted_content)),
|
||||
InferenceMode::Success,
|
||||
InferenceMode::Success,
|
||||
]);
|
||||
let provider = provider(StubSelector::success(), transport.clone()).await;
|
||||
|
||||
execute_successfully(
|
||||
&provider,
|
||||
reasoning_replay_operation(
|
||||
"session-a",
|
||||
json!([{"type":"message","role":"user","content":"first"}]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
execute_successfully(
|
||||
&provider,
|
||||
reasoning_replay_operation(
|
||||
"session-b",
|
||||
json!([{"type":"message","role":"user","content":"other"}]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
execute_successfully(
|
||||
&provider,
|
||||
reasoning_replay_operation(
|
||||
"session-a",
|
||||
json!([{"type":"message","role":"user","content":"next"}]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
|
||||
let requests = transport.requests.lock().expect("requests");
|
||||
let other: serde_json::Value =
|
||||
serde_json::from_slice(requests[1].body()).expect("other session body");
|
||||
let replayed: serde_json::Value =
|
||||
serde_json::from_slice(requests[2].body()).expect("replayed body");
|
||||
assert!(
|
||||
other["input"]
|
||||
.as_array()
|
||||
.is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning"))
|
||||
);
|
||||
assert_eq!(
|
||||
replayed.pointer("/input/0/encrypted_content"),
|
||||
Some(&json!(encrypted_content))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_replay_should_pair_context_scoped_custom_call_with_its_output() {
|
||||
let encrypted_content = replay_ciphertext(4);
|
||||
let transport = StubInferenceTransport::sequence([
|
||||
InferenceMode::SuccessBody(stateful_sse_with_custom_tool_call(
|
||||
&encrypted_content,
|
||||
"inspect skills",
|
||||
)),
|
||||
InferenceMode::Success,
|
||||
]);
|
||||
let provider = provider(StubSelector::success(), transport.clone()).await;
|
||||
|
||||
execute_successfully(
|
||||
&provider,
|
||||
contextual_reasoning_replay_operation(
|
||||
"conversation-from-header",
|
||||
json!([{"type":"message","role":"user","content":"find skills"}]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
execute_successfully(
|
||||
&provider,
|
||||
contextual_reasoning_replay_operation(
|
||||
"conversation-from-header",
|
||||
json!([{
|
||||
"type":"custom_tool_call_output",
|
||||
"call_id":"call_exec",
|
||||
"output":[{"type":"input_text","text":"skills found"}]
|
||||
}]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
|
||||
let requests = transport.requests.lock().expect("requests");
|
||||
let replayed: serde_json::Value =
|
||||
serde_json::from_slice(requests[1].body()).expect("replayed body");
|
||||
let item_types = replayed["input"]
|
||||
.as_array()
|
||||
.expect("replayed input")
|
||||
.iter()
|
||||
.filter_map(|item| item["type"].as_str())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(
|
||||
item_types,
|
||||
["reasoning", "function_call", "function_call_output"]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_replay_is_scoped_to_the_selected_account() {
|
||||
let encrypted_content = replay_ciphertext(1);
|
||||
let transport = StubInferenceTransport::sequence([
|
||||
InferenceMode::SuccessBody(stateful_sse(&encrypted_content)),
|
||||
InferenceMode::Success,
|
||||
InferenceMode::Success,
|
||||
]);
|
||||
let selector = SequencedAccountSelector::new([
|
||||
account_id("replay-a"),
|
||||
account_id("replay-b"),
|
||||
account_id("replay-a"),
|
||||
]);
|
||||
let provider = provider_with_catalog_transport(
|
||||
selector,
|
||||
transport.clone(),
|
||||
StubRecovery::new(GrokCredentialRecoveryOutcome::Unavailable),
|
||||
Arc::new(StaticCatalogTransport),
|
||||
)
|
||||
.await;
|
||||
|
||||
for prompt in ["first", "different account", "owner again"] {
|
||||
execute_successfully(
|
||||
&provider,
|
||||
reasoning_replay_operation(
|
||||
"shared-session",
|
||||
json!([{"type":"message","role":"user","content":prompt}]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let requests = transport.requests.lock().expect("requests");
|
||||
let different_account: serde_json::Value =
|
||||
serde_json::from_slice(requests[1].body()).expect("different account request body");
|
||||
let owner_again: serde_json::Value =
|
||||
serde_json::from_slice(requests[2].body()).expect("owner request body");
|
||||
assert!(
|
||||
different_account["input"]
|
||||
.as_array()
|
||||
.is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning"))
|
||||
);
|
||||
assert_eq!(
|
||||
owner_again.pointer("/input/0/encrypted_content"),
|
||||
Some(&json!(encrypted_content))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_replay_rejects_mismatched_assistant_history() {
|
||||
let encrypted_content = replay_ciphertext(2);
|
||||
let transport = StubInferenceTransport::sequence([
|
||||
InferenceMode::SuccessBody(stateful_sse_with_assistant(
|
||||
&encrypted_content,
|
||||
"cached answer",
|
||||
)),
|
||||
InferenceMode::Success,
|
||||
]);
|
||||
let provider = provider(StubSelector::success(), transport.clone()).await;
|
||||
|
||||
execute_successfully(
|
||||
&provider,
|
||||
reasoning_replay_operation(
|
||||
"assistant-session",
|
||||
json!([{"type":"message","role":"user","content":"first"}]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
execute_successfully(
|
||||
&provider,
|
||||
reasoning_replay_operation(
|
||||
"assistant-session",
|
||||
json!([
|
||||
{"type":"message","role":"user","content":"first"},
|
||||
{"type":"message","role":"assistant","content":[{"type":"output_text","text":"different answer"}]},
|
||||
{"type":"message","role":"user","content":"next"}
|
||||
]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
|
||||
let requests = transport.requests.lock().expect("requests");
|
||||
let body: serde_json::Value = serde_json::from_slice(requests[1].body()).expect("second body");
|
||||
assert!(
|
||||
body["input"]
|
||||
.as_array()
|
||||
.is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_compaction_clears_reasoning_replay_for_the_session() {
|
||||
let encrypted_content = replay_ciphertext(3);
|
||||
let summary = valid_compaction_summary("reasoning replay cleared");
|
||||
let transport = StubInferenceTransport::sequence([
|
||||
InferenceMode::SuccessBody(stateful_sse(&encrypted_content)),
|
||||
InferenceMode::SuccessBody(compaction_sse(&summary, None)),
|
||||
InferenceMode::Success,
|
||||
]);
|
||||
let provider = provider(StubSelector::success(), transport.clone()).await;
|
||||
|
||||
execute_successfully(
|
||||
&provider,
|
||||
reasoning_replay_operation(
|
||||
"compact-session",
|
||||
json!([{"type":"message","role":"user","content":"first"}]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
execute_successfully(
|
||||
&provider,
|
||||
compaction_operation_with_prompt_cache("compact-session"),
|
||||
)
|
||||
.await;
|
||||
execute_successfully(
|
||||
&provider,
|
||||
reasoning_replay_operation(
|
||||
"compact-session",
|
||||
json!([{"type":"message","role":"user","content":"after compact"}]),
|
||||
),
|
||||
)
|
||||
.await;
|
||||
|
||||
let requests = transport.requests.lock().expect("requests");
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_slice(requests[2].body()).expect("post-compaction body");
|
||||
assert!(
|
||||
body["input"]
|
||||
.as_array()
|
||||
.is_some_and(|input| input.iter().all(|item| item["type"] != "reasoning"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn connection_state_inherits_session_and_recovers_reasoning_on_pinned_account() {
|
||||
use base64::Engine as _;
|
||||
|
||||
@@ -642,6 +642,76 @@ fn decoder_should_preserve_image_and_hosted_tool_items_without_inventing_facts()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_hide_injected_search_calls_and_compact_later_output_indexes() {
|
||||
let request = tool_request(serde_json::json!({
|
||||
"model": "client",
|
||||
"prompt_cache_key": "cache-session",
|
||||
"input": "hello"
|
||||
}));
|
||||
let body = concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_cache\",\"tools\":[{\"type\":\"web_search\"},{\"type\":\"x_search\"}]}}\n\n",
|
||||
"event: response.output_item.added\n",
|
||||
"data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_xs\",\"call_id\":\"xs_call_1\",\"name\":\"x_semantic_search\",\"arguments\":\"{}\"}}\n\n",
|
||||
"event: response.function_call_arguments.delta\n",
|
||||
"data: {\"type\":\"response.function_call_arguments.delta\",\"output_index\":0,\"item_id\":\"item_xs\",\"call_id\":\"xs_call_1\",\"delta\":\"{}\"}\n\n",
|
||||
"event: response.output_item.done\n",
|
||||
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_xs\",\"call_id\":\"xs_call_1\",\"name\":\"x_semantic_search\",\"arguments\":\"{}\"}}\n\n",
|
||||
"event: response.output_item.added\n",
|
||||
"data: {\"type\":\"response.output_item.added\",\"output_index\":1,\"item\":{\"type\":\"image_generation_call\",\"id\":\"image_visible\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_cache\",\"status\":\"completed\",\"tools\":[{\"type\":\"web_search\"},{\"type\":\"x_search\"}],\"output\":[{\"type\":\"function_call\",\"call_id\":\"xs_call_1\",\"name\":\"x_semantic_search\",\"arguments\":\"{}\"},{\"type\":\"image_generation_call\",\"id\":\"image_visible\"}]}}\n\n",
|
||||
);
|
||||
|
||||
let events = GrokCanonicalDecoder::for_request("grok-4.5", &request)
|
||||
.push(body.as_bytes())
|
||||
.expect("filtered cache route response");
|
||||
let wire = wire_events(&events);
|
||||
|
||||
assert_eq!(wire.len(), 3);
|
||||
assert_eq!(
|
||||
wire[1].data().get("output_index"),
|
||||
Some(&serde_json::json!(0))
|
||||
);
|
||||
assert_eq!(
|
||||
wire[2].data().pointer("/response/output/0/id"),
|
||||
Some(&serde_json::json!("image_visible"))
|
||||
);
|
||||
assert_eq!(wire[2].data().pointer("/response/output/1"), None);
|
||||
assert_eq!(wire[2].data().pointer("/response/tools"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_preserve_an_internal_search_name_declared_by_the_client() {
|
||||
let request = tool_request(serde_json::json!({
|
||||
"model": "client",
|
||||
"prompt_cache_key": "cache-session",
|
||||
"input": "hello",
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"name": "x_user_search",
|
||||
"parameters": {"type": "object"}
|
||||
}]
|
||||
}));
|
||||
let body = concat!(
|
||||
"event: response.created\n",
|
||||
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_declared\"}}\n\n",
|
||||
"event: response.output_item.added\n",
|
||||
"data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"function_call\",\"id\":\"item_client\",\"call_id\":\"call_client\",\"name\":\"x_user_search\",\"arguments\":\"{}\"}}\n\n",
|
||||
"event: response.completed\n",
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_declared\",\"status\":\"completed\"}}\n\n",
|
||||
);
|
||||
|
||||
let events = GrokCanonicalDecoder::for_request("grok-4.5", &request)
|
||||
.push(body.as_bytes())
|
||||
.expect("client-declared search function");
|
||||
|
||||
assert!(wire_events(&events).iter().any(|event| {
|
||||
event.data().pointer("/item/name") == Some(&serde_json::json!("x_user_search"))
|
||||
}));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decoder_should_tolerate_mismatched_event_identity_and_preserve_the_wire_event() {
|
||||
let body = concat!(
|
||||
|
||||
@@ -91,7 +91,7 @@ fn headers_should_bind_identity_to_the_selected_oauth_account() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn headers_should_keep_openai_cache_requests_out_of_shell_session_mode() {
|
||||
fn headers_should_send_both_stable_session_headers_without_inventing_a_turn_index() {
|
||||
let headers = build_grok_headers(
|
||||
&crate::support::xai_wire_profile(),
|
||||
&selected_session(),
|
||||
@@ -109,7 +109,7 @@ fn headers_should_keep_openai_cache_requests_out_of_shell_session_mode() {
|
||||
};
|
||||
|
||||
assert_eq!(value("x-grok-conv-id"), Some("session-fixture"));
|
||||
assert_eq!(value("x-grok-session-id"), None);
|
||||
assert_eq!(value("x-grok-session-id"), Some("session-fixture"));
|
||||
assert_eq!(value("x-grok-turn-idx"), None);
|
||||
}
|
||||
|
||||
|
||||
@@ -256,7 +256,7 @@ fn explicit_session_should_enable_the_noop_native_cache_route() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_session_should_not_add_invokable_cache_tools_to_client_tool_requests() {
|
||||
fn explicit_session_should_add_only_internal_x_search_to_client_tool_requests() {
|
||||
let request = raw_request(json!({
|
||||
"model": "client",
|
||||
"prompt_cache_key": "conversation-42",
|
||||
@@ -273,15 +273,15 @@ fn explicit_session_should_not_add_invokable_cache_tools_to_client_tool_requests
|
||||
GrokResponsesRequest::encode(&request, "grok-4.5", &client_key()).expect("request");
|
||||
let body = Value::Object(encoded.body().clone());
|
||||
|
||||
// 缓存路由不得扩大模型可调用的工具集:带客户端工具时不注入原生搜索。
|
||||
assert_eq!(
|
||||
body.pointer("/tools"),
|
||||
Some(&json!([{
|
||||
body.pointer("/tools/0"),
|
||||
Some(&json!({
|
||||
"type": "function",
|
||||
"name": "read_file",
|
||||
"parameters": {"type": "object"}
|
||||
}]))
|
||||
}))
|
||||
);
|
||||
assert_eq!(body.pointer("/tools/1"), Some(&json!({"type": "x_search"})));
|
||||
assert_eq!(body.pointer("/tool_choice"), Some(&json!("auto")));
|
||||
}
|
||||
|
||||
@@ -384,7 +384,7 @@ fn function_parameters_should_report_the_invalid_nullable_root_field() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_session_should_not_inject_cache_tools_after_codex_additional_tools_normalization() {
|
||||
fn explicit_session_should_add_x_search_after_codex_additional_tools_normalization() {
|
||||
let request = raw_request(json!({
|
||||
"model": "client",
|
||||
"prompt_cache_key": "conversation-42",
|
||||
@@ -417,11 +417,38 @@ fn explicit_session_should_not_inject_cache_tools_after_codex_additional_tools_n
|
||||
Some(&json!(["patch"]))
|
||||
);
|
||||
assert_eq!(body.pointer("/tools/1/name"), Some(&json!("read_file")));
|
||||
// 归一化装载的会话工具视同客户端工具:不再追加可被调用的原生搜索。
|
||||
assert_eq!(body.pointer("/tools/2"), None);
|
||||
assert_eq!(body.pointer("/tools/2"), Some(&json!({"type": "x_search"})));
|
||||
assert_eq!(body.pointer("/tool_choice"), Some(&json!("auto")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_session_should_add_x_search_to_allowed_tools_choice() {
|
||||
let request = raw_request(json!({
|
||||
"model": "client",
|
||||
"prompt_cache_key": "conversation-42",
|
||||
"input": "first",
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"name": "read_file",
|
||||
"parameters": {"type": "object"}
|
||||
}],
|
||||
"tool_choice": {
|
||||
"type": "allowed_tools",
|
||||
"mode": "auto",
|
||||
"tools": [{"type": "function", "name": "read_file"}]
|
||||
}
|
||||
}));
|
||||
|
||||
let encoded =
|
||||
GrokResponsesRequest::encode(&request, "grok-4.5", &client_key()).expect("request");
|
||||
let body = Value::Object(encoded.body().clone());
|
||||
|
||||
assert_eq!(
|
||||
body.pointer("/tool_choice/tools/1"),
|
||||
Some(&json!({"type": "x_search"}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_session_should_not_duplicate_existing_native_cache_tools() {
|
||||
let request = raw_request(json!({
|
||||
|
||||
@@ -2,3 +2,4 @@
|
||||
0ae3945c02461347ab8dcfc64821027cd2366505778e0450d71bd82772c4e99d 0002_drop_opaque_response_id_index.sql
|
||||
e608ed89991f9d6317491714447a00a61ede90caca4fb030eb30695654afb664 0003_store_opaque_response_ids_as_bytes.sql
|
||||
df4e02cf2335733326c12429ced4cbd0b35b8a22b3fc24981c3c0549477d7c43 0004_add_model_request_service_tier.sql
|
||||
fe8f36913c6db7053148db450f05e008dab6cdb9ca8cdcfcb22d97ad9a87bf49 0005_allow_quota_exhausted_cooldown.sql
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
alter table provider_accounts
|
||||
drop constraint provider_accounts_cooldown_ck,
|
||||
add constraint provider_accounts_cooldown_ck check (
|
||||
(
|
||||
availability != 'cooldown'
|
||||
or cooldown_until is not null
|
||||
)
|
||||
and (
|
||||
cooldown_until is null
|
||||
or availability in ('cooldown', 'quota_exhausted')
|
||||
)
|
||||
);
|
||||
@@ -1 +1 @@
|
||||
version: "3.0.23"
|
||||
version: "3.0.26"
|
||||
|
||||
Reference in New Issue
Block a user