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

fix(providers): refresh Antigravity grouped quotas

Browse Source
This commit is contained in:
ZheFox
2026-09-03 10:46:44 +08:00
parent cc6f5e89b6
commit 45c840b8d3
20 changed files with 976 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-provider/pool/src/lib.rs
+71 -9
View File
@@ -15,10 +15,11 @@ pub use 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_codex_pool_reset_credit_consume_request,
build_codex_pool_reset_credits_request, build_gemini_cli_pool_quota_request,
build_kiro_pool_quota_request, build_windsurf_pool_model_configs_request,
build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request,
build_chatgpt_web_pool_quota_request, build_codex_pool_quota_request,
build_codex_pool_reset_credit_consume_request, build_codex_pool_reset_credits_request,
build_gemini_cli_pool_quota_request, build_kiro_pool_quota_request,
build_windsurf_pool_model_configs_request,
build_windsurf_pool_model_configs_request_with_base_url, build_windsurf_pool_quota_request,
build_windsurf_pool_quota_request_with_base_url, build_windsurf_pool_rate_limit_request,
build_windsurf_pool_rate_limit_request_with_base_url, enrich_chatgpt_web_quota_metadata,
@@ -27,11 +28,12 @@ pub use providers::{
AntigravityProviderPoolAdapter, ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter,
DefaultProviderPoolAdapter, GeminiCliProviderPoolAdapter, GrokProviderPoolAdapter,
KiroPoolQuotaAuthInput, KiroProviderPoolAdapter, UnsupportedQuotaProviderPoolAdapter,
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, CHATGPT_WEB_CONVERSATION_INIT_PATH,
CHATGPT_WEB_DEFAULT_BASE_URL, CODEX_WHAM_RESET_CREDITS_CONSUME_URL,
CODEX_WHAM_RESET_CREDITS_URL, CODEX_WHAM_USAGE_URL, GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH,
GEMINI_CLI_USER_AGENT, KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH, WINDSURF_USER_STATUS_PATH,
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH,
CHATGPT_WEB_CONVERSATION_INIT_PATH, CHATGPT_WEB_DEFAULT_BASE_URL,
CODEX_WHAM_RESET_CREDITS_CONSUME_URL, CODEX_WHAM_RESET_CREDITS_URL, CODEX_WHAM_USAGE_URL,
GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH, GEMINI_CLI_USER_AGENT, KIRO_USAGE_LIMITS_PATH,
KIRO_USAGE_SDK_VERSION, WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH,
WINDSURF_USER_STATUS_PATH,
};
pub use quota::{
provider_pool_key_account_quota_exhausted, provider_pool_key_model_quota_exhausted,
@@ -883,6 +885,66 @@ mod tests {
assert!(!available.quota_exhausted);
}
#[test]
fn antigravity_tiered_model_quota_wins_over_display_only_group_windows() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "antigravity",
"exhausted": true,
"windows": [{
"code": "model:gemini-3.7-flash-tiered",
"scope": "model",
"model": "gemini-3.7-flash-tiered",
"remaining_ratio": 0.906,
"used_ratio": 0.094,
"is_exhausted": false
}, {
"code": "group:0:3p-5h",
"scope": "quota_group",
"quota_group": "group:0",
"bucket_id": "3p-5h",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals =
service.member_signals("antigravity", &key, None, Some("gemini-3.7-flash-tiered"));
assert!(!signals.quota_exhausted);
assert!(!signals.quota_hard_blocked);
}
#[test]
fn antigravity_tiered_variant_does_not_match_another_model_family() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "antigravity",
"exhausted": false,
"windows": [{
"code": "model:gemini-3.7-pro-tiered",
"scope": "model",
"model": "gemini-3.7-pro-tiered",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals =
service.member_signals("antigravity", &key, None, Some("gemini-3.7-flash-tiered"));
assert!(!signals.quota_exhausted);
assert!(!signals.quota_hard_blocked);
}
#[test]
fn codex_standard_and_spark_quota_families_are_independent() {
let service = ProviderPoolService::with_builtin_adapters();
crates/aether-provider/pool/src/providers/antigravity.rs
+84
View File
@@ -12,6 +12,8 @@ use crate::quota::provider_pool_model_quota_exhausted;
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
pub const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels";
pub const ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH: &str =
"/v1internal:retrieveUserQuotaSummary";
#[derive(Debug, Clone, Default)]
pub struct AntigravityProviderPoolAdapter;
@@ -92,3 +94,85 @@ pub fn build_antigravity_pool_quota_request(
accept_invalid_certs: false,
}
}
pub fn build_antigravity_pool_quota_summary_request(
key_id: &str,
endpoint_base_url: &str,
authorization: (String, String),
project_id: Option<&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());
let json_body = project_id
.map(str::trim)
.filter(|project_id| !project_id.is_empty())
.map_or_else(|| json!({}), |project_id| json!({ "project": project_id }));
ProviderPoolQuotaRequestSpec {
request_id: format!("antigravity-quota-summary:{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_RETRIEVE_USER_QUOTA_SUMMARY_PATH
),
headers,
content_type: Some("application/json".to_string()),
json_body: Some(json_body),
client_api_format: "gemini:generate_content".to_string(),
provider_api_format: "antigravity:retrieve_user_quota_summary".to_string(),
model_name: Some("retrieveUserQuotaSummary".to_string()),
accept_invalid_certs: false,
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use super::{
build_antigravity_pool_quota_summary_request, ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH,
};
#[test]
fn grouped_quota_request_can_retry_without_project_on_the_same_endpoint() {
let with_project = build_antigravity_pool_quota_summary_request(
"key-1",
"https://daily-cloudcode-pa.googleapis.com/",
("authorization".to_string(), "Bearer token".to_string()),
Some("project-1"),
BTreeMap::new(),
);
let without_project = build_antigravity_pool_quota_summary_request(
"key-1",
"https://daily-cloudcode-pa.googleapis.com/",
("authorization".to_string(), "Bearer token".to_string()),
None,
BTreeMap::new(),
);
assert_eq!(
with_project.url,
format!(
"https://daily-cloudcode-pa.googleapis.com{ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH}"
)
);
assert_eq!(
with_project.json_body,
Some(json!({"project": "project-1"}))
);
assert_eq!(without_project.url, with_project.url);
assert_eq!(without_project.json_body, Some(json!({})));
}
}
crates/aether-provider/pool/src/providers/mod.rs
+2 -1
View File
@@ -10,7 +10,8 @@ pub mod windsurf;
pub use antigravity::AntigravityProviderPoolAdapter;
pub use antigravity::{
build_antigravity_pool_quota_request, ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH,
build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request,
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH,
};
pub use chatgpt_web::ChatGptWebProviderPoolAdapter;
pub use chatgpt_web::{
crates/aether-provider/pool/src/quota.rs
+1
View File
@@ -398,6 +398,7 @@ fn provider_pool_is_specific_model_token(token: &str) -> bool {
"team",
"think",
"thinking",
"tiered",
"vendor",
]
.contains(&token)