fix(pool): isolate dynamic model quota buckets and 429 scheduling

This commit is contained in:
zhefox
2026-09-02 15:23:23 +08:00
parent 2cd20da1ec
commit dbbe7b22ab
37 changed files with 1491 additions and 78 deletions
+282 -2
View File
@@ -35,6 +35,7 @@ pub use providers::{
};
pub use quota::{
provider_pool_key_account_quota_exhausted, provider_pool_key_quota_hard_blocked,
provider_pool_key_model_quota_exhausted, provider_pool_key_model_quota_hard_blocked,
provider_pool_key_scheduling_label, provider_pool_member_quota_snapshot,
provider_pool_quota_metadata_provider_type, provider_pool_quota_metadata_updated_at,
provider_pool_quota_snapshot_updated_at,
@@ -561,7 +562,7 @@ mod tests {
}
})));
let signals = service.member_signals("windsurf", &key, None);
let signals = service.member_signals("windsurf", &key, None, None);
assert!(!signals.quota_exhausted);
}
@@ -588,7 +589,7 @@ mod tests {
}
}));
let signals = service.member_signals("windsurf", &key, None);
let signals = service.member_signals("windsurf", &key, None, None);
assert!(signals.quota_exhausted);
}
@@ -876,6 +877,285 @@ mod tests {
));
}
#[test]
fn model_quota_windows_are_isolated_without_provider_specific_names() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [
{
"code": "alpha_short",
"quota_group": "alpha",
"model": "vendor-alpha-model",
"used_ratio": 1.0,
"is_exhausted": true
},
{
"code": "alpha_long",
"quota_group": "alpha",
"model": "vendor-alpha-model",
"used_ratio": 0.2,
"is_exhausted": false
},
{
"code": "beta_short",
"quota_group": "beta",
"model": "vendor-beta-model",
"used_ratio": 1.0,
"is_exhausted": true
}
]
}
}));
let alpha = service.member_signals(
"codex",
&key,
None,
Some("vendor-alpha-model"),
);
let beta = service.member_signals("codex", &key, None, Some("vendor-beta-model"));
assert!(!alpha.quota_exhausted, "one available alpha window must keep it usable");
assert!(beta.quota_exhausted);
assert!(!beta.quota_hard_blocked);
}
#[test]
fn legacy_family_prefix_is_matched_by_model_tokens() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [
{ "code": "alpha_weekly", "used_ratio": 1.0, "is_exhausted": true },
{ "code": "default_weekly", "used_ratio": 0.1, "is_exhausted": false }
]
}
}));
let alpha = service.member_signals("codex", &key, None, Some("vendor-alpha-v2"));
assert!(alpha.quota_exhausted);
// An unrelated model has no identifiable bucket and therefore falls
// back to the account-level snapshot rather than guessing a family.
let unknown = service.member_signals("codex", &key, None, Some("vendor-gamma-v2"));
assert!(unknown.quota_exhausted);
}
#[test]
fn compact_model_bucket_identity_matches_request_tokens() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [
{
"code": "additional_0_primary_window",
"scope": "model",
"model": "spark",
"used_ratio": 0.0,
"is_exhausted": false
}
]
}
}));
let signals = service.member_signals(
"codex",
&key,
None,
Some("gpt-5.3-codex-spark"),
);
assert!(
!signals.quota_exhausted,
"a compact bucket name should match a token in the selected model"
);
}
#[test]
fn model_family_token_matches_versioned_alias_without_hardcoding_name() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [{
"code": "additional_0_primary_window",
"scope": "model",
"model": "gpt-5.3-codex-spark",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals = service.member_signals(
"codex",
&key,
None,
Some("gpt-5.4-codex-spark"),
);
assert!(signals.quota_exhausted);
}
#[test]
fn model_only_snapshot_does_not_poison_account_fallback() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [{
"code": "model:vendor-alpha",
"scope": "model",
"model": "vendor-alpha",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals = service.member_signals("codex", &key, None, None);
assert!(
!signals.quota_exhausted,
"account-level inspection must ignore model-only buckets"
);
}
#[test]
fn model_map_quota_is_resolved_without_materialized_windows() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "grok",
"exhausted": false,
"quota_by_model": {
"vendor-alpha": {
"remaining": 0.0,
"total": 10.0
},
"vendor-beta": {
"remaining": 5.0,
"total": 10.0
}
}
}
}));
let alpha = service.member_signals("grok", &key, None, Some("vendor-alpha"));
let beta = service.member_signals("grok", &key, None, Some("vendor-beta"));
assert!(alpha.quota_exhausted);
assert!(!beta.quota_exhausted);
}
#[test]
fn raw_provider_metadata_model_bucket_overrides_account_snapshot() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(Some(json!({
"codex": {
"updated_at": 1_700_000_001u64,
"additional_quota_windows": [{
"scope": "model",
"model": "future-spark",
"used_ratio": 0.0,
"is_exhausted": false
}]
}
})));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"windows": [{
"code": "primary",
"scope": "account",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals = service.member_signals("codex", &key, None, Some("future-spark"));
assert!(
!signals.quota_exhausted,
"a newer model bucket in provider metadata must not inherit an exhausted account bucket"
);
}
#[test]
fn model_quota_availability_suppresses_account_hard_block() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"allowed": false,
"limit_reached": true,
"exhausted": true,
"windows": [{
"scope": "model",
"model": "future-model",
"used_ratio": 0.0,
"is_exhausted": false
}]
}
}));
let signals = service.member_signals("codex", &key, None, Some("future-model"));
assert!(!signals.quota_exhausted);
assert!(!signals.quota_hard_blocked);
}
#[test]
fn newer_model_quota_observation_wins_over_stale_snapshot() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(Some(json!({
"codex": {
"updated_at": 1_700_000_200u64,
"additional_quota_windows": [{
"scope": "model",
"model": "future-model",
"used_ratio": 0.0,
"is_exhausted": false
}]
}
})));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"observed_at": 1_700_000_100u64,
"exhausted": true,
"windows": [{
"scope": "model",
"model": "future-model",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals = service.member_signals("codex", &key, None, Some("future-model"));
assert!(!signals.quota_exhausted);
}
#[test]
fn codex_explicit_quota_block_is_hard_until_reset() {
let now = std::time::SystemTime::now()
+12 -1
View File
@@ -8,7 +8,8 @@ use crate::capability::{ProviderPoolCapabilities, ProviderPoolCapability};
use crate::plan::{derive_plan_tier, normalize_provider_plan_tier};
use crate::quota::{
provider_pool_account_blocked, provider_pool_quota_reset_seconds,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_quota_usage_ratio,
provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_quota_usage_ratio,
};
#[derive(Debug, Clone)]
@@ -16,6 +17,11 @@ pub struct ProviderPoolMemberInput<'a> {
pub provider_type: &'a str,
pub key: &'a StoredProviderCatalogKey,
pub auth_config: Option<&'a Map<String, Value>>,
/// The provider-side model selected for the current request. Quota
/// snapshots may contain several independent windows, so adapters use
/// this value to select the applicable bucket instead of treating the
/// whole account as one quota.
pub provider_model_name: Option<&'a str>,
}
pub trait ProviderPoolAdapter: Send + Sync {
@@ -70,6 +76,11 @@ pub trait ProviderPoolAdapter: Send + Sync {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
.unwrap_or(false)
}
@@ -6,7 +6,9 @@ use serde_json::json;
use crate::capability::ProviderPoolCapabilities;
use crate::provider::{
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
ProviderPoolMemberInput,
};
use crate::quota::provider_pool_model_quota_exhausted;
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
pub const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels";
@@ -26,6 +28,21 @@ impl ProviderPoolAdapter for AntigravityProviderPoolAdapter {
}
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
input
.provider_model_name
.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
})
.unwrap_or_else(|| {
crate::quota::provider_pool_quota_snapshot_exhausted_decision(
input.key,
input.provider_type,
)
.unwrap_or(false)
})
}
fn quota_refresh_endpoint(
&self,
endpoints: &[StoredProviderCatalogEndpoint],
@@ -12,7 +12,8 @@ use crate::provider::{
use crate::quota::{
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
provider_pool_model_quota_exhausted, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -40,6 +41,11 @@ impl ProviderPoolAdapter for ChatGptWebProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
@@ -12,8 +12,8 @@ use crate::provider::{
use crate::quota::{
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_member_quota_snapshot, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -49,6 +49,11 @@ impl ProviderPoolAdapter for CodexProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if let Some(quota_snapshot) =
provider_pool_member_quota_snapshot(input.key, input.provider_type)
{
@@ -76,6 +81,11 @@ impl ProviderPoolAdapter for CodexProviderPoolAdapter {
}
fn quota_hard_blocked(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if input.provider_model_name.is_some_and(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model).is_some()
}) {
return false;
}
codex_explicit_quota_block_active(input.key, input.provider_type)
}
@@ -9,7 +9,8 @@ use crate::provider::{
use crate::quota::{
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
provider_pool_model_quota_exhausted, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
};
pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")];
@@ -44,6 +45,11 @@ impl ProviderPoolAdapter for GrokProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
@@ -12,8 +12,8 @@ use crate::provider::{
};
use crate::quota::{
provider_pool_current_unix_secs, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
provider_pool_model_quota_exhausted, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -46,6 +46,11 @@ impl ProviderPoolAdapter for KiroProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
@@ -13,7 +13,8 @@ use crate::provider::{
};
use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_member_quota_snapshot,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_metadata_bucket, provider_pool_model_quota_exhausted,
provider_pool_quota_snapshot_exhausted_decision,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -50,6 +51,11 @@ impl ProviderPoolAdapter for WindsurfProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) = input.provider_model_name.and_then(|model| {
provider_pool_model_quota_exhausted(input.key, input.provider_type, model)
}) {
return exhausted;
}
if windsurf_quota_snapshot_hard_exhausted(input.key, input.provider_type) {
return true;
}
+415 -7
View File
@@ -15,6 +15,7 @@ pub fn provider_pool_key_account_quota_exhausted(
provider_type,
key,
auth_config: None,
provider_model_name: None,
})
}
@@ -27,6 +28,25 @@ pub fn provider_pool_key_quota_hard_blocked(
provider_type,
key,
auth_config: None,
provider_model_name: None,
})
}
/// Model-aware hard-block lookup for pre-scheduler candidate filtering. Some
/// providers expose permanent account flags alongside independent model
/// buckets; adapters can suppress the account flag when the selected model
/// has its own usable quota.
pub fn provider_pool_key_model_quota_hard_blocked(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: &str,
) -> bool {
let adapter = ProviderPoolService::with_builtin_adapters().adapter(provider_type);
adapter.quota_hard_blocked(&ProviderPoolMemberInput {
provider_type,
key,
auth_config: None,
provider_model_name: Some(provider_model_name),
})
}
@@ -44,6 +64,331 @@ pub fn provider_pool_member_quota_snapshot<'a>(
.then_some(quota_snapshot)
}
/// Resolve exhaustion for the quota bucket applicable to one provider model.
///
/// Providers are free to expose quota windows in different shapes. Newer
/// snapshots should put an explicit `model`/`models` (or `quota_group`) on a
/// window; legacy Codex snapshots use a family prefix in `code` (for example
/// `spark_5h`). We deliberately do not name any product or model here: the
/// resolver compares the metadata supplied by the provider with the selected
/// model and only falls back to account-level exhaustion when no model bucket
/// can be identified.
pub(crate) fn provider_pool_model_quota_exhausted(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: &str,
) -> Option<bool> {
let requested = provider_pool_identifier_tokens(provider_model_name);
if requested.is_empty() {
return None;
}
// Prefer the materialized status snapshot, but also inspect the raw
// provider metadata. A quota refresh and a request can race, leaving the
// latter newer than the snapshot; resolving both avoids falling back to an
// account-wide signal and incorrectly blocking an unrelated model bucket.
let sources = [
provider_pool_member_quota_snapshot(key, provider_type),
provider_pool_metadata_bucket(key.upstream_metadata.as_ref(), provider_type),
];
let mut resolved = None::<(Option<u64>, bool)>;
for source in sources.into_iter().flatten() {
let windows = provider_pool_collect_quota_windows(source);
if windows.is_empty() {
continue;
}
let observed_at = provider_pool_timestamp_unix_secs(source.get("observed_at"))
.or_else(|| provider_pool_timestamp_unix_secs(source.get("updated_at")));
let model_matches = windows
.iter()
.filter(|window| {
provider_pool_window_explicitly_matches_model(window, provider_model_name)
})
.collect::<Vec<_>>();
if !model_matches.is_empty() {
let exhausted = provider_pool_matching_windows_exhausted(
model_matches,
observed_at,
);
if resolved.is_none()
|| provider_pool_should_replace_model_quota_resolution(
resolved.as_ref().and_then(|(observed_at, _)| *observed_at),
observed_at,
)
{
resolved = Some((observed_at, exhausted));
}
continue;
}
// Legacy snapshots may not carry a model field. Match an opaque
// family token (the prefix before `_`/`:` in `code`, or an explicit
// family key) against the model's tokens. This keeps independent
// windows isolated without baking in names such as "spark".
let family_matches = windows
.iter()
.filter(|window| provider_pool_window_family_matches_model(window, &requested))
.collect::<Vec<_>>();
if !family_matches.is_empty() {
let exhausted = provider_pool_matching_windows_exhausted(
family_matches,
observed_at,
);
if resolved.is_none()
|| provider_pool_should_replace_model_quota_resolution(
resolved.as_ref().and_then(|(observed_at, _)| *observed_at),
observed_at,
)
{
resolved = Some((observed_at, exhausted));
}
}
}
resolved.map(|(_, exhausted)| exhausted)
}
fn provider_pool_should_replace_model_quota_resolution(
previous_observed_at: Option<u64>,
next_observed_at: Option<u64>,
) -> bool {
match (previous_observed_at, next_observed_at) {
(Some(previous), Some(next)) => next >= previous,
(None, Some(_)) => true,
(Some(_), None) => false,
// Preserve source order when neither side carries freshness metadata;
// the materialized status snapshot is preferred over raw metadata.
(None, None) => false,
}
}
/// Public adapter-independent model quota lookup used by schedulers that need
/// to prefilter candidates before constructing provider-pool signals.
pub fn provider_pool_key_model_quota_exhausted(
key: &StoredProviderCatalogKey,
provider_type: &str,
provider_model_name: &str,
) -> Option<bool> {
provider_pool_model_quota_exhausted(key, provider_type, provider_model_name)
}
fn provider_pool_matching_windows_exhausted(
windows: Vec<&Map<String, Value>>,
snapshot_observed_at: Option<u64>,
) -> bool {
let now_unix_secs = provider_pool_current_unix_secs();
// A request can use a bucket as long as at least one of its independent
// windows still has capacity (e.g. a short and a long rolling window).
windows.iter().all(|window| {
provider_pool_quota_window_is_exhausted(window)
&& !now_unix_secs.is_some_and(|now| {
provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now)
})
})
}
fn provider_pool_window_explicitly_matches_model(
window: &Map<String, Value>,
requested_model: &str,
) -> bool {
let requested = provider_pool_normalize_identifier(requested_model);
let requested_tokens = provider_pool_identifier_tokens(requested_model);
if requested.is_empty() {
return false;
}
[
"model",
"model_name",
"model_id",
"quota_model",
"quota_model_name",
"target_model",
"limit_name",
]
.iter()
.filter_map(|key| window.get(*key))
.any(|value| match value {
Value::String(value) => {
provider_pool_identifiers_match(&requested, value, &requested_tokens)
}
Value::Array(values) => values
.iter()
.filter_map(Value::as_str)
.any(|value| provider_pool_identifiers_match(&requested, value, &requested_tokens)),
_ => false,
}) || ["models", "model_ids"]
.iter()
.filter_map(|key| window.get(*key).and_then(Value::as_array))
.flatten()
.filter_map(Value::as_str)
.any(|value| provider_pool_identifiers_match(&requested, value, &requested_tokens))
}
fn provider_pool_window_family_matches_model(
window: &Map<String, Value>,
requested_tokens: &std::collections::BTreeSet<String>,
) -> bool {
let explicit_scope = window
.get("scope")
.and_then(Value::as_str)
.map(|scope| scope.trim().to_ascii_lowercase());
// A model-scoped window without an explicit model must not accidentally
// match a token from its opaque code.
if explicit_scope.as_deref() == Some("model")
|| provider_pool_window_has_explicit_model(window)
{
return false;
}
let mut families = Vec::new();
for key in ["quota_group", "quota_family", "family", "bucket"] {
if let Some(value) = window.get(key).and_then(Value::as_str) {
families.push(value.to_string());
}
}
if let Some(code) = window.get("code").and_then(Value::as_str) {
let code = code.trim();
if let Some((prefix, _)) = code.split_once(['_', ':', '/']) {
families.push(prefix.to_string());
}
}
families.into_iter().any(|family| {
let normalized = provider_pool_normalize_identifier(&family);
if normalized.is_empty()
|| [
"account",
"quota",
"window",
"primary",
"secondary",
"rate",
"reset",
]
.iter()
.any(|generic| normalized == *generic)
{
return false;
}
requested_tokens.iter().any(|token| {
token.len() >= 3
&& (token == &normalized
|| token.contains(&normalized)
|| normalized.contains(token))
})
})
}
fn provider_pool_identifiers_match(
requested: &str,
candidate: &str,
requested_tokens: &std::collections::BTreeSet<String>,
) -> bool {
let candidate_tokens = provider_pool_identifier_tokens(candidate);
let candidate = provider_pool_normalize_identifier(candidate);
if candidate.is_empty() {
return false;
}
requested == candidate
|| candidate_tokens
.iter()
.any(|token| {
requested_tokens.contains(token) && provider_pool_is_specific_model_token(token)
})
// Handle compact upstream identifiers such as `spark` embedded in a
// provider model name (`vendor-codex-spark`) while avoiding accidental
// one/two-character matches.
|| (candidate.len() >= 4
&& requested.len() >= 6
&& (requested.contains(&candidate) || candidate.contains(requested)))
}
fn provider_pool_is_specific_model_token(token: &str) -> bool {
token.len() >= 4
&& !token.chars().all(|character| character.is_ascii_digit())
&& ![
"auto", "base", "claude", "codex", "default", "fast", "flash", "free",
"gemini", "gpt", "latest", "mini", "model", "plus", "pro", "reasoning",
"team", "think", "thinking", "vendor",
]
.contains(&token)
}
fn provider_pool_window_has_explicit_model(window: &Map<String, Value>) -> bool {
[
"model",
"model_name",
"model_id",
"quota_model",
"quota_model_name",
"target_model",
"models",
"model_ids",
]
.iter()
.any(|key| match window.get(*key) {
Some(Value::String(value)) => !value.trim().is_empty(),
Some(Value::Array(values)) => values.iter().any(|value| {
value.as_str().is_some_and(|value| !value.trim().is_empty())
}),
_ => false,
})
}
fn provider_pool_normalize_identifier(value: &str) -> String {
value
.trim()
.to_ascii_lowercase()
.chars()
.filter(|character| character.is_ascii_alphanumeric())
.collect()
}
fn provider_pool_identifier_tokens(value: &str) -> std::collections::BTreeSet<String> {
value
.split(|character: char| !character.is_ascii_alphanumeric())
.map(|token| token.trim().to_ascii_lowercase())
.filter(|token| token.len() >= 3)
.collect()
}
/// Materialize quota windows from the small set of shapes emitted by current
/// and legacy adapters. Model maps (`quota_by_model`/`models`) are converted
/// to the same window representation used by status snapshots, with the map
/// key retained as the model identity. Keeping this normalization here means
/// provider adapters do not need to grow model-specific quota code whenever an
/// upstream introduces another independent bucket.
fn provider_pool_collect_quota_windows(source: &Map<String, Value>) -> Vec<Map<String, Value>> {
let mut windows = Vec::new();
for key in ["windows", "additional_quota_windows"] {
if let Some(values) = source.get(key).and_then(Value::as_array) {
windows.extend(values.iter().filter_map(Value::as_object).cloned());
}
}
for key in ["quota_by_model", "models", "model_quotas"] {
let Some(models) = source.get(key).and_then(Value::as_object) else {
continue;
};
for (model_name, item) in models {
let Some(item) = item.as_object() else {
continue;
};
let mut window = item.clone();
window
.entry("model".to_string())
.or_insert_with(|| json!(model_name));
window
.entry("scope".to_string())
.or_insert_with(|| json!("model"));
window
.entry("code".to_string())
.or_insert_with(|| json!(format!("model:{model_name}")));
windows.push(window);
}
}
windows
}
pub fn provider_pool_quota_snapshot_updated_at(
key: &StoredProviderCatalogKey,
provider_type: &str,
@@ -197,12 +542,51 @@ pub(crate) fn provider_pool_reset_deadline_elapsed(
fn provider_pool_quota_window_is_exhausted(window: &Map<String, Value>) -> bool {
provider_pool_json_bool(window.get("is_exhausted"))
.or_else(|| provider_pool_json_bool(window.get("exhausted")))
.or_else(|| {
provider_pool_json_f64(window.get("used_ratio")).map(|value| value >= 1.0 - 1e-6)
provider_pool_json_f64(window.get("used_ratio").or_else(|| window.get("usage_ratio")))
.map(|value| value >= 1.0 - 1e-6)
})
.or_else(|| {
provider_pool_json_f64(window.get("used_percent")).map(|value| value >= 100.0 - 1e-6)
})
.or_else(|| {
provider_pool_json_f64(
window
.get("remaining_ratio")
.or_else(|| window.get("remaining_fraction")),
)
.map(|value| value <= 1e-6)
})
.or_else(|| {
provider_pool_json_f64(window.get("remaining_percent"))
.map(|value| value <= 1e-6)
})
.or_else(|| {
let remaining = provider_pool_json_f64(
window
.get("remaining")
.or_else(|| window.get("remaining_value")),
)?;
let limit = provider_pool_json_f64(
window
.get("limit")
.or_else(|| window.get("limit_value"))
.or_else(|| window.get("total")),
)?;
(limit > 0.0).then_some(remaining <= 0.0)
})
.unwrap_or(false)
}
fn provider_pool_window_is_model_scoped(window: &Map<String, Value>) -> bool {
window
.get("scope")
.and_then(Value::as_str)
.is_some_and(|scope| scope.trim().eq_ignore_ascii_case("model"))
|| provider_pool_window_has_explicit_model(window)
}
fn provider_pool_quota_snapshot_matches_provider(
quota_snapshot: &Map<String, Value>,
provider_type: &str,
@@ -237,6 +621,18 @@ fn provider_pool_quota_snapshot_matches_provider(
.get("windows")
.and_then(Value::as_array)
.is_some_and(|windows| !windows.is_empty())
|| quota_snapshot
.get("additional_quota_windows")
.and_then(Value::as_array)
.is_some_and(|windows| !windows.is_empty())
|| ["quota_by_model", "models", "model_quotas"]
.iter()
.any(|key| {
quota_snapshot
.get(*key)
.and_then(Value::as_object)
.is_some_and(|models| !models.is_empty())
})
|| quota_snapshot
.get("credits")
.and_then(Value::as_object)
@@ -265,16 +661,28 @@ pub(crate) fn provider_pool_quota_snapshot_exhausted_decision(
provider_pool_timestamp_unix_secs(quota_snapshot.get("observed_at"))
.or_else(|| provider_pool_timestamp_unix_secs(quota_snapshot.get("updated_at")));
if let Some(windows) = quota_snapshot
.get("windows")
.and_then(Value::as_array)
.filter(|windows| !windows.is_empty())
{
let materialized_windows = provider_pool_collect_quota_windows(quota_snapshot);
if !materialized_windows.is_empty() {
// Model-scoped windows are evaluated only when the request model
// is known. They must not turn the account-level fallback into
// an exhausted state for unrelated models.
let account_scoped_windows = materialized_windows
.iter()
.filter(|window| !provider_pool_window_is_model_scoped(window))
.collect::<Vec<_>>();
// A snapshot containing only model-scoped buckets has no
// account-wide signal to apply when the caller did not provide a
// model name (for example, an admin status listing). Do not let a
// single exhausted model poison every sibling bucket.
if account_scoped_windows.is_empty() {
return Some(false);
}
let windows = account_scoped_windows;
let mut saw_exhausted_window = false;
let mut saw_active_exhausted_window = false;
let mut windows_max_ratio = None::<f64>;
for window in windows.iter().filter_map(Value::as_object) {
for window in windows.iter() {
if let Some(ratio) = provider_pool_json_f64(window.get("used_ratio")) {
windows_max_ratio =
Some(windows_max_ratio.map_or(ratio, |current| current.max(ratio)));
@@ -123,12 +123,14 @@ impl ProviderPoolService {
provider_type: &str,
key: &StoredProviderCatalogKey,
auth_config: Option<&Map<String, Value>>,
provider_model_name: Option<&str>,
) -> aether_pool_core::PoolMemberSignals {
let adapter = self.adapter(provider_type);
let input = ProviderPoolMemberInput {
provider_type,
key,
auth_config,
provider_model_name,
};
adapter.member_signals(&input)
}