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

feat(grok): add provider pool and transport support

Browse Source
This commit is contained in:
mayrain
2026-05-16 19:22:54 +08:00
parent e5f1f52759
commit cbfe1d378f
20 changed files with 2616 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

112
crates/aether-provider-pool/src/lib.rs
View File

@@ -17,9 +17,11 @@ 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,
enrich_chatgpt_web_quota_metadata, grok_mode_id_for_model, grok_pool_tier_from_quota_bucket,
grok_quota_window_key_for_model, grok_supported_quota_windows_for_tier,
normalize_chatgpt_web_image_quota_limit, AntigravityProviderPoolAdapter,
ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter, DefaultProviderPoolAdapter,
GrokProviderPoolAdapter, 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,
@@ -65,6 +67,7 @@ mod tests {
"claude_code",
"codex",
"gemini_cli",
"grok",
"kiro",
"vertex_ai"
]
@@ -82,10 +85,11 @@ mod tests {
assert_eq!(
service.provider_types_for_capability(ProviderPoolCapability::QuotaRefresh),
["antigravity", "chatgpt_web", "codex", "kiro"]
["antigravity", "chatgpt_web", "codex", "grok", "kiro"]
);
assert!(service.supports_quota_refresh("codex"));
assert!(service.supports_quota_refresh("antigravity"));
assert!(service.supports_quota_refresh("grok"));
assert!(!service.supports_quota_refresh("gemini_cli"));
assert_eq!(
service.quota_refresh_unsupported_message("claude_code"),
@@ -247,8 +251,11 @@ mod tests {
.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"]));
assert_eq!(free_first["providers"], json!(["codex", "grok", "kiro"]));
assert_eq!(
recent_refresh["providers"],
json!(["codex", "grok", "kiro"])
);
}
#[test]
@@ -331,6 +338,37 @@ mod tests {
}))),
"chatgpt_web",
));
assert!(provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"grok": {
"quota_by_model": {
"quota_fast": {
"is_exhausted": true,
"remaining": 0.0
}
}
}
}))),
"grok",
));
assert!(!provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"grok": {
"pool_tier": "basic",
"quota_by_model": {
"quota_fast": {
"is_exhausted": false,
"remaining": 1.0
},
"quota_heavy": {
"is_exhausted": true,
"remaining": 0.0
}
}
}
}))),
"grok",
));
assert!(!provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
@@ -342,6 +380,68 @@ mod tests {
));
}
#[test]
fn grok_quota_tier_boundaries_match_pool_modes() {
assert_eq!(
grok_supported_quota_windows_for_tier(Some("basic")),
[("quota_fast", "fast")]
);
assert_eq!(
grok_supported_quota_windows_for_tier(Some("super")),
[
("quota_auto", "auto"),
("quota_fast", "fast"),
("quota_expert", "expert"),
("quota_grok_4_3", "grok-420-computer-use-sa")
]
);
assert_eq!(
grok_supported_quota_windows_for_tier(Some("heavy")),
[
("quota_auto", "auto"),
("quota_fast", "fast"),
("quota_expert", "expert"),
("quota_heavy", "heavy"),
("quota_grok_4_3", "grok-420-computer-use-sa")
]
);
}
#[test]
fn grok_pool_tier_infers_from_live_quota_totals() {
let bucket = json!({
"quota_by_model": {
"quota_fast": {
"remaining": 20.0,
"total": 30.0
},
"quota_auto": {
"remaining": 7.0,
"total": 7.0
}
}
});
let bucket = bucket.as_object().expect("bucket should be object");
assert_eq!(grok_pool_tier_from_quota_bucket(bucket), Some("basic"));
}
#[test]
fn grok_model_name_maps_to_quota_window() {
assert_eq!(
grok_quota_window_key_for_model(Some("grok-4.20-fast")),
Some("quota_fast")
);
assert_eq!(
grok_quota_window_key_for_model(Some("grok-4.20-multi-agent-0309")),
Some("quota_heavy")
);
assert_eq!(
grok_quota_window_key_for_model(Some("grok-4.3-beta")),
Some("quota_grok_4_3")
);
}
#[test]
fn plan_tier_derivation_normalizes_provider_prefix() {
let key = sample_key(Some(json!({

220
crates/aether-provider-pool/src/providers/grok.rs Normal file
View File

@@ -0,0 +1,220 @@
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
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,
};
pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")];
pub const GROK_QUOTA_WINDOWS_SUPER: &[(&str, &str)] = &[
("quota_auto", "auto"),
("quota_fast", "fast"),
("quota_expert", "expert"),
("quota_grok_4_3", "grok-420-computer-use-sa"),
];
pub const GROK_QUOTA_WINDOWS_HEAVY: &[(&str, &str)] = &[
("quota_auto", "auto"),
("quota_fast", "fast"),
("quota_expert", "expert"),
("quota_heavy", "heavy"),
("quota_grok_4_3", "grok-420-computer-use-sa"),
];
#[derive(Debug, Clone, Default)]
pub struct GrokProviderPoolAdapter;
impl ProviderPoolAdapter for GrokProviderPoolAdapter {
fn provider_type(&self) -> &'static str {
"grok"
}
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, "openai:chat")
})
.or_else(|| provider_pool_matching_endpoint(endpoints, include_inactive, |_| true))
}
fn quota_refresh_missing_endpoint_message(&self) -> String {
"找不到有效的 Grok 端点".to_string()
}
}
pub fn grok_supported_quota_windows_for_tier(
tier: Option<&str>,
) -> &'static [(&'static str, &'static str)] {
match grok_normalize_pool_tier(tier) {
Some("basic") => GROK_QUOTA_WINDOWS_BASIC,
Some("super") => GROK_QUOTA_WINDOWS_SUPER,
Some("heavy") => GROK_QUOTA_WINDOWS_HEAVY,
_ => GROK_QUOTA_WINDOWS_HEAVY,
}
}
pub fn grok_pool_tier_from_quota_bucket(bucket: &Map<String, Value>) -> Option<&'static str> {
if let Some(tier) = grok_normalize_pool_tier(
grok_bucket_string(bucket, &["pool_tier", "tier", "plan_type", "plan"]).as_deref(),
) {
return Some(tier);
}
if let Some(auto_total) = grok_quota_total(bucket, "quota_auto") {
if (auto_total - 50.0).abs() < f64::EPSILON {
return Some("super");
}
if (auto_total - 150.0).abs() < f64::EPSILON {
return Some("heavy");
}
}
if let Some(fast_total) = grok_quota_total(bucket, "quota_fast") {
if (fast_total - 30.0).abs() < f64::EPSILON {
return Some("basic");
}
if (fast_total - 140.0).abs() < f64::EPSILON {
return Some("super");
}
if (fast_total - 400.0).abs() < f64::EPSILON {
return Some("heavy");
}
}
None
}
pub fn grok_quota_window_key_for_model(model: Option<&str>) -> Option<&'static str> {
Some(match grok_mode_id_for_model(model) {
"fast" => "quota_fast",
"auto" => "quota_auto",
"expert" => "quota_expert",
"heavy" => "quota_heavy",
"grok-420-computer-use-sa" => "quota_grok_4_3",
_ => return None,
})
}
pub fn grok_mode_id_for_model(model: Option<&str>) -> &'static str {
let model = model.unwrap_or_default().to_ascii_lowercase();
if model.contains("4.3") || model.contains("computer") {
"grok-420-computer-use-sa"
} else if model.contains("multi-agent") {
"heavy"
} else if model.contains("non-reasoning") || model.contains("fast") || model.contains("lite") {
"fast"
} else if model.contains("expert") || model.contains("reasoning") {
"expert"
} else if model.contains("0309-heavy") {
"auto"
} else if model.contains("heavy") {
"heavy"
} else {
"auto"
}
}
fn grok_normalize_pool_tier(value: Option<&str>) -> Option<&'static str> {
match value?.trim().to_ascii_lowercase().as_str() {
"basic" => Some("basic"),
"super" => Some("super"),
"heavy" => Some("heavy"),
_ => None,
}
}
fn grok_bucket_string(bucket: &Map<String, Value>, fields: &[&str]) -> Option<String> {
fields.iter().find_map(|field| {
bucket
.get(*field)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn grok_quota_total(quota_by_model: &Map<String, Value>, key: &str) -> Option<f64> {
let models = quota_by_model
.get("quota_by_model")
.or_else(|| quota_by_model.get("models"))
.and_then(Value::as_object)
.unwrap_or(quota_by_model);
models
.get(key)
.and_then(Value::as_object)
.and_then(|quota| quota.get("total"))
.and_then(Value::as_f64)
.filter(|value| value.is_finite() && *value > 0.0)
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
let models = bucket
.get("quota_by_model")
.or_else(|| bucket.get("models"))
.and_then(Value::as_object);
let Some(models) = models else {
return false;
};
let supported_mode_keys =
grok_supported_quota_windows_for_tier(grok_pool_tier_from_quota_bucket(bucket))
.iter()
.map(|(quota_key, _)| *quota_key)
.collect::<Vec<_>>();
let mut model_count = 0usize;
let mut exhausted_count = 0usize;
for (model_key, item) in models.iter() {
if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) {
continue;
}
let Some(item) = item.as_object() else {
continue;
};
let has_quota_data = provider_pool_json_bool(item.get("is_exhausted")).is_some()
|| provider_pool_json_f64(item.get("used_percent")).is_some()
|| provider_pool_json_f64(item.get("remaining")).is_some()
|| provider_pool_json_f64(item.get("remaining_fraction")).is_some();
if !has_quota_data {
continue;
}
model_count += 1;
if provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
|| provider_pool_json_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0)
|| provider_pool_json_f64(item.get("remaining")).is_some_and(|value| value <= 0.0)
|| provider_pool_json_f64(item.get("remaining_fraction"))
.is_some_and(|value| value <= 0.0)
{
exhausted_count += 1;
}
}
model_count > 0 && model_count == exhausted_count
}

5
crates/aether-provider-pool/src/providers/mod.rs
View File

@@ -2,6 +2,7 @@ pub mod antigravity;
pub mod chatgpt_web;
pub mod codex;
pub mod default;
pub mod grok;
pub mod kiro;
pub mod unsupported;
@@ -18,6 +19,10 @@ pub use chatgpt_web::{
pub use codex::CodexProviderPoolAdapter;
pub use codex::{build_codex_pool_quota_request, CODEX_WHAM_USAGE_URL};
pub use default::DefaultProviderPoolAdapter;
pub use grok::{
grok_mode_id_for_model, grok_pool_tier_from_quota_bucket, grok_quota_window_key_for_model,
grok_supported_quota_windows_for_tier, GrokProviderPoolAdapter,
};
pub use kiro::KiroProviderPoolAdapter;
pub use kiro::{
build_kiro_pool_quota_request, KiroPoolQuotaAuthInput, KIRO_USAGE_LIMITS_PATH,

6
crates/aether-provider-pool/src/service.rs
View File

@@ -12,8 +12,9 @@ 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,
DefaultProviderPoolAdapter, GrokProviderPoolAdapter, KiroProviderPoolAdapter,
CLAUDE_CODE_PROVIDER_POOL_ADAPTER, GEMINI_CLI_PROVIDER_POOL_ADAPTER,
VERTEX_AI_PROVIDER_POOL_ADAPTER,
};
#[derive(Clone)]
@@ -50,6 +51,7 @@ impl ProviderPoolService {
.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(GrokProviderPoolAdapter))
.with_adapter(Arc::new(KiroProviderPoolAdapter))
.with_adapter(Arc::new(ChatGptWebProviderPoolAdapter))
.with_adapter(Arc::new(VERTEX_AI_PROVIDER_POOL_ADAPTER))