mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-08 02:17:46 +08:00
feat: support Gemini CLI v1internal quota
This commit is contained in:
@@ -16,15 +16,16 @@ 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_kiro_pool_quota_request,
|
||||
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,
|
||||
build_codex_pool_quota_request, build_gemini_cli_pool_quota_request,
|
||||
build_kiro_pool_quota_request, 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, GeminiCliProviderPoolAdapter, GrokProviderPoolAdapter,
|
||||
KiroPoolQuotaAuthInput, KiroProviderPoolAdapter, UnsupportedQuotaProviderPoolAdapter,
|
||||
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, CHATGPT_WEB_CONVERSATION_INIT_PATH,
|
||||
CHATGPT_WEB_DEFAULT_BASE_URL, CODEX_WHAM_USAGE_URL, GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH,
|
||||
GEMINI_CLI_USER_AGENT, KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
|
||||
};
|
||||
pub use quota::{
|
||||
provider_pool_key_account_quota_exhausted, provider_pool_key_scheduling_label,
|
||||
@@ -85,12 +86,19 @@ mod tests {
|
||||
|
||||
assert_eq!(
|
||||
service.provider_types_for_capability(ProviderPoolCapability::QuotaRefresh),
|
||||
["antigravity", "chatgpt_web", "codex", "grok", "kiro"]
|
||||
[
|
||||
"antigravity",
|
||||
"chatgpt_web",
|
||||
"codex",
|
||||
"gemini_cli",
|
||||
"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!(service.supports_quota_refresh("gemini_cli"));
|
||||
assert_eq!(
|
||||
service.quota_refresh_unsupported_message("claude_code"),
|
||||
"Claude Code 暂不支持自动刷新额度:上游没有稳定可用的账号额度查询接口"
|
||||
@@ -120,6 +128,32 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_cli_quota_request_uses_v1internal_retrieve_user_quota() {
|
||||
let spec = build_gemini_cli_pool_quota_request(
|
||||
"key-1",
|
||||
"https://cloudcode-pa.googleapis.com/",
|
||||
("authorization".to_string(), "Bearer access".to_string()),
|
||||
"project-1",
|
||||
);
|
||||
|
||||
assert_eq!(spec.method, "POST");
|
||||
assert_eq!(
|
||||
spec.url,
|
||||
"https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota"
|
||||
);
|
||||
assert_eq!(
|
||||
spec.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer access")
|
||||
);
|
||||
assert_eq!(
|
||||
spec.json_body.as_ref().and_then(|body| body.get("project")),
|
||||
Some(&json!("project-1"))
|
||||
);
|
||||
assert_eq!(spec.client_api_format, "gemini:generate_content");
|
||||
assert_eq!(spec.provider_api_format, "gemini_cli:retrieve_user_quota");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_quota_request_uses_wham_usage_endpoint() {
|
||||
let spec = build_codex_pool_quota_request(
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
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 GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH: &str = "/v1internal:retrieveUserQuota";
|
||||
pub const GEMINI_CLI_USER_AGENT: &str = "GeminiCLI/0.1.5 (Windows; AMD64)";
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct GeminiCliProviderPoolAdapter;
|
||||
|
||||
impl ProviderPoolAdapter for GeminiCliProviderPoolAdapter {
|
||||
fn provider_type(&self) -> &'static str {
|
||||
"gemini_cli"
|
||||
}
|
||||
|
||||
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_gemini_cli_pool_quota_request(
|
||||
key_id: &str,
|
||||
endpoint_base_url: &str,
|
||||
authorization: (String, String),
|
||||
project_id: &str,
|
||||
) -> ProviderPoolQuotaRequestSpec {
|
||||
let headers = BTreeMap::from([
|
||||
("authorization".to_string(), authorization.1),
|
||||
("content-type".to_string(), "application/json".to_string()),
|
||||
("accept".to_string(), "application/json".to_string()),
|
||||
("user-agent".to_string(), GEMINI_CLI_USER_AGENT.to_string()),
|
||||
]);
|
||||
|
||||
ProviderPoolQuotaRequestSpec {
|
||||
request_id: format!("gemini-cli-quota:{key_id}"),
|
||||
provider_name: "gemini_cli".to_string(),
|
||||
quota_kind: "gemini_cli".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: format!(
|
||||
"{}{}",
|
||||
endpoint_base_url.trim_end_matches('/'),
|
||||
GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH
|
||||
),
|
||||
headers,
|
||||
content_type: Some("application/json".to_string()),
|
||||
json_body: Some(json!({
|
||||
"project": project_id,
|
||||
"userAgent": GEMINI_CLI_USER_AGENT,
|
||||
})),
|
||||
client_api_format: "gemini:generate_content".to_string(),
|
||||
provider_api_format: "gemini_cli:retrieve_user_quota".to_string(),
|
||||
model_name: Some("retrieveUserQuota".to_string()),
|
||||
accept_invalid_certs: false,
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@ pub mod antigravity;
|
||||
pub mod chatgpt_web;
|
||||
pub mod codex;
|
||||
pub mod default;
|
||||
pub mod gemini_cli;
|
||||
pub mod grok;
|
||||
pub mod kiro;
|
||||
pub mod unsupported;
|
||||
@@ -19,6 +20,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 gemini_cli::GeminiCliProviderPoolAdapter;
|
||||
pub use gemini_cli::{
|
||||
build_gemini_cli_pool_quota_request, GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH, GEMINI_CLI_USER_AGENT,
|
||||
};
|
||||
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,
|
||||
@@ -30,5 +35,5 @@ pub use kiro::{
|
||||
};
|
||||
pub use unsupported::{
|
||||
UnsupportedQuotaProviderPoolAdapter, CLAUDE_CODE_PROVIDER_POOL_ADAPTER,
|
||||
GEMINI_CLI_PROVIDER_POOL_ADAPTER, VERTEX_AI_PROVIDER_POOL_ADAPTER,
|
||||
VERTEX_AI_PROVIDER_POOL_ADAPTER,
|
||||
};
|
||||
|
||||
@@ -34,12 +34,6 @@ pub const CLAUDE_CODE_PROVIDER_POOL_ADAPTER: UnsupportedQuotaProviderPoolAdapter
|
||||
"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",
|
||||
|
||||
@@ -12,9 +12,8 @@ use crate::presets::normalize_provider_scheduling_presets;
|
||||
use crate::provider::{ProviderPoolAdapter, ProviderPoolMemberInput};
|
||||
use crate::providers::{
|
||||
AntigravityProviderPoolAdapter, ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter,
|
||||
DefaultProviderPoolAdapter, GrokProviderPoolAdapter, KiroProviderPoolAdapter,
|
||||
CLAUDE_CODE_PROVIDER_POOL_ADAPTER, GEMINI_CLI_PROVIDER_POOL_ADAPTER,
|
||||
VERTEX_AI_PROVIDER_POOL_ADAPTER,
|
||||
DefaultProviderPoolAdapter, GeminiCliProviderPoolAdapter, GrokProviderPoolAdapter,
|
||||
KiroProviderPoolAdapter, CLAUDE_CODE_PROVIDER_POOL_ADAPTER, VERTEX_AI_PROVIDER_POOL_ADAPTER,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -50,7 +49,7 @@ impl ProviderPoolService {
|
||||
.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(GeminiCliProviderPoolAdapter))
|
||||
.with_adapter(Arc::new(GrokProviderPoolAdapter))
|
||||
.with_adapter(Arc::new(KiroProviderPoolAdapter))
|
||||
.with_adapter(Arc::new(ChatGptWebProviderPoolAdapter))
|
||||
|
||||
Reference in New Issue
Block a user