i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: extract provider pool abstractions

Browse Source
This commit is contained in:
fawney19
2026-05-13 18:19:15 +08:00
parent 3c2497f019
commit 5d1460e051
55 changed files with 3469 additions and 2184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/aether-admin/Cargo.toml
View File

@@ -12,6 +12,7 @@ aether-billing.workspace = true
aether-contracts.workspace = true
aether-data.workspace = true
aether-data-contracts.workspace = true
aether-provider-pool.workspace = true
axum.workspace = true
base64.workspace = true
chrono.workspace = true

314
crates/aether-admin/src/provider/pool.rs
View File

@@ -92,180 +92,11 @@ fn admin_pool_reason_indicates_ban(reason: &str) -> bool {
.any(|hint| normalized.contains(hint))
}
fn admin_pool_metadata_bucket<'a>(
upstream_metadata: Option<&'a Value>,
provider_type: &str,
) -> Option<&'a serde_json::Map<String, Value>> {
upstream_metadata
.and_then(Value::as_object)
.and_then(|metadata| metadata.get(&provider_type.trim().to_ascii_lowercase()))
.and_then(Value::as_object)
}
fn admin_pool_json_bool(value: Option<&Value>) -> Option<bool> {
match value {
Some(Value::Bool(value)) => Some(*value),
Some(Value::String(value)) => match value.trim().to_ascii_lowercase().as_str() {
"true" | "1" => Some(true),
"false" | "0" => Some(false),
_ => None,
},
_ => None,
}
}
fn admin_pool_json_f64(value: Option<&Value>) -> Option<f64> {
match value {
Some(Value::Number(number)) => number.as_f64(),
Some(Value::String(value)) => value.trim().parse::<f64>().ok(),
_ => None,
}
.filter(|value| value.is_finite())
}
fn admin_pool_quota_snapshot_matches_provider(
quota_snapshot: &serde_json::Map<String, Value>,
provider_type: &str,
) -> bool {
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
match quota_snapshot
.get("provider_type")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
Some(provider_type) => provider_type.eq_ignore_ascii_case(&normalized_provider_type),
None => {
admin_pool_json_bool(quota_snapshot.get("exhausted")) == Some(true)
|| quota_snapshot
.get("code")
.and_then(Value::as_str)
.is_some_and(|code| !code.trim().eq_ignore_ascii_case("unknown"))
|| quota_snapshot
.get("updated_at")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("observed_at")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("usage_ratio")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("reset_seconds")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("windows")
.and_then(Value::as_array)
.is_some_and(|windows| !windows.is_empty())
|| quota_snapshot
.get("credits")
.and_then(Value::as_object)
.is_some_and(|credits| !credits.is_empty())
}
}
}
fn admin_pool_key_quota_snapshot<'a>(
key: &'a StoredProviderCatalogKey,
provider_type: &str,
) -> Option<&'a serde_json::Map<String, Value>> {
let quota_snapshot = key
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)?;
admin_pool_quota_snapshot_matches_provider(quota_snapshot, provider_type)
.then_some(quota_snapshot)
}
pub fn admin_pool_key_account_quota_exhausted(
key: &StoredProviderCatalogKey,
provider_type: &str,
) -> bool {
if let Some(exhausted) =
admin_pool_key_quota_snapshot(key, provider_type).and_then(|quota_snapshot| {
let exhausted = admin_pool_json_bool(quota_snapshot.get("exhausted"))?;
if exhausted {
let windows_max_ratio = quota_snapshot
.get("windows")
.and_then(Value::as_array)
.filter(|w| !w.is_empty())
.and_then(|windows| {
windows
.iter()
.filter_map(Value::as_object)
.filter_map(|w| w.get("used_ratio"))
.filter_map(Value::as_f64)
.max_by(f64::total_cmp)
});
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
return Some(false);
}
}
Some(exhausted)
})
{
return exhausted;
}
let provider_type = provider_type.trim().to_ascii_lowercase();
let Some(bucket) = admin_pool_metadata_bucket(key.upstream_metadata.as_ref(), &provider_type)
else {
return false;
};
match provider_type.as_str() {
"codex" => {
if admin_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
return false;
}
let has_window_data = admin_pool_json_f64(bucket.get("primary_used_percent")).is_some()
|| admin_pool_json_f64(bucket.get("secondary_used_percent")).is_some();
if !has_window_data && admin_pool_json_bool(bucket.get("has_credits")) == Some(false) {
return true;
}
admin_pool_json_f64(bucket.get("primary_used_percent"))
.is_some_and(|value| value >= 100.0)
|| admin_pool_json_f64(bucket.get("secondary_used_percent"))
.is_some_and(|value| value >= 100.0)
}
"kiro" => {
if admin_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) {
return true;
}
if admin_pool_json_f64(bucket.get("usage_percentage"))
.is_some_and(|value| value >= 100.0)
{
return true;
}
match (
admin_pool_json_f64(bucket.get("usage_limit")),
admin_pool_json_f64(bucket.get("current_usage")),
) {
(Some(limit), Some(current)) if limit > 0.0 => current >= limit,
_ => false,
}
}
"chatgpt_web" => {
if admin_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) {
return true;
}
if admin_pool_json_f64(bucket.get("image_quota_remaining"))
.is_some_and(|value| value <= 0.0)
{
return true;
}
match (
admin_pool_json_f64(bucket.get("image_quota_total")),
admin_pool_json_f64(bucket.get("image_quota_used")),
) {
(Some(limit), Some(used)) if limit > 0.0 => used >= limit,
_ => false,
}
}
_ => false,
}
aether_provider_pool::provider_pool_key_account_quota_exhausted(key, provider_type)
}
fn admin_pool_has_proxy(key: &StoredProviderCatalogKey) -> bool {
@@ -937,148 +768,7 @@ pub fn build_admin_pool_key_payload(
}
pub fn build_admin_pool_scheduling_presets_payload() -> Value {
json!([
{
"name": "lru",
"label": "LRU 轮转",
"description": "最久未使用的 Key 优先",
"providers": [],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": "distribution_mode",
"evidence_hint": "依据 LRU 时间戳（最近未使用优先）",
},
{
"name": "cache_affinity",
"label": "缓存亲和",
"description": "优先复用最近使用过的 Key，利用 Prompt Caching",
"providers": [],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": "distribution_mode",
"evidence_hint": "依据 LRU 时间戳（最近使用优先，与 LRU 轮转相反）",
},
{
"name": "cost_first",
"label": "成本优先",
"description": "优先选择窗口消耗更低的账号",
"providers": [],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据窗口成本/Token 用量，缺失时回退配额使用率",
},
{
"name": "free_first",
"label": "Free 优先",
"description": "优先消耗 Free 账号（依赖 plan_type）",
"providers": ["codex", "kiro"],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据 plan_type（Free 账号优先调度）",
},
{
"name": "health_first",
"label": "健康优先",
"description": "优先选择健康分更高、失败更少的账号",
"providers": [],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据 health_by_format 聚合分（含熔断/失败衰减）",
},
{
"name": "latency_first",
"label": "延迟优先",
"description": "优先选择最近延迟更低的账号",
"providers": [],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据号池延迟窗口均值（latency_window_seconds）",
},
{
"name": "load_balance",
"label": "负载均衡",
"description": "随机分散 Key 使用，均匀分摊负载",
"providers": [],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": "distribution_mode",
"evidence_hint": "每次随机分值，实现完全均匀分散",
},
{
"name": "plus_first",
"label": "Plus 优先",
"description": "优先消耗 Plus 账号（依赖 plan_type）",
"providers": ["codex", "kiro"],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据 plan_type（Plus 账号优先调度）",
},
{
"name": "pro_first",
"label": "Pro 优先",
"description": "优先消耗 Pro 账号（依赖 plan_type）",
"providers": ["codex", "kiro"],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据 plan_type（Pro 账号优先调度）",
},
{
"name": "priority_first",
"label": "优先级优先",
"description": "按账号优先级顺序调度（数字越小越优先）",
"providers": [],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据 internal_priority（支持拖拽/手工编辑）",
},
{
"name": "quota_balanced",
"label": "额度平均",
"description": "优先选额度消耗最少的账号",
"providers": [],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据账号配额使用率；无配额时回退到窗口成本使用",
},
{
"name": "recent_refresh",
"label": "额度刷新优先",
"description": "优先选即将刷新额度的账号",
"providers": ["codex", "kiro"],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据账号额度重置倒计时（next_reset / reset_seconds）",
},
{
"name": "single_account",
"label": "单号优先",
"description": "集中使用同一账号（反向 LRU）",
"providers": [],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": "distribution_mode",
"evidence_hint": "先按账号优先级（internal_priority），同级再按反向 LRU 集中",
},
{
"name": "team_first",
"label": "Team 优先",
"description": "优先消耗 Team 账号（依赖 plan_type）",
"providers": ["codex", "kiro"],
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": Value::Null,
"evidence_hint": "依据 plan_type（Team 账号优先调度）",
}
])
aether_provider_pool::build_admin_pool_scheduling_presets_payload()
}
pub fn admin_pool_batch_delete_task_parts(request_path: &str) -> Option<(String, String)> {

1
crates/aether-ai-serving/Cargo.toml
View File

@@ -10,6 +10,7 @@ description = "AI serving application contracts and ports for Aether"
aether-ai-formats.workspace = true
aether-contracts.workspace = true
aether-data-contracts.workspace = true
aether-pool-core.workspace = true
aether-scheduler-core.workspace = true
async-trait.workspace = true
http.workspace = true

45
crates/aether-ai-serving/src/lib.rs
View File

@@ -15,8 +15,6 @@ pub mod dto;
pub mod execution_path;
pub mod failure_diagnostic;
pub mod plan_payload;
pub mod pool_scheduler;
pub mod pool_scores;
pub mod ports;
pub mod ranking_metadata;
pub mod report_context;
@@ -24,6 +22,34 @@ pub mod request_body_diagnostics;
pub mod runtime_miss;
pub mod surface_spec;
pub use aether_pool_core::{
normalize_enabled_pool_presets, run_pool_scheduler, PoolCandidateFacts, PoolCandidateInput,
PoolCandidateOrchestration, PoolMemberSignals, PoolRuntimeState, PoolScheduledCandidate,
PoolSchedulerOutcome, PoolSchedulingConfig, PoolSchedulingPreset, PoolSkippedCandidate,
POOL_ACCOUNT_BLOCKED_SKIP_REASON, POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
POOL_COOLDOWN_SKIP_REASON, POOL_COST_LIMIT_REACHED_SKIP_REASON,
};
pub use aether_pool_core::{
normalize_enabled_pool_presets as normalize_enabled_ai_pool_presets,
run_pool_scheduler as run_ai_pool_scheduler, PoolCandidateFacts as AiPoolCandidateFacts,
PoolCandidateInput as AiPoolCandidateInput,
PoolCandidateOrchestration as AiPoolCandidateOrchestration,
PoolMemberSignals as AiPoolCatalogKeyContext, PoolRuntimeState as AiPoolRuntimeState,
PoolScheduledCandidate as AiPoolScheduledCandidate,
PoolSchedulerOutcome as AiPoolSchedulerOutcome, PoolSchedulingConfig as AiPoolSchedulingConfig,
PoolSchedulingPreset as AiPoolSchedulingPreset, PoolSkippedCandidate as AiPoolSkippedCandidate,
POOL_ACCOUNT_BLOCKED_SKIP_REASON as AI_POOL_ACCOUNT_BLOCKED_SKIP_REASON,
POOL_ACCOUNT_EXHAUSTED_SKIP_REASON as AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
POOL_COOLDOWN_SKIP_REASON as AI_POOL_COOLDOWN_SKIP_REASON,
POOL_COST_LIMIT_REACHED_SKIP_REASON as AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
};
pub use aether_pool_core::{
probe_freshness_score, probe_freshness_score_with_ttl, score_pool_member,
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreOutput,
PoolMemberScoreRules, PoolMemberScoreWeights, POOL_SCORE_VERSION,
PROBE_FAILURE_COOLDOWN_THRESHOLD, PROBE_FAILURE_PENALTY, PROBE_FRESHNESS_TTL_SECONDS,
REQUEST_FAILURE_PENALTY, UNSCHEDULABLE_SCORE_CAP,
};
pub use attempt_loop::{
run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt,
};
@@ -92,21 +118,6 @@ pub use failure_diagnostic::{CandidateFailureDiagnostic, CandidateFailureDiagnos
pub use plan_payload::{
build_ai_stream_execution_plan_payload, build_ai_sync_execution_plan_payload,
};
pub use pool_scheduler::{
normalize_enabled_ai_pool_presets, run_ai_pool_scheduler, AiPoolCandidateFacts,
AiPoolCandidateInput, AiPoolCandidateOrchestration, AiPoolCatalogKeyContext,
AiPoolRuntimeState, AiPoolScheduledCandidate, AiPoolSchedulerOutcome, AiPoolSchedulingConfig,
AiPoolSchedulingPreset, AiPoolSkippedCandidate, AI_POOL_ACCOUNT_BLOCKED_SKIP_REASON,
AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON, AI_POOL_COOLDOWN_SKIP_REASON,
AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
};
pub use pool_scores::{
probe_freshness_score, probe_freshness_score_with_ttl, score_pool_member,
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreOutput,
PoolMemberScoreRules, PoolMemberScoreWeights, POOL_SCORE_VERSION,
PROBE_FAILURE_COOLDOWN_THRESHOLD, PROBE_FAILURE_PENALTY, PROBE_FRESHNESS_TTL_SECONDS,
REQUEST_FAILURE_PENALTY, UNSCHEDULABLE_SCORE_CAP,
};
pub use ranking_metadata::append_ai_ranking_metadata_to_object;
pub use report_context::{
build_ai_execution_report_context, build_ai_report_context_original_request_echo,

345
crates/aether-oauth/src/provider/providers/kiro.rs
View File

@@ -11,6 +11,10 @@ use sha2::{Digest, Sha256};
use std::collections::BTreeMap;
pub const KIRO_PROVIDER_TYPE: &str = "kiro";
pub const DEFAULT_REGION: &str = "us-east-1";
pub const DEFAULT_KIRO_VERSION: &str = "0.3.210";
pub const DEFAULT_NODE_VERSION: &str = "22.21.1";
pub const DEFAULT_SYSTEM_VERSION: &str = "other#unknown";
const 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";
@@ -39,7 +43,8 @@ impl KiroAuthConfig {
auth_method: string_field(
object,
&["auth_method", "authMethod", "auth_type", "authType"],
),
)
.map(|value| normalize_kiro_auth_method(&value)),
refresh_token: string_field(object, &["refresh_token", "refreshToken"]),
expires_at: u64_field(object.get("expires_at"))
.or_else(|| u64_field(object.get("expiresAt"))),
@@ -93,7 +98,7 @@ impl KiroAuthConfig {
.or(self.region.as_deref())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("us-east-1")
.unwrap_or(DEFAULT_REGION)
}
pub fn effective_api_region(&self) -> &str {
@@ -101,7 +106,7 @@ impl KiroAuthConfig {
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("us-east-1")
.unwrap_or(DEFAULT_REGION)
}
pub fn effective_kiro_version(&self) -> &str {
@@ -109,40 +114,111 @@ impl KiroAuthConfig {
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("0.3.210")
.unwrap_or(DEFAULT_KIRO_VERSION)
}
pub fn effective_system_version(&self) -> &str {
self.system_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_SYSTEM_VERSION)
}
pub fn effective_node_version(&self) -> &str {
self.node_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_NODE_VERSION)
}
pub fn cached_access_token(&self) -> Option<&str> {
self.access_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn cached_access_token_requires_refresh(&self, skew_seconds: u64) -> bool {
let Some(expires_at) = self.expires_at else {
return self.can_refresh_access_token();
};
let now = current_unix_secs();
now >= expires_at.saturating_sub(skew_seconds)
}
pub fn is_idc_auth(&self) -> bool {
self.auth_method
let explicit_method = self
.auth_method
.as_deref()
.map(normalize_kiro_auth_method)
.unwrap_or_else(|| "social".to_string());
if explicit_method != "social" {
return matches!(explicit_method.as_str(), "idc" | "external_idp");
}
self.client_id
.as_deref()
.map(str::trim)
.map(str::to_ascii_lowercase)
.is_some_and(|value| matches!(value.as_str(), "idc" | "external_idp"))
|| (self
.client_id
.filter(|value| !value.is_empty())
.is_some()
&& self
.client_secret
.as_deref()
.is_some_and(|value| !value.trim().is_empty())
&& self
.client_secret
.as_deref()
.is_some_and(|value| !value.trim().is_empty()))
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
}
pub fn uses_external_idp_token_type(&self) -> bool {
self.auth_method
.as_deref()
.map(normalize_kiro_auth_method)
.as_deref()
== Some("external_idp")
}
pub fn profile_arn_for_payload(&self) -> Option<&str> {
if self.is_idc_auth() {
return None;
}
self.profile_arn
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn profile_arn_for_mcp(&self) -> Option<&str> {
self.profile_arn
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn can_refresh_access_token(&self) -> bool {
self.refresh_token
let refresh_token = self
.refresh_token
.as_deref()
.map(str::trim)
.filter(|value| value.len() >= 100 && !value.contains("..."))
.filter(|value| !value.is_empty())
.filter(|value| value.len() >= 100 && !value.contains("..."));
if refresh_token.is_none() {
return false;
}
if !self.is_idc_auth() {
return true;
}
self.client_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
&& (!self.is_idc_auth()
|| (self
.client_id
.as_deref()
.is_some_and(|value| !value.trim().is_empty())
&& self
.client_secret
.as_deref()
.is_some_and(|value| !value.trim().is_empty())))
&& self
.client_secret
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
}
}
@@ -163,7 +239,7 @@ impl KiroProviderOAuthAdapter {
self
}
async fn refresh_auth_config(
pub async fn refresh_auth_config(
&self,
executor: &dyn OAuthHttpExecutor,
ctx: &ProviderOAuthTransportContext,
@@ -436,7 +512,7 @@ pub fn generate_kiro_machine_id(
if let Some(machine_id) = auth_config
.machine_id
.as_deref()
.and_then(normalize_machine_id)
.and_then(normalize_kiro_machine_id)
{
return Some(machine_id);
}
@@ -497,7 +573,7 @@ fn resolve_expires_at(payload: &Value) -> u64 {
current_unix_secs().saturating_add(expires_in)
}
fn normalize_machine_id(raw: &str) -> Option<String> {
pub fn normalize_kiro_machine_id(raw: &str) -> Option<String> {
let raw = raw.trim();
if raw.len() == 64 && raw.bytes().all(|byte| byte.is_ascii_hexdigit()) {
return Some(raw.to_ascii_lowercase());
@@ -514,6 +590,26 @@ fn normalize_machine_id(raw: &str) -> Option<String> {
None
}
fn normalize_kiro_auth_method(raw: &str) -> String {
let value = raw.trim().to_ascii_lowercase();
match value.as_str() {
"" => "social".to_string(),
"builder-id"
| "builder_id"
| "builderid"
| "device"
| "device-auth"
| "device_authorization"
| "iam"
| "identity-center"
| "identity_center"
| "identitycenter"
| "idc" => "idc".to_string(),
"external-idp" | "external_idp" | "externalidp" => "external_idp".to_string(),
_ => value,
}
}
fn string_field(object: &serde_json::Map<String, Value>, keys: &[&str]) -> Option<String> {
keys.iter()
.find_map(|key| object.get(*key))
@@ -548,7 +644,51 @@ fn secret_fingerprint(value: &str) -> String {
#[cfg(test)]
mod tests {
use super::{generate_kiro_machine_id, KiroAuthConfig};
use super::{
generate_kiro_machine_id, KiroAuthConfig, KiroProviderOAuthAdapter, IDC_AMZ_USER_AGENT,
};
use crate::network::{OAuthHttpExecutor, OAuthHttpRequest, OAuthHttpResponse};
use crate::provider::ProviderOAuthTransportContext;
use async_trait::async_trait;
use serde_json::json;
use std::sync::{Arc, Mutex};
#[derive(Debug, Clone)]
struct StaticExecutor {
seen_request: Arc<Mutex<Option<OAuthHttpRequest>>>,
response: serde_json::Value,
}
#[async_trait]
impl OAuthHttpExecutor for StaticExecutor {
async fn execute(
&self,
request: OAuthHttpRequest,
) -> Result<OAuthHttpResponse, crate::core::OAuthError> {
*self.seen_request.lock().expect("mutex should lock") = Some(request);
Ok(OAuthHttpResponse {
status_code: 200,
body_text: self.response.to_string(),
json_body: Some(self.response.clone()),
})
}
}
fn test_ctx() -> ProviderOAuthTransportContext {
ProviderOAuthTransportContext {
provider_id: "provider-1".to_string(),
provider_type: "kiro".to_string(),
endpoint_id: None,
key_id: Some("key-1".to_string()),
auth_type: Some("oauth".to_string()),
decrypted_api_key: None,
decrypted_auth_config: None,
provider_config: None,
endpoint_config: None,
key_config: None,
network: crate::network::OAuthNetworkContext::provider_operation(None),
}
}
#[test]
fn normalizes_kiro_uuid_machine_id() {
@@ -573,4 +713,151 @@ mod tests {
Some("123e4567e89b12d3a456426614174000123e4567e89b12d3a456426614174000")
);
}
#[tokio::test]
async fn refreshes_social_auth_config_with_provider_adapter() {
let seen_request = Arc::new(Mutex::new(None));
let executor = StaticExecutor {
seen_request: Arc::clone(&seen_request),
response: json!({
"accessToken": "new-social-access-token",
"refreshToken": "s".repeat(120),
"expiresIn": 3600,
"profileArn": "arn:aws:bedrock:demo"
}),
};
let adapter = KiroProviderOAuthAdapter::default()
.with_refresh_base_urls(Some("https://auth.example.test".to_string()), None);
let auth_config = KiroAuthConfig {
auth_method: Some("social".to_string()),
refresh_token: Some("r".repeat(120)),
expires_at: Some(1),
profile_arn: None,
region: None,
auth_region: None,
api_region: None,
client_id: None,
client_secret: None,
machine_id: Some("123e4567-e89b-12d3-a456-426614174000".to_string()),
kiro_version: Some("1.2.3".to_string()),
system_version: None,
node_version: None,
access_token: None,
};
let refreshed = adapter
.refresh_auth_config(&executor, &test_ctx(), &auth_config)
.await
.expect("social refresh should succeed");
assert_eq!(
refreshed.access_token.as_deref(),
Some("new-social-access-token")
);
let expected_rotated_refresh_token = "s".repeat(120);
assert_eq!(
refreshed.refresh_token.as_deref(),
Some(expected_rotated_refresh_token.as_str())
);
assert_eq!(
refreshed.profile_arn.as_deref(),
Some("arn:aws:bedrock:demo")
);
assert!(refreshed.expires_at.is_some());
let seen = seen_request
.lock()
.expect("mutex should lock")
.clone()
.expect("request should be captured");
assert_eq!(seen.request_id, "provider-oauth:kiro-social-refresh");
assert_eq!(seen.method, reqwest::Method::POST);
assert_eq!(seen.url, "https://auth.example.test/refreshToken");
assert_eq!(
seen.headers.get("user-agent").map(String::as_str),
Some("KiroIDE-1.2.3-123e4567e89b12d3a456426614174000123e4567e89b12d3a456426614174000")
);
let expected_refresh_token = "r".repeat(120);
assert_eq!(
seen.json_body
.as_ref()
.and_then(|body| body.get("refreshToken"))
.and_then(serde_json::Value::as_str),
Some(expected_refresh_token.as_str())
);
}
#[tokio::test]
async fn refreshes_idc_auth_config_with_provider_adapter() {
let seen_request = Arc::new(Mutex::new(None));
let executor = StaticExecutor {
seen_request: Arc::clone(&seen_request),
response: json!({
"accessToken": "new-idc-access-token",
"refreshToken": "i".repeat(120),
"expiresIn": 1800
}),
};
let adapter = KiroProviderOAuthAdapter::default()
.with_refresh_base_urls(None, Some("https://idc.example.test".to_string()));
let auth_config = KiroAuthConfig {
auth_method: Some("identity_center".to_string()),
refresh_token: Some("r".repeat(120)),
expires_at: Some(1),
profile_arn: Some("arn:aws:bedrock:demo".to_string()),
region: None,
auth_region: None,
api_region: None,
client_id: Some("client-id".to_string()),
client_secret: Some("client-secret".to_string()),
machine_id: None,
kiro_version: None,
system_version: None,
node_version: None,
access_token: None,
};
let refreshed = adapter
.refresh_auth_config(&executor, &test_ctx(), &auth_config)
.await
.expect("idc refresh should succeed");
assert_eq!(
refreshed.access_token.as_deref(),
Some("new-idc-access-token")
);
let expected_rotated_refresh_token = "i".repeat(120);
assert_eq!(
refreshed.refresh_token.as_deref(),
Some(expected_rotated_refresh_token.as_str())
);
assert!(refreshed.machine_id.is_some());
assert!(refreshed.expires_at.is_some());
let seen = seen_request
.lock()
.expect("mutex should lock")
.clone()
.expect("request should be captured");
assert_eq!(seen.request_id, "provider-oauth:kiro-idc-refresh");
assert_eq!(seen.method, reqwest::Method::POST);
assert_eq!(seen.url, "https://idc.example.test/token");
assert_eq!(
seen.headers.get("user-agent").map(String::as_str),
Some("node")
);
assert_eq!(
seen.headers.get("x-amz-user-agent").map(String::as_str),
Some(IDC_AMZ_USER_AGENT)
);
let body = seen.json_body.expect("json body should exist");
assert_eq!(body.get("grantType"), Some(&json!("refresh_token")));
assert_eq!(body.get("clientId"), Some(&json!("client-id")));
assert_eq!(body.get("clientSecret"), Some(&json!("client-secret")));
let expected_refresh_token = "r".repeat(120);
assert_eq!(
body.get("refreshToken").and_then(serde_json::Value::as_str),
Some(expected_refresh_token.as_str())
);
}
}

4
crates/aether-oauth/src/provider/providers/mod.rs
View File

@@ -9,5 +9,7 @@ pub use generic::{
GenericProviderOAuthAdapter, GenericProviderOAuthTemplate, GENERIC_PROVIDER_OAUTH_TEMPLATES,
};
pub use kiro::{
generate_kiro_machine_id, KiroAuthConfig, KiroProviderOAuthAdapter, KIRO_PROVIDER_TYPE,
generate_kiro_machine_id, normalize_kiro_machine_id, KiroAuthConfig, KiroProviderOAuthAdapter,
DEFAULT_KIRO_VERSION, DEFAULT_NODE_VERSION, DEFAULT_REGION, DEFAULT_SYSTEM_VERSION,
KIRO_PROVIDER_TYPE,
};

11
crates/aether-pool-core/Cargo.toml Normal file
View File

@@ -0,0 +1,11 @@
[package]
name = "aether-pool-core"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Provider-independent pool scheduling and scoring primitives for Aether"
[dependencies]
aether-data-contracts.workspace = true
serde_json.workspace = true

17
crates/aether-pool-core/src/lib.rs Normal file
View File

@@ -0,0 +1,17 @@
mod scheduler;
mod scoring;
pub use scheduler::{
normalize_enabled_pool_presets, run_pool_scheduler, PoolCandidateFacts, PoolCandidateInput,
PoolCandidateOrchestration, PoolMemberSignals, PoolRuntimeState, PoolScheduledCandidate,
PoolSchedulerOutcome, PoolSchedulingConfig, PoolSchedulingPreset, PoolSkippedCandidate,
POOL_ACCOUNT_BLOCKED_SKIP_REASON, POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
POOL_COOLDOWN_SKIP_REASON, POOL_COST_LIMIT_REACHED_SKIP_REASON,
};
pub use scoring::{
probe_freshness_score, probe_freshness_score_with_ttl, score_pool_member,
score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreOutput,
PoolMemberScoreRules, PoolMemberScoreWeights, POOL_SCORE_VERSION,
PROBE_FAILURE_COOLDOWN_THRESHOLD, PROBE_FAILURE_PENALTY, PROBE_FRESHNESS_TTL_SECONDS,
REQUEST_FAILURE_PENALTY, UNSCHEDULABLE_SCORE_CAP,
};

344
crates/aether-ai-serving/src/pool_scheduler.rs → crates/aether-pool-core/src/scheduler.rs
View File

@@ -2,28 +2,28 @@ use std::cmp::Ordering;
use std::collections::{btree_map::Entry, BTreeMap, BTreeSet};
use std::hash::{Hash, Hasher};
pub const AI_POOL_ACCOUNT_BLOCKED_SKIP_REASON: &str = "pool_account_blocked";
pub const AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON: &str = "pool_account_exhausted";
pub const AI_POOL_COOLDOWN_SKIP_REASON: &str = "pool_cooldown";
pub const AI_POOL_COST_LIMIT_REACHED_SKIP_REASON: &str = "pool_cost_limit_reached";
pub const POOL_ACCOUNT_BLOCKED_SKIP_REASON: &str = "pool_account_blocked";
pub const POOL_ACCOUNT_EXHAUSTED_SKIP_REASON: &str = "pool_account_exhausted";
pub const POOL_COOLDOWN_SKIP_REASON: &str = "pool_cooldown";
pub const POOL_COST_LIMIT_REACHED_SKIP_REASON: &str = "pool_cost_limit_reached";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AiPoolSchedulingPreset {
pub struct PoolSchedulingPreset {
pub preset: String,
pub enabled: bool,
pub mode: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AiPoolSchedulingConfig {
pub scheduling_presets: Vec<AiPoolSchedulingPreset>,
pub struct PoolSchedulingConfig {
pub scheduling_presets: Vec<PoolSchedulingPreset>,
pub lru_enabled: bool,
pub skip_exhausted_accounts: bool,
pub cost_limit_per_key_tokens: Option<u64>,
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct AiPoolRuntimeState {
pub struct PoolRuntimeState {
pub sticky_bound_key_id: Option<String>,
pub cooldown_reason_by_key: BTreeMap<String, String>,
pub cost_window_usage_by_key: BTreeMap<String, u64>,
@@ -32,8 +32,8 @@ pub struct AiPoolRuntimeState {
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct AiPoolCatalogKeyContext {
pub oauth_plan_type: Option<String>,
pub struct PoolMemberSignals {
pub plan_tier: Option<String>,
pub quota_usage_ratio: Option<f64>,
pub quota_reset_seconds: Option<f64>,
pub account_blocked: bool,
@@ -44,47 +44,46 @@ pub struct AiPoolCatalogKeyContext {
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AiPoolCandidateFacts {
pub struct PoolCandidateFacts {
pub provider_id: String,
pub endpoint_id: String,
pub model_id: String,
pub selected_provider_model_name: String,
pub provider_api_format: String,
pub provider_type: String,
pub key_id: String,
pub key_internal_priority: i32,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AiPoolCandidateOrchestration {
pub struct PoolCandidateOrchestration {
pub candidate_group_id: Option<String>,
pub pool_key_index: Option<u32>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct AiPoolCandidateInput<Candidate> {
pub struct PoolCandidateInput<Candidate> {
pub candidate: Candidate,
pub facts: AiPoolCandidateFacts,
pub pool_config: Option<AiPoolSchedulingConfig>,
pub key_context: AiPoolCatalogKeyContext,
pub facts: PoolCandidateFacts,
pub pool_config: Option<PoolSchedulingConfig>,
pub key_context: PoolMemberSignals,
}
#[derive(Debug, Clone, PartialEq)]
pub struct AiPoolScheduledCandidate<Candidate> {
pub struct PoolScheduledCandidate<Candidate> {
pub candidate: Candidate,
pub orchestration: AiPoolCandidateOrchestration,
pub orchestration: PoolCandidateOrchestration,
}
#[derive(Debug, Clone, PartialEq)]
pub struct AiPoolSkippedCandidate<Candidate> {
pub struct PoolSkippedCandidate<Candidate> {
pub candidate: Candidate,
pub skip_reason: &'static str,
}
#[derive(Debug, Clone, PartialEq)]
pub struct AiPoolSchedulerOutcome<Candidate> {
pub candidates: Vec<AiPoolScheduledCandidate<Candidate>>,
pub skipped_candidates: Vec<AiPoolSkippedCandidate<Candidate>>,
pub struct PoolSchedulerOutcome<Candidate> {
pub candidates: Vec<PoolScheduledCandidate<Candidate>>,
pub skipped_candidates: Vec<PoolSkippedCandidate<Candidate>>,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
@@ -103,13 +102,13 @@ struct NormalizedPoolPreset {
mode: Option<String>,
}
pub fn run_ai_pool_scheduler<Candidate>(
candidates: Vec<AiPoolCandidateInput<Candidate>>,
runtime_by_provider: &BTreeMap<String, AiPoolRuntimeState>,
pub fn run_pool_scheduler<Candidate>(
candidates: Vec<PoolCandidateInput<Candidate>>,
runtime_by_provider: &BTreeMap<String, PoolRuntimeState>,
load_balance_seed_nonce: &str,
) -> AiPoolSchedulerOutcome<Candidate> {
) -> PoolSchedulerOutcome<Candidate> {
let mut group_order = Vec::new();
let mut groups = BTreeMap::<PoolGroupKey, Vec<AiPoolCandidateInput<Candidate>>>::new();
let mut groups = BTreeMap::<PoolGroupKey, Vec<PoolCandidateInput<Candidate>>>::new();
for candidate in candidates {
let pool_enabled = candidate.pool_config.is_some();
@@ -127,20 +126,20 @@ pub fn run_ai_pool_scheduler<Candidate>(
let mut reordered = Vec::new();
let mut skipped = Vec::new();
let default_runtime = AiPoolRuntimeState::default();
let default_runtime = PoolRuntimeState::default();
for group_key in group_order {
let Some(group) = groups.remove(&group_key) else {
continue;
};
let candidate_group_id = ai_pool_candidate_group_id(&group_key);
let candidate_group_id = pool_candidate_group_id(&group_key);
let Some(pool_config) = group
.first()
.expect("group should exist")
.pool_config
.clone()
else {
reordered.extend(annotate_ai_pool_candidates(
reordered.extend(annotate_pool_candidates(
group,
candidate_group_id.as_str(),
false,
@@ -161,14 +160,14 @@ pub fn run_ai_pool_scheduler<Candidate>(
skipped.extend(outcome.skipped_candidates);
}
AiPoolSchedulerOutcome {
PoolSchedulerOutcome {
candidates: reordered,
skipped_candidates: skipped,
}
}
fn pool_group_key<Candidate>(
candidate: &AiPoolCandidateInput<Candidate>,
candidate: &PoolCandidateInput<Candidate>,
pool_enabled: bool,
) -> PoolGroupKey {
PoolGroupKey {
@@ -181,7 +180,7 @@ fn pool_group_key<Candidate>(
}
}
fn ai_pool_candidate_group_id(group_key: &PoolGroupKey) -> String {
fn pool_candidate_group_id(group_key: &PoolGroupKey) -> String {
format!(
"provider={}|endpoint={}|model={}|selected_model={}|api_format={}|singleton_key={}",
group_key.provider_id,
@@ -194,18 +193,13 @@ fn ai_pool_candidate_group_id(group_key: &PoolGroupKey) -> String {
}
fn schedule_pool_group<Candidate>(
group: Vec<AiPoolCandidateInput<Candidate>>,
pool_config: &AiPoolSchedulingConfig,
runtime: &AiPoolRuntimeState,
group: Vec<PoolCandidateInput<Candidate>>,
pool_config: &PoolSchedulingConfig,
runtime: &PoolRuntimeState,
candidate_group_id: &str,
load_balance_seed_nonce: &str,
) -> AiPoolSchedulerOutcome<Candidate> {
let provider_type = group
.first()
.map(|candidate| candidate.facts.provider_type.trim().to_ascii_lowercase())
.unwrap_or_default();
let active_presets =
normalize_enabled_pool_presets(&pool_config.scheduling_presets, provider_type.as_str());
) -> PoolSchedulerOutcome<Candidate> {
let active_presets = normalize_enabled_pool_preset_entries(&pool_config.scheduling_presets);
let lru_distribution_enabled = pool_config.lru_enabled
&& !active_presets
.iter()
@@ -223,25 +217,25 @@ fn schedule_pool_group<Candidate>(
.or(item.key_context.latency_avg_ms);
if item.key_context.account_blocked {
skipped.push(AiPoolSkippedCandidate {
skipped.push(PoolSkippedCandidate {
candidate: item.candidate,
skip_reason: AI_POOL_ACCOUNT_BLOCKED_SKIP_REASON,
skip_reason: POOL_ACCOUNT_BLOCKED_SKIP_REASON,
});
continue;
}
if pool_config.skip_exhausted_accounts && item.key_context.quota_exhausted {
skipped.push(AiPoolSkippedCandidate {
skipped.push(PoolSkippedCandidate {
candidate: item.candidate,
skip_reason: AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
skip_reason: POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
});
continue;
}
if runtime.cooldown_reason_by_key.contains_key(&key_id) {
skipped.push(AiPoolSkippedCandidate {
skipped.push(PoolSkippedCandidate {
candidate: item.candidate,
skip_reason: AI_POOL_COOLDOWN_SKIP_REASON,
skip_reason: POOL_COOLDOWN_SKIP_REASON,
});
continue;
}
@@ -250,9 +244,9 @@ fn schedule_pool_group<Candidate>(
.cost_limit_per_key_tokens
.is_some_and(|limit| runtime_cost_usage(runtime, key_id.as_str()) >= limit)
{
skipped.push(AiPoolSkippedCandidate {
skipped.push(PoolSkippedCandidate {
candidate: item.candidate,
skip_reason: AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
skip_reason: POOL_COST_LIMIT_REACHED_SKIP_REASON,
});
continue;
}
@@ -269,7 +263,7 @@ fn schedule_pool_group<Candidate>(
}
if available.is_empty() {
return AiPoolSchedulerOutcome {
return PoolSchedulerOutcome {
candidates: Vec::new(),
skipped_candidates: skipped,
};
@@ -295,7 +289,6 @@ fn schedule_pool_group<Candidate>(
&active_presets,
lru_distribution_enabled,
group_sort_seed(
provider_type.as_str(),
available.first().map(|item| &item.item.facts),
load_balance_seed_nonce,
)
@@ -324,23 +317,23 @@ fn schedule_pool_group<Candidate>(
}
ordered.extend(available.into_iter().map(|item| item.item));
AiPoolSchedulerOutcome {
candidates: annotate_ai_pool_candidates(ordered, candidate_group_id, true),
PoolSchedulerOutcome {
candidates: annotate_pool_candidates(ordered, candidate_group_id, true),
skipped_candidates: skipped,
}
}
fn annotate_ai_pool_candidates<Candidate>(
candidates: Vec<AiPoolCandidateInput<Candidate>>,
fn annotate_pool_candidates<Candidate>(
candidates: Vec<PoolCandidateInput<Candidate>>,
candidate_group_id: &str,
pool_enabled: bool,
) -> Vec<AiPoolScheduledCandidate<Candidate>> {
) -> Vec<PoolScheduledCandidate<Candidate>> {
candidates
.into_iter()
.enumerate()
.map(|(index, item)| AiPoolScheduledCandidate {
.map(|(index, item)| PoolScheduledCandidate {
candidate: item.candidate,
orchestration: AiPoolCandidateOrchestration {
orchestration: PoolCandidateOrchestration {
candidate_group_id: Some(candidate_group_id.to_string()),
pool_key_index: pool_enabled.then_some(index as u32),
},
@@ -350,7 +343,7 @@ fn annotate_ai_pool_candidates<Candidate>(
#[derive(Debug)]
struct PoolGroupCandidateOrdering<Candidate> {
item: AiPoolCandidateInput<Candidate>,
item: PoolCandidateInput<Candidate>,
original_index: usize,
lru_score: Option<f64>,
cost_usage: u64,
@@ -473,7 +466,7 @@ fn plan_ranks<Candidate>(
(
item.item.facts.key_id.clone(),
Some(plan_priority_score(
item.item.key_context.oauth_plan_type.as_deref(),
item.item.key_context.plan_tier.as_deref(),
mode,
)),
)
@@ -574,19 +567,18 @@ fn load_balance_ranks<Candidate>(
}
fn group_sort_seed(
provider_type: &str,
candidate: Option<&AiPoolCandidateFacts>,
candidate: Option<&PoolCandidateFacts>,
load_balance_seed_nonce: &str,
) -> String {
match candidate {
Some(candidate) => format!(
"{provider_type}:{}:{}:{}:{}:{load_balance_seed_nonce}",
"{}:{}:{}:{}:{load_balance_seed_nonce}",
candidate.provider_id,
candidate.endpoint_id,
candidate.model_id,
candidate.selected_provider_model_name,
),
None => format!("{provider_type}:{load_balance_seed_nonce}"),
None => load_balance_seed_nonce.to_string(),
}
}
@@ -736,21 +728,16 @@ fn plan_priority_score(plan_type: Option<&str>, mode: Option<&str>) -> f64 {
}
}
pub fn normalize_enabled_ai_pool_presets(
scheduling_presets: &[AiPoolSchedulingPreset],
provider_type: &str,
) -> Vec<String> {
normalize_enabled_pool_presets(scheduling_presets, provider_type)
pub fn normalize_enabled_pool_presets(scheduling_presets: &[PoolSchedulingPreset]) -> Vec<String> {
normalize_enabled_pool_preset_entries(scheduling_presets)
.into_iter()
.map(|preset| preset.preset)
.collect()
}
fn normalize_enabled_pool_presets(
scheduling_presets: &[AiPoolSchedulingPreset],
provider_type: &str,
fn normalize_enabled_pool_preset_entries(
scheduling_presets: &[PoolSchedulingPreset],
) -> Vec<NormalizedPoolPreset> {
let provider_type = provider_type.trim().to_ascii_lowercase();
let mut entries = Vec::<(usize, String, bool, Option<String>)>::new();
let mut seen = BTreeSet::new();
@@ -762,20 +749,11 @@ fn normalize_enabled_pool_presets(
entries.push((index, preset, item.enabled, item.mode.clone()));
}
if provider_type == "codex"
&& !entries.is_empty()
&& entries
.iter()
.all(|(_, preset, _, _)| preset != "recent_refresh")
{
entries.push((entries.len(), "recent_refresh".to_string(), true, None));
}
let mut distribution_mode = None::<(usize, String, Option<String>)>;
let mut strategy_presets = Vec::<(usize, String, Option<String>)>::new();
for (index, preset, enabled, mode) in entries {
if !enabled || !pool_preset_supported_for_provider(&preset, &provider_type) {
if !enabled {
continue;
}
@@ -809,15 +787,6 @@ fn normalize_enabled_pool_presets(
normalized
}
fn pool_preset_supported_for_provider(preset: &str, provider_type: &str) -> bool {
match preset {
"free_first" | "plus_first" | "pro_first" | "recent_refresh" | "team_first" => {
matches!(provider_type, "codex" | "kiro")
}
_ => true,
}
}
fn pool_preset_mutex_group(preset: &str) -> Option<&'static str> {
match preset {
"lru" | "cache_affinity" | "load_balance" | "single_account" => Some("distribution_mode"),
@@ -825,11 +794,11 @@ fn pool_preset_mutex_group(preset: &str) -> Option<&'static str> {
}
}
fn runtime_lru_score(runtime: &AiPoolRuntimeState, key_id: &str) -> Option<f64> {
fn runtime_lru_score(runtime: &PoolRuntimeState, key_id: &str) -> Option<f64> {
runtime.lru_score_by_key.get(key_id).copied()
}
fn runtime_cost_usage(runtime: &AiPoolRuntimeState, key_id: &str) -> u64 {
fn runtime_cost_usage(runtime: &PoolRuntimeState, key_id: &str) -> u64 {
runtime
.cost_window_usage_by_key
.get(key_id)
@@ -849,16 +818,16 @@ mod tests {
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AiPoolRuntimeState {
PoolRuntimeState {
lru_score_by_key: BTreeMap::from([
("key-pool-a".to_string(), 20.0),
("key-pool-b".to_string(), 10.0),
]),
..AiPoolRuntimeState::default()
..PoolRuntimeState::default()
},
)]);
let outcome = run_ai_pool_scheduler(
let outcome = run_pool_scheduler(
vec![pool_first, other, pool_second],
&runtime_by_provider,
"seed",
@@ -887,17 +856,17 @@ mod tests {
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AiPoolRuntimeState {
PoolRuntimeState {
cooldown_reason_by_key: BTreeMap::from([(
"key-cooldown".to_string(),
"429".to_string(),
)]),
cost_window_usage_by_key: BTreeMap::from([("key-cost".to_string(), 100)]),
..AiPoolRuntimeState::default()
..PoolRuntimeState::default()
},
)]);
let outcome = run_ai_pool_scheduler(
let outcome = run_pool_scheduler(
vec![key_ready, key_cooldown, key_cost],
&runtime_by_provider,
"seed",
@@ -927,13 +896,13 @@ mod tests {
#[test]
fn pool_scheduler_promotes_sticky_hit_before_other_sorted_keys() {
let key_a = sample_candidate("provider-pool", "endpoint-1", "key-a", 10, true)
.with_presets(vec![AiPoolSchedulingPreset {
.with_presets(vec![PoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
}]);
let key_b = sample_candidate("provider-pool", "endpoint-1", "key-b", 10, true)
.with_presets(vec![AiPoolSchedulingPreset {
.with_presets(vec![PoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
@@ -941,17 +910,17 @@ mod tests {
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AiPoolRuntimeState {
PoolRuntimeState {
sticky_bound_key_id: Some("key-a".to_string()),
lru_score_by_key: BTreeMap::from([
("key-a".to_string(), 50.0),
("key-b".to_string(), 10.0),
]),
..AiPoolRuntimeState::default()
..PoolRuntimeState::default()
},
)]);
let outcome = run_ai_pool_scheduler(vec![key_a, key_b], &runtime_by_provider, "seed");
let outcome = run_pool_scheduler(vec![key_a, key_b], &runtime_by_provider, "seed");
assert!(outcome.skipped_candidates.is_empty());
assert_eq!(
@@ -967,13 +936,13 @@ mod tests {
#[test]
fn load_balance_distribution_ignores_sticky_hit() {
let key_a = sample_candidate("provider-pool", "endpoint-1", "key-a", 10, true)
.with_presets(vec![AiPoolSchedulingPreset {
.with_presets(vec![PoolSchedulingPreset {
preset: "load_balance".to_string(),
enabled: true,
mode: None,
}]);
let key_b = sample_candidate("provider-pool", "endpoint-1", "key-b", 10, true)
.with_presets(vec![AiPoolSchedulingPreset {
.with_presets(vec![PoolSchedulingPreset {
preset: "load_balance".to_string(),
enabled: true,
mode: None,
@@ -981,20 +950,20 @@ mod tests {
let nonce = (0..1000)
.map(|index| format!("seed-{index}"))
.find(|nonce| {
let group_seed = format!("codex:provider-pool:endpoint-1:model-1:gpt-5:{nonce}");
let group_seed = format!("provider-pool:endpoint-1:model-1:gpt-5:{nonce}");
stable_hash_score(format!("{group_seed}:key-b").as_str())
< stable_hash_score(format!("{group_seed}:key-a").as_str())
})
.expect("test seed should exist");
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AiPoolRuntimeState {
PoolRuntimeState {
sticky_bound_key_id: Some("key-a".to_string()),
..AiPoolRuntimeState::default()
..PoolRuntimeState::default()
},
)]);
let outcome = run_ai_pool_scheduler(vec![key_a, key_b], &runtime_by_provider, &nonce);
let outcome = run_pool_scheduler(vec![key_a, key_b], &runtime_by_provider, &nonce);
assert!(outcome.skipped_candidates.is_empty());
assert_eq!(
@@ -1010,21 +979,21 @@ mod tests {
#[test]
fn pool_scheduler_uses_plan_preset_with_catalog_context() {
let key_free = sample_candidate("provider-pool", "endpoint-1", "key-free", 10, true)
.with_presets(vec![AiPoolSchedulingPreset {
.with_presets(vec![PoolSchedulingPreset {
preset: "plus_first".to_string(),
enabled: true,
mode: None,
}])
.with_plan("free");
let key_plus = sample_candidate("provider-pool", "endpoint-1", "key-plus", 10, true)
.with_presets(vec![AiPoolSchedulingPreset {
.with_presets(vec![PoolSchedulingPreset {
preset: "plus_first".to_string(),
enabled: true,
mode: None,
}])
.with_plan("plus");
let outcome = run_ai_pool_scheduler(vec![key_free, key_plus], &BTreeMap::new(), "seed");
let outcome = run_pool_scheduler(vec![key_free, key_plus], &BTreeMap::new(), "seed");
assert!(outcome.skipped_candidates.is_empty());
assert_eq!(
@@ -1042,12 +1011,12 @@ mod tests {
let key_cache_hit =
sample_candidate("provider-pool", "endpoint-1", "key-cache-hit", 50, true)
.with_presets(vec![
AiPoolSchedulingPreset {
PoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
PoolSchedulingPreset {
preset: "priority_first".to_string(),
enabled: true,
mode: None,
@@ -1056,12 +1025,12 @@ mod tests {
let key_high_priority =
sample_candidate("provider-pool", "endpoint-1", "key-high-priority", 10, true)
.with_presets(vec![
AiPoolSchedulingPreset {
PoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
PoolSchedulingPreset {
preset: "priority_first".to_string(),
enabled: true,
mode: None,
@@ -1070,16 +1039,16 @@ mod tests {
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AiPoolRuntimeState {
PoolRuntimeState {
lru_score_by_key: BTreeMap::from([
("key-cache-hit".to_string(), 200.0),
("key-high-priority".to_string(), 10.0),
]),
..AiPoolRuntimeState::default()
..PoolRuntimeState::default()
},
)]);
let outcome = run_ai_pool_scheduler(
let outcome = run_pool_scheduler(
vec![key_cache_hit, key_high_priority],
&runtime_by_provider,
"seed",
@@ -1101,12 +1070,12 @@ mod tests {
let key_random_first =
sample_candidate("provider-pool", "endpoint-1", "key-random-first", 50, true)
.with_presets(vec![
AiPoolSchedulingPreset {
PoolSchedulingPreset {
preset: "load_balance".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
PoolSchedulingPreset {
preset: "priority_first".to_string(),
enabled: true,
mode: None,
@@ -1115,12 +1084,12 @@ mod tests {
let key_high_priority =
sample_candidate("provider-pool", "endpoint-1", "key-high-priority", 10, true)
.with_presets(vec![
AiPoolSchedulingPreset {
PoolSchedulingPreset {
preset: "load_balance".to_string(),
enabled: true,
mode: None,
},
AiPoolSchedulingPreset {
PoolSchedulingPreset {
preset: "priority_first".to_string(),
enabled: true,
mode: None,
@@ -1129,13 +1098,13 @@ mod tests {
let nonce = (0..1000)
.map(|index| format!("seed-{index}"))
.find(|nonce| {
let group_seed = format!("codex:provider-pool:endpoint-1:model-1:gpt-5:{nonce}");
let group_seed = format!("provider-pool:endpoint-1:model-1:gpt-5:{nonce}");
stable_hash_score(format!("{group_seed}:key-random-first").as_str())
< stable_hash_score(format!("{group_seed}:key-high-priority").as_str())
})
.expect("test seed should exist");
let outcome = run_ai_pool_scheduler(
let outcome = run_pool_scheduler(
vec![key_random_first, key_high_priority],
&BTreeMap::new(),
nonce.as_str(),
@@ -1156,7 +1125,7 @@ mod tests {
fn single_account_distribution_orders_by_priority_then_reverse_lru() {
let key_priority_old =
sample_candidate("provider-pool", "endpoint-1", "key-priority-old", 10, true)
.with_presets(vec![AiPoolSchedulingPreset {
.with_presets(vec![PoolSchedulingPreset {
preset: "single_account".to_string(),
enabled: true,
mode: None,
@@ -1168,7 +1137,7 @@ mod tests {
10,
true,
)
.with_presets(vec![AiPoolSchedulingPreset {
.with_presets(vec![PoolSchedulingPreset {
preset: "single_account".to_string(),
enabled: true,
mode: None,
@@ -1180,7 +1149,7 @@ mod tests {
50,
true,
)
.with_presets(vec![AiPoolSchedulingPreset {
.with_presets(vec![PoolSchedulingPreset {
preset: "single_account".to_string(),
enabled: true,
mode: None,
@@ -1188,17 +1157,17 @@ mod tests {
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AiPoolRuntimeState {
PoolRuntimeState {
lru_score_by_key: BTreeMap::from([
("key-priority-old".to_string(), 10.0),
("key-priority-recent".to_string(), 200.0),
("key-lower-priority-recent".to_string(), 500.0),
]),
..AiPoolRuntimeState::default()
..PoolRuntimeState::default()
},
)]);
let outcome = run_ai_pool_scheduler(
let outcome = run_pool_scheduler(
vec![
key_priority_old,
key_lower_priority_recent,
@@ -1225,57 +1194,51 @@ 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 = normalize_enabled_pool_presets(&[
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,
},
]);
assert_eq!(presets, ["single_account", "priority_first"]);
}
#[test]
fn normalizes_lru_as_mutually_exclusive_distribution_mode() {
let presets = normalize_enabled_ai_pool_presets(
&[
AiPoolSchedulingPreset {
preset: "lru".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 = normalize_enabled_pool_presets(&[
PoolSchedulingPreset {
preset: "lru".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,
},
]);
assert_eq!(presets, ["priority_first"]);
}
@@ -1286,37 +1249,36 @@ mod tests {
key_id: &str,
internal_priority: i32,
pool_enabled: bool,
) -> AiPoolCandidateInput<String> {
let pool_config = pool_enabled.then(|| AiPoolSchedulingConfig {
) -> PoolCandidateInput<String> {
let pool_config = pool_enabled.then(|| PoolSchedulingConfig {
scheduling_presets: Vec::new(),
lru_enabled: true,
skip_exhausted_accounts: false,
cost_limit_per_key_tokens: None,
});
AiPoolCandidateInput {
PoolCandidateInput {
candidate: key_id.to_string(),
facts: AiPoolCandidateFacts {
facts: PoolCandidateFacts {
provider_id: provider_id.to_string(),
endpoint_id: endpoint_id.to_string(),
model_id: "model-1".to_string(),
selected_provider_model_name: "gpt-5".to_string(),
provider_api_format: "openai:chat".to_string(),
provider_type: "codex".to_string(),
key_id: key_id.to_string(),
key_internal_priority: internal_priority,
},
pool_config,
key_context: AiPoolCatalogKeyContext::default(),
key_context: PoolMemberSignals::default(),
}
}
trait TestCandidateExt {
fn with_cost_limit(self, limit: u64) -> Self;
fn with_presets(self, presets: Vec<AiPoolSchedulingPreset>) -> Self;
fn with_presets(self, presets: Vec<PoolSchedulingPreset>) -> Self;
fn with_plan(self, plan: &str) -> Self;
}
impl TestCandidateExt for AiPoolCandidateInput<String> {
impl TestCandidateExt for PoolCandidateInput<String> {
fn with_cost_limit(mut self, limit: u64) -> Self {
if let Some(config) = self.pool_config.as_mut() {
config.cost_limit_per_key_tokens = Some(limit);
@@ -1324,7 +1286,7 @@ mod tests {
self
}
fn with_presets(mut self, presets: Vec<AiPoolSchedulingPreset>) -> Self {
fn with_presets(mut self, presets: Vec<PoolSchedulingPreset>) -> Self {
if let Some(config) = self.pool_config.as_mut() {
config.scheduling_presets = presets;
}
@@ -1332,7 +1294,7 @@ mod tests {
}
fn with_plan(mut self, plan: &str) -> Self {
self.key_context.oauth_plan_type = Some(plan.to_string());
self.key_context.plan_tier = Some(plan.to_string());
self
}
}

0
crates/aether-ai-serving/src/pool_scores.rs → crates/aether-pool-core/src/scoring.rs
View File

14
crates/aether-provider-pool/Cargo.toml Normal file
View File

@@ -0,0 +1,14 @@
[package]
name = "aether-provider-pool"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Provider-specific pool behavior adapters for Aether"
[dependencies]
aether-data-contracts.workspace = true
aether-pool-core.workspace = true
serde_json.workspace = true
url.workspace = true
uuid.workspace = true

23
crates/aether-provider-pool/src/capability.rs Normal file
View File

@@ -0,0 +1,23 @@
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProviderPoolCapability {
PlanTier,
QuotaReset,
QuotaRefresh,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct ProviderPoolCapabilities {
pub plan_tier: bool,
pub quota_reset: bool,
pub quota_refresh: bool,
}
impl ProviderPoolCapabilities {
pub fn supports(self, capability: ProviderPoolCapability) -> bool {
match capability {
ProviderPoolCapability::PlanTier => self.plan_tier,
ProviderPoolCapability::QuotaReset => self.quota_reset,
ProviderPoolCapability::QuotaRefresh => self.quota_refresh,
}
}
}

350
crates/aether-provider-pool/src/lib.rs Normal file
View File

@@ -0,0 +1,350 @@
mod capability;
mod plan;
mod presets;
mod provider;
mod quota;
mod quota_refresh;
mod service;
pub mod providers;
pub use capability::{ProviderPoolCapabilities, ProviderPoolCapability};
pub use plan::{derive_oauth_plan_type, derive_plan_tier, normalize_provider_plan_tier};
pub use presets::{
build_admin_pool_scheduling_presets_payload, normalize_provider_scheduling_presets,
};
pub use provider::{ProviderPoolAdapter, ProviderPoolMemberInput};
pub use providers::{
build_antigravity_pool_quota_request, build_chatgpt_web_pool_quota_request,
build_codex_pool_quota_request, build_kiro_pool_quota_request,
enrich_chatgpt_web_quota_metadata, normalize_chatgpt_web_image_quota_limit,
AntigravityProviderPoolAdapter, ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter,
DefaultProviderPoolAdapter, KiroPoolQuotaAuthInput, KiroProviderPoolAdapter,
UnsupportedQuotaProviderPoolAdapter, ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH,
CHATGPT_WEB_CONVERSATION_INIT_PATH, CHATGPT_WEB_DEFAULT_BASE_URL, CODEX_WHAM_USAGE_URL,
KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
};
pub use quota::{
provider_pool_key_account_quota_exhausted, provider_pool_key_scheduling_label,
provider_pool_member_quota_snapshot, provider_pool_quota_metadata_provider_type,
provider_pool_quota_metadata_updated_at, provider_pool_quota_snapshot_updated_at,
};
pub use quota_refresh::ProviderPoolQuotaRequestSpec;
pub use service::ProviderPoolService;
#[cfg(test)]
mod tests {
use super::*;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_pool_core::PoolSchedulingPreset;
use serde_json::{json, Value};
fn sample_key(upstream_metadata: Option<Value>) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key-1".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.upstream_metadata = upstream_metadata;
key
}
#[test]
fn builtin_service_registers_provider_pool_adapters() {
let service = ProviderPoolService::with_builtin_adapters();
assert_eq!(
service.provider_types().collect::<Vec<_>>(),
[
"antigravity",
"chatgpt_web",
"claude_code",
"codex",
"gemini_cli",
"kiro",
"vertex_ai"
]
);
assert!(service
.adapter("codex")
.capabilities()
.supports(ProviderPoolCapability::PlanTier));
assert_eq!(service.adapter("unknown").provider_type(), "default");
}
#[test]
fn builtin_service_owns_quota_refresh_support_and_endpoint_selection() {
let service = ProviderPoolService::with_builtin_adapters();
assert_eq!(
service.provider_types_for_capability(ProviderPoolCapability::QuotaRefresh),
["antigravity", "chatgpt_web", "codex", "kiro"]
);
assert!(service.supports_quota_refresh("codex"));
assert!(service.supports_quota_refresh("antigravity"));
assert!(!service.supports_quota_refresh("gemini_cli"));
assert_eq!(
service.quota_refresh_unsupported_message("claude_code"),
"Claude Code 暂不支持自动刷新额度：上游没有稳定可用的账号额度查询接口"
);
assert_eq!(
service.quota_refresh_unsupported_message("vertex_ai"),
"Vertex AI 暂不支持自动刷新额度：额度属于 Google Cloud 项目/区域配额"
);
}
#[test]
fn codex_quota_request_adds_account_header_for_paid_accounts() {
let spec = build_codex_pool_quota_request(
"key-1",
Some(("authorization".to_string(), "Bearer access".to_string())),
None,
Some(&json!({
"plan_type": "plus",
"account_id": "acct-1"
})),
)
.expect("spec should build");
assert_eq!(
spec.headers.get("chatgpt-account-id").map(String::as_str),
Some("acct-1")
);
}
#[test]
fn codex_quota_request_skips_account_header_for_free_accounts() {
let spec = build_codex_pool_quota_request(
"key-1",
Some(("authorization".to_string(), "Bearer access".to_string())),
None,
Some(&json!({
"plan_type": "codex:free",
"account_id": "acct-1"
})),
)
.expect("spec should build");
assert!(!spec.headers.contains_key("chatgpt-account-id"));
}
#[test]
fn kiro_quota_request_includes_profile_arn_when_present() {
let spec = build_kiro_pool_quota_request(
"key-1",
&KiroPoolQuotaAuthInput {
authorization_value: "Bearer access".to_string(),
api_region: "us-west-2".to_string(),
kiro_version: "0.3.210".to_string(),
machine_id: "machine".to_string(),
profile_arn: Some("arn:aws:sso:::profile/p-1".to_string()),
},
);
assert!(spec.url.contains("q.us-west-2.amazonaws.com"));
assert!(spec
.url
.contains("profileArn=arn%3Aaws%3Asso%3A%3A%3Aprofile%2Fp-1"));
}
#[test]
fn chatgpt_web_quota_request_uses_default_base_url_when_empty() {
let spec = build_chatgpt_web_pool_quota_request(
"key-1",
"",
("authorization".to_string(), "Bearer access".to_string()),
);
assert_eq!(
spec.url,
"https://chatgpt.com/backend-api/conversation/init"
);
assert_eq!(
spec.headers.get("origin").map(String::as_str),
Some("https://chatgpt.com")
);
assert!(spec.accept_invalid_certs);
}
#[test]
fn chatgpt_web_quota_metadata_enriches_auth_and_normalizes_free_limit() {
let mut metadata = json!({
"image_quota_remaining": 12,
});
enrich_chatgpt_web_quota_metadata(
&mut metadata,
Some(&json!({
"plan": "free",
"email": "user@example.com",
"accountId": "acct-1"
})),
);
normalize_chatgpt_web_image_quota_limit(&mut metadata, None);
assert_eq!(metadata["plan_type"], json!("free"));
assert_eq!(metadata["email"], json!("user@example.com"));
assert_eq!(metadata["account_id"], json!("acct-1"));
assert_eq!(metadata["image_quota_total"], json!(25.0));
assert_eq!(metadata["image_quota_used"], json!(13.0));
}
#[test]
fn chatgpt_web_quota_metadata_preserves_existing_paid_limit() {
let mut metadata = json!({
"plan_type": "plus",
"image_quota_remaining": 7,
});
normalize_chatgpt_web_image_quota_limit(
&mut metadata,
Some(&json!({
"chatgpt_web": {
"image_quota_total": 40
}
})),
);
assert_eq!(metadata["image_quota_total"], json!(40.0));
assert_eq!(metadata["image_quota_used"], json!(33.0));
}
#[test]
fn preset_payload_derives_provider_support_from_capabilities() {
let payload = build_admin_pool_scheduling_presets_payload();
let items = payload.as_array().expect("payload should be array");
let free_first = items
.iter()
.find(|item| item["name"] == "free_first")
.expect("free_first should exist");
let recent_refresh = items
.iter()
.find(|item| item["name"] == "recent_refresh")
.expect("recent_refresh should exist");
assert_eq!(free_first["providers"], json!(["codex", "kiro"]));
assert_eq!(recent_refresh["providers"], json!(["codex", "kiro"]));
}
#[test]
fn quota_metadata_provider_type_comes_from_pool_registry() {
assert_eq!(
provider_pool_quota_metadata_provider_type(&json!({
"gemini_cli": {
"updated_at": 1_700_000_000u64
}
}))
.as_deref(),
Some("gemini_cli")
);
assert_eq!(
provider_pool_quota_metadata_provider_type(&json!({
"custom_provider": {
"updated_at": 1_700_000_000u64
}
}))
.as_deref(),
Some("custom_provider")
);
}
#[test]
fn codex_adapter_injects_recent_refresh_and_filters_by_capability() {
let service = ProviderPoolService::with_builtin_adapters();
let normalized = service.normalize_scheduling_presets(
"codex",
&[PoolSchedulingPreset {
preset: "cache_affinity".to_string(),
enabled: true,
mode: None,
}],
);
assert_eq!(
normalized
.iter()
.map(|preset| preset.preset.as_str())
.collect::<Vec<_>>(),
["cache_affinity", "recent_refresh"]
);
let unsupported = service.normalize_scheduling_presets(
"chatgpt_web",
&[PoolSchedulingPreset {
preset: "plus_first".to_string(),
enabled: true,
mode: None,
}],
);
assert!(unsupported.is_empty());
}
#[test]
fn provider_quota_exhaustion_is_adapter_owned() {
assert!(provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
"has_credits": false,
"credits_unlimited": false
}
}))),
"codex",
));
assert!(provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"kiro": {
"remaining": 0
}
}))),
"kiro",
));
assert!(provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"chatgpt_web": {
"image_quota_blocked": true
}
}))),
"chatgpt_web",
));
assert!(!provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
"has_credits": false,
"credits_unlimited": true
}
}))),
"codex",
));
}
#[test]
fn plan_tier_derivation_normalizes_provider_prefix() {
let key = sample_key(Some(json!({
"codex": {
"plan_type": "codex:Plus"
}
})));
assert_eq!(
derive_oauth_plan_type("codex", &key, None).as_deref(),
Some("plus")
);
}
#[test]
fn plan_tier_derivation_reads_quota_snapshot() {
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"plan_type": "team"
}
}));
assert_eq!(
derive_oauth_plan_type("codex", &key, None).as_deref(),
Some("team")
);
}
}

104
crates/aether-provider-pool/src/plan.rs Normal file
View File

@@ -0,0 +1,104 @@
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::{Map, Value};
pub fn derive_plan_tier(
provider_type: &str,
key: &StoredProviderCatalogKey,
auth_config: Option<&Map<String, Value>>,
) -> Option<String> {
let has_auth_config = auth_config.is_some()
|| key
.encrypted_auth_config
.as_deref()
.is_some_and(|value| !value.trim().is_empty());
if !provider_pool_auth_managed(key, provider_type, has_auth_config) {
return None;
}
if let Some(quota_snapshot) = key
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)
{
if let Some(normalized) = derive_plan_tier_from_map(quota_snapshot, provider_type) {
return Some(normalized);
}
}
if let Some(upstream_metadata) = key.upstream_metadata.as_ref().and_then(Value::as_object) {
let provider_bucket = upstream_metadata
.get(&provider_type.trim().to_ascii_lowercase())
.and_then(Value::as_object);
for source in provider_bucket
.into_iter()
.chain(std::iter::once(upstream_metadata))
{
if let Some(normalized) = derive_plan_tier_from_map(source, provider_type) {
return Some(normalized);
}
}
}
if let Some(config) = auth_config {
if let Some(normalized) = derive_plan_tier_from_map(config, provider_type) {
return Some(normalized);
}
}
None
}
fn derive_plan_tier_from_map(source: &Map<String, Value>, provider_type: &str) -> Option<String> {
for field in [
"plan_type",
"tier",
"plan",
"subscription_title",
"subscription_plan",
] {
if let Some(value) = source.get(field).and_then(Value::as_str) {
if let Some(normalized) = normalize_provider_plan_tier(value, provider_type) {
return Some(normalized);
}
}
}
None
}
pub fn derive_oauth_plan_type(
provider_type: &str,
key: &StoredProviderCatalogKey,
auth_config: Option<&Map<String, Value>>,
) -> Option<String> {
derive_plan_tier(provider_type, key, auth_config)
}
pub fn normalize_provider_plan_tier(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 provider_pool_auth_managed(
key: &StoredProviderCatalogKey,
provider_type: &str,
has_auth_config: bool,
) -> bool {
key.auth_type.trim().eq_ignore_ascii_case("oauth")
|| (provider_type.trim().eq_ignore_ascii_case("kiro")
&& key.auth_type.trim().eq_ignore_ascii_case("bearer")
&& has_auth_config)
}

244
crates/aether-provider-pool/src/presets.rs Normal file
View File

@@ -0,0 +1,244 @@
use std::collections::BTreeSet;
use aether_pool_core::PoolSchedulingPreset;
use serde_json::{json, Value};
use crate::capability::ProviderPoolCapability;
use crate::provider::ProviderPoolAdapter;
use crate::service::ProviderPoolService;
pub fn normalize_provider_scheduling_presets(
adapter: &dyn ProviderPoolAdapter,
scheduling_presets: &[PoolSchedulingPreset],
) -> Vec<PoolSchedulingPreset> {
let mut entries = Vec::<(usize, PoolSchedulingPreset)>::new();
let mut seen = BTreeSet::new();
for (index, item) in scheduling_presets.iter().enumerate() {
let preset = item.preset.trim().to_ascii_lowercase();
if preset.is_empty() || !seen.insert(preset.clone()) {
continue;
}
if !provider_pool_supports_preset(adapter, &preset) {
continue;
}
entries.push((
index,
PoolSchedulingPreset {
preset,
enabled: item.enabled,
mode: item.mode.clone(),
},
));
}
if !entries.is_empty() {
for preset in adapter.default_scheduling_presets() {
let preset_name = preset.preset.trim().to_ascii_lowercase();
if preset_name.is_empty() || seen.contains(&preset_name) {
continue;
}
if !provider_pool_supports_preset(adapter, &preset_name) {
continue;
}
seen.insert(preset_name.clone());
entries.push((
entries.len(),
PoolSchedulingPreset {
preset: preset_name,
enabled: preset.enabled,
mode: preset.mode,
},
));
}
}
let mut distribution_mode = None::<(usize, PoolSchedulingPreset)>;
let mut strategy_presets = Vec::<(usize, PoolSchedulingPreset)>::new();
for (index, preset) in entries {
if !preset.enabled {
continue;
}
if let Some(mutex_group) = provider_pool_preset_mutex_group(&preset.preset) {
if mutex_group == "distribution_mode"
&& distribution_mode
.as_ref()
.is_none_or(|current| index < current.0)
{
distribution_mode = Some((index, preset));
}
} else {
strategy_presets.push((index, preset));
}
}
let mut normalized = Vec::new();
if let Some((_, preset)) = distribution_mode.filter(|(_, preset)| preset.preset != "lru") {
normalized.push(preset);
}
strategy_presets.sort_by_key(|left| left.0);
normalized.extend(strategy_presets.into_iter().map(|(_, preset)| preset));
normalized
}
pub fn build_admin_pool_scheduling_presets_payload() -> Value {
let service = ProviderPoolService::with_builtin_adapters();
json!([
provider_pool_preset_payload(
"lru",
"LRU 轮转",
"最久未使用的 Key 优先",
None,
"依据 LRU 时间戳（最近未使用优先）",
&service,
),
provider_pool_preset_payload(
"cache_affinity",
"缓存亲和",
"优先复用最近使用过的 Key，利用 Prompt Caching",
None,
"依据 LRU 时间戳（最近使用优先，与 LRU 轮转相反）",
&service,
),
provider_pool_preset_payload(
"cost_first",
"成本优先",
"优先选择窗口消耗更低的账号",
None,
"依据窗口成本/Token 用量，缺失时回退配额使用率",
&service,
),
provider_pool_preset_payload(
"free_first",
"Free 优先",
"优先消耗 Free 账号（依赖 plan_type）",
Some(ProviderPoolCapability::PlanTier),
"依据 plan_type（Free 账号优先调度）",
&service,
),
provider_pool_preset_payload(
"health_first",
"健康优先",
"优先选择健康分更高、失败更少的账号",
None,
"依据 health_by_format 聚合分（含熔断/失败衰减）",
&service,
),
provider_pool_preset_payload(
"latency_first",
"延迟优先",
"优先选择最近延迟更低的账号",
None,
"依据号池延迟窗口均值（latency_window_seconds）",
&service,
),
provider_pool_preset_payload(
"load_balance",
"负载均衡",
"随机分散 Key 使用，均匀分摊负载",
None,
"每次随机分值，实现完全均匀分散",
&service,
),
provider_pool_preset_payload(
"plus_first",
"Plus 优先",
"优先消耗 Plus 账号（依赖 plan_type）",
Some(ProviderPoolCapability::PlanTier),
"依据 plan_type（Plus 账号优先调度）",
&service,
),
provider_pool_preset_payload(
"pro_first",
"Pro 优先",
"优先消耗 Pro 账号（依赖 plan_type）",
Some(ProviderPoolCapability::PlanTier),
"依据 plan_type（Pro 账号优先调度）",
&service,
),
provider_pool_preset_payload(
"priority_first",
"优先级优先",
"按账号优先级顺序调度（数字越小越优先）",
None,
"依据 internal_priority（支持拖拽/手工编辑）",
&service,
),
provider_pool_preset_payload(
"quota_balanced",
"额度平均",
"优先选额度消耗最少的账号",
None,
"依据账号配额使用率；无配额时回退到窗口成本使用",
&service,
),
provider_pool_preset_payload(
"recent_refresh",
"额度刷新优先",
"优先选即将刷新额度的账号",
Some(ProviderPoolCapability::QuotaReset),
"依据账号额度重置倒计时（next_reset / reset_seconds）",
&service,
),
provider_pool_preset_payload(
"single_account",
"单号优先",
"集中使用同一账号（反向 LRU）",
None,
"先按账号优先级（internal_priority），同级再按反向 LRU 集中",
&service,
),
provider_pool_preset_payload(
"team_first",
"Team 优先",
"优先消耗 Team 账号（依赖 plan_type）",
Some(ProviderPoolCapability::PlanTier),
"依据 plan_type（Team 账号优先调度）",
&service,
),
])
}
fn provider_pool_preset_payload(
name: &'static str,
label: &'static str,
description: &'static str,
capability: Option<ProviderPoolCapability>,
evidence_hint: &'static str,
service: &ProviderPoolService,
) -> Value {
let providers = capability
.map(|capability| service.provider_types_for_capability(capability))
.unwrap_or_default();
json!({
"name": name,
"label": label,
"description": description,
"providers": providers,
"modes": Value::Null,
"default_mode": Value::Null,
"mutex_group": provider_pool_preset_mutex_group(name),
"evidence_hint": evidence_hint,
})
}
fn provider_pool_supports_preset(adapter: &dyn ProviderPoolAdapter, preset: &str) -> bool {
match preset {
"free_first" | "plus_first" | "pro_first" | "team_first" => adapter
.capabilities()
.supports(ProviderPoolCapability::PlanTier),
"recent_refresh" => adapter
.capabilities()
.supports(ProviderPoolCapability::QuotaReset),
_ => true,
}
}
fn provider_pool_preset_mutex_group(preset: &str) -> Option<&'static str> {
match preset {
"lru" | "cache_affinity" | "load_balance" | "single_account" => Some("distribution_mode"),
_ => None,
}
}

107
crates/aether-provider-pool/src/provider.rs Normal file
View File

@@ -0,0 +1,107 @@
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_pool_core::{PoolMemberSignals, PoolSchedulingPreset};
use serde_json::{Map, Value};
use crate::capability::{ProviderPoolCapabilities, ProviderPoolCapability};
use crate::plan::{derive_plan_tier, normalize_provider_plan_tier};
use crate::quota::{
provider_pool_account_blocked, provider_pool_quota_reset_seconds,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_quota_usage_ratio,
};
#[derive(Debug, Clone)]
pub struct ProviderPoolMemberInput<'a> {
pub provider_type: &'a str,
pub key: &'a StoredProviderCatalogKey,
pub auth_config: Option<&'a Map<String, Value>>,
}
pub trait ProviderPoolAdapter: Send + Sync {
fn provider_type(&self) -> &'static str;
fn capabilities(&self) -> ProviderPoolCapabilities {
ProviderPoolCapabilities::default()
}
fn default_scheduling_presets(&self) -> Vec<PoolSchedulingPreset> {
Vec::new()
}
fn supports_quota_refresh(&self) -> bool {
self.capabilities()
.supports(ProviderPoolCapability::QuotaRefresh)
}
fn quota_refresh_endpoint(
&self,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
if !self.supports_quota_refresh() {
return None;
}
provider_pool_matching_endpoint(endpoints, include_inactive, |_| true)
}
fn quota_refresh_unsupported_message(&self) -> String {
"该 Provider 暂不支持自动刷新额度".to_string()
}
fn quota_refresh_missing_endpoint_message(&self) -> String {
"找不到有效端点".to_string()
}
fn normalize_plan_tier(&self, value: &str) -> Option<String> {
normalize_provider_plan_tier(value, self.provider_type())
}
fn member_signals(&self, input: &ProviderPoolMemberInput<'_>) -> PoolMemberSignals {
PoolMemberSignals {
plan_tier: derive_plan_tier(input.provider_type, input.key, input.auth_config),
quota_usage_ratio: provider_pool_quota_usage_ratio(input.key),
quota_reset_seconds: provider_pool_quota_reset_seconds(input.key),
account_blocked: provider_pool_account_blocked(input.key),
quota_exhausted: self.quota_exhausted(input),
..PoolMemberSignals::default()
}
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
.unwrap_or(false)
}
}
pub(crate) fn provider_pool_matching_endpoint<F>(
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
predicate: F,
) -> Option<StoredProviderCatalogEndpoint>
where
F: Fn(&StoredProviderCatalogEndpoint) -> bool,
{
endpoints
.iter()
.find(|endpoint| endpoint.is_active && predicate(endpoint))
.cloned()
.or_else(|| {
include_inactive.then(|| {
endpoints
.iter()
.find(|endpoint| !endpoint.is_active && predicate(endpoint))
.cloned()
})?
})
}
pub(crate) fn provider_pool_endpoint_format_matches(
endpoint: &StoredProviderCatalogEndpoint,
expected: &str,
) -> bool {
endpoint
.api_format
.trim()
.eq_ignore_ascii_case(expected.trim())
}

77
crates/aether-provider-pool/src/providers/antigravity.rs Normal file
View File

@@ -0,0 +1,77 @@
use std::collections::BTreeMap;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
use serde_json::json;
use crate::capability::ProviderPoolCapabilities;
use crate::provider::{
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
pub const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels";
#[derive(Debug, Clone, Default)]
pub struct AntigravityProviderPoolAdapter;
impl ProviderPoolAdapter for AntigravityProviderPoolAdapter {
fn provider_type(&self) -> &'static str {
"antigravity"
}
fn capabilities(&self) -> ProviderPoolCapabilities {
ProviderPoolCapabilities {
quota_refresh: true,
..ProviderPoolCapabilities::default()
}
}
fn quota_refresh_endpoint(
&self,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
provider_pool_endpoint_format_matches(endpoint, "gemini:generate_content")
})
}
fn quota_refresh_missing_endpoint_message(&self) -> String {
"找不到有效的 gemini:generate_content 端点".to_string()
}
}
pub fn build_antigravity_pool_quota_request(
key_id: &str,
endpoint_base_url: &str,
authorization: (String, String),
project_id: &str,
mut identity_headers: BTreeMap<String, String>,
) -> ProviderPoolQuotaRequestSpec {
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());
ProviderPoolQuotaRequestSpec {
request_id: format!("antigravity-quota:{key_id}"),
provider_name: "antigravity".to_string(),
quota_kind: "antigravity".to_string(),
method: "POST".to_string(),
url: format!(
"{}{}",
endpoint_base_url.trim_end_matches('/'),
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH
),
headers,
content_type: Some("application/json".to_string()),
json_body: Some(json!({ "project": project_id })),
client_api_format: "gemini:generate_content".to_string(),
provider_api_format: "antigravity:fetch_available_models".to_string(),
model_name: Some("fetchAvailableModels".to_string()),
accept_invalid_certs: false,
}
}

272
crates/aether-provider-pool/src/providers/chatgpt_web.rs Normal file
View File

@@ -0,0 +1,272 @@
use std::collections::BTreeMap;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
use serde_json::{json, Map, Value};
use uuid::Uuid;
use crate::capability::ProviderPoolCapabilities;
use crate::provider::{
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
ProviderPoolMemberInput,
};
use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
pub const CHATGPT_WEB_DEFAULT_BASE_URL: &str = "https://chatgpt.com";
pub 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 CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT: f64 = 25.0;
#[derive(Debug, Clone, Default)]
pub struct ChatGptWebProviderPoolAdapter;
impl ProviderPoolAdapter for ChatGptWebProviderPoolAdapter {
fn provider_type(&self) -> &'static str {
"chatgpt_web"
}
fn capabilities(&self) -> ProviderPoolCapabilities {
ProviderPoolCapabilities {
quota_refresh: true,
..ProviderPoolCapabilities::default()
}
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
return exhausted;
}
provider_pool_metadata_bucket(input.key.upstream_metadata.as_ref(), input.provider_type)
.is_some_and(quota_exhausted_from_bucket)
}
fn quota_refresh_endpoint(
&self,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
provider_pool_endpoint_format_matches(endpoint, "openai:image")
})
}
fn quota_refresh_missing_endpoint_message(&self) -> String {
"找不到有效的 openai:image 端点".to_string()
}
}
pub fn build_chatgpt_web_pool_quota_request(
key_id: &str,
endpoint_base_url: &str,
authorization: (String, String),
) -> ProviderPoolQuotaRequestSpec {
let base_url = chatgpt_web_base_url(endpoint_base_url);
let device_id = Uuid::new_v4().to_string();
let session_id = 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.clone()),
("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(),
),
]);
headers.insert(authorization.0.to_ascii_lowercase(), authorization.1);
ProviderPoolQuotaRequestSpec {
request_id: format!("chatgpt-web-quota:{key_id}"),
provider_name: "chatgpt_web".to_string(),
quota_kind: "chatgpt_web".to_string(),
method: "POST".to_string(),
url: format!("{base_url}{CHATGPT_WEB_CONVERSATION_INIT_PATH}"),
headers,
content_type: Some("application/json".to_string()),
json_body: Some(json!({
"gizmo_id": Value::Null,
"requested_default_model": Value::Null,
"conversation_id": Value::Null,
"timezone_offset_min": -480,
"system_hints": ["picture_v2"],
})),
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()),
accept_invalid_certs: true,
}
}
fn chatgpt_web_base_url(endpoint_base_url: &str) -> 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()
}
}
pub fn enrich_chatgpt_web_quota_metadata(metadata: &mut Value, auth_config: Option<&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));
}
}
}
pub fn normalize_chatgpt_web_image_quota_limit(
metadata: &mut Value,
upstream_metadata: Option<&Value>,
) {
let existing_limit = existing_chatgpt_web_image_quota_limit(upstream_metadata);
let Some(object) = metadata.as_object_mut() else {
return;
};
let remaining = provider_pool_json_f64(object.get("image_quota_remaining"));
let explicit_limit =
provider_pool_json_f64(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(Value::as_bool) == Some(true) {
object.insert("image_quota_used".to_string(), json!(limit));
}
}
}
}
fn chatgpt_web_auth_config_string(auth_config: Option<&Value>, fields: &[&str]) -> Option<String> {
let object = auth_config.and_then(Value::as_object)?;
fields.iter().find_map(|field| {
object
.get(*field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn chatgpt_web_json_string(value: Option<&Value>) -> Option<&str> {
value
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn existing_chatgpt_web_image_quota_limit(upstream_metadata: Option<&Value>) -> Option<f64> {
upstream_metadata
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("chatgpt_web"))
.and_then(Value::as_object)
.and_then(|bucket| provider_pool_json_f64(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)
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if provider_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) {
return true;
}
if provider_pool_json_f64(bucket.get("image_quota_remaining")).is_some_and(|value| value <= 0.0)
{
return true;
}
match (
provider_pool_json_f64(bucket.get("image_quota_total")),
provider_pool_json_f64(bucket.get("image_quota_used")),
) {
(Some(limit), Some(used)) if limit > 0.0 => used >= limit,
_ => false,
}
}

136
crates/aether-provider-pool/src/providers/codex.rs Normal file
View File

@@ -0,0 +1,136 @@
use std::collections::BTreeMap;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
use aether_pool_core::PoolSchedulingPreset;
use serde_json::{Map, Value};
use crate::capability::ProviderPoolCapabilities;
use crate::provider::{
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
ProviderPoolMemberInput,
};
use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
pub const CODEX_WHAM_USAGE_URL: &str = "https://chatgpt.com/backend-api/wham/usage";
const PLACEHOLDER_API_KEY: &str = "__placeholder__";
#[derive(Debug, Clone, Default)]
pub struct CodexProviderPoolAdapter;
impl ProviderPoolAdapter for CodexProviderPoolAdapter {
fn provider_type(&self) -> &'static str {
"codex"
}
fn capabilities(&self) -> ProviderPoolCapabilities {
ProviderPoolCapabilities {
plan_tier: true,
quota_reset: true,
quota_refresh: true,
}
}
fn default_scheduling_presets(&self) -> Vec<PoolSchedulingPreset> {
vec![PoolSchedulingPreset {
preset: "recent_refresh".to_string(),
enabled: true,
mode: None,
}]
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
return exhausted;
}
provider_pool_metadata_bucket(input.key.upstream_metadata.as_ref(), input.provider_type)
.is_some_and(quota_exhausted_from_bucket)
}
fn quota_refresh_endpoint(
&self,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
provider_pool_endpoint_format_matches(endpoint, "openai:responses")
})
}
fn quota_refresh_missing_endpoint_message(&self) -> String {
"找不到有效的 openai:responses 端点".to_string()
}
}
pub fn build_codex_pool_quota_request(
key_id: &str,
resolved_oauth_auth: Option<(String, String)>,
decrypted_api_key: Option<&str>,
auth_config: Option<&Value>,
) -> Result<ProviderPoolQuotaRequestSpec, 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 = decrypted_api_key.unwrap_or_default().trim();
if decrypted_key.is_empty() || decrypted_key == PLACEHOLDER_API_KEY {
return Err("缺少 OAuth 认证信息，请先授权/刷新 Token".to_string());
}
headers.insert(
"authorization".to_string(),
format!("Bearer {decrypted_key}"),
);
}
let oauth_plan_type = auth_config
.and_then(|value| value.get("plan_type"))
.and_then(Value::as_str)
.and_then(|value| crate::plan::normalize_provider_plan_tier(value, "codex"));
let oauth_account_id = auth_config
.and_then(|value| value.get("account_id"))
.and_then(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(ProviderPoolQuotaRequestSpec {
request_id: format!("codex-quota:{key_id}"),
provider_name: "codex".to_string(),
quota_kind: "codex".to_string(),
method: "GET".to_string(),
url: CODEX_WHAM_USAGE_URL.to_string(),
headers,
content_type: None,
json_body: None,
client_api_format: "openai:responses".to_string(),
provider_api_format: "openai:responses".to_string(),
model_name: Some("codex-wham-usage".to_string()),
accept_invalid_certs: false,
})
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
return false;
}
let has_window_data = provider_pool_json_f64(bucket.get("primary_used_percent")).is_some()
|| provider_pool_json_f64(bucket.get("secondary_used_percent")).is_some();
if !has_window_data && provider_pool_json_bool(bucket.get("has_credits")) == Some(false) {
return true;
}
provider_pool_json_f64(bucket.get("primary_used_percent")).is_some_and(|value| value >= 100.0)
|| provider_pool_json_f64(bucket.get("secondary_used_percent"))
.is_some_and(|value| value >= 100.0)
}

10
crates/aether-provider-pool/src/providers/default.rs Normal file
View File

@@ -0,0 +1,10 @@
use crate::provider::ProviderPoolAdapter;
#[derive(Debug, Clone, Default)]
pub struct DefaultProviderPoolAdapter;
impl ProviderPoolAdapter for DefaultProviderPoolAdapter {
fn provider_type(&self) -> &'static str {
"default"
}
}

167
crates/aether-provider-pool/src/providers/kiro.rs Normal file
View File

@@ -0,0 +1,167 @@
use std::collections::BTreeMap;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
use serde_json::{Map, Value};
use url::form_urlencoded;
use uuid::Uuid;
use crate::capability::ProviderPoolCapabilities;
use crate::provider::{
provider_pool_endpoint_format_matches, provider_pool_matching_endpoint, ProviderPoolAdapter,
ProviderPoolMemberInput,
};
use crate::quota::{
provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
pub const KIRO_USAGE_LIMITS_PATH: &str = "/getUsageLimits";
pub const KIRO_USAGE_SDK_VERSION: &str = "1.0.0";
#[derive(Debug, Clone, PartialEq)]
pub struct KiroPoolQuotaAuthInput {
pub authorization_value: String,
pub api_region: String,
pub kiro_version: String,
pub machine_id: String,
pub profile_arn: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct KiroProviderPoolAdapter;
impl ProviderPoolAdapter for KiroProviderPoolAdapter {
fn provider_type(&self) -> &'static str {
"kiro"
}
fn capabilities(&self) -> ProviderPoolCapabilities {
ProviderPoolCapabilities {
plan_tier: true,
quota_reset: true,
quota_refresh: true,
}
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
return exhausted;
}
provider_pool_metadata_bucket(input.key.upstream_metadata.as_ref(), input.provider_type)
.is_some_and(quota_exhausted_from_bucket)
}
fn quota_refresh_endpoint(
&self,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
provider_pool_matching_endpoint(endpoints, include_inactive, |endpoint| {
provider_pool_endpoint_format_matches(endpoint, "claude:messages")
})
.or_else(|| provider_pool_matching_endpoint(endpoints, include_inactive, |_| true))
}
fn quota_refresh_missing_endpoint_message(&self) -> String {
"找不到有效的 Kiro 端点".to_string()
}
}
pub fn build_kiro_pool_quota_request(
key_id: &str,
auth: &KiroPoolQuotaAuthInput,
) -> ProviderPoolQuotaRequestSpec {
let host = format!("q.{}.amazonaws.com", normalize_region(&auth.api_region));
let machine_id = auth.machine_id.trim();
let ide_tag = if machine_id.is_empty() {
format!("KiroIDE-{}", normalize_kiro_version(&auth.kiro_version))
} else {
format!(
"KiroIDE-{}-{machine_id}",
normalize_kiro_version(&auth.kiro_version)
)
};
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
.profile_arn
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
serializer.append_pair("profileArn", profile_arn);
}
ProviderPoolQuotaRequestSpec {
request_id: format!("kiro-quota:{key_id}"),
provider_name: "kiro".to_string(),
quota_kind: "kiro".to_string(),
method: "GET".to_string(),
url: format!("https://{host}{KIRO_USAGE_LIMITS_PATH}?{}", serializer.finish()),
headers: 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.authorization_value.clone(),
),
("connection".to_string(), "close".to_string()),
]),
content_type: None,
json_body: None,
client_api_format: "claude:messages".to_string(),
provider_api_format: "kiro:usage".to_string(),
model_name: Some("kiro-usage-limits".to_string()),
accept_invalid_certs: false,
}
}
fn normalize_region(value: &str) -> &str {
let value = value.trim();
if value.is_empty() {
"us-east-1"
} else {
value
}
}
fn normalize_kiro_version(value: &str) -> &str {
let value = value.trim();
if value.is_empty() {
"0.3.210"
} else {
value
}
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if provider_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) {
return true;
}
if provider_pool_json_f64(bucket.get("usage_percentage")).is_some_and(|value| value >= 100.0) {
return true;
}
match (
provider_pool_json_f64(bucket.get("usage_limit")),
provider_pool_json_f64(bucket.get("current_usage")),
) {
(Some(limit), Some(current)) if limit > 0.0 => current >= limit,
_ => false,
}
}

29
crates/aether-provider-pool/src/providers/mod.rs Normal file
View File

@@ -0,0 +1,29 @@
pub mod antigravity;
pub mod chatgpt_web;
pub mod codex;
pub mod default;
pub mod kiro;
pub mod unsupported;
pub use antigravity::AntigravityProviderPoolAdapter;
pub use antigravity::{
build_antigravity_pool_quota_request, ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH,
};
pub use chatgpt_web::ChatGptWebProviderPoolAdapter;
pub use chatgpt_web::{
build_chatgpt_web_pool_quota_request, enrich_chatgpt_web_quota_metadata,
normalize_chatgpt_web_image_quota_limit, CHATGPT_WEB_CONVERSATION_INIT_PATH,
CHATGPT_WEB_DEFAULT_BASE_URL,
};
pub use codex::CodexProviderPoolAdapter;
pub use codex::{build_codex_pool_quota_request, CODEX_WHAM_USAGE_URL};
pub use default::DefaultProviderPoolAdapter;
pub use kiro::KiroProviderPoolAdapter;
pub use kiro::{
build_kiro_pool_quota_request, KiroPoolQuotaAuthInput, KIRO_USAGE_LIMITS_PATH,
KIRO_USAGE_SDK_VERSION,
};
pub use unsupported::{
UnsupportedQuotaProviderPoolAdapter, CLAUDE_CODE_PROVIDER_POOL_ADAPTER,
GEMINI_CLI_PROVIDER_POOL_ADAPTER, VERTEX_AI_PROVIDER_POOL_ADAPTER,
};

47
crates/aether-provider-pool/src/providers/unsupported.rs Normal file
View File

@@ -0,0 +1,47 @@
use crate::provider::ProviderPoolAdapter;
#[derive(Debug, Clone, Copy)]
pub struct UnsupportedQuotaProviderPoolAdapter {
provider_type: &'static str,
quota_refresh_unsupported_message: &'static str,
}
impl UnsupportedQuotaProviderPoolAdapter {
pub const fn new(
provider_type: &'static str,
quota_refresh_unsupported_message: &'static str,
) -> Self {
Self {
provider_type,
quota_refresh_unsupported_message,
}
}
}
impl ProviderPoolAdapter for UnsupportedQuotaProviderPoolAdapter {
fn provider_type(&self) -> &'static str {
self.provider_type
}
fn quota_refresh_unsupported_message(&self) -> String {
self.quota_refresh_unsupported_message.to_string()
}
}
pub const CLAUDE_CODE_PROVIDER_POOL_ADAPTER: UnsupportedQuotaProviderPoolAdapter =
UnsupportedQuotaProviderPoolAdapter::new(
"claude_code",
"Claude Code 暂不支持自动刷新额度：上游没有稳定可用的账号额度查询接口",
);
pub const GEMINI_CLI_PROVIDER_POOL_ADAPTER: UnsupportedQuotaProviderPoolAdapter =
UnsupportedQuotaProviderPoolAdapter::new(
"gemini_cli",
"Gemini CLI 暂不支持自动刷新额度：当前只能通过模型同步/缓存快照展示已知配额信息",
);
pub const VERTEX_AI_PROVIDER_POOL_ADAPTER: UnsupportedQuotaProviderPoolAdapter =
UnsupportedQuotaProviderPoolAdapter::new(
"vertex_ai",
"Vertex AI 暂不支持自动刷新额度：额度属于 Google Cloud 项目/区域配额",
);

263
crates/aether-provider-pool/src/quota.rs Normal file
View File

@@ -0,0 +1,263 @@
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::{json, Map, Value};
use crate::provider::ProviderPoolMemberInput;
use crate::service::ProviderPoolService;
pub fn provider_pool_key_account_quota_exhausted(
key: &StoredProviderCatalogKey,
provider_type: &str,
) -> bool {
let adapter = ProviderPoolService::with_builtin_adapters().adapter(provider_type);
adapter.quota_exhausted(&ProviderPoolMemberInput {
provider_type,
key,
auth_config: None,
})
}
pub fn provider_pool_member_quota_snapshot<'a>(
key: &'a StoredProviderCatalogKey,
provider_type: &str,
) -> Option<&'a Map<String, Value>> {
let quota_snapshot = key
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)?;
provider_pool_quota_snapshot_matches_provider(quota_snapshot, provider_type)
.then_some(quota_snapshot)
}
pub fn provider_pool_quota_snapshot_updated_at(
key: &StoredProviderCatalogKey,
provider_type: &str,
) -> Option<u64> {
let quota_snapshot = provider_pool_member_quota_snapshot(key, provider_type)?;
provider_pool_timestamp_unix_secs(quota_snapshot.get("updated_at"))
}
pub fn provider_pool_quota_metadata_updated_at(
upstream_metadata: Option<&Value>,
provider_type: &str,
) -> Option<u64> {
let bucket = provider_pool_metadata_bucket(upstream_metadata, provider_type)?;
provider_pool_timestamp_unix_secs(bucket.get("updated_at"))
}
pub fn provider_pool_quota_metadata_provider_type(metadata_update: &Value) -> Option<String> {
let object = metadata_update.as_object()?;
let service = ProviderPoolService::with_builtin_adapters();
let known_provider_type = service
.provider_types()
.find(|provider_type| object.contains_key(*provider_type))
.map(ToOwned::to_owned);
known_provider_type.or_else(|| {
object
.iter()
.find(|(_, value)| value.is_object())
.map(|(provider_type, _)| provider_type.clone())
})
}
pub fn provider_pool_key_scheduling_label(
is_active: bool,
cooldown_reason: Option<&str>,
cooldown_ttl_seconds: Option<u64>,
) -> (String, String, String, Vec<Value>) {
if !is_active {
return (
"blocked".to_string(),
"inactive".to_string(),
"已禁用".to_string(),
vec![json!({
"code": "inactive",
"label": "已禁用",
"blocking": true,
"source": "manual",
"ttl_seconds": Value::Null,
"detail": Value::Null,
})],
);
}
if let Some(reason) = cooldown_reason {
return (
"degraded".to_string(),
"cooldown".to_string(),
"冷却中".to_string(),
vec![json!({
"code": "cooldown",
"label": "冷却中",
"blocking": true,
"source": "pool",
"ttl_seconds": cooldown_ttl_seconds,
"detail": reason,
})],
);
}
(
"available".to_string(),
"available".to_string(),
"可用".to_string(),
Vec::new(),
)
}
pub(crate) fn provider_pool_metadata_bucket<'a>(
upstream_metadata: Option<&'a Value>,
provider_type: &str,
) -> Option<&'a Map<String, Value>> {
upstream_metadata
.and_then(Value::as_object)
.and_then(|metadata| metadata.get(&provider_type.trim().to_ascii_lowercase()))
.and_then(Value::as_object)
}
pub(crate) fn provider_pool_json_bool(value: Option<&Value>) -> Option<bool> {
match value {
Some(Value::Bool(value)) => Some(*value),
Some(Value::String(value)) => match value.trim().to_ascii_lowercase().as_str() {
"true" | "1" => Some(true),
"false" | "0" => Some(false),
_ => None,
},
_ => None,
}
}
pub(crate) fn provider_pool_json_f64(value: Option<&Value>) -> Option<f64> {
match value {
Some(Value::Number(number)) => number.as_f64(),
Some(Value::String(value)) => value.trim().parse::<f64>().ok(),
_ => None,
}
.filter(|value| value.is_finite())
}
fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
let mut timestamp = provider_pool_json_f64(value)?;
if timestamp <= 0.0 {
return None;
}
if timestamp > 1_000_000_000_000.0 {
timestamp /= 1000.0;
}
Some(timestamp as u64)
}
fn provider_pool_quota_snapshot_matches_provider(
quota_snapshot: &Map<String, Value>,
provider_type: &str,
) -> bool {
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
match quota_snapshot
.get("provider_type")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
Some(provider_type) => provider_type.eq_ignore_ascii_case(&normalized_provider_type),
None => {
provider_pool_json_bool(quota_snapshot.get("exhausted")) == Some(true)
|| quota_snapshot
.get("code")
.and_then(Value::as_str)
.is_some_and(|code| !code.trim().eq_ignore_ascii_case("unknown"))
|| quota_snapshot
.get("updated_at")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("observed_at")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("usage_ratio")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("reset_seconds")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("windows")
.and_then(Value::as_array)
.is_some_and(|windows| !windows.is_empty())
|| quota_snapshot
.get("credits")
.and_then(Value::as_object)
.is_some_and(|credits| !credits.is_empty())
}
}
}
pub(crate) fn provider_pool_quota_snapshot_exhausted_decision(
key: &StoredProviderCatalogKey,
provider_type: &str,
) -> Option<bool> {
let quota_snapshot = key
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)?;
if !provider_pool_quota_snapshot_matches_provider(quota_snapshot, provider_type) {
return None;
}
let exhausted = provider_pool_json_bool(quota_snapshot.get("exhausted"))?;
if exhausted {
let windows_max_ratio = quota_snapshot
.get("windows")
.and_then(Value::as_array)
.filter(|w| !w.is_empty())
.and_then(|windows| {
windows
.iter()
.filter_map(Value::as_object)
.filter_map(|w| w.get("used_ratio"))
.filter_map(Value::as_f64)
.max_by(f64::total_cmp)
});
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
return Some(false);
}
}
Some(exhausted)
}
pub(crate) fn provider_pool_quota_usage_ratio(key: &StoredProviderCatalogKey) -> Option<f64> {
key.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)
.and_then(|quota| provider_pool_json_f64(quota.get("usage_ratio")))
}
pub(crate) fn provider_pool_quota_reset_seconds(key: &StoredProviderCatalogKey) -> Option<f64> {
key.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)
.and_then(|quota| provider_pool_json_f64(quota.get("reset_seconds")))
}
pub(crate) fn provider_pool_account_blocked(key: &StoredProviderCatalogKey) -> bool {
key.oauth_invalid_reason.as_deref().is_some_and(|reason| {
let normalized = reason.trim().to_ascii_lowercase();
!normalized.is_empty()
&& [
"banned",
"forbidden",
"blocked",
"suspend",
"deactivated",
"disabled",
"verification",
"workspace",
"受限",
"封",
"禁",
]
.iter()
.any(|hint| normalized.contains(hint))
})
}

19
crates/aether-provider-pool/src/quota_refresh.rs Normal file
View File

@@ -0,0 +1,19 @@
use std::collections::BTreeMap;
use serde_json::Value;
#[derive(Debug, Clone, PartialEq)]
pub struct ProviderPoolQuotaRequestSpec {
pub request_id: String,
pub provider_name: String,
pub quota_kind: String,
pub method: String,
pub url: String,
pub headers: BTreeMap<String, String>,
pub content_type: Option<String>,
pub json_body: Option<Value>,
pub client_api_format: String,
pub provider_api_format: String,
pub model_name: Option<String>,
pub accept_invalid_certs: bool,
}

132
crates/aether-provider-pool/src/service.rs Normal file
View File

@@ -0,0 +1,132 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_pool_core::PoolSchedulingPreset;
use serde_json::{Map, Value};
use crate::capability::ProviderPoolCapability;
use crate::presets::normalize_provider_scheduling_presets;
use crate::provider::{ProviderPoolAdapter, ProviderPoolMemberInput};
use crate::providers::{
AntigravityProviderPoolAdapter, ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter,
DefaultProviderPoolAdapter, KiroProviderPoolAdapter, CLAUDE_CODE_PROVIDER_POOL_ADAPTER,
GEMINI_CLI_PROVIDER_POOL_ADAPTER, VERTEX_AI_PROVIDER_POOL_ADAPTER,
};
#[derive(Clone)]
pub struct ProviderPoolService {
adapters: BTreeMap<String, Arc<dyn ProviderPoolAdapter>>,
default_adapter: Arc<dyn ProviderPoolAdapter>,
}
impl std::fmt::Debug for ProviderPoolService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ProviderPoolService")
.field("provider_types", &self.adapters.keys().collect::<Vec<_>>())
.finish()
}
}
impl Default for ProviderPoolService {
fn default() -> Self {
Self {
adapters: BTreeMap::new(),
default_adapter: Arc::new(DefaultProviderPoolAdapter),
}
}
}
impl ProviderPoolService {
pub fn new() -> Self {
Self::default()
}
pub fn with_builtin_adapters() -> Self {
Self::new()
.with_adapter(Arc::new(AntigravityProviderPoolAdapter))
.with_adapter(Arc::new(CLAUDE_CODE_PROVIDER_POOL_ADAPTER))
.with_adapter(Arc::new(CodexProviderPoolAdapter))
.with_adapter(Arc::new(GEMINI_CLI_PROVIDER_POOL_ADAPTER))
.with_adapter(Arc::new(KiroProviderPoolAdapter))
.with_adapter(Arc::new(ChatGptWebProviderPoolAdapter))
.with_adapter(Arc::new(VERTEX_AI_PROVIDER_POOL_ADAPTER))
}
pub fn with_adapter(mut self, adapter: Arc<dyn ProviderPoolAdapter>) -> Self {
self.adapters
.insert(adapter.provider_type().trim().to_ascii_lowercase(), adapter);
self
}
pub fn adapter(&self, provider_type: &str) -> Arc<dyn ProviderPoolAdapter> {
self.adapters
.get(provider_type.trim().to_ascii_lowercase().as_str())
.cloned()
.unwrap_or_else(|| self.default_adapter.clone())
}
pub fn provider_types(&self) -> impl Iterator<Item = &str> {
self.adapters.keys().map(String::as_str)
}
pub fn provider_types_for_capability(&self, capability: ProviderPoolCapability) -> Vec<String> {
self.adapters
.iter()
.filter(|(_, adapter)| adapter.capabilities().supports(capability))
.map(|(provider_type, _)| provider_type.clone())
.collect()
}
pub fn supports_quota_refresh(&self, provider_type: &str) -> bool {
self.adapter(provider_type).supports_quota_refresh()
}
pub fn quota_refresh_endpoint_for_provider(
&self,
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
self.adapter(provider_type)
.quota_refresh_endpoint(endpoints, include_inactive)
}
pub fn quota_refresh_unsupported_message(&self, provider_type: &str) -> String {
self.adapter(provider_type)
.quota_refresh_unsupported_message()
}
pub fn quota_refresh_missing_endpoint_message(&self, provider_type: &str) -> String {
self.adapter(provider_type)
.quota_refresh_missing_endpoint_message()
}
pub fn normalize_scheduling_presets(
&self,
provider_type: &str,
scheduling_presets: &[PoolSchedulingPreset],
) -> Vec<PoolSchedulingPreset> {
normalize_provider_scheduling_presets(
self.adapter(provider_type).as_ref(),
scheduling_presets,
)
}
pub fn member_signals(
&self,
provider_type: &str,
key: &StoredProviderCatalogKey,
auth_config: Option<&Map<String, Value>>,
) -> aether_pool_core::PoolMemberSignals {
let adapter = self.adapter(provider_type);
let input = ProviderPoolMemberInput {
provider_type,
key,
auth_config,
};
adapter.member_signals(&input)
}
}

334
crates/aether-provider-transport/src/kiro/credentials.rs
View File

@@ -1,333 +1,7 @@
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::time::{SystemTime, UNIX_EPOCH};
pub const DEFAULT_REGION: &str = "us-east-1";
pub const DEFAULT_KIRO_VERSION: &str = "0.3.210";
pub const DEFAULT_NODE_VERSION: &str = "22.21.1";
pub const DEFAULT_SYSTEM_VERSION: &str = "other#unknown";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KiroAuthConfig {
pub auth_method: Option<String>,
pub refresh_token: Option<String>,
pub expires_at: Option<u64>,
pub profile_arn: Option<String>,
pub region: Option<String>,
pub auth_region: Option<String>,
pub api_region: Option<String>,
pub client_id: Option<String>,
pub client_secret: Option<String>,
pub machine_id: Option<String>,
pub kiro_version: Option<String>,
pub system_version: Option<String>,
pub node_version: Option<String>,
pub access_token: Option<String>,
}
impl KiroAuthConfig {
pub fn from_raw_json(raw: Option<&str>) -> Option<Self> {
let raw = raw?.trim();
if raw.is_empty() {
return None;
}
let parsed: Value = serde_json::from_str(raw).ok()?;
Self::from_json_value(&parsed)
}
pub fn from_json_value(raw: &Value) -> Option<Self> {
let object = raw.as_object()?;
Some(Self {
auth_method: get_nonempty_string(
object,
&["auth_method", "authMethod", "auth_type", "authType"],
)
.map(|value| normalize_auth_method(&value)),
refresh_token: get_nonempty_string(object, &["refresh_token", "refreshToken"]),
expires_at: get_epoch_seconds(object.get("expires_at"))
.or_else(|| get_epoch_seconds(object.get("expiresAt"))),
profile_arn: get_nonempty_string(object, &["profile_arn", "profileArn"]),
region: get_nonempty_string(object, &["region"]),
auth_region: get_nonempty_string(object, &["auth_region", "authRegion"]),
api_region: get_nonempty_string(object, &["api_region", "apiRegion"]),
client_id: get_nonempty_string(object, &["client_id", "clientId"]),
client_secret: get_nonempty_string(object, &["client_secret", "clientSecret"]),
machine_id: get_nonempty_string(object, &["machine_id", "machineId"]),
kiro_version: get_nonempty_string(object, &["kiro_version", "kiroVersion"]),
system_version: get_nonempty_string(object, &["system_version", "systemVersion"]),
node_version: get_nonempty_string(object, &["node_version", "nodeVersion"]),
access_token: get_nonempty_string(object, &["access_token", "accessToken"]),
})
}
pub fn to_json_value(&self) -> Value {
let mut object = serde_json::Map::new();
insert_optional_string(&mut object, "auth_method", self.auth_method.as_deref());
insert_optional_string(&mut object, "refresh_token", self.refresh_token.as_deref());
if let Some(expires_at) = self.expires_at {
object.insert("expires_at".to_string(), Value::from(expires_at));
}
insert_optional_string(&mut object, "profile_arn", self.profile_arn.as_deref());
insert_optional_string(&mut object, "region", self.region.as_deref());
insert_optional_string(&mut object, "auth_region", self.auth_region.as_deref());
insert_optional_string(&mut object, "api_region", self.api_region.as_deref());
insert_optional_string(&mut object, "client_id", self.client_id.as_deref());
insert_optional_string(&mut object, "client_secret", self.client_secret.as_deref());
insert_optional_string(&mut object, "machine_id", self.machine_id.as_deref());
insert_optional_string(&mut object, "kiro_version", self.kiro_version.as_deref());
insert_optional_string(
&mut object,
"system_version",
self.system_version.as_deref(),
);
insert_optional_string(&mut object, "node_version", self.node_version.as_deref());
insert_optional_string(&mut object, "access_token", self.access_token.as_deref());
Value::Object(object)
}
pub fn effective_api_region(&self) -> &str {
self.api_region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_REGION)
}
pub fn effective_auth_region(&self) -> &str {
self.auth_region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
self.region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
})
.unwrap_or(DEFAULT_REGION)
}
pub fn effective_kiro_version(&self) -> &str {
self.kiro_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_KIRO_VERSION)
}
pub fn effective_system_version(&self) -> &str {
self.system_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_SYSTEM_VERSION)
}
pub fn effective_node_version(&self) -> &str {
self.node_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_NODE_VERSION)
}
pub fn cached_access_token(&self) -> Option<&str> {
self.access_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn cached_access_token_requires_refresh(&self, skew_seconds: u64) -> bool {
let Some(expires_at) = self.expires_at else {
return self.can_refresh_access_token();
};
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|value| value.as_secs())
.unwrap_or_default();
now >= expires_at.saturating_sub(skew_seconds)
}
pub fn is_idc_auth(&self) -> bool {
let explicit_method = self
.auth_method
.as_deref()
.map(normalize_auth_method)
.unwrap_or_else(|| "social".to_string());
if explicit_method != "social" {
return matches!(explicit_method.as_str(), "idc" | "external_idp");
}
self.client_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
&& self
.client_secret
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
}
pub fn uses_external_idp_token_type(&self) -> bool {
self.auth_method
.as_deref()
.map(normalize_auth_method)
.as_deref()
== Some("external_idp")
}
pub fn profile_arn_for_payload(&self) -> Option<&str> {
if self.is_idc_auth() {
return None;
}
self.profile_arn
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn profile_arn_for_mcp(&self) -> Option<&str> {
self.profile_arn
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn can_refresh_access_token(&self) -> bool {
let refresh_token = self
.refresh_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| value.len() >= 100 && !value.contains("..."));
if refresh_token.is_none() {
return false;
}
if !self.is_idc_auth() {
return true;
}
self.client_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
&& self
.client_secret
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
}
}
pub fn normalize_machine_id(raw: &str) -> Option<String> {
let raw = raw.trim();
if raw.is_empty() {
return None;
}
if raw.len() == 64 && raw.bytes().all(|byte| byte.is_ascii_hexdigit()) {
return Some(raw.to_ascii_lowercase());
}
if raw.len() == 36
&& raw.chars().enumerate().all(|(idx, ch)| match idx {
8 | 13 | 18 | 23 => ch == '-',
_ => ch.is_ascii_hexdigit(),
})
{
let normalized = raw.replace('-', "").to_ascii_lowercase();
return Some(format!("{normalized}{normalized}"));
}
None
}
pub fn generate_machine_id(
auth_config: &KiroAuthConfig,
fallback_secret: Option<&str>,
) -> Option<String> {
if let Some(machine_id) = auth_config
.machine_id
.as_deref()
.and_then(normalize_machine_id)
{
return Some(machine_id);
}
let seed = auth_config
.refresh_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
fallback_secret
.map(str::trim)
.filter(|value| !value.is_empty())
})?;
let mut hasher = Sha256::new();
hasher.update(b"KotlinNativeAPI/");
hasher.update(seed.as_bytes());
Some(format!("{:x}", hasher.finalize()))
}
fn get_nonempty_string(object: &serde_json::Map<String, Value>, keys: &[&str]) -> Option<String> {
keys.iter()
.find_map(|key| object.get(*key))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn insert_optional_string(
object: &mut serde_json::Map<String, Value>,
key: &str,
value: Option<&str>,
) {
let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else {
return;
};
object.insert(key.to_string(), Value::String(value.to_string()));
}
fn get_epoch_seconds(value: Option<&Value>) -> Option<u64> {
match value? {
Value::Number(number) => number.as_u64().or_else(|| {
number
.as_i64()
.and_then(|value| (value >= 0).then_some(value as u64))
}),
Value::String(text) => text.trim().parse::<u64>().ok(),
_ => None,
}
}
fn normalize_auth_method(raw: &str) -> String {
let value = raw.trim().to_ascii_lowercase();
match value.as_str() {
"" => "social".to_string(),
"builder-id"
| "builder_id"
| "builderid"
| "device"
| "device-auth"
| "device_authorization"
| "iam"
| "identity-center"
| "identity_center"
| "identitycenter"
| "idc" => "idc".to_string(),
"external-idp" | "external_idp" | "externalidp" => "external_idp".to_string(),
_ => value,
}
}
pub use aether_oauth::provider::providers::{
generate_kiro_machine_id as generate_machine_id,
normalize_kiro_machine_id as normalize_machine_id, KiroAuthConfig, DEFAULT_REGION,
};
#[cfg(test)]
mod tests {

25
crates/aether-provider-transport/src/kiro/refresh.rs
View File

@@ -1,7 +1,4 @@
use std::collections::BTreeMap;
use aether_oauth::provider::providers::KiroProviderOAuthAdapter as CoreKiroProviderOAuthAdapter;
use aether_oauth::provider::{ProviderOAuthAccount, ProviderOAuthAdapter};
use async_trait::async_trait;
use super::super::oauth_refresh::{
@@ -48,26 +45,10 @@ impl KiroOAuthRefreshAdapter {
let oauth_executor =
ProviderOAuthLocalHttpExecutor::new(PROVIDER_TYPE, transport, executor);
let ctx = provider_oauth_transport_context_from_snapshot(transport);
let account = ProviderOAuthAccount {
provider_type: PROVIDER_TYPE.to_string(),
access_token: auth_config
.cached_access_token()
.map(ToOwned::to_owned)
.unwrap_or_default(),
auth_config: auth_config.to_json_value(),
expires_at_unix_secs: auth_config.expires_at,
identity: BTreeMap::new(),
};
let refreshed = adapter
.refresh(&oauth_executor, &ctx, &account)
adapter
.refresh_auth_config(&oauth_executor, &ctx, auth_config)
.await
.map_err(|error| oauth_error_to_local_refresh_error(PROVIDER_TYPE, error))?;
KiroAuthConfig::from_json_value(&refreshed.auth_config).ok_or_else(|| {
LocalOAuthRefreshError::InvalidResponse {
provider_type: PROVIDER_TYPE,
message: "kiro refresh returned invalid auth_config".to_string(),
}
})
.map_err(|error| oauth_error_to_local_refresh_error(PROVIDER_TYPE, error))
}
fn auth_config_from_entry(entry: &CachedOAuthEntry) -> Option<KiroAuthConfig> {