mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
Fix provider pool exhaustion scheduling
This commit is contained in:
@@ -709,31 +709,39 @@ impl<'a> PoolKeyCursor<'a> {
|
||||
}
|
||||
|
||||
async fn refill_queued_candidates(&mut self) -> bool {
|
||||
let mut candidates = Vec::new();
|
||||
let refill_target = self.window_size.max(1) as usize;
|
||||
// Keep pool expansion bounded; the cursor freezes one small window at a time.
|
||||
while candidates.len() < refill_target {
|
||||
let Some(mut page_candidates) = self.next_page_candidates().await else {
|
||||
break;
|
||||
};
|
||||
candidates.append(&mut page_candidates);
|
||||
}
|
||||
|
||||
if candidates.is_empty() {
|
||||
return false;
|
||||
}
|
||||
loop {
|
||||
let mut candidates = Vec::new();
|
||||
// Keep pool expansion bounded; the cursor freezes one small window at a time.
|
||||
while candidates.len() < refill_target {
|
||||
let Some(mut page_candidates) = self.next_page_candidates().await else {
|
||||
break;
|
||||
};
|
||||
candidates.append(&mut page_candidates);
|
||||
}
|
||||
|
||||
let (mut scheduled, mut skipped) = schedule_pool_page_candidates(
|
||||
self.state,
|
||||
candidates,
|
||||
self.sticky_session_token.as_deref(),
|
||||
)
|
||||
.await;
|
||||
scheduled.truncate(refill_target);
|
||||
self.record_skipped_candidates(&skipped);
|
||||
self.queued_candidates.extend(scheduled.drain(..));
|
||||
self.skipped_candidates.append(&mut skipped);
|
||||
!self.queued_candidates.is_empty()
|
||||
if candidates.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let (mut scheduled, mut skipped) = schedule_pool_page_candidates(
|
||||
self.state,
|
||||
candidates,
|
||||
self.sticky_session_token.as_deref(),
|
||||
)
|
||||
.await;
|
||||
self.record_skipped_candidates(&skipped);
|
||||
self.skipped_candidates.append(&mut skipped);
|
||||
|
||||
if scheduled.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
scheduled.truncate(refill_target);
|
||||
self.queued_candidates.extend(scheduled.drain(..));
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
|
||||
@@ -2798,6 +2806,78 @@ mod tests {
|
||||
assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_key_cursor_continues_after_exhausted_window() {
|
||||
let provider_config = Some(json!({
|
||||
"pool_advanced": {
|
||||
"skip_exhausted_accounts": true
|
||||
}
|
||||
}));
|
||||
let (provider, endpoint, mut keys, rows) = large_pool_fixture(3, provider_config.clone());
|
||||
for key in keys.iter_mut().take(2) {
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"provider_type": "openai",
|
||||
"exhausted": true,
|
||||
"usage_ratio": 1.0,
|
||||
"windows": [
|
||||
{
|
||||
"code": "daily",
|
||||
"used_ratio": 1.0,
|
||||
"remaining_ratio": 0.0
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
let data_state =
|
||||
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
|
||||
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
keys,
|
||||
)),
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows)),
|
||||
)
|
||||
.with_encryption_key_for_tests(aether_crypto::DEVELOPMENT_ENCRYPTION_KEY);
|
||||
let app = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
let group = sample_eligible_candidate(
|
||||
"provider-pool",
|
||||
"endpoint-1",
|
||||
"pool-group",
|
||||
10,
|
||||
provider_config,
|
||||
);
|
||||
|
||||
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
|
||||
cursor.window_size = 2;
|
||||
cursor.page_size = 2;
|
||||
cursor.max_scanned_keys = 4;
|
||||
|
||||
let candidate = cursor
|
||||
.next_key()
|
||||
.await
|
||||
.expect("cursor should scan past an exhausted window");
|
||||
assert_eq!(candidate.candidate.key_id, "key-00002");
|
||||
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
|
||||
assert!(candidate.orchestration.pool_key_lease.is_none());
|
||||
assert_eq!(
|
||||
cursor
|
||||
.skip_reason_counts
|
||||
.get(aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON),
|
||||
Some(&2)
|
||||
);
|
||||
|
||||
let skipped = cursor.take_skipped_candidates();
|
||||
assert_eq!(skipped.len(), 2);
|
||||
assert!(skipped.iter().all(|candidate| {
|
||||
candidate.skip_reason == aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
|
||||
}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_key_cursor_simulates_large_lru_pool_with_lazy_pages_and_dynamic_skips() {
|
||||
const KEY_COUNT: usize = 2048;
|
||||
|
||||
@@ -2,7 +2,9 @@ use super::state::{
|
||||
decode_jwt_claims, enrich_admin_provider_oauth_auth_config, json_non_empty_string,
|
||||
json_u64_value,
|
||||
};
|
||||
use crate::handlers::admin::admin_provider_pool_config;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
|
||||
use crate::provider_key_auth::provider_active_api_formats;
|
||||
use crate::GatewayError;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
@@ -168,6 +170,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
|
||||
.app()
|
||||
.invalidate_local_oauth_refresh_entry(&key.id)
|
||||
.await;
|
||||
seed_provider_oauth_pool_score(state, provider_id, key, now_unix_secs).await;
|
||||
}
|
||||
Ok(created)
|
||||
}
|
||||
@@ -221,10 +224,75 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
|
||||
.app()
|
||||
.invalidate_local_oauth_refresh_entry(&key.id)
|
||||
.await;
|
||||
seed_provider_oauth_pool_score(state, &existing_key.provider_id, key, now_unix_secs).await;
|
||||
}
|
||||
Ok(persisted)
|
||||
}
|
||||
|
||||
async fn seed_provider_oauth_pool_score(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
key: &StoredProviderCatalogKey,
|
||||
now_unix_secs: u64,
|
||||
) {
|
||||
let provider_id = provider_id.to_string();
|
||||
let provider = match state
|
||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
{
|
||||
Ok(mut providers) => providers.pop(),
|
||||
Err(err) => {
|
||||
tracing::debug!(
|
||||
provider_id = %provider_id,
|
||||
key_id = %key.id,
|
||||
error = ?err,
|
||||
"gateway provider oauth provisioning: failed to read provider for pool score seed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
let Some(provider) = provider else {
|
||||
return;
|
||||
};
|
||||
let Some(pool_config) = admin_provider_pool_config(&provider) else {
|
||||
return;
|
||||
};
|
||||
let endpoints = match state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
{
|
||||
Ok(endpoints) => endpoints,
|
||||
Err(err) => {
|
||||
tracing::debug!(
|
||||
provider_id = %provider_id,
|
||||
key_id = %key.id,
|
||||
error = ?err,
|
||||
"gateway provider oauth provisioning: failed to read endpoints for pool score seed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize).clamp(1, 50_000);
|
||||
if let Err(err) = ensure_provider_key_pool_scores_for_keys(
|
||||
state.as_ref(),
|
||||
&provider,
|
||||
&pool_config,
|
||||
&endpoints,
|
||||
std::slice::from_ref(key),
|
||||
now_unix_secs,
|
||||
score_ensure_budget,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::debug!(
|
||||
provider_id = %provider_id,
|
||||
key_id = %key.id,
|
||||
error = ?err,
|
||||
"gateway provider oauth provisioning: failed to seed pool score row"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn provider_oauth_catalog_key_api_formats(
|
||||
provider_type: &str,
|
||||
api_formats: &[String],
|
||||
|
||||
@@ -12,8 +12,12 @@ use aether_data::repository::management_tokens::{
|
||||
use aether_data::repository::oauth_providers::{
|
||||
InMemoryOAuthProviderRepository, OAuthProviderReadRepository,
|
||||
};
|
||||
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
|
||||
use aether_data_contracts::repository::pool_scores::{
|
||||
GetPoolMemberScoresByIdsQuery, PoolMemberHardState, PoolMemberIdentity, PoolScoreReadRepository,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
|
||||
};
|
||||
@@ -32,6 +36,7 @@ use super::super::{
|
||||
use crate::admin_api::{
|
||||
maybe_build_local_admin_provider_oauth_response, AdminAppState, AdminRequestContext,
|
||||
};
|
||||
use crate::ai_serving::{provider_key_pool_score_id, provider_key_pool_score_scope};
|
||||
use crate::audit::AdminAuditEvent;
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
|
||||
@@ -1930,6 +1935,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
|
||||
|
||||
let mut provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 10);
|
||||
provider.provider_type = "chatgpt_web".to_string();
|
||||
provider.config = Some(json!({"pool_advanced": {}}));
|
||||
let endpoint = sample_endpoint(
|
||||
"endpoint-chatgpt-web-image",
|
||||
"provider-chatgpt-web",
|
||||
@@ -1941,6 +1947,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
|
||||
vec![endpoint],
|
||||
vec![],
|
||||
));
|
||||
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::default());
|
||||
|
||||
let (token_url, token_handle) = start_server(token_server).await;
|
||||
let gateway = build_router_with_state(
|
||||
@@ -1950,6 +1957,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository))
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
)
|
||||
.with_provider_oauth_token_url_for_tests(
|
||||
@@ -2030,6 +2038,20 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
|
||||
assert_eq!(auth_config["plan_type"], "plus");
|
||||
assert_eq!(auth_config["user_id"], "user-pool-image");
|
||||
|
||||
let score_scope = provider_key_pool_score_scope();
|
||||
let score_identity =
|
||||
PoolMemberIdentity::provider_api_key("provider-chatgpt-web", persisted.id.clone());
|
||||
let scores = pool_score_repository
|
||||
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
|
||||
ids: vec![provider_key_pool_score_id(&score_identity, &score_scope)],
|
||||
})
|
||||
.await
|
||||
.expect("pool score should load");
|
||||
assert_eq!(scores.len(), 1);
|
||||
assert_eq!(scores[0].member_id, persisted.id);
|
||||
assert_eq!(scores[0].hard_state, PoolMemberHardState::Unknown);
|
||||
assert!(scores[0].score > 0.0);
|
||||
|
||||
gateway_handle.abort();
|
||||
token_handle.abort();
|
||||
}
|
||||
|
||||
@@ -380,6 +380,81 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_quota_exhaustion_snapshot_expires_after_reset_at() {
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("system time should be after unix epoch")
|
||||
.as_secs();
|
||||
|
||||
let mut expired = sample_key(None);
|
||||
expired.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"code": "exhausted",
|
||||
"exhausted": true,
|
||||
"updated_at": now.saturating_sub(600),
|
||||
"windows": [{
|
||||
"code": "5h",
|
||||
"used_ratio": 1.0,
|
||||
"reset_at": now.saturating_sub(60),
|
||||
"is_exhausted": true
|
||||
}]
|
||||
}
|
||||
}));
|
||||
assert!(!provider_pool_key_account_quota_exhausted(
|
||||
&expired, "codex"
|
||||
));
|
||||
|
||||
let mut active = sample_key(None);
|
||||
active.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"code": "exhausted",
|
||||
"exhausted": true,
|
||||
"updated_at": now,
|
||||
"windows": [{
|
||||
"code": "5h",
|
||||
"used_ratio": 1.0,
|
||||
"reset_at": now.saturating_add(3600),
|
||||
"is_exhausted": true
|
||||
}]
|
||||
}
|
||||
}));
|
||||
assert!(provider_pool_key_account_quota_exhausted(&active, "codex"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_quota_exhaustion_metadata_expires_after_reset_at() {
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("system time should be after unix epoch")
|
||||
.as_secs();
|
||||
|
||||
assert!(!provider_pool_key_account_quota_exhausted(
|
||||
&sample_key(Some(json!({
|
||||
"codex": {
|
||||
"updated_at": now.saturating_sub(600),
|
||||
"primary_used_percent": 100.0,
|
||||
"primary_reset_at": now.saturating_sub(60)
|
||||
}
|
||||
}))),
|
||||
"codex",
|
||||
));
|
||||
assert!(provider_pool_key_account_quota_exhausted(
|
||||
&sample_key(Some(json!({
|
||||
"codex": {
|
||||
"updated_at": now,
|
||||
"primary_used_percent": 100.0,
|
||||
"primary_reset_at": now.saturating_add(3600)
|
||||
}
|
||||
}))),
|
||||
"codex",
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn grok_quota_tier_boundaries_match_pool_modes() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -10,8 +10,9 @@ use crate::provider::{
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision,
|
||||
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,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
@@ -255,6 +256,19 @@ fn infer_chatgpt_web_image_quota_limit(
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_current_unix_secs().is_some_and(|now| {
|
||||
let mut image_quota = Map::new();
|
||||
if let Some(value) = bucket.get("image_quota_reset_at") {
|
||||
image_quota.insert("reset_at".to_string(), value.clone());
|
||||
}
|
||||
provider_pool_reset_deadline_elapsed(
|
||||
&image_quota,
|
||||
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
|
||||
now,
|
||||
)
|
||||
}) {
|
||||
return false;
|
||||
}
|
||||
if provider_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) {
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -10,8 +10,9 @@ use crate::provider::{
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision,
|
||||
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,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
@@ -121,6 +122,37 @@ pub fn build_codex_pool_quota_request(
|
||||
})
|
||||
}
|
||||
|
||||
fn codex_window_reset_elapsed(bucket: &Map<String, Value>, prefix: &str) -> bool {
|
||||
let Some(now_unix_secs) = provider_pool_current_unix_secs() else {
|
||||
return false;
|
||||
};
|
||||
let mut window = Map::new();
|
||||
for (target, source) in [
|
||||
("reset_at", format!("{prefix}_reset_at")),
|
||||
("next_reset_at", format!("{prefix}_next_reset_at")),
|
||||
("reset_seconds", format!("{prefix}_reset_seconds")),
|
||||
(
|
||||
"reset_after_seconds",
|
||||
format!("{prefix}_reset_after_seconds"),
|
||||
),
|
||||
] {
|
||||
if let Some(value) = bucket.get(source.as_str()) {
|
||||
window.insert(target.to_string(), value.clone());
|
||||
}
|
||||
}
|
||||
provider_pool_reset_deadline_elapsed(
|
||||
&window,
|
||||
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
|
||||
now_unix_secs,
|
||||
)
|
||||
}
|
||||
|
||||
fn codex_window_used_percent_exhausted(bucket: &Map<String, Value>, prefix: &str) -> bool {
|
||||
let used_percent_key = format!("{prefix}_used_percent");
|
||||
provider_pool_json_f64(bucket.get(used_percent_key.as_str()))
|
||||
.is_some_and(|value| value >= 100.0 && !codex_window_reset_elapsed(bucket, prefix))
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
|
||||
return false;
|
||||
@@ -130,7 +162,6 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if !has_window_data && provider_pool_json_bool(bucket.get("has_credits")) == Some(false) {
|
||||
return true;
|
||||
}
|
||||
provider_pool_json_f64(bucket.get("primary_used_percent")).is_some_and(|value| value >= 100.0)
|
||||
|| provider_pool_json_f64(bucket.get("secondary_used_percent"))
|
||||
.is_some_and(|value| value >= 100.0)
|
||||
codex_window_used_percent_exhausted(bucket, "primary")
|
||||
|| codex_window_used_percent_exhausted(bucket, "secondary")
|
||||
}
|
||||
|
||||
@@ -7,8 +7,9 @@ use crate::provider::{
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision,
|
||||
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,
|
||||
};
|
||||
|
||||
pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")];
|
||||
@@ -191,6 +192,8 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
|
||||
let mut model_count = 0usize;
|
||||
let mut exhausted_count = 0usize;
|
||||
let now_unix_secs = provider_pool_current_unix_secs();
|
||||
let bucket_updated_at = provider_pool_timestamp_unix_secs(bucket.get("updated_at"));
|
||||
for (model_key, item) in models.iter() {
|
||||
if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) {
|
||||
continue;
|
||||
@@ -207,12 +210,14 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
continue;
|
||||
}
|
||||
model_count += 1;
|
||||
if provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
|
||||
let quota_exhausted = provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
|
||||
|| provider_pool_json_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0)
|
||||
|| provider_pool_json_f64(item.get("remaining")).is_some_and(|value| value <= 0.0)
|
||||
|| provider_pool_json_f64(item.get("remaining_fraction"))
|
||||
.is_some_and(|value| value <= 0.0)
|
||||
{
|
||||
.is_some_and(|value| value <= 0.0);
|
||||
let reset_elapsed = now_unix_secs
|
||||
.is_some_and(|now| provider_pool_reset_deadline_elapsed(item, bucket_updated_at, now));
|
||||
if quota_exhausted && !reset_elapsed {
|
||||
exhausted_count += 1;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,8 +11,9 @@ use crate::provider::{
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision,
|
||||
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,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
@@ -151,6 +152,15 @@ fn normalize_kiro_version(value: &str) -> &str {
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_current_unix_secs().is_some_and(|now| {
|
||||
provider_pool_reset_deadline_elapsed(
|
||||
bucket,
|
||||
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
|
||||
now,
|
||||
)
|
||||
}) {
|
||||
return false;
|
||||
}
|
||||
if provider_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) {
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
@@ -135,7 +137,7 @@ pub(crate) fn provider_pool_json_f64(value: Option<&Value>) -> Option<f64> {
|
||||
.filter(|value| value.is_finite())
|
||||
}
|
||||
|
||||
fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
|
||||
pub(crate) fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
|
||||
let mut timestamp = provider_pool_json_f64(value)?;
|
||||
if timestamp <= 0.0 {
|
||||
return None;
|
||||
@@ -146,6 +148,49 @@ fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
|
||||
Some(timestamp as u64)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_pool_current_unix_secs() -> Option<u64> {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
}
|
||||
|
||||
fn provider_pool_reset_deadline_unix_secs(
|
||||
item: &Map<String, Value>,
|
||||
fallback_observed_at: Option<u64>,
|
||||
) -> Option<u64> {
|
||||
provider_pool_timestamp_unix_secs(item.get("reset_at"))
|
||||
.or_else(|| provider_pool_timestamp_unix_secs(item.get("next_reset_at")))
|
||||
.or_else(|| {
|
||||
let reset_seconds = provider_pool_json_f64(item.get("reset_seconds"))
|
||||
.or_else(|| provider_pool_json_f64(item.get("reset_after_seconds")))?;
|
||||
if reset_seconds < 0.0 {
|
||||
return None;
|
||||
}
|
||||
let base = provider_pool_timestamp_unix_secs(item.get("observed_at"))
|
||||
.or_else(|| provider_pool_timestamp_unix_secs(item.get("updated_at")))
|
||||
.or(fallback_observed_at)?;
|
||||
Some(base.saturating_add(reset_seconds.ceil() as u64))
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn provider_pool_reset_deadline_elapsed(
|
||||
item: &Map<String, Value>,
|
||||
fallback_observed_at: Option<u64>,
|
||||
now_unix_secs: u64,
|
||||
) -> bool {
|
||||
provider_pool_reset_deadline_unix_secs(item, fallback_observed_at)
|
||||
.is_some_and(|reset_at| reset_at <= now_unix_secs)
|
||||
}
|
||||
|
||||
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_f64(window.get("used_ratio")).map(|value| value >= 1.0 - 1e-6)
|
||||
})
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn provider_pool_quota_snapshot_matches_provider(
|
||||
quota_snapshot: &Map<String, Value>,
|
||||
provider_type: &str,
|
||||
@@ -203,19 +248,45 @@ pub(crate) fn provider_pool_quota_snapshot_exhausted_decision(
|
||||
}
|
||||
let exhausted = provider_pool_json_bool(quota_snapshot.get("exhausted"))?;
|
||||
if exhausted {
|
||||
let windows_max_ratio = quota_snapshot
|
||||
let now_unix_secs = provider_pool_current_unix_secs();
|
||||
let snapshot_observed_at =
|
||||
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(|w| !w.is_empty())
|
||||
.and_then(|windows| {
|
||||
windows
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.filter_map(|w| w.get("used_ratio"))
|
||||
.filter_map(Value::as_f64)
|
||||
.max_by(f64::total_cmp)
|
||||
});
|
||||
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
|
||||
.filter(|windows| !windows.is_empty())
|
||||
{
|
||||
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) {
|
||||
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)));
|
||||
}
|
||||
if provider_pool_quota_window_is_exhausted(window) {
|
||||
saw_exhausted_window = true;
|
||||
let reset_elapsed = now_unix_secs.is_some_and(|now| {
|
||||
provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now)
|
||||
});
|
||||
if !reset_elapsed {
|
||||
saw_active_exhausted_window = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if saw_exhausted_window {
|
||||
return Some(saw_active_exhausted_window);
|
||||
}
|
||||
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
|
||||
return Some(false);
|
||||
}
|
||||
} else if now_unix_secs.is_some_and(|now| {
|
||||
provider_pool_reset_deadline_elapsed(quota_snapshot, snapshot_observed_at, now)
|
||||
}) {
|
||||
return Some(false);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user