mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor: extract provider pool abstractions
This commit is contained in:
@@ -20,6 +20,8 @@ aether-dispatch-core.workspace = true
|
||||
aether-http.workspace = true
|
||||
aether-model-fetch.workspace = true
|
||||
aether-oauth.workspace = true
|
||||
aether-pool-core.workspace = true
|
||||
aether-provider-pool.workspace = true
|
||||
aether-provider-transport.workspace = true
|
||||
aether-scheduler-core.workspace = true
|
||||
aether-runtime.workspace = true
|
||||
|
||||
@@ -3,10 +3,9 @@ pub(crate) use crate::handlers::admin::{
|
||||
build_internal_control_error_response, create_provider_oauth_catalog_key,
|
||||
find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response,
|
||||
maybe_build_local_admin_response, provider_oauth_maintenance_endpoint_for_provider,
|
||||
provider_oauth_runtime_endpoint_for_provider, provider_type_supports_quota_refresh,
|
||||
reconcile_admin_fixed_provider_template_endpoints, refresh_antigravity_provider_quota_locally,
|
||||
refresh_chatgpt_web_provider_quota_locally, refresh_codex_provider_quota_locally,
|
||||
refresh_kiro_provider_quota_locally, refresh_provider_oauth_account_state_after_update,
|
||||
provider_oauth_runtime_endpoint_for_provider, provider_quota_refresh_endpoint_for_provider,
|
||||
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
|
||||
refresh_provider_oauth_account_state_after_update, refresh_provider_pool_quota_locally,
|
||||
update_existing_provider_oauth_catalog_key, AdminAppState,
|
||||
AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError, AdminRequestContext,
|
||||
AdminRouteRequest, AdminRouteResponse, AdminRouteResult, AdminStatsTimeRange,
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
use aether_ai_serving::{
|
||||
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreRules, POOL_SCORE_VERSION,
|
||||
};
|
||||
use aether_data_contracts::repository::pool_scores::{
|
||||
PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope, UpsertPoolMemberScore,
|
||||
POOL_SCORE_CAPABILITY_ACCOUNT, POOL_SCORE_SCOPE_KIND_ACCOUNT,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_pool_core::{
|
||||
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreRules, POOL_SCORE_VERSION,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::handlers::shared::{provider_key_health_summary, provider_key_status_snapshot_payload};
|
||||
|
||||
@@ -2,11 +2,6 @@ use std::collections::{btree_map::Entry, BTreeMap, BTreeSet, VecDeque};
|
||||
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
|
||||
|
||||
use aether_admin::provider::pool as admin_provider_pool_pure;
|
||||
use aether_ai_serving::{
|
||||
normalize_enabled_ai_pool_presets, run_ai_pool_scheduler, AiPoolCandidateFacts,
|
||||
AiPoolCandidateInput, AiPoolCandidateOrchestration, AiPoolCatalogKeyContext,
|
||||
AiPoolRuntimeState, AiPoolSchedulingConfig, AiPoolSchedulingPreset,
|
||||
};
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
|
||||
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
|
||||
@@ -16,7 +11,11 @@ use aether_data_contracts::repository::pool_scores::{
|
||||
PoolMemberScheduleFeedback, PoolScoreScope, StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use serde_json::{Map, Value};
|
||||
use aether_pool_core::{
|
||||
run_pool_scheduler, PoolCandidateFacts, PoolCandidateInput, PoolCandidateOrchestration,
|
||||
PoolMemberSignals, PoolRuntimeState, PoolSchedulingConfig, PoolSchedulingPreset,
|
||||
};
|
||||
use aether_provider_pool::ProviderPoolService;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::{
|
||||
@@ -35,16 +34,12 @@ use crate::handlers::shared::provider_pool::{
|
||||
read_admin_provider_pool_key_cooldown_reason, AdminProviderPoolConfig,
|
||||
AdminProviderPoolRuntimeState,
|
||||
};
|
||||
use crate::handlers::shared::{
|
||||
parse_catalog_auth_config_json, provider_key_health_summary,
|
||||
provider_key_status_snapshot_payload,
|
||||
};
|
||||
use crate::handlers::shared::{parse_catalog_auth_config_json, provider_key_health_summary};
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
use crate::provider_key_auth::provider_key_auth_semantics;
|
||||
|
||||
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
type PoolCatalogKeyContext = AiPoolCatalogKeyContext;
|
||||
type PoolCatalogKeyContext = PoolMemberSignals;
|
||||
|
||||
pub(crate) async fn apply_local_execution_pool_scheduler(
|
||||
state: PlannerAppState<'_>,
|
||||
@@ -796,6 +791,8 @@ async fn read_pool_catalog_key_contexts_by_id(
|
||||
}
|
||||
};
|
||||
|
||||
let provider_pool_service = ProviderPoolService::with_builtin_adapters();
|
||||
|
||||
keys.into_iter()
|
||||
.map(|key| {
|
||||
let provider_type = provider_type_by_key_id
|
||||
@@ -804,7 +801,7 @@ async fn read_pool_catalog_key_contexts_by_id(
|
||||
.unwrap_or_default();
|
||||
(
|
||||
key.id.clone(),
|
||||
build_pool_catalog_key_context(state, &key, provider_type),
|
||||
build_pool_catalog_key_context(state, &provider_pool_service, &key, provider_type),
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
@@ -812,24 +809,15 @@ async fn read_pool_catalog_key_contexts_by_id(
|
||||
|
||||
fn build_pool_catalog_key_context(
|
||||
state: PlannerAppState<'_>,
|
||||
provider_pool_service: &ProviderPoolService,
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> PoolCatalogKeyContext {
|
||||
let status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
|
||||
let quota_snapshot = status_snapshot
|
||||
.as_object()
|
||||
.and_then(|snapshot| snapshot.get("quota"))
|
||||
.and_then(Value::as_object);
|
||||
let account_snapshot = status_snapshot
|
||||
.as_object()
|
||||
.and_then(|snapshot| snapshot.get("account"))
|
||||
.and_then(Value::as_object);
|
||||
|
||||
let (health_score, _, _, _, _) = provider_key_health_summary(key);
|
||||
let health_score = key
|
||||
.health_by_format
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.filter(|payload| !payload.is_empty())
|
||||
.map(|_| health_score);
|
||||
let latency_avg_ms = key
|
||||
@@ -841,131 +829,14 @@ fn build_pool_catalog_key_context(
|
||||
})
|
||||
.filter(|value| value.is_finite() && *value >= 0.0);
|
||||
|
||||
PoolCatalogKeyContext {
|
||||
oauth_plan_type: quota_snapshot
|
||||
.and_then(|quota| quota.get("plan_type"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|value| normalize_pool_plan_type(value, provider_type))
|
||||
.or_else(|| derive_pool_oauth_plan_type(state, key, provider_type)),
|
||||
quota_usage_ratio: quota_snapshot
|
||||
.and_then(|quota| quota.get("usage_ratio"))
|
||||
.and_then(json_f64)
|
||||
.map(|value| value.clamp(0.0, 1.0)),
|
||||
quota_reset_seconds: quota_snapshot
|
||||
.and_then(|quota| quota.get("reset_seconds"))
|
||||
.and_then(json_f64)
|
||||
.filter(|value| *value >= 0.0),
|
||||
account_blocked: account_snapshot
|
||||
.and_then(|account| account.get("blocked"))
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
|| admin_provider_pool_pure::admin_pool_key_is_known_banned(key),
|
||||
quota_exhausted: pool_catalog_key_quota_exhausted(key, provider_type, quota_snapshot),
|
||||
health_score,
|
||||
latency_avg_ms,
|
||||
catalog_lru_score: Some(key.last_used_at_unix_secs.unwrap_or(0) as f64),
|
||||
}
|
||||
}
|
||||
|
||||
fn pool_catalog_key_quota_exhausted(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
quota_snapshot: Option<&Map<String, Value>>,
|
||||
) -> bool {
|
||||
match provider_type.trim().to_ascii_lowercase().as_str() {
|
||||
"codex" | "kiro" | "chatgpt_web" => {
|
||||
admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type)
|
||||
}
|
||||
_ => quota_snapshot
|
||||
.and_then(|quota| quota.get("exhausted"))
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn derive_pool_oauth_plan_type(
|
||||
state: PlannerAppState<'_>,
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> Option<String> {
|
||||
if !provider_key_auth_semantics(key, provider_type).oauth_managed() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let provider_type_key = provider_type.trim().to_ascii_lowercase();
|
||||
if let Some(upstream_metadata) = key.upstream_metadata.as_ref().and_then(Value::as_object) {
|
||||
let provider_bucket = upstream_metadata
|
||||
.get(&provider_type_key)
|
||||
.and_then(Value::as_object);
|
||||
for source in provider_bucket
|
||||
.into_iter()
|
||||
.chain(std::iter::once(upstream_metadata))
|
||||
{
|
||||
if let Some(plan_type) = pool_plan_type_from_source(
|
||||
source,
|
||||
provider_type,
|
||||
&[
|
||||
"plan_type",
|
||||
"tier",
|
||||
"subscription_title",
|
||||
"subscription_plan",
|
||||
"plan",
|
||||
],
|
||||
) {
|
||||
return Some(plan_type);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
parse_catalog_auth_config_json(state.app(), key).and_then(|auth_config| {
|
||||
pool_plan_type_from_source(
|
||||
&auth_config,
|
||||
provider_type,
|
||||
&["plan_type", "tier", "plan", "subscription_plan"],
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn pool_plan_type_from_source(
|
||||
source: &Map<String, Value>,
|
||||
provider_type: &str,
|
||||
fields: &[&str],
|
||||
) -> Option<String> {
|
||||
for field in fields {
|
||||
let Some(value) = source.get(*field).and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
if let Some(normalized) = normalize_pool_plan_type(value, provider_type) {
|
||||
return Some(normalized);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn normalize_pool_plan_type(value: &str, provider_type: &str) -> Option<String> {
|
||||
let mut normalized = value.trim().to_string();
|
||||
if normalized.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let provider_type = provider_type.trim().to_ascii_lowercase();
|
||||
if !provider_type.is_empty() && normalized.to_ascii_lowercase().starts_with(&provider_type) {
|
||||
normalized = normalized[provider_type.len()..]
|
||||
.trim_matches(|ch: char| [' ', ':', '-', '_'].contains(&ch))
|
||||
.to_string();
|
||||
}
|
||||
|
||||
let normalized = normalized.trim().to_ascii_lowercase();
|
||||
(!normalized.is_empty()).then_some(normalized)
|
||||
}
|
||||
|
||||
fn json_f64(value: &Value) -> Option<f64> {
|
||||
match value {
|
||||
Value::Number(number) => number.as_f64(),
|
||||
Value::String(text) => text.trim().parse::<f64>().ok(),
|
||||
_ => None,
|
||||
}
|
||||
.filter(|value| value.is_finite())
|
||||
let auth_config = parse_catalog_auth_config_json(state.app(), key);
|
||||
let mut signals =
|
||||
provider_pool_service.member_signals(provider_type, key, auth_config.as_ref());
|
||||
signals.account_blocked |= admin_provider_pool_pure::admin_pool_key_is_known_banned(key);
|
||||
signals.health_score = health_score;
|
||||
signals.latency_avg_ms = latency_avg_ms;
|
||||
signals.catalog_lru_score = Some(key.last_used_at_unix_secs.unwrap_or(0) as f64);
|
||||
signals
|
||||
}
|
||||
|
||||
fn apply_local_execution_pool_scheduler_with_runtime_map(
|
||||
@@ -978,7 +849,7 @@ fn apply_local_execution_pool_scheduler_with_runtime_map(
|
||||
) {
|
||||
let runtime_by_provider = runtime_by_provider
|
||||
.iter()
|
||||
.map(|(provider_id, runtime)| (provider_id.clone(), ai_pool_runtime_state(runtime)))
|
||||
.map(|(provider_id, runtime)| (provider_id.clone(), pool_runtime_state(runtime)))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let inputs = candidates
|
||||
.into_iter()
|
||||
@@ -987,20 +858,25 @@ fn apply_local_execution_pool_scheduler_with_runtime_map(
|
||||
.get(&candidate.candidate.key_id)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
AiPoolCandidateInput {
|
||||
facts: ai_pool_candidate_facts(&candidate),
|
||||
pool_config: pool_config_for_candidate(&candidate).map(ai_pool_scheduling_config),
|
||||
PoolCandidateInput {
|
||||
facts: pool_candidate_facts(&candidate),
|
||||
pool_config: pool_config_for_candidate(&candidate).map(|config| {
|
||||
pool_scheduling_config(
|
||||
config,
|
||||
candidate.transport.provider.provider_type.as_str(),
|
||||
)
|
||||
}),
|
||||
key_context,
|
||||
candidate,
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let outcome = run_ai_pool_scheduler(inputs, &runtime_by_provider, pool_sort_seed().as_str());
|
||||
let outcome = run_pool_scheduler(inputs, &runtime_by_provider, pool_sort_seed().as_str());
|
||||
|
||||
let candidates = outcome
|
||||
.candidates
|
||||
.into_iter()
|
||||
.map(|scheduled| apply_ai_pool_orchestration(scheduled.candidate, scheduled.orchestration))
|
||||
.map(|scheduled| apply_pool_orchestration(scheduled.candidate, scheduled.orchestration))
|
||||
.collect::<Vec<_>>();
|
||||
let skipped_candidates = outcome
|
||||
.skipped_candidates
|
||||
@@ -1032,16 +908,17 @@ fn pool_key_candidate_order_for_group(
|
||||
let presets = pool_config
|
||||
.scheduling_presets
|
||||
.iter()
|
||||
.map(|preset| AiPoolSchedulingPreset {
|
||||
.map(|preset| PoolSchedulingPreset {
|
||||
preset: preset.preset.clone(),
|
||||
enabled: preset.enabled,
|
||||
mode: preset.mode.clone(),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let active_presets = normalize_enabled_ai_pool_presets(
|
||||
&presets,
|
||||
group.transport.provider.provider_type.as_str(),
|
||||
);
|
||||
let active_presets = ProviderPoolService::with_builtin_adapters()
|
||||
.normalize_scheduling_presets(group.transport.provider.provider_type.as_str(), &presets)
|
||||
.into_iter()
|
||||
.map(|preset| preset.preset)
|
||||
.collect::<Vec<_>>();
|
||||
if let Some(distribution_mode) = active_presets
|
||||
.iter()
|
||||
.find(|preset| pool_distribution_mode_preset(preset.as_str()))
|
||||
@@ -1072,38 +949,43 @@ fn pool_sort_seed() -> String {
|
||||
format!("{now_ms}:{sequence}")
|
||||
}
|
||||
|
||||
fn ai_pool_candidate_facts(candidate: &EligibleLocalExecutionCandidate) -> AiPoolCandidateFacts {
|
||||
AiPoolCandidateFacts {
|
||||
fn pool_candidate_facts(candidate: &EligibleLocalExecutionCandidate) -> PoolCandidateFacts {
|
||||
PoolCandidateFacts {
|
||||
provider_id: candidate.candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.candidate.endpoint_id.clone(),
|
||||
model_id: candidate.candidate.model_id.clone(),
|
||||
selected_provider_model_name: candidate.candidate.selected_provider_model_name.clone(),
|
||||
provider_api_format: candidate.provider_api_format.clone(),
|
||||
provider_type: candidate.transport.provider.provider_type.clone(),
|
||||
key_id: candidate.candidate.key_id.clone(),
|
||||
key_internal_priority: candidate.candidate.key_internal_priority,
|
||||
}
|
||||
}
|
||||
|
||||
fn ai_pool_scheduling_config(config: AdminProviderPoolConfig) -> AiPoolSchedulingConfig {
|
||||
AiPoolSchedulingConfig {
|
||||
scheduling_presets: config
|
||||
.scheduling_presets
|
||||
.into_iter()
|
||||
.map(|preset| AiPoolSchedulingPreset {
|
||||
preset: preset.preset,
|
||||
enabled: preset.enabled,
|
||||
mode: preset.mode,
|
||||
})
|
||||
.collect(),
|
||||
fn pool_scheduling_config(
|
||||
config: AdminProviderPoolConfig,
|
||||
provider_type: &str,
|
||||
) -> PoolSchedulingConfig {
|
||||
let service = ProviderPoolService::with_builtin_adapters();
|
||||
let scheduling_presets = config
|
||||
.scheduling_presets
|
||||
.into_iter()
|
||||
.map(|preset| PoolSchedulingPreset {
|
||||
preset: preset.preset,
|
||||
enabled: preset.enabled,
|
||||
mode: preset.mode,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
PoolSchedulingConfig {
|
||||
scheduling_presets: service
|
||||
.normalize_scheduling_presets(provider_type, &scheduling_presets),
|
||||
lru_enabled: config.lru_enabled,
|
||||
skip_exhausted_accounts: config.skip_exhausted_accounts,
|
||||
cost_limit_per_key_tokens: config.cost_limit_per_key_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
fn ai_pool_runtime_state(runtime: &AdminProviderPoolRuntimeState) -> AiPoolRuntimeState {
|
||||
AiPoolRuntimeState {
|
||||
fn pool_runtime_state(runtime: &AdminProviderPoolRuntimeState) -> PoolRuntimeState {
|
||||
PoolRuntimeState {
|
||||
sticky_bound_key_id: runtime.sticky_bound_key_id.clone(),
|
||||
cooldown_reason_by_key: runtime.cooldown_reason_by_key.clone(),
|
||||
cost_window_usage_by_key: runtime.cost_window_usage_by_key.clone(),
|
||||
@@ -1112,9 +994,9 @@ fn ai_pool_runtime_state(runtime: &AdminProviderPoolRuntimeState) -> AiPoolRunti
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_ai_pool_orchestration(
|
||||
fn apply_pool_orchestration(
|
||||
mut candidate: EligibleLocalExecutionCandidate,
|
||||
orchestration: AiPoolCandidateOrchestration,
|
||||
orchestration: PoolCandidateOrchestration,
|
||||
) -> EligibleLocalExecutionCandidate {
|
||||
let scheduler_affinity_epoch = candidate.orchestration.scheduler_affinity_epoch;
|
||||
candidate.orchestration = LocalExecutionCandidateMetadata {
|
||||
@@ -1142,7 +1024,6 @@ mod tests {
|
||||
};
|
||||
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
use crate::{AppState, LocalExecutionRuntimeMissDiagnostic};
|
||||
use aether_ai_serving::{normalize_enabled_ai_pool_presets, AiPoolSchedulingPreset};
|
||||
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::candidate_selection::{
|
||||
@@ -1151,6 +1032,8 @@ mod tests {
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_pool_core::PoolSchedulingPreset;
|
||||
use aether_provider_pool::ProviderPoolService;
|
||||
use aether_provider_transport::snapshot::{
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider,
|
||||
@@ -1592,14 +1475,14 @@ mod tests {
|
||||
(
|
||||
"key-free".to_string(),
|
||||
PoolCatalogKeyContext {
|
||||
oauth_plan_type: Some("free".to_string()),
|
||||
plan_tier: Some("free".to_string()),
|
||||
..PoolCatalogKeyContext::default()
|
||||
},
|
||||
),
|
||||
(
|
||||
"key-plus".to_string(),
|
||||
PoolCatalogKeyContext {
|
||||
oauth_plan_type: Some("plus".to_string()),
|
||||
plan_tier: Some("plus".to_string()),
|
||||
..PoolCatalogKeyContext::default()
|
||||
},
|
||||
),
|
||||
@@ -1661,7 +1544,7 @@ mod tests {
|
||||
(
|
||||
"key-plus".to_string(),
|
||||
PoolCatalogKeyContext {
|
||||
oauth_plan_type: Some("plus".to_string()),
|
||||
plan_tier: Some("plus".to_string()),
|
||||
catalog_lru_score: Some(300.0),
|
||||
..PoolCatalogKeyContext::default()
|
||||
},
|
||||
@@ -1669,7 +1552,7 @@ mod tests {
|
||||
(
|
||||
"key-pro".to_string(),
|
||||
PoolCatalogKeyContext {
|
||||
oauth_plan_type: Some("pro".to_string()),
|
||||
plan_tier: Some("pro".to_string()),
|
||||
catalog_lru_score: Some(100.0),
|
||||
..PoolCatalogKeyContext::default()
|
||||
},
|
||||
@@ -1677,7 +1560,7 @@ mod tests {
|
||||
(
|
||||
"key-team".to_string(),
|
||||
PoolCatalogKeyContext {
|
||||
oauth_plan_type: Some("team".to_string()),
|
||||
plan_tier: Some("team".to_string()),
|
||||
catalog_lru_score: Some(50.0),
|
||||
..PoolCatalogKeyContext::default()
|
||||
},
|
||||
@@ -1740,21 +1623,21 @@ mod tests {
|
||||
(
|
||||
"key-plus".to_string(),
|
||||
PoolCatalogKeyContext {
|
||||
oauth_plan_type: Some("plus".to_string()),
|
||||
plan_tier: Some("plus".to_string()),
|
||||
..PoolCatalogKeyContext::default()
|
||||
},
|
||||
),
|
||||
(
|
||||
"key-pro".to_string(),
|
||||
PoolCatalogKeyContext {
|
||||
oauth_plan_type: Some("pro".to_string()),
|
||||
plan_tier: Some("pro".to_string()),
|
||||
..PoolCatalogKeyContext::default()
|
||||
},
|
||||
),
|
||||
(
|
||||
"key-team".to_string(),
|
||||
PoolCatalogKeyContext {
|
||||
oauth_plan_type: Some("team".to_string()),
|
||||
plan_tier: Some("team".to_string()),
|
||||
..PoolCatalogKeyContext::default()
|
||||
},
|
||||
),
|
||||
@@ -1822,31 +1705,35 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn normalizes_distribution_mode_before_strategy_presets() {
|
||||
let presets = normalize_enabled_ai_pool_presets(
|
||||
&[
|
||||
AiPoolSchedulingPreset {
|
||||
preset: "lru".to_string(),
|
||||
enabled: false,
|
||||
mode: None,
|
||||
},
|
||||
AiPoolSchedulingPreset {
|
||||
preset: "single_account".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
},
|
||||
AiPoolSchedulingPreset {
|
||||
preset: "cache_affinity".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
},
|
||||
AiPoolSchedulingPreset {
|
||||
preset: "priority_first".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
},
|
||||
],
|
||||
"openai",
|
||||
);
|
||||
let presets = ProviderPoolService::with_builtin_adapters()
|
||||
.normalize_scheduling_presets(
|
||||
"openai",
|
||||
&[
|
||||
PoolSchedulingPreset {
|
||||
preset: "lru".to_string(),
|
||||
enabled: false,
|
||||
mode: None,
|
||||
},
|
||||
PoolSchedulingPreset {
|
||||
preset: "single_account".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
},
|
||||
PoolSchedulingPreset {
|
||||
preset: "cache_affinity".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
},
|
||||
PoolSchedulingPreset {
|
||||
preset: "priority_first".to_string(),
|
||||
enabled: true,
|
||||
mode: None,
|
||||
},
|
||||
],
|
||||
)
|
||||
.into_iter()
|
||||
.map(|preset| preset.preset)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(presets, ["single_account", "priority_first"]);
|
||||
}
|
||||
@@ -2286,9 +2173,14 @@ mod tests {
|
||||
)),
|
||||
));
|
||||
|
||||
let context = build_pool_catalog_key_context(PlannerAppState::new(&app), &key, "codex");
|
||||
let context = build_pool_catalog_key_context(
|
||||
PlannerAppState::new(&app),
|
||||
&ProviderPoolService::with_builtin_adapters(),
|
||||
&key,
|
||||
"codex",
|
||||
);
|
||||
|
||||
assert_eq!(context.oauth_plan_type.as_deref(), Some("team"));
|
||||
assert_eq!(context.plan_tier.as_deref(), Some("team"));
|
||||
assert_eq!(context.quota_usage_ratio, Some(0.25));
|
||||
assert_eq!(context.quota_reset_seconds, Some(3600.0));
|
||||
assert_eq!(context.latency_avg_ms, Some(50.0));
|
||||
@@ -2326,7 +2218,12 @@ mod tests {
|
||||
}));
|
||||
|
||||
let app = app_state_with_catalog_key(key.clone());
|
||||
let context = build_pool_catalog_key_context(PlannerAppState::new(&app), &key, "codex");
|
||||
let context = build_pool_catalog_key_context(
|
||||
PlannerAppState::new(&app),
|
||||
&ProviderPoolService::with_builtin_adapters(),
|
||||
&key,
|
||||
"codex",
|
||||
);
|
||||
|
||||
assert!(!context.quota_exhausted);
|
||||
}
|
||||
@@ -2341,7 +2238,12 @@ mod tests {
|
||||
}));
|
||||
|
||||
let app = app_state_with_catalog_key(key.clone());
|
||||
let context = build_pool_catalog_key_context(PlannerAppState::new(&app), &key, "codex");
|
||||
let context = build_pool_catalog_key_context(
|
||||
PlannerAppState::new(&app),
|
||||
&ProviderPoolService::with_builtin_adapters(),
|
||||
&key,
|
||||
"codex",
|
||||
);
|
||||
|
||||
assert!(context.quota_exhausted);
|
||||
}
|
||||
@@ -2366,8 +2268,12 @@ mod tests {
|
||||
}));
|
||||
|
||||
let app = app_state_with_catalog_key(key.clone());
|
||||
let context =
|
||||
build_pool_catalog_key_context(PlannerAppState::new(&app), &key, "antigravity");
|
||||
let context = build_pool_catalog_key_context(
|
||||
PlannerAppState::new(&app),
|
||||
&ProviderPoolService::with_builtin_adapters(),
|
||||
&key,
|
||||
"antigravity",
|
||||
);
|
||||
|
||||
assert!(context.quota_exhausted);
|
||||
}
|
||||
@@ -2383,7 +2289,12 @@ mod tests {
|
||||
}));
|
||||
|
||||
let app = app_state_with_catalog_key(key.clone());
|
||||
let context = build_pool_catalog_key_context(PlannerAppState::new(&app), &key, "codex");
|
||||
let context = build_pool_catalog_key_context(
|
||||
PlannerAppState::new(&app),
|
||||
&ProviderPoolService::with_builtin_adapters(),
|
||||
&key,
|
||||
"codex",
|
||||
);
|
||||
|
||||
assert!(context.account_blocked);
|
||||
}
|
||||
|
||||
@@ -25,10 +25,8 @@ pub(crate) use self::provider::oauth::errors::build_internal_control_error_respo
|
||||
pub(crate) use self::provider::oauth::provisioning::{
|
||||
create_provider_oauth_catalog_key, update_existing_provider_oauth_catalog_key,
|
||||
};
|
||||
pub(crate) use self::provider::oauth::quota::antigravity::refresh_antigravity_provider_quota_locally;
|
||||
pub(crate) use self::provider::oauth::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
|
||||
pub(crate) use self::provider::oauth::quota::codex::refresh_codex_provider_quota_locally;
|
||||
pub(crate) use self::provider::oauth::quota::kiro::refresh_kiro_provider_quota_locally;
|
||||
pub(crate) use self::provider::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
|
||||
pub(crate) use self::provider::oauth::quota::shared::provider_quota_refresh_endpoint_for_provider;
|
||||
pub(crate) use self::provider::oauth::quota::shared::provider_type_supports_quota_refresh;
|
||||
pub(crate) use self::provider::oauth::runtime::{
|
||||
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
|
||||
|
||||
@@ -13,15 +13,12 @@ use axum::{
|
||||
use serde_json::json;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use super::super::oauth::quota::antigravity::refresh_antigravity_provider_quota_locally;
|
||||
use super::super::oauth::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
|
||||
use super::super::oauth::quota::codex::refresh_codex_provider_quota_locally;
|
||||
use super::super::oauth::quota::kiro::refresh_kiro_provider_quota_locally;
|
||||
use super::super::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
|
||||
use super::super::oauth::quota::shared::normalize_string_id_list;
|
||||
use super::super::oauth::quota::shared::{
|
||||
provider_quota_refresh_endpoint_for_provider, provider_quota_refresh_missing_endpoint_message,
|
||||
provider_type_supports_quota_refresh, unsupported_provider_quota_refresh_message,
|
||||
};
|
||||
use super::super::oauth::runtime::provider_oauth_maintenance_endpoint_for_provider;
|
||||
use super::super::write::provider::reconcile_admin_fixed_provider_template_endpoints;
|
||||
|
||||
fn unsupported_provider_quota_refresh_response(provider_type: &str) -> Response<Body> {
|
||||
@@ -115,7 +112,7 @@ pub(super) async fn maybe_handle(
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
|
||||
.await?;
|
||||
let mut endpoint =
|
||||
provider_oauth_maintenance_endpoint_for_provider(&normalized_provider_type, &endpoints);
|
||||
provider_quota_refresh_endpoint_for_provider(&normalized_provider_type, &endpoints, true);
|
||||
|
||||
if endpoint.is_none() && is_fixed_provider {
|
||||
if !state.has_provider_catalog_data_writer() {
|
||||
@@ -136,8 +133,11 @@ pub(super) async fn maybe_handle(
|
||||
endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
|
||||
.await?;
|
||||
endpoint =
|
||||
provider_oauth_maintenance_endpoint_for_provider(&normalized_provider_type, &endpoints);
|
||||
endpoint = provider_quota_refresh_endpoint_for_provider(
|
||||
&normalized_provider_type,
|
||||
&endpoints,
|
||||
true,
|
||||
);
|
||||
}
|
||||
|
||||
if !provider_type_supports_quota_refresh(&normalized_provider_type) {
|
||||
@@ -147,15 +147,7 @@ pub(super) async fn maybe_handle(
|
||||
}
|
||||
|
||||
let Some(endpoint) = endpoint else {
|
||||
let detail = match normalized_provider_type.as_str() {
|
||||
"codex" => "找不到有效的 openai:responses 端点",
|
||||
"antigravity" => "找不到有效的 gemini:generate_content 端点",
|
||||
"kiro" => "找不到有效的 Kiro 端点",
|
||||
"chatgpt_web" => "找不到有效的 openai:image 端点",
|
||||
"claude_code" => "找不到有效的 claude:messages 端点",
|
||||
"gemini_cli" | "vertex_ai" => "找不到有效的 gemini:generate_content 端点",
|
||||
_ => "找不到有效端点",
|
||||
};
|
||||
let detail = provider_quota_refresh_missing_endpoint_message(&normalized_provider_type);
|
||||
return Ok(Some(
|
||||
(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
@@ -231,23 +223,16 @@ pub(super) async fn maybe_handle(
|
||||
));
|
||||
}
|
||||
|
||||
let Some(payload) = (match normalized_provider_type.as_str() {
|
||||
"codex" => {
|
||||
refresh_codex_provider_quota_locally(state, &provider, &endpoint, keys, None).await?
|
||||
}
|
||||
"kiro" => {
|
||||
refresh_kiro_provider_quota_locally(state, &provider, &endpoint, keys, None).await?
|
||||
}
|
||||
"antigravity" => {
|
||||
refresh_antigravity_provider_quota_locally(state, &provider, &endpoint, keys, None)
|
||||
.await?
|
||||
}
|
||||
"chatgpt_web" => {
|
||||
refresh_chatgpt_web_provider_quota_locally(state, &provider, &endpoint, keys, None)
|
||||
.await?
|
||||
}
|
||||
_ => None,
|
||||
}) else {
|
||||
let Some(payload) = refresh_provider_pool_quota_locally(
|
||||
state,
|
||||
&provider,
|
||||
&endpoint,
|
||||
&normalized_provider_type,
|
||||
keys,
|
||||
None,
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(Some(Json(payload).into_response()))
|
||||
|
||||
@@ -5,13 +5,12 @@ use crate::handlers::admin::provider::shared::payloads::{
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig};
|
||||
use crate::provider_transport::kiro::{build_kiro_request_auth_from_config, KiroRequestAuth};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use serde_json::{json, Value};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use aether_oauth::core::OAuthError;
|
||||
use aether_oauth::provider::providers::KiroProviderOAuthAdapter;
|
||||
use aether_oauth::provider::ProviderOAuthTransportContext;
|
||||
use serde_json::Value;
|
||||
use url::form_urlencoded;
|
||||
|
||||
const KIRO_IDC_AMZ_USER_AGENT: &str =
|
||||
"aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE";
|
||||
|
||||
pub(super) fn admin_provider_oauth_kiro_refresh_base_url_override(
|
||||
state: &AdminAppState<'_>,
|
||||
override_key: &str,
|
||||
@@ -21,29 +20,6 @@ pub(super) fn admin_provider_oauth_kiro_refresh_base_url_override(
|
||||
(!normalized.is_empty()).then(|| normalized.to_string())
|
||||
}
|
||||
|
||||
fn admin_provider_oauth_kiro_build_refresh_url(
|
||||
auth_config: &AdminKiroAuthConfig,
|
||||
override_base_url: Option<&str>,
|
||||
path: &str,
|
||||
default_host: impl FnOnce(&str) -> String,
|
||||
) -> String {
|
||||
if let Some(base_url) = override_base_url
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return format!("{}/{}", base_url.trim_end_matches('/'), path);
|
||||
}
|
||||
let region = auth_config.effective_auth_region();
|
||||
default_host(region)
|
||||
}
|
||||
|
||||
fn admin_provider_oauth_kiro_effective_host(url: &str, fallback_host: String) -> String {
|
||||
reqwest::Url::parse(url)
|
||||
.ok()
|
||||
.and_then(|value| value.host_str().map(ToOwned::to_owned))
|
||||
.unwrap_or(fallback_host)
|
||||
}
|
||||
|
||||
fn admin_provider_oauth_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> String {
|
||||
if machine_id.trim().is_empty() {
|
||||
format!("KiroIDE-{kiro_version}")
|
||||
@@ -52,41 +28,49 @@ fn admin_provider_oauth_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> St
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_provider_oauth_kiro_refresh_expires_at(payload: &Value) -> u64 {
|
||||
let expires_in = payload
|
||||
.get("expiresIn")
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_str()?.parse::<u64>().ok())
|
||||
})
|
||||
.unwrap_or(3600);
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|value| value.as_secs())
|
||||
.unwrap_or_default()
|
||||
.saturating_add(expires_in)
|
||||
fn admin_provider_oauth_kiro_refresh_context(
|
||||
proxy: Option<ProxySnapshot>,
|
||||
) -> ProviderOAuthTransportContext {
|
||||
ProviderOAuthTransportContext {
|
||||
provider_id: String::new(),
|
||||
provider_type: "kiro".to_string(),
|
||||
endpoint_id: None,
|
||||
key_id: None,
|
||||
auth_type: Some("oauth".to_string()),
|
||||
decrypted_api_key: None,
|
||||
decrypted_auth_config: None,
|
||||
provider_config: None,
|
||||
endpoint_config: None,
|
||||
key_config: None,
|
||||
network: aether_oauth::network::OAuthNetworkContext::provider_operation(proxy),
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_provider_oauth_kiro_refresh_response_json(
|
||||
body_text: &str,
|
||||
json_body: Option<Value>,
|
||||
) -> Result<Value, String> {
|
||||
json_body
|
||||
.or_else(|| serde_json::from_str::<Value>(body_text).ok())
|
||||
.ok_or_else(|| "refresh 接口返回了非 JSON 响应".to_string())
|
||||
}
|
||||
|
||||
fn admin_provider_oauth_kiro_refresh_error_detail(
|
||||
status: http::StatusCode,
|
||||
body_text: &str,
|
||||
fn admin_provider_oauth_kiro_refresh_error(
|
||||
auth_config: &AdminKiroAuthConfig,
|
||||
error: OAuthError,
|
||||
) -> String {
|
||||
let detail = body_text.trim();
|
||||
if detail.is_empty() {
|
||||
format!("HTTP {}", status.as_u16())
|
||||
let prefix = if auth_config.is_idc_auth() {
|
||||
"IDC refresh"
|
||||
} else {
|
||||
detail.to_string()
|
||||
"social refresh"
|
||||
};
|
||||
match error {
|
||||
OAuthError::HttpStatus {
|
||||
status_code,
|
||||
body_excerpt,
|
||||
} => {
|
||||
let detail = body_excerpt.trim();
|
||||
if detail.is_empty() {
|
||||
format!("{prefix} 失败: HTTP {status_code}")
|
||||
} else {
|
||||
format!("{prefix} 失败: {detail}")
|
||||
}
|
||||
}
|
||||
OAuthError::Transport(message) => format!("{prefix} 请求失败: {message}"),
|
||||
OAuthError::InvalidRequest(message) => format!("{prefix} 参数无效: {message}"),
|
||||
OAuthError::InvalidResponse(message) => format!("{prefix} 返回无效响应: {message}"),
|
||||
error => format!("{prefix} 失败: {error}"),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -97,216 +81,25 @@ pub(super) async fn refresh_admin_provider_oauth_kiro_auth_config(
|
||||
social_refresh_base_url: Option<&str>,
|
||||
idc_refresh_base_url: Option<&str>,
|
||||
) -> Result<AdminKiroAuthConfig, String> {
|
||||
if auth_config.is_idc_auth() {
|
||||
let fallback_host = format!("oidc.{}.amazonaws.com", auth_config.effective_auth_region());
|
||||
let url = admin_provider_oauth_kiro_build_refresh_url(
|
||||
let adapter = KiroProviderOAuthAdapter::default().with_refresh_base_urls(
|
||||
social_refresh_base_url
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
idc_refresh_base_url
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
);
|
||||
let ctx = admin_provider_oauth_kiro_refresh_context(proxy);
|
||||
adapter
|
||||
.refresh_auth_config(
|
||||
&crate::oauth::GatewayOAuthHttpExecutor::new(*state),
|
||||
&ctx,
|
||||
auth_config,
|
||||
idc_refresh_base_url,
|
||||
"token",
|
||||
|region| format!("https://oidc.{region}.amazonaws.com/token"),
|
||||
);
|
||||
let host = admin_provider_oauth_kiro_effective_host(&url, fallback_host);
|
||||
let headers = reqwest::header::HeaderMap::from_iter([
|
||||
(
|
||||
reqwest::header::CONTENT_TYPE,
|
||||
reqwest::header::HeaderValue::from_static("application/json"),
|
||||
),
|
||||
(
|
||||
reqwest::header::HOST,
|
||||
reqwest::header::HeaderValue::from_str(&host)
|
||||
.map_err(|_| "IDC host 无效".to_string())?,
|
||||
),
|
||||
(
|
||||
reqwest::header::HeaderName::from_static("x-amz-user-agent"),
|
||||
reqwest::header::HeaderValue::from_static(KIRO_IDC_AMZ_USER_AGENT),
|
||||
),
|
||||
(
|
||||
reqwest::header::USER_AGENT,
|
||||
reqwest::header::HeaderValue::from_static("node"),
|
||||
),
|
||||
(
|
||||
reqwest::header::ACCEPT,
|
||||
reqwest::header::HeaderValue::from_static("*/*"),
|
||||
),
|
||||
]);
|
||||
let response = state
|
||||
.execute_admin_provider_oauth_http_request(
|
||||
"kiro_batch_refresh:idc",
|
||||
reqwest::Method::POST,
|
||||
&url,
|
||||
&headers,
|
||||
Some("application/json"),
|
||||
Some(json!({
|
||||
"clientId": auth_config
|
||||
.client_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default(),
|
||||
"clientSecret": auth_config
|
||||
.client_secret
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default(),
|
||||
"refreshToken": auth_config
|
||||
.refresh_token
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default(),
|
||||
"grantType": "refresh_token",
|
||||
})),
|
||||
None,
|
||||
proxy.clone(),
|
||||
)
|
||||
.await
|
||||
.map_err(|err| format!("IDC refresh 请求失败: {err}"))?;
|
||||
if !response.status.is_success() {
|
||||
return Err(format!(
|
||||
"IDC refresh 失败: {}",
|
||||
admin_provider_oauth_kiro_refresh_error_detail(
|
||||
response.status,
|
||||
&response.body_text
|
||||
)
|
||||
));
|
||||
}
|
||||
let payload = admin_provider_oauth_kiro_refresh_response_json(
|
||||
&response.body_text,
|
||||
response.json_body,
|
||||
)?;
|
||||
let access_token = payload
|
||||
.get("accessToken")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| "IDC refresh 返回了空 accessToken".to_string())?;
|
||||
|
||||
let mut refreshed = auth_config.clone();
|
||||
refreshed.access_token = Some(access_token.to_string());
|
||||
refreshed.expires_at = Some(admin_provider_oauth_kiro_refresh_expires_at(&payload));
|
||||
if refreshed
|
||||
.machine_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_none_or(|value| value.is_empty())
|
||||
{
|
||||
refreshed.machine_id =
|
||||
crate::provider_transport::kiro::generate_machine_id(auth_config, None);
|
||||
}
|
||||
if let Some(refresh_token) = payload
|
||||
.get("refreshToken")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
refreshed.refresh_token = Some(refresh_token.to_string());
|
||||
}
|
||||
return Ok(refreshed);
|
||||
}
|
||||
|
||||
let machine_id = crate::provider_transport::kiro::generate_machine_id(auth_config, None)
|
||||
.ok_or_else(|| "缺少 machine_id 种子,无法刷新 social token".to_string())?;
|
||||
let fallback_host = format!(
|
||||
"prod.{}.auth.desktop.kiro.dev",
|
||||
auth_config.effective_auth_region()
|
||||
);
|
||||
let url = admin_provider_oauth_kiro_build_refresh_url(
|
||||
auth_config,
|
||||
social_refresh_base_url,
|
||||
"refreshToken",
|
||||
|region| format!("https://prod.{region}.auth.desktop.kiro.dev/refreshToken"),
|
||||
);
|
||||
let host = admin_provider_oauth_kiro_effective_host(&url, fallback_host);
|
||||
let user_agent =
|
||||
admin_provider_oauth_kiro_ide_tag(auth_config.effective_kiro_version(), &machine_id);
|
||||
let headers = reqwest::header::HeaderMap::from_iter([
|
||||
(
|
||||
reqwest::header::USER_AGENT,
|
||||
reqwest::header::HeaderValue::from_str(&user_agent)
|
||||
.map_err(|_| "Kiro User-Agent 无效".to_string())?,
|
||||
),
|
||||
(
|
||||
reqwest::header::HOST,
|
||||
reqwest::header::HeaderValue::from_str(&host)
|
||||
.map_err(|_| "Kiro host 无效".to_string())?,
|
||||
),
|
||||
(
|
||||
reqwest::header::ACCEPT,
|
||||
reqwest::header::HeaderValue::from_static("application/json, text/plain, */*"),
|
||||
),
|
||||
(
|
||||
reqwest::header::CONTENT_TYPE,
|
||||
reqwest::header::HeaderValue::from_static("application/json"),
|
||||
),
|
||||
(
|
||||
reqwest::header::CONNECTION,
|
||||
reqwest::header::HeaderValue::from_static("close"),
|
||||
),
|
||||
(
|
||||
reqwest::header::ACCEPT_ENCODING,
|
||||
reqwest::header::HeaderValue::from_static("gzip, compress, deflate, br"),
|
||||
),
|
||||
]);
|
||||
let response = state
|
||||
.execute_admin_provider_oauth_http_request(
|
||||
"kiro_batch_refresh:social",
|
||||
reqwest::Method::POST,
|
||||
&url,
|
||||
&headers,
|
||||
Some("application/json"),
|
||||
Some(json!({
|
||||
"refreshToken": auth_config
|
||||
.refresh_token
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default(),
|
||||
})),
|
||||
None,
|
||||
proxy,
|
||||
)
|
||||
.await
|
||||
.map_err(|err| format!("social refresh 请求失败: {err}"))?;
|
||||
if !response.status.is_success() {
|
||||
return Err(format!(
|
||||
"social refresh 失败: {}",
|
||||
admin_provider_oauth_kiro_refresh_error_detail(response.status, &response.body_text)
|
||||
));
|
||||
}
|
||||
let payload =
|
||||
admin_provider_oauth_kiro_refresh_response_json(&response.body_text, response.json_body)?;
|
||||
let access_token = payload
|
||||
.get("accessToken")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| "social refresh 返回了空 accessToken".to_string())?;
|
||||
|
||||
let mut refreshed = auth_config.clone();
|
||||
refreshed.access_token = Some(access_token.to_string());
|
||||
refreshed.expires_at = Some(admin_provider_oauth_kiro_refresh_expires_at(&payload));
|
||||
if refreshed
|
||||
.machine_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_none_or(|value| value.is_empty())
|
||||
{
|
||||
refreshed.machine_id = Some(machine_id);
|
||||
}
|
||||
if let Some(refresh_token) = payload
|
||||
.get("refreshToken")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
refreshed.refresh_token = Some(refresh_token.to_string());
|
||||
}
|
||||
if let Some(profile_arn) = payload
|
||||
.get("profileArn")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
refreshed.profile_arn = Some(profile_arn.to_string());
|
||||
}
|
||||
Ok(refreshed)
|
||||
.map_err(|error| admin_provider_oauth_kiro_refresh_error(auth_config, error))
|
||||
}
|
||||
|
||||
fn build_kiro_usage_url(auth: &KiroRequestAuth) -> String {
|
||||
|
||||
@@ -1,17 +1,17 @@
|
||||
use super::shared::{
|
||||
build_quota_snapshot_payload, coerce_json_f64, coerce_json_string,
|
||||
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
|
||||
build_provider_quota_execution_plan, build_quota_snapshot_payload, coerce_json_f64,
|
||||
coerce_json_string, default_provider_quota_execution_timeouts, execute_provider_quota_plan,
|
||||
extract_execution_error_message, persist_provider_quota_refresh_state,
|
||||
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
|
||||
};
|
||||
use crate::handlers::admin::provider::shared::payloads::ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||
use crate::GatewayError;
|
||||
use aether_admin::provider::quota::parse_antigravity_usage_response;
|
||||
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_pool::build_antigravity_pool_quota_request;
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
@@ -21,18 +21,9 @@ async fn execute_antigravity_quota_plan(
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
authorization: (String, String),
|
||||
project_id: &str,
|
||||
mut identity_headers: BTreeMap<String, String>,
|
||||
identity_headers: BTreeMap<String, String>,
|
||||
proxy_override: Option<&ProxySnapshot>,
|
||||
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
|
||||
let mut headers = std::mem::take(&mut identity_headers);
|
||||
headers.insert("authorization".to_string(), authorization.1);
|
||||
headers.insert("content-type".to_string(), "application/json".to_string());
|
||||
headers.insert("accept".to_string(), "application/json".to_string());
|
||||
headers
|
||||
.entry("user-agent".to_string())
|
||||
.or_insert_with(|| "antigravity".to_string());
|
||||
|
||||
let body = json!({ "project": project_id });
|
||||
let proxy = match proxy_override {
|
||||
Some(proxy) => Some(proxy.clone()),
|
||||
None => {
|
||||
@@ -46,35 +37,20 @@ async fn execute_antigravity_quota_plan(
|
||||
.or(Some(default_provider_quota_execution_timeouts(
|
||||
proxy.as_ref(),
|
||||
)));
|
||||
let plan = ExecutionPlan {
|
||||
request_id: format!("antigravity-quota:{}", transport.key.id),
|
||||
candidate_id: None,
|
||||
provider_name: Some("antigravity".to_string()),
|
||||
provider_id: transport.provider.id.clone(),
|
||||
endpoint_id: transport.endpoint.id.clone(),
|
||||
key_id: transport.key.id.clone(),
|
||||
method: "POST".to_string(),
|
||||
url: format!(
|
||||
"{}{}",
|
||||
transport.endpoint.base_url.trim_end_matches('/'),
|
||||
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH
|
||||
),
|
||||
headers,
|
||||
content_type: Some("application/json".to_string()),
|
||||
content_encoding: None,
|
||||
body: RequestBody {
|
||||
json_body: Some(body),
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: false,
|
||||
client_api_format: "gemini:generate_content".to_string(),
|
||||
provider_api_format: "antigravity:fetch_available_models".to_string(),
|
||||
model_name: Some("fetchAvailableModels".to_string()),
|
||||
let spec = build_antigravity_pool_quota_request(
|
||||
&transport.key.id,
|
||||
&transport.endpoint.base_url,
|
||||
authorization,
|
||||
project_id,
|
||||
identity_headers,
|
||||
);
|
||||
let plan = build_provider_quota_execution_plan(
|
||||
transport,
|
||||
spec,
|
||||
proxy,
|
||||
transport_profile: state.resolve_transport_profile(transport),
|
||||
state.resolve_transport_profile(transport),
|
||||
timeouts,
|
||||
};
|
||||
);
|
||||
|
||||
execute_provider_quota_plan(state, transport, plan, "antigravity").await
|
||||
}
|
||||
|
||||
@@ -10,94 +10,18 @@ use crate::handlers::admin::provider::shared::payloads::{
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||
use crate::GatewayError;
|
||||
use aether_admin::provider::quota::parse_chatgpt_web_conversation_init_response;
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ProxySnapshot, RequestBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
|
||||
};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_pool::{
|
||||
build_chatgpt_web_pool_quota_request, enrich_chatgpt_web_quota_metadata,
|
||||
normalize_chatgpt_web_image_quota_limit,
|
||||
};
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
const CHATGPT_WEB_DEFAULT_BASE_URL: &str = "https://chatgpt.com";
|
||||
const CHATGPT_WEB_CONVERSATION_INIT_PATH: &str = "/backend-api/conversation/init";
|
||||
const CHATGPT_WEB_USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36 Edg/143.0.0.0";
|
||||
const CHATGPT_WEB_CLIENT_VERSION: &str = "prod-be885abbfcfe7b1f511e88b3003d9ee44757fbad";
|
||||
const CHATGPT_WEB_BUILD_NUMBER: &str = "5955942";
|
||||
const CHATGPT_WEB_SEC_CH_UA: &str =
|
||||
r#""Microsoft Edge";v="143", "Chromium";v="143", "Not A(Brand";v="24""#;
|
||||
const PLACEHOLDER_API_KEY: &str = "__placeholder__";
|
||||
const CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT: f64 = 25.0;
|
||||
|
||||
fn chatgpt_web_base_url(endpoint: &StoredProviderCatalogEndpoint) -> String {
|
||||
let base_url = endpoint.base_url.trim().trim_end_matches('/');
|
||||
if base_url.is_empty() {
|
||||
CHATGPT_WEB_DEFAULT_BASE_URL.to_string()
|
||||
} else {
|
||||
base_url.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn build_chatgpt_web_quota_headers(
|
||||
authorization: (String, String),
|
||||
base_url: &str,
|
||||
) -> BTreeMap<String, String> {
|
||||
let device_id = uuid::Uuid::new_v4().to_string();
|
||||
let session_id = uuid::Uuid::new_v4().to_string();
|
||||
let mut headers = BTreeMap::from([
|
||||
("accept".to_string(), "application/json".to_string()),
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
("user-agent".to_string(), CHATGPT_WEB_USER_AGENT.to_string()),
|
||||
("origin".to_string(), base_url.to_string()),
|
||||
("referer".to_string(), format!("{base_url}/")),
|
||||
(
|
||||
"accept-language".to_string(),
|
||||
"zh-CN,zh;q=0.9,en;q=0.8,en-US;q=0.7".to_string(),
|
||||
),
|
||||
("cache-control".to_string(), "no-cache".to_string()),
|
||||
("pragma".to_string(), "no-cache".to_string()),
|
||||
("priority".to_string(), "u=1, i".to_string()),
|
||||
("sec-ch-ua".to_string(), CHATGPT_WEB_SEC_CH_UA.to_string()),
|
||||
("sec-ch-ua-arch".to_string(), r#""x86""#.to_string()),
|
||||
("sec-ch-ua-bitness".to_string(), r#""64""#.to_string()),
|
||||
("sec-ch-ua-mobile".to_string(), "?0".to_string()),
|
||||
("sec-ch-ua-model".to_string(), r#""""#.to_string()),
|
||||
("sec-ch-ua-platform".to_string(), r#""Windows""#.to_string()),
|
||||
(
|
||||
"sec-ch-ua-platform-version".to_string(),
|
||||
r#""19.0.0""#.to_string(),
|
||||
),
|
||||
("sec-fetch-dest".to_string(), "empty".to_string()),
|
||||
("sec-fetch-mode".to_string(), "cors".to_string()),
|
||||
("sec-fetch-site".to_string(), "same-origin".to_string()),
|
||||
("oai-device-id".to_string(), device_id),
|
||||
("oai-session-id".to_string(), session_id),
|
||||
("oai-language".to_string(), "zh-CN".to_string()),
|
||||
(
|
||||
"oai-client-version".to_string(),
|
||||
CHATGPT_WEB_CLIENT_VERSION.to_string(),
|
||||
),
|
||||
(
|
||||
"oai-client-build-number".to_string(),
|
||||
CHATGPT_WEB_BUILD_NUMBER.to_string(),
|
||||
),
|
||||
(
|
||||
"x-openai-target-path".to_string(),
|
||||
CHATGPT_WEB_CONVERSATION_INIT_PATH.to_string(),
|
||||
),
|
||||
(
|
||||
"x-openai-target-route".to_string(),
|
||||
CHATGPT_WEB_CONVERSATION_INIT_PATH.to_string(),
|
||||
),
|
||||
(
|
||||
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER.to_string(),
|
||||
"true".to_string(),
|
||||
),
|
||||
]);
|
||||
headers.insert(authorization.0.to_ascii_lowercase(), authorization.1);
|
||||
headers
|
||||
}
|
||||
|
||||
fn chatgpt_web_auth_config(
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
@@ -111,133 +35,6 @@ fn chatgpt_web_auth_config(
|
||||
.and_then(|value| serde_json::from_str::<serde_json::Value>(value).ok())
|
||||
}
|
||||
|
||||
fn chatgpt_web_auth_config_string(
|
||||
auth_config: Option<&serde_json::Value>,
|
||||
fields: &[&str],
|
||||
) -> Option<String> {
|
||||
let object = auth_config.and_then(serde_json::Value::as_object)?;
|
||||
fields.iter().find_map(|field| {
|
||||
object
|
||||
.get(*field)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
})
|
||||
}
|
||||
|
||||
fn enrich_chatgpt_web_quota_metadata(
|
||||
metadata: &mut serde_json::Value,
|
||||
auth_config: Option<&serde_json::Value>,
|
||||
) {
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for (target, fields) in [
|
||||
("plan_type", &["plan_type", "tier", "plan"][..]),
|
||||
("email", &["email"][..]),
|
||||
("account_id", &["account_id", "accountId"][..]),
|
||||
("account_user_id", &["account_user_id", "accountUserId"][..]),
|
||||
("user_id", &["user_id", "userId"][..]),
|
||||
] {
|
||||
if object.contains_key(target) {
|
||||
continue;
|
||||
}
|
||||
if let Some(value) = chatgpt_web_auth_config_string(auth_config, fields) {
|
||||
object.insert(target.to_string(), json!(value));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn chatgpt_web_json_number(value: Option<&serde_json::Value>) -> Option<f64> {
|
||||
let value = value?;
|
||||
if let Some(number) = value.as_f64() {
|
||||
return number.is_finite().then_some(number);
|
||||
}
|
||||
value
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(|value| value.parse::<f64>().ok())
|
||||
.filter(|value| value.is_finite())
|
||||
}
|
||||
|
||||
fn chatgpt_web_json_string(value: Option<&serde_json::Value>) -> Option<&str> {
|
||||
value
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn existing_chatgpt_web_image_quota_limit(
|
||||
upstream_metadata: Option<&serde_json::Value>,
|
||||
) -> Option<f64> {
|
||||
upstream_metadata
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|metadata| metadata.get("chatgpt_web"))
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|bucket| chatgpt_web_json_number(bucket.get("image_quota_total")))
|
||||
.filter(|value| *value > 0.0)
|
||||
}
|
||||
|
||||
fn infer_chatgpt_web_image_quota_limit(
|
||||
plan_type: Option<&str>,
|
||||
remaining: Option<f64>,
|
||||
existing_limit: Option<f64>,
|
||||
) -> Option<f64> {
|
||||
let normalized_plan = plan_type.unwrap_or_default().trim().to_ascii_lowercase();
|
||||
if normalized_plan == "free" {
|
||||
return Some(CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT);
|
||||
}
|
||||
|
||||
if let Some(existing_limit) = existing_limit.filter(|value| *value > 0.0) {
|
||||
return Some(existing_limit);
|
||||
}
|
||||
|
||||
remaining.filter(|value| *value > 0.0)
|
||||
}
|
||||
|
||||
fn normalize_chatgpt_web_image_quota_limit(
|
||||
metadata: &mut serde_json::Value,
|
||||
upstream_metadata: Option<&serde_json::Value>,
|
||||
) {
|
||||
let existing_limit = existing_chatgpt_web_image_quota_limit(upstream_metadata);
|
||||
let Some(object) = metadata.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
|
||||
let remaining = chatgpt_web_json_number(object.get("image_quota_remaining"));
|
||||
let explicit_limit =
|
||||
chatgpt_web_json_number(object.get("image_quota_total")).filter(|value| *value > 0.0);
|
||||
let plan_type = chatgpt_web_json_string(object.get("plan_type"));
|
||||
let is_free_plan = plan_type.is_some_and(|value| value.trim().eq_ignore_ascii_case("free"));
|
||||
let limit = if is_free_plan {
|
||||
Some(CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT)
|
||||
} else {
|
||||
explicit_limit
|
||||
.or_else(|| infer_chatgpt_web_image_quota_limit(plan_type, remaining, existing_limit))
|
||||
};
|
||||
|
||||
if let Some(limit) = limit {
|
||||
object.insert("image_quota_total".to_string(), json!(limit));
|
||||
|
||||
if !object.contains_key("image_quota_used") {
|
||||
if let Some(remaining) = remaining {
|
||||
object.insert(
|
||||
"image_quota_used".to_string(),
|
||||
json!((limit - remaining).max(0.0)),
|
||||
);
|
||||
} else if object
|
||||
.get("image_quota_blocked")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
== Some(true)
|
||||
{
|
||||
object.insert("image_quota_used".to_string(), json!(limit));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn resolve_chatgpt_web_quota_auth(
|
||||
state: &AdminAppState<'_>,
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
@@ -262,7 +59,6 @@ async fn execute_chatgpt_web_quota_plan(
|
||||
authorization: (String, String),
|
||||
proxy_override: Option<&ProxySnapshot>,
|
||||
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
|
||||
let base_url = chatgpt_web_base_url(endpoint);
|
||||
let proxy = match proxy_override {
|
||||
Some(proxy) => Some(proxy.clone()),
|
||||
None => {
|
||||
@@ -276,33 +72,15 @@ async fn execute_chatgpt_web_quota_plan(
|
||||
.or(Some(default_provider_quota_execution_timeouts(
|
||||
proxy.as_ref(),
|
||||
)));
|
||||
let plan = ExecutionPlan {
|
||||
request_id: format!("chatgpt-web-quota:{}", transport.key.id),
|
||||
candidate_id: None,
|
||||
provider_name: Some("chatgpt_web".to_string()),
|
||||
provider_id: transport.provider.id.clone(),
|
||||
endpoint_id: transport.endpoint.id.clone(),
|
||||
key_id: transport.key.id.clone(),
|
||||
method: "POST".to_string(),
|
||||
url: format!("{base_url}{CHATGPT_WEB_CONVERSATION_INIT_PATH}"),
|
||||
headers: build_chatgpt_web_quota_headers(authorization, base_url.as_str()),
|
||||
content_type: Some("application/json".to_string()),
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({
|
||||
"gizmo_id": serde_json::Value::Null,
|
||||
"requested_default_model": serde_json::Value::Null,
|
||||
"conversation_id": serde_json::Value::Null,
|
||||
"timezone_offset_min": -480,
|
||||
"system_hints": ["picture_v2"],
|
||||
})),
|
||||
stream: false,
|
||||
client_api_format: "openai:image".to_string(),
|
||||
provider_api_format: "chatgpt_web:conversation_init".to_string(),
|
||||
model_name: Some("chatgpt-web-conversation-init".to_string()),
|
||||
let spec =
|
||||
build_chatgpt_web_pool_quota_request(&transport.key.id, &endpoint.base_url, authorization);
|
||||
let plan = super::shared::build_provider_quota_execution_plan(
|
||||
transport,
|
||||
spec,
|
||||
proxy,
|
||||
transport_profile: state.resolve_transport_profile(transport),
|
||||
state.resolve_transport_profile(transport),
|
||||
timeouts,
|
||||
};
|
||||
);
|
||||
|
||||
execute_provider_quota_plan(state, transport, plan, "chatgpt_web").await
|
||||
}
|
||||
|
||||
@@ -11,7 +11,7 @@ use self::parse::{
|
||||
build_codex_quota_exhausted_fallback_metadata, parse_codex_usage_headers,
|
||||
parse_codex_wham_usage_response,
|
||||
};
|
||||
use self::plan::{build_codex_refresh_headers, execute_codex_quota_plan};
|
||||
use self::plan::{build_codex_quota_request_spec, execute_codex_quota_plan};
|
||||
use super::shared::{
|
||||
build_quota_snapshot_payload, extract_execution_error_message,
|
||||
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
|
||||
@@ -82,8 +82,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
None
|
||||
};
|
||||
|
||||
let headers = match build_codex_refresh_headers(&transport, resolved_oauth_auth) {
|
||||
Ok(headers) => headers,
|
||||
let request_spec = match build_codex_quota_request_spec(&transport, resolved_oauth_auth) {
|
||||
Ok(request_spec) => request_spec,
|
||||
Err(message) => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
@@ -96,23 +96,27 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
}
|
||||
};
|
||||
|
||||
let result =
|
||||
match execute_codex_quota_plan(state, &transport, headers, proxy_override.as_ref())
|
||||
.await?
|
||||
{
|
||||
ProviderQuotaExecutionOutcome::Response(result) => result,
|
||||
ProviderQuotaExecutionOutcome::Failure(detail) => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": format!("wham/usage 请求执行失败: {detail}"),
|
||||
"status_code": 502,
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let result = match execute_codex_quota_plan(
|
||||
state,
|
||||
&transport,
|
||||
request_spec,
|
||||
proxy_override.as_ref(),
|
||||
)
|
||||
.await?
|
||||
{
|
||||
ProviderQuotaExecutionOutcome::Response(result) => result,
|
||||
ProviderQuotaExecutionOutcome::Failure(detail) => {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
"key_name": key.name,
|
||||
"status": "error",
|
||||
"message": format!("wham/usage 请求执行失败: {detail}"),
|
||||
"status_code": 502,
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
};
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
|
||||
@@ -1,65 +1,33 @@
|
||||
use super::super::shared::{
|
||||
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
|
||||
ProviderQuotaExecutionOutcome,
|
||||
build_provider_quota_execution_plan, default_provider_quota_execution_timeouts,
|
||||
execute_provider_quota_plan, ProviderQuotaExecutionOutcome,
|
||||
};
|
||||
use super::parse::normalize_codex_plan_type;
|
||||
use crate::handlers::admin::provider::shared::payloads::CODEX_WHAM_USAGE_URL;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||
use crate::GatewayError;
|
||||
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
|
||||
use std::collections::BTreeMap;
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_provider_pool::{build_codex_pool_quota_request, ProviderPoolQuotaRequestSpec};
|
||||
|
||||
pub(super) fn build_codex_refresh_headers(
|
||||
pub(super) fn build_codex_quota_request_spec(
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
resolved_oauth_auth: Option<(String, String)>,
|
||||
) -> Result<BTreeMap<String, String>, String> {
|
||||
let mut headers = BTreeMap::new();
|
||||
headers.insert("accept".to_string(), "application/json".to_string());
|
||||
|
||||
if let Some((name, value)) = resolved_oauth_auth {
|
||||
headers.insert(name.to_ascii_lowercase(), value);
|
||||
} else {
|
||||
let decrypted_key = transport.key.decrypted_api_key.trim();
|
||||
if decrypted_key.is_empty() || decrypted_key == "__placeholder__" {
|
||||
return Err("缺少 OAuth 认证信息,请先授权/刷新 Token".to_string());
|
||||
}
|
||||
headers.insert(
|
||||
"authorization".to_string(),
|
||||
format!("Bearer {decrypted_key}"),
|
||||
);
|
||||
}
|
||||
|
||||
) -> Result<ProviderPoolQuotaRequestSpec, String> {
|
||||
let auth_config = transport
|
||||
.key
|
||||
.decrypted_auth_config
|
||||
.as_deref()
|
||||
.and_then(|raw| serde_json::from_str::<serde_json::Value>(raw).ok());
|
||||
let oauth_plan_type = normalize_codex_plan_type(
|
||||
auth_config
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("plan_type"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
);
|
||||
let oauth_account_id = auth_config
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("account_id"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
if oauth_account_id.is_some() && oauth_plan_type.as_deref() != Some("free") {
|
||||
headers.insert(
|
||||
"chatgpt-account-id".to_string(),
|
||||
oauth_account_id.unwrap_or_default().to_string(),
|
||||
);
|
||||
}
|
||||
|
||||
Ok(headers)
|
||||
build_codex_pool_quota_request(
|
||||
&transport.key.id,
|
||||
resolved_oauth_auth,
|
||||
Some(transport.key.decrypted_api_key.as_str()),
|
||||
auth_config.as_ref(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) async fn execute_codex_quota_plan(
|
||||
state: &AdminAppState<'_>,
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
headers: BTreeMap<String, String>,
|
||||
spec: ProviderPoolQuotaRequestSpec,
|
||||
proxy_override: Option<&ProxySnapshot>,
|
||||
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
|
||||
let proxy = match proxy_override {
|
||||
@@ -75,30 +43,12 @@ pub(super) async fn execute_codex_quota_plan(
|
||||
.or(Some(default_provider_quota_execution_timeouts(
|
||||
proxy.as_ref(),
|
||||
)));
|
||||
let plan = ExecutionPlan {
|
||||
request_id: format!("codex-quota:{}", transport.key.id),
|
||||
candidate_id: None,
|
||||
provider_name: Some("codex".to_string()),
|
||||
provider_id: transport.provider.id.clone(),
|
||||
endpoint_id: transport.endpoint.id.clone(),
|
||||
key_id: transport.key.id.clone(),
|
||||
method: "GET".to_string(),
|
||||
url: CODEX_WHAM_USAGE_URL.to_string(),
|
||||
headers,
|
||||
content_type: None,
|
||||
content_encoding: None,
|
||||
body: RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
},
|
||||
stream: false,
|
||||
client_api_format: "openai:responses".to_string(),
|
||||
provider_api_format: "openai:responses".to_string(),
|
||||
model_name: Some("codex-wham-usage".to_string()),
|
||||
let plan = build_provider_quota_execution_plan(
|
||||
transport,
|
||||
spec,
|
||||
proxy,
|
||||
transport_profile: state.resolve_transport_profile(transport),
|
||||
state.resolve_transport_profile(transport),
|
||||
timeouts,
|
||||
};
|
||||
);
|
||||
execute_provider_quota_plan(state, transport, plan, "codex").await
|
||||
}
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
|
||||
use super::antigravity::refresh_antigravity_provider_quota_locally;
|
||||
use super::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
|
||||
use super::codex::refresh_codex_provider_quota_locally;
|
||||
use super::kiro::refresh_kiro_provider_quota_locally;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::GatewayError;
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
|
||||
type ProviderQuotaRefreshFuture<'a> =
|
||||
Pin<Box<dyn Future<Output = Result<Option<serde_json::Value>, GatewayError>> + Send + 'a>>;
|
||||
|
||||
type ProviderQuotaRefreshHandler = for<'a> fn(
|
||||
&'a AdminAppState<'a>,
|
||||
&'a StoredProviderCatalogProvider,
|
||||
&'a StoredProviderCatalogEndpoint,
|
||||
Vec<StoredProviderCatalogKey>,
|
||||
Option<ProxySnapshot>,
|
||||
) -> ProviderQuotaRefreshFuture<'a>;
|
||||
|
||||
const PROVIDER_QUOTA_REFRESH_HANDLERS: &[(&str, ProviderQuotaRefreshHandler)] = &[
|
||||
(
|
||||
"antigravity",
|
||||
refresh_antigravity_provider_quota_locally_boxed,
|
||||
),
|
||||
(
|
||||
"chatgpt_web",
|
||||
refresh_chatgpt_web_provider_quota_locally_boxed,
|
||||
),
|
||||
("codex", refresh_codex_provider_quota_locally_boxed),
|
||||
("kiro", refresh_kiro_provider_quota_locally_boxed),
|
||||
];
|
||||
|
||||
pub(crate) async fn refresh_provider_pool_quota_locally(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
provider_type: &str,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> Result<Option<serde_json::Value>, GatewayError> {
|
||||
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
|
||||
let Some((_, handler)) = PROVIDER_QUOTA_REFRESH_HANDLERS
|
||||
.iter()
|
||||
.find(|(supported_provider_type, _)| *supported_provider_type == normalized_provider_type)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
handler(state, provider, endpoint, keys, proxy_override).await
|
||||
}
|
||||
|
||||
fn refresh_antigravity_provider_quota_locally_boxed<'a>(
|
||||
state: &'a AdminAppState<'a>,
|
||||
provider: &'a StoredProviderCatalogProvider,
|
||||
endpoint: &'a StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> ProviderQuotaRefreshFuture<'a> {
|
||||
Box::pin(refresh_antigravity_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
endpoint,
|
||||
keys,
|
||||
proxy_override,
|
||||
))
|
||||
}
|
||||
|
||||
fn refresh_chatgpt_web_provider_quota_locally_boxed<'a>(
|
||||
state: &'a AdminAppState<'a>,
|
||||
provider: &'a StoredProviderCatalogProvider,
|
||||
endpoint: &'a StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> ProviderQuotaRefreshFuture<'a> {
|
||||
Box::pin(refresh_chatgpt_web_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
endpoint,
|
||||
keys,
|
||||
proxy_override,
|
||||
))
|
||||
}
|
||||
|
||||
fn refresh_codex_provider_quota_locally_boxed<'a>(
|
||||
state: &'a AdminAppState<'a>,
|
||||
provider: &'a StoredProviderCatalogProvider,
|
||||
endpoint: &'a StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> ProviderQuotaRefreshFuture<'a> {
|
||||
Box::pin(refresh_codex_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
endpoint,
|
||||
keys,
|
||||
proxy_override,
|
||||
))
|
||||
}
|
||||
|
||||
fn refresh_kiro_provider_quota_locally_boxed<'a>(
|
||||
state: &'a AdminAppState<'a>,
|
||||
provider: &'a StoredProviderCatalogProvider,
|
||||
endpoint: &'a StoredProviderCatalogEndpoint,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
proxy_override: Option<ProxySnapshot>,
|
||||
) -> ProviderQuotaRefreshFuture<'a> {
|
||||
Box::pin(refresh_kiro_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
endpoint,
|
||||
keys,
|
||||
proxy_override,
|
||||
))
|
||||
}
|
||||
@@ -1,66 +1,13 @@
|
||||
use super::super::shared::default_provider_quota_execution_timeouts;
|
||||
use super::super::shared::{execute_provider_quota_plan, ProviderQuotaExecutionOutcome};
|
||||
use crate::handlers::admin::provider::shared::payloads::{
|
||||
KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
|
||||
use super::super::shared::{
|
||||
build_provider_quota_execution_plan, default_provider_quota_execution_timeouts,
|
||||
execute_provider_quota_plan, ProviderQuotaExecutionOutcome,
|
||||
};
|
||||
use crate::handlers::admin::request::{
|
||||
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminKiroRequestAuth,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
|
||||
use std::collections::BTreeMap;
|
||||
use url::form_urlencoded;
|
||||
use uuid::Uuid;
|
||||
|
||||
fn build_kiro_usage_headers(auth: &AdminKiroRequestAuth) -> BTreeMap<String, String> {
|
||||
let kiro_version = auth.auth_config.effective_kiro_version();
|
||||
let machine_id = auth.machine_id.trim();
|
||||
let ide_tag = if machine_id.is_empty() {
|
||||
format!("KiroIDE-{kiro_version}")
|
||||
} else {
|
||||
format!("KiroIDE-{kiro_version}-{machine_id}")
|
||||
};
|
||||
let host = format!(
|
||||
"q.{}.amazonaws.com",
|
||||
auth.auth_config.effective_api_region()
|
||||
);
|
||||
|
||||
BTreeMap::from([
|
||||
(
|
||||
"x-amz-user-agent".to_string(),
|
||||
format!("aws-sdk-js/{KIRO_USAGE_SDK_VERSION} {ide_tag}"),
|
||||
),
|
||||
(
|
||||
"user-agent".to_string(),
|
||||
format!(
|
||||
"aws-sdk-js/{KIRO_USAGE_SDK_VERSION} ua/2.1 os/other#unknown lang/js md/nodejs#22.21.1 api/codewhispererruntime#1.0.0 m/N,E {ide_tag}"
|
||||
),
|
||||
),
|
||||
("host".to_string(), host),
|
||||
("amz-sdk-invocation-id".to_string(), Uuid::new_v4().to_string()),
|
||||
("amz-sdk-request".to_string(), "attempt=1; max=1".to_string()),
|
||||
("authorization".to_string(), auth.value.clone()),
|
||||
("connection".to_string(), "close".to_string()),
|
||||
])
|
||||
}
|
||||
|
||||
fn build_kiro_usage_url(auth: &AdminKiroRequestAuth) -> String {
|
||||
let host = format!(
|
||||
"q.{}.amazonaws.com",
|
||||
auth.auth_config.effective_api_region()
|
||||
);
|
||||
let mut serializer = form_urlencoded::Serializer::new(String::new());
|
||||
serializer.append_pair("origin", "AI_EDITOR");
|
||||
serializer.append_pair("resourceType", "AGENTIC_REQUEST");
|
||||
serializer.append_pair("isEmailRequired", "true");
|
||||
if let Some(profile_arn) = auth.auth_config.profile_arn_for_payload() {
|
||||
serializer.append_pair("profileArn", profile_arn);
|
||||
}
|
||||
format!(
|
||||
"https://{host}{KIRO_USAGE_LIMITS_PATH}?{}",
|
||||
serializer.finish()
|
||||
)
|
||||
}
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_provider_pool::{build_kiro_pool_quota_request, KiroPoolQuotaAuthInput};
|
||||
|
||||
pub(super) async fn execute_kiro_quota_plan(
|
||||
state: &AdminAppState<'_>,
|
||||
@@ -81,31 +28,26 @@ pub(super) async fn execute_kiro_quota_plan(
|
||||
.or(Some(default_provider_quota_execution_timeouts(
|
||||
proxy.as_ref(),
|
||||
)));
|
||||
let plan = ExecutionPlan {
|
||||
request_id: format!("kiro-quota:{}", transport.key.id),
|
||||
candidate_id: None,
|
||||
provider_name: Some("kiro".to_string()),
|
||||
provider_id: transport.provider.id.clone(),
|
||||
endpoint_id: transport.endpoint.id.clone(),
|
||||
key_id: transport.key.id.clone(),
|
||||
method: "GET".to_string(),
|
||||
url: build_kiro_usage_url(auth),
|
||||
headers: build_kiro_usage_headers(auth),
|
||||
content_type: None,
|
||||
content_encoding: None,
|
||||
body: RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
let spec = build_kiro_pool_quota_request(
|
||||
&transport.key.id,
|
||||
&KiroPoolQuotaAuthInput {
|
||||
authorization_value: auth.value.clone(),
|
||||
api_region: auth.auth_config.effective_api_region().to_string(),
|
||||
kiro_version: auth.auth_config.effective_kiro_version().to_string(),
|
||||
machine_id: auth.machine_id.clone(),
|
||||
profile_arn: auth
|
||||
.auth_config
|
||||
.profile_arn_for_payload()
|
||||
.map(str::to_string),
|
||||
},
|
||||
stream: false,
|
||||
client_api_format: "claude:messages".to_string(),
|
||||
provider_api_format: "kiro:usage".to_string(),
|
||||
model_name: Some("kiro-usage-limits".to_string()),
|
||||
);
|
||||
let plan = build_provider_quota_execution_plan(
|
||||
transport,
|
||||
spec,
|
||||
proxy,
|
||||
transport_profile: state.resolve_transport_profile(transport),
|
||||
state.resolve_transport_profile(transport),
|
||||
timeouts,
|
||||
};
|
||||
);
|
||||
|
||||
execute_provider_quota_plan(state, transport, plan, "kiro").await
|
||||
}
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
pub(crate) mod antigravity;
|
||||
pub(crate) mod chatgpt_web;
|
||||
pub(crate) mod codex;
|
||||
pub(crate) mod dispatch;
|
||||
pub(crate) mod kiro;
|
||||
pub(crate) mod shared;
|
||||
|
||||
@@ -7,8 +7,14 @@ use crate::handlers::shared::{
|
||||
};
|
||||
use crate::GatewayError;
|
||||
use aether_admin::provider::quota as admin_provider_quota_pure;
|
||||
use aether_contracts::{ExecutionPlan, ExecutionResult, ExecutionTimeouts, ProxySnapshot};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use aether_contracts::{
|
||||
ExecutionPlan, ExecutionResult, ExecutionTimeouts, ProxySnapshot, RequestBody,
|
||||
ResolvedTransportProfile, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
};
|
||||
use aether_provider_pool::{ProviderPoolQuotaRequestSpec, ProviderPoolService};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use tracing::warn;
|
||||
|
||||
@@ -51,22 +57,28 @@ pub(crate) fn normalize_string_id_list(values: Option<Vec<String>>) -> Option<Ve
|
||||
}
|
||||
|
||||
pub(crate) fn provider_type_supports_quota_refresh(provider_type: &str) -> bool {
|
||||
matches!(
|
||||
provider_type.trim().to_ascii_lowercase().as_str(),
|
||||
"codex" | "kiro" | "antigravity" | "chatgpt_web"
|
||||
)
|
||||
ProviderPoolService::with_builtin_adapters().supports_quota_refresh(provider_type)
|
||||
}
|
||||
|
||||
pub(crate) fn unsupported_provider_quota_refresh_message(provider_type: &str) -> String {
|
||||
match provider_type.trim().to_ascii_lowercase().as_str() {
|
||||
"claude_code" => "Claude Code 暂不支持自动刷新额度:上游没有稳定可用的账号额度查询接口",
|
||||
"gemini_cli" => {
|
||||
"Gemini CLI 暂不支持自动刷新额度:当前只能通过模型同步/缓存快照展示已知配额信息"
|
||||
}
|
||||
"vertex_ai" => "Vertex AI 暂不支持自动刷新额度:额度属于 Google Cloud 项目/区域配额",
|
||||
_ => "该 Provider 暂不支持自动刷新额度",
|
||||
}
|
||||
.to_string()
|
||||
ProviderPoolService::with_builtin_adapters().quota_refresh_unsupported_message(provider_type)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_quota_refresh_endpoint_for_provider(
|
||||
provider_type: &str,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
include_inactive: bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
ProviderPoolService::with_builtin_adapters().quota_refresh_endpoint_for_provider(
|
||||
provider_type,
|
||||
endpoints,
|
||||
include_inactive,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_quota_refresh_missing_endpoint_message(provider_type: &str) -> String {
|
||||
ProviderPoolService::with_builtin_adapters()
|
||||
.quota_refresh_missing_endpoint_message(provider_type)
|
||||
}
|
||||
|
||||
pub(super) fn coerce_json_u64(value: &serde_json::Value) -> Option<u64> {
|
||||
@@ -125,6 +137,63 @@ pub(super) fn build_quota_snapshot_payload(
|
||||
updated_snapshot.get("quota").cloned()
|
||||
}
|
||||
|
||||
pub(super) fn build_provider_quota_execution_plan(
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
spec: ProviderPoolQuotaRequestSpec,
|
||||
proxy: Option<ProxySnapshot>,
|
||||
transport_profile: Option<ResolvedTransportProfile>,
|
||||
timeouts: Option<ExecutionTimeouts>,
|
||||
) -> ExecutionPlan {
|
||||
let ProviderPoolQuotaRequestSpec {
|
||||
request_id,
|
||||
provider_name,
|
||||
quota_kind: _,
|
||||
method,
|
||||
url,
|
||||
mut headers,
|
||||
content_type,
|
||||
json_body,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name,
|
||||
accept_invalid_certs,
|
||||
} = spec;
|
||||
if accept_invalid_certs {
|
||||
headers.insert(
|
||||
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER.to_string(),
|
||||
"true".to_string(),
|
||||
);
|
||||
}
|
||||
let body = json_body
|
||||
.map(RequestBody::from_json)
|
||||
.unwrap_or(RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
});
|
||||
ExecutionPlan {
|
||||
request_id,
|
||||
candidate_id: None,
|
||||
provider_name: Some(provider_name),
|
||||
provider_id: transport.provider.id.clone(),
|
||||
endpoint_id: transport.endpoint.id.clone(),
|
||||
key_id: transport.key.id.clone(),
|
||||
method,
|
||||
url,
|
||||
headers,
|
||||
content_type,
|
||||
content_encoding: None,
|
||||
body,
|
||||
stream: false,
|
||||
client_api_format,
|
||||
provider_api_format,
|
||||
model_name,
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn persist_provider_quota_refresh_state(
|
||||
state: &AdminAppState<'_>,
|
||||
key_id: &str,
|
||||
@@ -142,24 +211,21 @@ pub(crate) async fn persist_provider_quota_refresh_state(
|
||||
return Ok(false);
|
||||
};
|
||||
|
||||
let mut quota_snapshot_provider_type = None::<&str>;
|
||||
let mut quota_snapshot_provider_type = None::<String>;
|
||||
if let Some(metadata_update) = metadata_update {
|
||||
latest_key.upstream_metadata = Some(merge_upstream_metadata(
|
||||
latest_key.upstream_metadata.as_ref(),
|
||||
metadata_update,
|
||||
));
|
||||
quota_snapshot_provider_type = metadata_update.as_object().and_then(|object| {
|
||||
["codex", "kiro", "antigravity", "gemini_cli", "chatgpt_web"]
|
||||
.into_iter()
|
||||
.find(|provider_type| object.contains_key(*provider_type))
|
||||
});
|
||||
quota_snapshot_provider_type =
|
||||
aether_provider_pool::provider_pool_quota_metadata_provider_type(metadata_update);
|
||||
}
|
||||
if let Some(encrypted_auth_config) = encrypted_auth_config {
|
||||
latest_key.encrypted_auth_config = Some(encrypted_auth_config);
|
||||
}
|
||||
latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs;
|
||||
latest_key.oauth_invalid_reason = oauth_invalid_reason;
|
||||
if let Some(provider_type) = quota_snapshot_provider_type {
|
||||
if let Some(provider_type) = quota_snapshot_provider_type.as_deref() {
|
||||
latest_key.status_snapshot = sync_provider_key_quota_status_snapshot(
|
||||
latest_key.status_snapshot.as_ref(),
|
||||
provider_type,
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
use super::quota::antigravity::refresh_antigravity_provider_quota_locally;
|
||||
use super::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
|
||||
use super::quota::codex::refresh_codex_provider_quota_locally;
|
||||
use super::quota::kiro::refresh_kiro_provider_quota_locally;
|
||||
use super::quota::dispatch::refresh_provider_pool_quota_locally;
|
||||
use super::quota::shared::provider_quota_refresh_endpoint_for_provider;
|
||||
use super::quota::shared::provider_type_supports_quota_refresh;
|
||||
use crate::handlers::admin::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::provider_key_auth::provider_key_is_oauth_managed;
|
||||
@@ -113,16 +112,20 @@ pub(crate) struct ProviderOAuthRuntimeEndpoints {
|
||||
pub(crate) runtime_endpoint: Option<StoredProviderCatalogEndpoint>,
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
|
||||
async fn resolve_provider_runtime_endpoints_with_selector(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
provider_type: &str,
|
||||
endpoint_selector: fn(
|
||||
&str,
|
||||
&[StoredProviderCatalogEndpoint],
|
||||
bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint>,
|
||||
) -> Result<ProviderOAuthRuntimeEndpoints, GatewayError> {
|
||||
let mut endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await?;
|
||||
let mut runtime_endpoint =
|
||||
provider_oauth_maintenance_endpoint_for_provider(provider_type, &endpoints);
|
||||
let mut runtime_endpoint = endpoint_selector(provider_type, &endpoints, true);
|
||||
if runtime_endpoint.is_none()
|
||||
&& state
|
||||
.fixed_provider_template(&provider.provider_type)
|
||||
@@ -133,8 +136,7 @@ pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
|
||||
endpoints = state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await?;
|
||||
runtime_endpoint =
|
||||
provider_oauth_maintenance_endpoint_for_provider(provider_type, &endpoints);
|
||||
runtime_endpoint = endpoint_selector(provider_type, &endpoints, true);
|
||||
}
|
||||
|
||||
Ok(ProviderOAuthRuntimeEndpoints {
|
||||
@@ -143,6 +145,34 @@ pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
provider_type: &str,
|
||||
) -> Result<ProviderOAuthRuntimeEndpoints, GatewayError> {
|
||||
resolve_provider_runtime_endpoints_with_selector(
|
||||
state,
|
||||
provider,
|
||||
provider_type,
|
||||
select_provider_oauth_runtime_endpoint,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn resolve_provider_quota_runtime_endpoints(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
provider_type: &str,
|
||||
) -> Result<ProviderOAuthRuntimeEndpoints, GatewayError> {
|
||||
resolve_provider_runtime_endpoints_with_selector(
|
||||
state,
|
||||
provider,
|
||||
provider_type,
|
||||
provider_quota_refresh_endpoint_for_provider,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn refresh_provider_oauth_account_state_after_update(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
@@ -150,16 +180,13 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
|
||||
proxy_override: Option<&ProxySnapshot>,
|
||||
) -> Result<(bool, Option<String>), GatewayError> {
|
||||
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
|
||||
if !matches!(
|
||||
provider_type.as_str(),
|
||||
"codex" | "kiro" | "antigravity" | "chatgpt_web"
|
||||
) {
|
||||
if !provider_type_supports_quota_refresh(&provider_type) {
|
||||
return Ok((false, None));
|
||||
}
|
||||
|
||||
let ProviderOAuthRuntimeEndpoints {
|
||||
runtime_endpoint, ..
|
||||
} = resolve_provider_oauth_runtime_endpoints(state, provider, &provider_type).await?;
|
||||
} = resolve_provider_quota_runtime_endpoints(state, provider, &provider_type).await?;
|
||||
let Some(endpoint) = runtime_endpoint else {
|
||||
return Ok((false, None));
|
||||
};
|
||||
@@ -175,50 +202,15 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
|
||||
return Ok((false, None));
|
||||
}
|
||||
|
||||
let proxy_override = proxy_override.cloned();
|
||||
let payload = match provider_type.as_str() {
|
||||
"codex" => {
|
||||
refresh_codex_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
&endpoint,
|
||||
vec![key],
|
||||
proxy_override.clone(),
|
||||
)
|
||||
.await?
|
||||
}
|
||||
"kiro" => {
|
||||
refresh_kiro_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
&endpoint,
|
||||
vec![key],
|
||||
proxy_override.clone(),
|
||||
)
|
||||
.await?
|
||||
}
|
||||
"antigravity" => {
|
||||
refresh_antigravity_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
&endpoint,
|
||||
vec![key],
|
||||
proxy_override,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
"chatgpt_web" => {
|
||||
refresh_chatgpt_web_provider_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
&endpoint,
|
||||
vec![key],
|
||||
proxy_override,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let payload = refresh_provider_pool_quota_locally(
|
||||
state,
|
||||
provider,
|
||||
&endpoint,
|
||||
&provider_type,
|
||||
vec![key],
|
||||
proxy_override.cloned(),
|
||||
)
|
||||
.await?;
|
||||
let Some(payload) = payload else {
|
||||
return Ok((false, None));
|
||||
};
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::handlers::admin::provider::shared::support::{
|
||||
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule,
|
||||
};
|
||||
use aether_ai_serving::{PoolMemberScoreRules, PoolMemberScoreWeights};
|
||||
use aether_pool_core::{PoolMemberScoreRules, PoolMemberScoreWeights};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
const POOL_ALLOWED_SCHEDULING_PRESETS: &[&str] = &[
|
||||
|
||||
@@ -618,7 +618,7 @@ mod tests {
|
||||
probe_concurrency: 4,
|
||||
score_top_n: 128,
|
||||
score_fallback_scan_limit: 1024,
|
||||
score_rules: aether_ai_serving::PoolMemberScoreRules::default(),
|
||||
score_rules: aether_pool_core::PoolMemberScoreRules::default(),
|
||||
stream_timeout_threshold: 3,
|
||||
stream_timeout_window_seconds: 1800,
|
||||
stream_timeout_cooldown_seconds: 300,
|
||||
|
||||
@@ -95,27 +95,6 @@ fn admin_pool_oauth_organizations(
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn admin_pool_normalize_oauth_plan_type(value: &str, provider_type: &str) -> Option<String> {
|
||||
let mut normalized = value.trim().to_string();
|
||||
if normalized.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let provider_type = provider_type.trim().to_ascii_lowercase();
|
||||
if !provider_type.is_empty() && normalized.to_ascii_lowercase().starts_with(&provider_type) {
|
||||
normalized = normalized[provider_type.len()..]
|
||||
.trim_matches(|ch: char| [' ', ':', '-', '_'].contains(&ch))
|
||||
.to_string();
|
||||
}
|
||||
|
||||
let normalized = normalized.trim().to_ascii_lowercase();
|
||||
if normalized.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(normalized)
|
||||
}
|
||||
}
|
||||
|
||||
fn admin_pool_derive_oauth_expires_at(
|
||||
provider_type: &str,
|
||||
key: &StoredProviderCatalogKey,
|
||||
@@ -139,57 +118,6 @@ fn admin_pool_derive_oauth_expires_at(
|
||||
None
|
||||
}
|
||||
|
||||
fn admin_pool_derive_oauth_plan_type(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
auth_config: Option<&serde_json::Map<String, serde_json::Value>>,
|
||||
) -> Option<String> {
|
||||
if !provider_key_auth_semantics(key, provider_type).oauth_managed() {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(upstream_metadata) = key
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_object)
|
||||
{
|
||||
let provider_bucket = upstream_metadata
|
||||
.get(&provider_type.trim().to_ascii_lowercase())
|
||||
.and_then(serde_json::Value::as_object);
|
||||
for source in provider_bucket
|
||||
.into_iter()
|
||||
.chain(std::iter::once(upstream_metadata))
|
||||
{
|
||||
for field in [
|
||||
"plan_type",
|
||||
"tier",
|
||||
"subscription_title",
|
||||
"subscription_plan",
|
||||
] {
|
||||
if let Some(value) = source.get(field).and_then(serde_json::Value::as_str) {
|
||||
let normalized = admin_pool_normalize_oauth_plan_type(value, provider_type);
|
||||
if normalized.is_some() {
|
||||
return normalized;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(config) = auth_config {
|
||||
for field in ["plan_type", "tier", "plan", "subscription_plan"] {
|
||||
if let Some(value) = config.get(field).and_then(serde_json::Value::as_str) {
|
||||
let normalized = admin_pool_normalize_oauth_plan_type(value, provider_type);
|
||||
if normalized.is_some() {
|
||||
return normalized;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn admin_pool_format_percent(value: f64) -> String {
|
||||
format!("{:.1}%", value.clamp(0.0, 100.0))
|
||||
}
|
||||
@@ -925,8 +853,11 @@ pub(super) fn build_admin_pool_key_payload(
|
||||
let auth_config = state.parse_catalog_auth_config_json(key);
|
||||
let oauth_expires_at =
|
||||
admin_pool_derive_oauth_expires_at(provider_type, key, auth_config.as_ref());
|
||||
let oauth_plan_type =
|
||||
admin_pool_derive_oauth_plan_type(key, provider_type, auth_config.as_ref());
|
||||
let oauth_plan_type = if auth_semantics.oauth_managed() {
|
||||
aether_provider_pool::derive_plan_tier(provider_type, key, auth_config.as_ref())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let mut status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
|
||||
if provider_type.trim().eq_ignore_ascii_case("codex") {
|
||||
admin_pool_apply_codex_window_usage_summaries(
|
||||
|
||||
@@ -22,74 +22,17 @@ fn admin_pool_parse_auth_config_json(
|
||||
.cloned()
|
||||
}
|
||||
|
||||
fn admin_pool_derive_oauth_plan_type(
|
||||
fn admin_pool_derive_plan_tier(
|
||||
state: &AdminAppState<'_>,
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> Option<String> {
|
||||
let normalize = |value: &str| {
|
||||
let mut text = value.trim().to_string();
|
||||
if text.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let provider_type = provider_type.trim().to_ascii_lowercase();
|
||||
if !provider_type.is_empty() && text.to_ascii_lowercase().starts_with(&provider_type) {
|
||||
text = text[provider_type.len()..]
|
||||
.trim_matches(|ch: char| [' ', ':', '-', '_'].contains(&ch))
|
||||
.to_string();
|
||||
}
|
||||
if text.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(text.to_ascii_lowercase())
|
||||
}
|
||||
};
|
||||
|
||||
if !provider_key_is_oauth_managed(key, provider_type) {
|
||||
return None;
|
||||
}
|
||||
|
||||
if let Some(upstream_metadata) = key
|
||||
.upstream_metadata
|
||||
.as_ref()
|
||||
.and_then(serde_json::Value::as_object)
|
||||
{
|
||||
let provider_bucket = upstream_metadata
|
||||
.get(&provider_type.trim().to_ascii_lowercase())
|
||||
.and_then(serde_json::Value::as_object);
|
||||
for source in provider_bucket
|
||||
.into_iter()
|
||||
.chain(std::iter::once(upstream_metadata))
|
||||
{
|
||||
for plan_key in [
|
||||
"plan_type",
|
||||
"tier",
|
||||
"subscription_title",
|
||||
"subscription_plan",
|
||||
] {
|
||||
if let Some(value) = source.get(plan_key).and_then(serde_json::Value::as_str) {
|
||||
if let Some(normalized) = normalize(value) {
|
||||
return Some(normalized);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(auth_config) = admin_pool_parse_auth_config_json(state, key) {
|
||||
for plan_key in ["plan_type", "tier", "plan", "subscription_plan"] {
|
||||
if let Some(value) = auth_config
|
||||
.get(plan_key)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
{
|
||||
if let Some(normalized) = normalize(value) {
|
||||
return Some(normalized);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
let auth_config = admin_pool_parse_auth_config_json(state, key);
|
||||
aether_provider_pool::derive_plan_tier(provider_type, key, auth_config.as_ref())
|
||||
}
|
||||
|
||||
pub(super) fn admin_pool_matches_quick_selector(
|
||||
@@ -98,7 +41,7 @@ pub(super) fn admin_pool_matches_quick_selector(
|
||||
provider_type: &str,
|
||||
selector: &str,
|
||||
) -> bool {
|
||||
let oauth_plan_type = admin_pool_derive_oauth_plan_type(state, key, provider_type);
|
||||
let oauth_plan_type = admin_pool_derive_plan_tier(state, key, provider_type);
|
||||
admin_provider_pool_pure::admin_pool_matches_quick_selector(
|
||||
key,
|
||||
selector,
|
||||
@@ -113,7 +56,7 @@ pub(super) fn admin_pool_matches_search(
|
||||
provider_type: &str,
|
||||
search: Option<&str>,
|
||||
) -> bool {
|
||||
let oauth_plan_type = admin_pool_derive_oauth_plan_type(state, key, provider_type);
|
||||
let oauth_plan_type = admin_pool_derive_plan_tier(state, key, provider_type);
|
||||
admin_provider_pool_pure::admin_pool_matches_search(key, search, oauth_plan_type.as_deref())
|
||||
}
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::LocalProviderDeleteTaskState;
|
||||
use aether_ai_serving::PoolMemberScoreRules;
|
||||
use aether_pool_core::PoolMemberScoreRules;
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
|
||||
@@ -9,16 +9,16 @@ use aether_data_contracts::repository::pool_scores::{
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_provider_pool::provider_pool_quota_metadata_updated_at;
|
||||
use aether_runtime_state::{RuntimeLockLease, RuntimeState};
|
||||
use futures_util::{stream, StreamExt};
|
||||
use serde_json::Value;
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use crate::admin_api::{
|
||||
admin_provider_pool_config, provider_oauth_maintenance_endpoint_for_provider,
|
||||
admin_provider_pool_config, provider_quota_refresh_endpoint_for_provider,
|
||||
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
|
||||
refresh_antigravity_provider_quota_locally, refresh_chatgpt_web_provider_quota_locally,
|
||||
refresh_codex_provider_quota_locally, refresh_kiro_provider_quota_locally, AdminAppState,
|
||||
refresh_provider_pool_quota_locally, AdminAppState,
|
||||
};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
@@ -112,38 +112,6 @@ fn provider_supports_quota_probe(provider_type: &str) -> bool {
|
||||
provider_type_supports_quota_refresh(provider_type)
|
||||
}
|
||||
|
||||
fn json_number(value: Option<&Value>) -> Option<f64> {
|
||||
let value = value?;
|
||||
if let Some(number) = value.as_f64() {
|
||||
return Some(number);
|
||||
}
|
||||
value
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.and_then(|value| value.parse::<f64>().ok())
|
||||
}
|
||||
|
||||
fn extract_quota_updated_at(provider_type: &str, upstream_metadata: Option<&Value>) -> Option<u64> {
|
||||
let metadata = upstream_metadata?.as_object()?;
|
||||
let bucket_name = match provider_type.trim().to_ascii_lowercase().as_str() {
|
||||
"codex" => "codex",
|
||||
"kiro" => "kiro",
|
||||
"antigravity" => "antigravity",
|
||||
"chatgpt_web" => "chatgpt_web",
|
||||
_ => return None,
|
||||
};
|
||||
let bucket = metadata.get(bucket_name)?.as_object()?;
|
||||
let mut updated_at = json_number(bucket.get("updated_at"))?;
|
||||
if updated_at <= 0.0 {
|
||||
return None;
|
||||
}
|
||||
if updated_at > 1_000_000_000_000.0 {
|
||||
updated_at /= 1000.0;
|
||||
}
|
||||
Some(updated_at as u64)
|
||||
}
|
||||
|
||||
fn parse_probe_stamp(raw_value: Option<&str>) -> Option<u64> {
|
||||
let parsed = raw_value
|
||||
.map(str::trim)
|
||||
@@ -169,7 +137,7 @@ pub(crate) fn select_pool_quota_probe_key_ids(
|
||||
continue;
|
||||
}
|
||||
let quota_updated_ts =
|
||||
extract_quota_updated_at(provider_type, key.upstream_metadata.as_ref());
|
||||
provider_pool_quota_metadata_updated_at(key.upstream_metadata.as_ref(), provider_type);
|
||||
let last_probe_ts = last_probe_timestamps.get(&key.id).copied();
|
||||
let anchor_ts = quota_updated_ts
|
||||
.unwrap_or(0)
|
||||
@@ -421,7 +389,7 @@ fn endpoint_for_probe(
|
||||
provider_type: &str,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
provider_oauth_maintenance_endpoint_for_provider(provider_type, endpoints)
|
||||
provider_quota_refresh_endpoint_for_provider(provider_type, endpoints, true)
|
||||
}
|
||||
|
||||
async fn endpoint_for_probe_with_reconcile(
|
||||
@@ -462,23 +430,8 @@ async fn refresh_provider_probe_keys(
|
||||
provider_type: &str,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
) -> Result<Option<Value>, GatewayError> {
|
||||
match provider_type {
|
||||
"codex" => {
|
||||
refresh_codex_provider_quota_locally(admin_state, provider, endpoint, keys, None).await
|
||||
}
|
||||
"kiro" => {
|
||||
refresh_kiro_provider_quota_locally(admin_state, provider, endpoint, keys, None).await
|
||||
}
|
||||
"antigravity" => {
|
||||
refresh_antigravity_provider_quota_locally(admin_state, provider, endpoint, keys, None)
|
||||
.await
|
||||
}
|
||||
"chatgpt_web" => {
|
||||
refresh_chatgpt_web_provider_quota_locally(admin_state, provider, endpoint, keys, None)
|
||||
.await
|
||||
}
|
||||
_ => Ok(None),
|
||||
}
|
||||
refresh_provider_pool_quota_locally(admin_state, provider, endpoint, provider_type, keys, None)
|
||||
.await
|
||||
}
|
||||
|
||||
fn update_summary_from_payload(
|
||||
@@ -971,16 +924,16 @@ mod tests {
|
||||
#[test]
|
||||
fn parses_quota_updated_at_seconds_and_milliseconds() {
|
||||
assert_eq!(
|
||||
extract_quota_updated_at(
|
||||
"codex",
|
||||
Some(&json!({ "codex": { "updated_at": 1_700_000_000 } }))
|
||||
provider_pool_quota_metadata_updated_at(
|
||||
Some(&json!({ "codex": { "updated_at": 1_700_000_000 } })),
|
||||
"codex"
|
||||
),
|
||||
Some(1_700_000_000)
|
||||
);
|
||||
assert_eq!(
|
||||
extract_quota_updated_at(
|
||||
"kiro",
|
||||
Some(&json!({ "kiro": { "updated_at": 1_700_000_000_000_u64 } }))
|
||||
provider_pool_quota_metadata_updated_at(
|
||||
Some(&json!({ "kiro": { "updated_at": 1_700_000_000_000_u64 } })),
|
||||
"kiro"
|
||||
),
|
||||
Some(1_700_000_000)
|
||||
);
|
||||
|
||||
@@ -1704,7 +1704,9 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
|
||||
read_workspace_file("apps/aether-gateway/src/handlers/admin/provider/oauth/quota/mod.rs");
|
||||
for pattern in [
|
||||
"pub(crate) mod antigravity;",
|
||||
"pub(crate) mod chatgpt_web;",
|
||||
"pub(crate) mod codex;",
|
||||
"pub(crate) mod dispatch;",
|
||||
"pub(crate) mod kiro;",
|
||||
"pub(crate) mod shared;",
|
||||
] {
|
||||
@@ -1730,9 +1732,7 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
|
||||
"apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/quota.rs",
|
||||
);
|
||||
for pattern in [
|
||||
"use super::super::oauth::quota::antigravity::refresh_antigravity_provider_quota_locally;",
|
||||
"use super::super::oauth::quota::codex::refresh_codex_provider_quota_locally;",
|
||||
"use super::super::oauth::quota::kiro::refresh_kiro_provider_quota_locally;",
|
||||
"use super::super::oauth::quota::dispatch::refresh_provider_pool_quota_locally;",
|
||||
"use super::super::oauth::quota::shared::normalize_string_id_list;",
|
||||
] {
|
||||
assert!(
|
||||
@@ -1744,15 +1744,51 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
|
||||
let oauth_runtime =
|
||||
read_workspace_file("apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs");
|
||||
for pattern in [
|
||||
"use super::quota::antigravity::refresh_antigravity_provider_quota_locally;",
|
||||
"use super::quota::codex::refresh_codex_provider_quota_locally;",
|
||||
"use super::quota::kiro::refresh_kiro_provider_quota_locally;",
|
||||
"use super::quota::dispatch::refresh_provider_pool_quota_locally;",
|
||||
"use super::quota::shared::provider_type_supports_quota_refresh;",
|
||||
] {
|
||||
assert!(
|
||||
oauth_runtime.contains(pattern),
|
||||
"handlers/admin/provider/oauth/runtime.rs should import quota helper via explicit owner {pattern}"
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
!oauth_runtime.contains("\"codex\" | \"kiro\" | \"antigravity\" | \"chatgpt_web\""),
|
||||
"handlers/admin/provider/oauth/runtime.rs should not hardcode quota refresh provider allow-list"
|
||||
);
|
||||
|
||||
let quota_shared = read_workspace_file(
|
||||
"apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs",
|
||||
);
|
||||
assert!(
|
||||
quota_shared.contains("aether_provider_pool::provider_pool_quota_metadata_provider_type("),
|
||||
"handlers/admin/provider/oauth/quota/shared.rs should delegate quota metadata provider detection to aether-provider-pool"
|
||||
);
|
||||
assert!(
|
||||
!quota_shared.contains("[\"codex\", \"kiro\", \"antigravity\", \"gemini_cli\", \"chatgpt_web\"]"),
|
||||
"handlers/admin/provider/oauth/quota/shared.rs should not hardcode quota metadata provider list"
|
||||
);
|
||||
|
||||
let quota_dispatch = read_workspace_file(
|
||||
"apps/aether-gateway/src/handlers/admin/provider/oauth/quota/dispatch.rs",
|
||||
);
|
||||
for pattern in [
|
||||
"pub(crate) async fn refresh_provider_pool_quota_locally(",
|
||||
"const PROVIDER_QUOTA_REFRESH_HANDLERS:",
|
||||
"refresh_codex_provider_quota_locally",
|
||||
"refresh_kiro_provider_quota_locally",
|
||||
"refresh_antigravity_provider_quota_locally",
|
||||
"refresh_chatgpt_web_provider_quota_locally",
|
||||
] {
|
||||
assert!(
|
||||
quota_dispatch.contains(pattern),
|
||||
"handlers/admin/provider/oauth/quota/dispatch.rs should centralize quota refresh dispatch {pattern}"
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
!quota_dispatch.contains("match provider_type.trim().to_ascii_lowercase().as_str()"),
|
||||
"handlers/admin/provider/oauth/quota/dispatch.rs should use provider handler registration instead of provider_type match"
|
||||
);
|
||||
|
||||
let quota_codex_mod = read_workspace_file(
|
||||
"apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs",
|
||||
@@ -1799,14 +1835,13 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
|
||||
"apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/plan.rs",
|
||||
);
|
||||
for pattern in [
|
||||
"use super::parse::normalize_codex_plan_type;",
|
||||
"use crate::handlers::admin::provider::shared::payloads::CODEX_WHAM_USAGE_URL;",
|
||||
"pub(super) fn build_codex_refresh_headers(",
|
||||
"use aether_provider_pool::{build_codex_pool_quota_request, ProviderPoolQuotaRequestSpec};",
|
||||
"pub(super) fn build_codex_quota_request_spec(",
|
||||
"pub(super) async fn execute_codex_quota_plan(",
|
||||
] {
|
||||
assert!(
|
||||
quota_codex_plan.contains(pattern),
|
||||
"handlers/admin/provider/oauth/quota/codex/plan.rs should own codex quota execution helper {pattern}"
|
||||
"handlers/admin/provider/oauth/quota/codex/plan.rs should delegate codex quota request construction and own execution helper {pattern}"
|
||||
);
|
||||
}
|
||||
let quota_kiro_mod = read_workspace_file(
|
||||
@@ -1843,13 +1878,13 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
|
||||
"apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/plan.rs",
|
||||
);
|
||||
for pattern in [
|
||||
"use super::super::shared::{execute_provider_quota_plan, ProviderQuotaExecutionOutcome};",
|
||||
"use crate::handlers::admin::provider::shared::payloads::{",
|
||||
"build_provider_quota_execution_plan",
|
||||
"use aether_provider_pool::{build_kiro_pool_quota_request, KiroPoolQuotaAuthInput};",
|
||||
"pub(super) async fn execute_kiro_quota_plan(",
|
||||
] {
|
||||
assert!(
|
||||
quota_kiro_plan.contains(pattern),
|
||||
"handlers/admin/provider/oauth/quota/kiro/plan.rs should own kiro quota execution helper {pattern}"
|
||||
"handlers/admin/provider/oauth/quota/kiro/plan.rs should delegate kiro quota request construction and own execution helper {pattern}"
|
||||
);
|
||||
}
|
||||
let quota_antigravity = read_workspace_file(
|
||||
@@ -1859,6 +1894,30 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
|
||||
quota_antigravity.contains("use super::shared::{"),
|
||||
"handlers/admin/provider/oauth/quota/antigravity.rs should import common quota helpers from shared.rs"
|
||||
);
|
||||
assert!(
|
||||
quota_antigravity.contains("use aether_provider_pool::build_antigravity_pool_quota_request;"),
|
||||
"handlers/admin/provider/oauth/quota/antigravity.rs should delegate antigravity quota request construction to aether-provider-pool"
|
||||
);
|
||||
let quota_chatgpt_web = read_workspace_file(
|
||||
"apps/aether-gateway/src/handlers/admin/provider/oauth/quota/chatgpt_web.rs",
|
||||
);
|
||||
assert!(
|
||||
quota_chatgpt_web.contains("use aether_provider_pool::{")
|
||||
&& quota_chatgpt_web.contains("build_chatgpt_web_pool_quota_request")
|
||||
&& quota_chatgpt_web.contains("enrich_chatgpt_web_quota_metadata")
|
||||
&& quota_chatgpt_web.contains("normalize_chatgpt_web_image_quota_limit"),
|
||||
"handlers/admin/provider/oauth/quota/chatgpt_web.rs should delegate chatgpt_web quota request and metadata behavior to aether-provider-pool"
|
||||
);
|
||||
for forbidden in [
|
||||
"fn enrich_chatgpt_web_quota_metadata(",
|
||||
"fn normalize_chatgpt_web_image_quota_limit(",
|
||||
"fn chatgpt_web_auth_config_string(",
|
||||
] {
|
||||
assert!(
|
||||
!quota_chatgpt_web.contains(forbidden),
|
||||
"handlers/admin/provider/oauth/quota/chatgpt_web.rs should not own provider-pool chatgpt_web quota metadata helper {forbidden}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1926,6 +1985,34 @@ fn admin_provider_oauth_refresh_helpers_use_specific_local_owners() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_provider_oauth_kiro_token_refresh_delegates_to_oauth_adapter() {
|
||||
let kiro_dispatch = read_workspace_file(
|
||||
"apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/kiro.rs",
|
||||
);
|
||||
assert!(
|
||||
kiro_dispatch.contains("use aether_oauth::provider::providers::KiroProviderOAuthAdapter;"),
|
||||
"gateway Kiro OAuth dispatch should depend on the shared provider OAuth adapter"
|
||||
);
|
||||
assert!(
|
||||
kiro_dispatch.contains(".refresh_auth_config("),
|
||||
"gateway Kiro OAuth dispatch should delegate token refresh to aether-oauth"
|
||||
);
|
||||
for forbidden in [
|
||||
"fn admin_provider_oauth_kiro_build_refresh_url(",
|
||||
"fn admin_provider_oauth_kiro_refresh_response_json(",
|
||||
"\"kiro_batch_refresh:social\"",
|
||||
"\"kiro_batch_refresh:idc\"",
|
||||
"\"refreshToken\": auth_config",
|
||||
"\"grantType\": \"refresh_token\"",
|
||||
] {
|
||||
assert!(
|
||||
!kiro_dispatch.contains(forbidden),
|
||||
"gateway Kiro OAuth dispatch should not own provider-specific token refresh detail {forbidden}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_provider_oauth_dispatch_batch_mod_stays_thin() {
|
||||
let batch_mod = read_workspace_file(
|
||||
|
||||
@@ -1280,10 +1280,10 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
|
||||
}
|
||||
for forbidden in [
|
||||
"pub(crate) async fn apply_local_execution_pool_scheduler(",
|
||||
"run_ai_pool_scheduler(",
|
||||
"fn ai_pool_candidate_facts(",
|
||||
"fn ai_pool_scheduling_config(",
|
||||
"fn ai_pool_runtime_state(",
|
||||
"run_pool_scheduler(",
|
||||
"fn pool_candidate_facts(",
|
||||
"fn pool_scheduling_config(",
|
||||
"fn pool_runtime_state(",
|
||||
] {
|
||||
assert!(
|
||||
!planner_pool_scheduler.contains(forbidden),
|
||||
@@ -1295,10 +1295,10 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
|
||||
for pattern in [
|
||||
"pub(crate) async fn apply_local_execution_pool_scheduler(",
|
||||
"pub(crate) struct PoolKeyCursor",
|
||||
"run_ai_pool_scheduler(",
|
||||
"fn ai_pool_candidate_facts(",
|
||||
"fn ai_pool_scheduling_config(",
|
||||
"fn ai_pool_runtime_state(",
|
||||
"run_pool_scheduler(",
|
||||
"fn pool_candidate_facts(",
|
||||
"fn pool_scheduling_config(",
|
||||
"fn pool_runtime_state(",
|
||||
"DEFAULT_POOL_WINDOW_SIZE",
|
||||
"DEFAULT_POOL_PAGE_SIZE",
|
||||
"DEFAULT_POOL_MAX_SCAN",
|
||||
@@ -1351,21 +1351,218 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
|
||||
);
|
||||
}
|
||||
|
||||
let serving_pool_scheduler =
|
||||
read_workspace_file("crates/aether-ai-serving/src/pool_scheduler.rs");
|
||||
let pool_core_lib = read_workspace_file("crates/aether-pool-core/src/lib.rs");
|
||||
for pattern in [
|
||||
"pub fn run_ai_pool_scheduler",
|
||||
"pub struct AiPoolCandidateInput",
|
||||
"pub struct AiPoolRuntimeState",
|
||||
"pub struct AiPoolSchedulingConfig",
|
||||
"run_pool_scheduler",
|
||||
"PoolCandidateInput",
|
||||
"PoolRuntimeState",
|
||||
"PoolSchedulingConfig",
|
||||
] {
|
||||
assert!(
|
||||
pool_core_lib.contains(pattern),
|
||||
"aether-pool-core lib.rs should expose pool scheduling primitive {pattern}"
|
||||
);
|
||||
}
|
||||
|
||||
let pool_core_scheduler = read_workspace_file("crates/aether-pool-core/src/scheduler.rs");
|
||||
for pattern in [
|
||||
"pub fn run_pool_scheduler",
|
||||
"fn schedule_pool_group",
|
||||
"fn build_pool_sort_vectors",
|
||||
"fn plan_priority_score(",
|
||||
"fn stable_hash_score(",
|
||||
] {
|
||||
assert!(
|
||||
serving_pool_scheduler.contains(pattern),
|
||||
"aether-ai-serving pool_scheduler.rs should own pool scheduling use-case primitive {pattern}"
|
||||
pool_core_scheduler.contains(pattern),
|
||||
"aether-pool-core scheduler.rs should own pool scheduling use-case primitive {pattern}"
|
||||
);
|
||||
}
|
||||
for forbidden in ["codex", "kiro", "chatgpt_web", "provider_type"] {
|
||||
assert!(
|
||||
!pool_core_scheduler.contains(forbidden) && !pool_core_lib.contains(forbidden),
|
||||
"aether-pool-core should stay provider-agnostic and not embed provider behavior {forbidden}"
|
||||
);
|
||||
}
|
||||
|
||||
let serving_lib = read_workspace_file("crates/aether-ai-serving/src/lib.rs");
|
||||
for forbidden in ["pub mod pool_scheduler;", "pub mod pool_scores;"] {
|
||||
assert!(
|
||||
!serving_lib.contains(forbidden),
|
||||
"aether-ai-serving should not own pool core module {forbidden}"
|
||||
);
|
||||
}
|
||||
|
||||
let provider_pool_lib = read_workspace_file("crates/aether-provider-pool/src/lib.rs");
|
||||
for pattern in [
|
||||
"mod capability;",
|
||||
"mod plan;",
|
||||
"mod presets;",
|
||||
"mod provider;",
|
||||
"mod quota;",
|
||||
"mod service;",
|
||||
"pub mod providers;",
|
||||
"pub use provider::{ProviderPoolAdapter, ProviderPoolMemberInput};",
|
||||
"pub use service::ProviderPoolService;",
|
||||
] {
|
||||
assert!(
|
||||
provider_pool_lib.contains(pattern),
|
||||
"aether-provider-pool lib.rs should stay a thin module/re-export root through {pattern}"
|
||||
);
|
||||
}
|
||||
|
||||
let provider_pool_provider = read_workspace_file("crates/aether-provider-pool/src/provider.rs");
|
||||
for pattern in [
|
||||
"pub trait ProviderPoolAdapter",
|
||||
"ProviderPoolMemberInput",
|
||||
"supports_quota_refresh",
|
||||
"quota_refresh_endpoint",
|
||||
] {
|
||||
assert!(
|
||||
provider_pool_provider.contains(pattern),
|
||||
"aether-provider-pool provider.rs should own adapter contract {pattern}"
|
||||
);
|
||||
}
|
||||
|
||||
let provider_pool_service = read_workspace_file("crates/aether-provider-pool/src/service.rs");
|
||||
for pattern in [
|
||||
"pub struct ProviderPoolService",
|
||||
"with_builtin_adapters",
|
||||
"AntigravityProviderPoolAdapter",
|
||||
"CodexProviderPoolAdapter",
|
||||
"KiroProviderPoolAdapter",
|
||||
"ChatGptWebProviderPoolAdapter",
|
||||
"CLAUDE_CODE_PROVIDER_POOL_ADAPTER",
|
||||
"GEMINI_CLI_PROVIDER_POOL_ADAPTER",
|
||||
"VERTEX_AI_PROVIDER_POOL_ADAPTER",
|
||||
"provider_types_for_capability",
|
||||
"supports_quota_refresh",
|
||||
"quota_refresh_endpoint_for_provider",
|
||||
] {
|
||||
assert!(
|
||||
provider_pool_service.contains(pattern),
|
||||
"aether-provider-pool service.rs should own adapter registry/service primitive {pattern}"
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
!provider_pool_service.contains("match provider_type.trim().to_ascii_lowercase().as_str()"),
|
||||
"aether-provider-pool service.rs should delegate provider-specific behavior to adapters"
|
||||
);
|
||||
|
||||
let provider_pool_providers =
|
||||
read_workspace_file("crates/aether-provider-pool/src/providers/mod.rs");
|
||||
for pattern in [
|
||||
"pub mod default;",
|
||||
"pub mod unsupported;",
|
||||
"pub mod antigravity;",
|
||||
"pub mod codex;",
|
||||
"pub mod kiro;",
|
||||
"pub mod chatgpt_web;",
|
||||
] {
|
||||
assert!(
|
||||
provider_pool_providers.contains(pattern),
|
||||
"aether-provider-pool providers/mod.rs should expose provider-specific module {pattern}"
|
||||
);
|
||||
}
|
||||
for (path, patterns) in [
|
||||
(
|
||||
"crates/aether-provider-pool/src/providers/default.rs",
|
||||
vec!["DefaultProviderPoolAdapter"],
|
||||
),
|
||||
(
|
||||
"crates/aether-provider-pool/src/providers/antigravity.rs",
|
||||
vec!["AntigravityProviderPoolAdapter"],
|
||||
),
|
||||
(
|
||||
"crates/aether-provider-pool/src/providers/codex.rs",
|
||||
vec![
|
||||
"CodexProviderPoolAdapter",
|
||||
"recent_refresh",
|
||||
"quota_exhausted_from_bucket",
|
||||
],
|
||||
),
|
||||
(
|
||||
"crates/aether-provider-pool/src/providers/kiro.rs",
|
||||
vec!["KiroProviderPoolAdapter", "quota_exhausted_from_bucket"],
|
||||
),
|
||||
(
|
||||
"crates/aether-provider-pool/src/providers/chatgpt_web.rs",
|
||||
vec![
|
||||
"ChatGptWebProviderPoolAdapter",
|
||||
"build_chatgpt_web_pool_quota_request",
|
||||
"enrich_chatgpt_web_quota_metadata",
|
||||
"normalize_chatgpt_web_image_quota_limit",
|
||||
"quota_exhausted_from_bucket",
|
||||
],
|
||||
),
|
||||
(
|
||||
"crates/aether-provider-pool/src/providers/unsupported.rs",
|
||||
vec![
|
||||
"UnsupportedQuotaProviderPoolAdapter",
|
||||
"CLAUDE_CODE_PROVIDER_POOL_ADAPTER",
|
||||
"GEMINI_CLI_PROVIDER_POOL_ADAPTER",
|
||||
"VERTEX_AI_PROVIDER_POOL_ADAPTER",
|
||||
],
|
||||
),
|
||||
] {
|
||||
let source = read_workspace_file(path);
|
||||
for pattern in patterns {
|
||||
assert!(
|
||||
source.contains(pattern),
|
||||
"{path} should own provider-specific pool behavior {pattern}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let provider_pool_plan = read_workspace_file("crates/aether-provider-pool/src/plan.rs");
|
||||
for pattern in ["normalize_provider_plan_tier", "derive_plan_tier"] {
|
||||
assert!(
|
||||
provider_pool_plan.contains(pattern),
|
||||
"aether-provider-pool plan.rs should own provider plan-tier normalization primitive {pattern}"
|
||||
);
|
||||
}
|
||||
|
||||
let provider_pool_quota = read_workspace_file("crates/aether-provider-pool/src/quota.rs");
|
||||
for pattern in [
|
||||
"provider_pool_key_account_quota_exhausted",
|
||||
"provider_pool_member_quota_snapshot",
|
||||
"provider_pool_quota_metadata_updated_at",
|
||||
"provider_pool_quota_metadata_provider_type",
|
||||
"provider_pool_key_scheduling_label",
|
||||
"provider_pool_quota_snapshot_updated_at",
|
||||
] {
|
||||
assert!(
|
||||
provider_pool_quota.contains(pattern),
|
||||
"aether-provider-pool quota.rs should own provider quota/scheduling signal primitive {pattern}"
|
||||
);
|
||||
}
|
||||
|
||||
let provider_pool_presets = read_workspace_file("crates/aether-provider-pool/src/presets.rs");
|
||||
for pattern in [
|
||||
"normalize_provider_scheduling_presets",
|
||||
"build_admin_pool_scheduling_presets_payload",
|
||||
] {
|
||||
assert!(
|
||||
provider_pool_presets.contains(pattern),
|
||||
"aether-provider-pool presets.rs should own provider preset adaptation primitive {pattern}"
|
||||
);
|
||||
}
|
||||
for forbidden in [
|
||||
"run_pool_scheduler",
|
||||
"PoolSchedulerOutcome",
|
||||
"schedule_pool_group",
|
||||
"plan_priority_score(",
|
||||
] {
|
||||
let mut violations = Vec::new();
|
||||
for file in collect_workspace_rust_files("crates/aether-provider-pool/src") {
|
||||
let source = std::fs::read_to_string(&file).expect("source file should be readable");
|
||||
if source.contains(forbidden) {
|
||||
violations.push(file.display().to_string());
|
||||
}
|
||||
}
|
||||
assert!(
|
||||
violations.is_empty(),
|
||||
"aether-provider-pool should not own generic pool scheduler primitive {forbidden}:\n{}",
|
||||
violations.join("\n")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4302,7 +4302,7 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_p
|
||||
.and_then(|proxy| proxy.node_id.as_deref()),
|
||||
Some("proxy-node-kiro-batch-runtime")
|
||||
);
|
||||
if plan.request_id == "kiro_batch_refresh:social" {
|
||||
if plan.request_id == "provider-oauth:kiro-social-refresh" {
|
||||
assert_eq!(plan.url, "https://oauth.example/refreshToken");
|
||||
assert_eq!(
|
||||
plan.headers
|
||||
|
||||
Reference in New Issue
Block a user