mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +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 {
|
async fn refill_queued_candidates(&mut self) -> bool {
|
||||||
let mut candidates = Vec::new();
|
|
||||||
let refill_target = self.window_size.max(1) as usize;
|
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() {
|
loop {
|
||||||
return false;
|
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(
|
if candidates.is_empty() {
|
||||||
self.state,
|
return false;
|
||||||
candidates,
|
}
|
||||||
self.sticky_session_token.as_deref(),
|
|
||||||
)
|
let (mut scheduled, mut skipped) = schedule_pool_page_candidates(
|
||||||
.await;
|
self.state,
|
||||||
scheduled.truncate(refill_target);
|
candidates,
|
||||||
self.record_skipped_candidates(&skipped);
|
self.sticky_session_token.as_deref(),
|
||||||
self.queued_candidates.extend(scheduled.drain(..));
|
)
|
||||||
self.skipped_candidates.append(&mut skipped);
|
.await;
|
||||||
!self.queued_candidates.is_empty()
|
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> {
|
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));
|
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]
|
#[tokio::test]
|
||||||
async fn pool_key_cursor_simulates_large_lru_pool_with_lazy_pages_and_dynamic_skips() {
|
async fn pool_key_cursor_simulates_large_lru_pool_with_lazy_pages_and_dynamic_skips() {
|
||||||
const KEY_COUNT: usize = 2048;
|
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,
|
decode_jwt_claims, enrich_admin_provider_oauth_auth_config, json_non_empty_string,
|
||||||
json_u64_value,
|
json_u64_value,
|
||||||
};
|
};
|
||||||
|
use crate::handlers::admin::admin_provider_pool_config;
|
||||||
use crate::handlers::admin::request::AdminAppState;
|
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::provider_key_auth::provider_active_api_formats;
|
||||||
use crate::GatewayError;
|
use crate::GatewayError;
|
||||||
use aether_data_contracts::repository::provider_catalog::{
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
@@ -168,6 +170,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
|
|||||||
.app()
|
.app()
|
||||||
.invalidate_local_oauth_refresh_entry(&key.id)
|
.invalidate_local_oauth_refresh_entry(&key.id)
|
||||||
.await;
|
.await;
|
||||||
|
seed_provider_oauth_pool_score(state, provider_id, key, now_unix_secs).await;
|
||||||
}
|
}
|
||||||
Ok(created)
|
Ok(created)
|
||||||
}
|
}
|
||||||
@@ -221,10 +224,75 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
|
|||||||
.app()
|
.app()
|
||||||
.invalidate_local_oauth_refresh_entry(&key.id)
|
.invalidate_local_oauth_refresh_entry(&key.id)
|
||||||
.await;
|
.await;
|
||||||
|
seed_provider_oauth_pool_score(state, &existing_key.provider_id, key, now_unix_secs).await;
|
||||||
}
|
}
|
||||||
Ok(persisted)
|
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(
|
fn provider_oauth_catalog_key_api_formats(
|
||||||
provider_type: &str,
|
provider_type: &str,
|
||||||
api_formats: &[String],
|
api_formats: &[String],
|
||||||
|
|||||||
@@ -12,8 +12,12 @@ use aether_data::repository::management_tokens::{
|
|||||||
use aether_data::repository::oauth_providers::{
|
use aether_data::repository::oauth_providers::{
|
||||||
InMemoryOAuthProviderRepository, OAuthProviderReadRepository,
|
InMemoryOAuthProviderRepository, OAuthProviderReadRepository,
|
||||||
};
|
};
|
||||||
|
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
|
||||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||||
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
|
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::{
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
|
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
|
||||||
};
|
};
|
||||||
@@ -32,6 +36,7 @@ use super::super::{
|
|||||||
use crate::admin_api::{
|
use crate::admin_api::{
|
||||||
maybe_build_local_admin_provider_oauth_response, AdminAppState, AdminRequestContext,
|
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::audit::AdminAuditEvent;
|
||||||
use crate::constants::{
|
use crate::constants::{
|
||||||
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
|
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);
|
let mut provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 10);
|
||||||
provider.provider_type = "chatgpt_web".to_string();
|
provider.provider_type = "chatgpt_web".to_string();
|
||||||
|
provider.config = Some(json!({"pool_advanced": {}}));
|
||||||
let endpoint = sample_endpoint(
|
let endpoint = sample_endpoint(
|
||||||
"endpoint-chatgpt-web-image",
|
"endpoint-chatgpt-web-image",
|
||||||
"provider-chatgpt-web",
|
"provider-chatgpt-web",
|
||||||
@@ -1941,6 +1947,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
|
|||||||
vec![endpoint],
|
vec![endpoint],
|
||||||
vec![],
|
vec![],
|
||||||
));
|
));
|
||||||
|
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::default());
|
||||||
|
|
||||||
let (token_url, token_handle) = start_server(token_server).await;
|
let (token_url, token_handle) = start_server(token_server).await;
|
||||||
let gateway = build_router_with_state(
|
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(
|
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||||
provider_catalog_repository.clone(),
|
provider_catalog_repository.clone(),
|
||||||
)
|
)
|
||||||
|
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository))
|
||||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||||
)
|
)
|
||||||
.with_provider_oauth_token_url_for_tests(
|
.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["plan_type"], "plus");
|
||||||
assert_eq!(auth_config["user_id"], "user-pool-image");
|
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();
|
gateway_handle.abort();
|
||||||
token_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]
|
#[test]
|
||||||
fn grok_quota_tier_boundaries_match_pool_modes() {
|
fn grok_quota_tier_boundaries_match_pool_modes() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
|
|||||||
@@ -10,8 +10,9 @@ use crate::provider::{
|
|||||||
ProviderPoolMemberInput,
|
ProviderPoolMemberInput,
|
||||||
};
|
};
|
||||||
use crate::quota::{
|
use crate::quota::{
|
||||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
|
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
|
||||||
provider_pool_quota_snapshot_exhausted_decision,
|
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;
|
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 {
|
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) {
|
if provider_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) {
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -10,8 +10,9 @@ use crate::provider::{
|
|||||||
ProviderPoolMemberInput,
|
ProviderPoolMemberInput,
|
||||||
};
|
};
|
||||||
use crate::quota::{
|
use crate::quota::{
|
||||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
|
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
|
||||||
provider_pool_quota_snapshot_exhausted_decision,
|
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;
|
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 {
|
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||||
if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
|
if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
|
||||||
return false;
|
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) {
|
if !has_window_data && provider_pool_json_bool(bucket.get("has_credits")) == Some(false) {
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
provider_pool_json_f64(bucket.get("primary_used_percent")).is_some_and(|value| value >= 100.0)
|
codex_window_used_percent_exhausted(bucket, "primary")
|
||||||
|| provider_pool_json_f64(bucket.get("secondary_used_percent"))
|
|| codex_window_used_percent_exhausted(bucket, "secondary")
|
||||||
.is_some_and(|value| value >= 100.0)
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -7,8 +7,9 @@ use crate::provider::{
|
|||||||
ProviderPoolMemberInput,
|
ProviderPoolMemberInput,
|
||||||
};
|
};
|
||||||
use crate::quota::{
|
use crate::quota::{
|
||||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
|
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
|
||||||
provider_pool_quota_snapshot_exhausted_decision,
|
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")];
|
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 model_count = 0usize;
|
||||||
let mut exhausted_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() {
|
for (model_key, item) in models.iter() {
|
||||||
if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) {
|
if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) {
|
||||||
continue;
|
continue;
|
||||||
@@ -207,12 +210,14 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
|||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
model_count += 1;
|
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("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")).is_some_and(|value| value <= 0.0)
|
||||||
|| provider_pool_json_f64(item.get("remaining_fraction"))
|
|| 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;
|
exhausted_count += 1;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,8 +11,9 @@ use crate::provider::{
|
|||||||
ProviderPoolMemberInput,
|
ProviderPoolMemberInput,
|
||||||
};
|
};
|
||||||
use crate::quota::{
|
use crate::quota::{
|
||||||
provider_pool_json_f64, provider_pool_metadata_bucket,
|
provider_pool_current_unix_secs, provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||||
provider_pool_quota_snapshot_exhausted_decision,
|
provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
|
||||||
|
provider_pool_timestamp_unix_secs,
|
||||||
};
|
};
|
||||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
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 {
|
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) {
|
if provider_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) {
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,3 +1,5 @@
|
|||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||||
use serde_json::{json, Map, Value};
|
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())
|
.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)?;
|
let mut timestamp = provider_pool_json_f64(value)?;
|
||||||
if timestamp <= 0.0 {
|
if timestamp <= 0.0 {
|
||||||
return None;
|
return None;
|
||||||
@@ -146,6 +148,49 @@ fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
|
|||||||
Some(timestamp as 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(
|
fn provider_pool_quota_snapshot_matches_provider(
|
||||||
quota_snapshot: &Map<String, Value>,
|
quota_snapshot: &Map<String, Value>,
|
||||||
provider_type: &str,
|
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"))?;
|
let exhausted = provider_pool_json_bool(quota_snapshot.get("exhausted"))?;
|
||||||
if 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")
|
.get("windows")
|
||||||
.and_then(Value::as_array)
|
.and_then(Value::as_array)
|
||||||
.filter(|w| !w.is_empty())
|
.filter(|windows| !windows.is_empty())
|
||||||
.and_then(|windows| {
|
{
|
||||||
windows
|
let mut saw_exhausted_window = false;
|
||||||
.iter()
|
let mut saw_active_exhausted_window = false;
|
||||||
.filter_map(Value::as_object)
|
let mut windows_max_ratio = None::<f64>;
|
||||||
.filter_map(|w| w.get("used_ratio"))
|
|
||||||
.filter_map(Value::as_f64)
|
for window in windows.iter().filter_map(Value::as_object) {
|
||||||
.max_by(f64::total_cmp)
|
if let Some(ratio) = provider_pool_json_f64(window.get("used_ratio")) {
|
||||||
});
|
windows_max_ratio =
|
||||||
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
|
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);
|
return Some(false);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user