mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 08:27:46 +08:00
fix(pool): isolate dynamic model quota buckets and 429 scheduling
This commit is contained in:
@@ -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()
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user