feat: support Gemini CLI v1internal quota

This commit is contained in:
Mas0nShi
2026-05-21 15:38:10 +08:00
parent b84e4a96e2
commit e53d5f07e8
67 changed files with 1898 additions and 79 deletions
+45 -11
View File
@@ -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",
+3 -4
View File
@@ -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))