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
@@ -370,6 +370,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: None,
},
@@ -90,6 +90,7 @@ mod tests {
expires_at_unix_secs: Some(1),
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "sk-test".to_string(),
decrypted_auth_config: Some("{\"token\":\"x\"}".to_string()),
},
@@ -130,6 +130,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: Some(json!({"transport_profile":"chrome_136"})),
upstream_metadata: None,
decrypted_api_key: "sk-ant-123".to_string(),
decrypted_auth_config: None,
},
@@ -494,6 +494,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
@@ -357,6 +357,7 @@ mod tests {
}
}
})),
upstream_metadata: None,
decrypted_api_key: "sk-test".to_string(),
decrypted_auth_config: None,
},
@@ -412,6 +413,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: None,
},
@@ -0,0 +1,25 @@
mod request;
mod url;
pub use request::{
build_gemini_cli_v1internal_request, resolve_gemini_cli_project_id,
GeminiCliRequestEnvelopeSupport, GeminiCliRequestEnvelopeUnsupportedReason,
};
pub use url::{
build_gemini_cli_v1internal_url, GeminiCliRequestUrlAction,
GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH, GEMINI_CLI_USER_AGENT,
GEMINI_CLI_V1INTERNAL_PATH_TEMPLATE,
};
use crate::snapshot::GatewayProviderTransportSnapshot;
pub const GEMINI_CLI_PROVIDER_TYPE: &str = "gemini_cli";
pub const GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME: &str = "gemini_cli:v1internal";
pub fn is_gemini_cli_provider_transport(transport: &GatewayProviderTransportSnapshot) -> bool {
transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(GEMINI_CLI_PROVIDER_TYPE)
}
@@ -0,0 +1,258 @@
use serde_json::{Map, Value};
use crate::snapshot::GatewayProviderTransportSnapshot;
#[derive(Debug, Clone, PartialEq)]
pub enum GeminiCliRequestEnvelopeSupport {
Supported(Value),
Unsupported(GeminiCliRequestEnvelopeUnsupportedReason),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GeminiCliRequestEnvelopeUnsupportedReason {
NonObjectBody,
MissingContents,
MissingProjectId,
MissingUserPromptId,
MissingModel,
}
pub fn resolve_gemini_cli_project_id(
transport: &GatewayProviderTransportSnapshot,
) -> Option<String> {
transport
.key
.upstream_metadata
.as_ref()
.and_then(|metadata| {
find_string_by_paths(
metadata,
&[
&["gemini_cli", "project_id"],
&["gemini_cli", "projectId"],
&["gemini_cli", "cloudaicompanionProject"],
&["gemini_cli", "cloudaicompanionProject", "id"],
&["project_id"],
&["projectId"],
],
)
})
.or_else(|| {
transport
.key
.decrypted_auth_config
.as_deref()
.and_then(parse_project_id_from_auth_config)
})
}
fn parse_project_id_from_auth_config(raw_auth_config: &str) -> Option<String> {
let auth_config = serde_json::from_str::<Value>(raw_auth_config).ok()?;
find_string_by_paths(
&auth_config,
&[
&["project_id"],
&["projectId"],
&["project", "id"],
&["project", "project_id"],
&["project", "projectId"],
&["gemini_cli", "project_id"],
&["gemini_cli", "projectId"],
&["metadata", "project_id"],
&["metadata", "projectId"],
],
)
}
pub fn build_gemini_cli_v1internal_request(
project_id: &str,
user_prompt_id: &str,
model: &str,
request_body: &Value,
) -> GeminiCliRequestEnvelopeSupport {
let project_id = project_id.trim();
if project_id.is_empty() {
return GeminiCliRequestEnvelopeSupport::Unsupported(
GeminiCliRequestEnvelopeUnsupportedReason::MissingProjectId,
);
}
let user_prompt_id = user_prompt_id.trim();
if user_prompt_id.is_empty() {
return GeminiCliRequestEnvelopeSupport::Unsupported(
GeminiCliRequestEnvelopeUnsupportedReason::MissingUserPromptId,
);
}
let model = model.trim();
if model.is_empty() {
return GeminiCliRequestEnvelopeSupport::Unsupported(
GeminiCliRequestEnvelopeUnsupportedReason::MissingModel,
);
}
let Value::Object(source) = request_body else {
return GeminiCliRequestEnvelopeSupport::Unsupported(
GeminiCliRequestEnvelopeUnsupportedReason::NonObjectBody,
);
};
if !source.contains_key("contents") {
return GeminiCliRequestEnvelopeSupport::Unsupported(
GeminiCliRequestEnvelopeUnsupportedReason::MissingContents,
);
}
let mut inner_request: Map<String, Value> = source.clone();
inner_request.remove("model");
inner_request.remove("stream");
GeminiCliRequestEnvelopeSupport::Supported(serde_json::json!({
"model": model,
"project": project_id,
"user_prompt_id": user_prompt_id,
"request": Value::Object(inner_request),
}))
}
fn find_string_by_paths(value: &Value, paths: &[&[&str]]) -> Option<String> {
for path in paths {
let mut current = value;
let mut matched = true;
for segment in *path {
let Some(next) = current.get(*segment) else {
matched = false;
break;
};
current = next;
}
if !matched {
continue;
}
if let Some(string) = current
.as_str()
.map(str::trim)
.filter(|item| !item.is_empty())
{
return Some(string.to_string());
}
}
None
}
#[cfg(test)]
mod tests {
use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use serde_json::json;
use super::{
build_gemini_cli_v1internal_request, resolve_gemini_cli_project_id,
GeminiCliRequestEnvelopeSupport,
};
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Gemini CLI".to_string(),
provider_type: "gemini_cli".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "gemini:generate_content".to_string(),
api_family: None,
endpoint_kind: None,
is_active: true,
base_url: "https://cloudcode-pa.googleapis.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "oauth".to_string(),
is_active: true,
api_formats: None,
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: Some(json!({
"gemini_cli": {
"project_id": "metadata-project"
}
})),
decrypted_api_key: String::new(),
decrypted_auth_config: Some(r#"{"project_id":"auth-project"}"#.to_string()),
},
}
}
#[test]
fn project_id_prefers_upstream_metadata_then_auth_config() {
let mut transport = sample_transport();
assert_eq!(
resolve_gemini_cli_project_id(&transport).as_deref(),
Some("metadata-project")
);
transport.key.upstream_metadata = None;
assert_eq!(
resolve_gemini_cli_project_id(&transport).as_deref(),
Some("auth-project")
);
}
#[test]
fn wraps_gemini_body_in_code_assist_envelope() {
let body = json!({
"model": "ignored",
"stream": true,
"contents": [{"role": "user", "parts": [{"text": "hi"}]}],
"generationConfig": {"temperature": 0.2}
});
let envelope = match build_gemini_cli_v1internal_request(
"project-1",
"trace-1",
"gemini-2.5-pro",
&body,
) {
GeminiCliRequestEnvelopeSupport::Supported(value) => value,
other => panic!("expected supported envelope, got {other:?}"),
};
assert_eq!(envelope["project"], json!("project-1"));
assert_eq!(envelope["user_prompt_id"], json!("trace-1"));
assert_eq!(envelope["model"], json!("gemini-2.5-pro"));
assert_eq!(
envelope["request"]["generationConfig"]["temperature"],
json!(0.2)
);
assert!(envelope["request"].get("model").is_none());
assert!(envelope["request"].get("stream").is_none());
}
}
@@ -0,0 +1,120 @@
use std::collections::BTreeMap;
use url::form_urlencoded;
pub const GEMINI_CLI_USER_AGENT: &str = "GeminiCLI/0.1.5 (Windows; AMD64)";
pub const GEMINI_CLI_V1INTERNAL_PATH_TEMPLATE: &str = "/v1internal:{action}";
pub const GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH: &str = "/v1internal:retrieveUserQuota";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GeminiCliRequestUrlAction {
GenerateContent,
StreamGenerateContent,
}
impl GeminiCliRequestUrlAction {
fn as_str(self) -> &'static str {
match self {
Self::GenerateContent => "generateContent",
Self::StreamGenerateContent => "streamGenerateContent",
}
}
fn is_stream(self) -> bool {
matches!(self, Self::StreamGenerateContent)
}
}
pub fn build_gemini_cli_v1internal_url(
base_url: &str,
action: GeminiCliRequestUrlAction,
query: Option<&BTreeMap<String, String>>,
) -> Option<String> {
let trimmed_base = base_url.trim();
if trimmed_base.is_empty() {
return None;
}
let path = GEMINI_CLI_V1INTERNAL_PATH_TEMPLATE.replace("{action}", action.as_str());
let mut url = format!("{}{}", trimmed_base.trim_end_matches('/'), path);
let mut params = BTreeMap::new();
if let Some(query) = query {
for (key, value) in query {
let key = key.trim();
let value = value.trim();
if key.is_empty()
|| value.is_empty()
|| key.eq_ignore_ascii_case("beta")
|| key.eq_ignore_ascii_case("key")
{
continue;
}
params.insert(key.to_string(), value.to_string());
}
}
if action.is_stream() {
params
.entry(String::from("alt"))
.or_insert_with(|| String::from("sse"));
}
if !params.is_empty() {
let mut serializer = form_urlencoded::Serializer::new(String::new());
for (key, value) in params {
serializer.append_pair(key.as_str(), value.as_str());
}
let query_string = serializer.finish();
if !query_string.is_empty() {
url.push('?');
url.push_str(&query_string);
}
}
Some(url)
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::{
build_gemini_cli_v1internal_url, GeminiCliRequestUrlAction,
GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH,
};
#[test]
fn builds_gemini_cli_stream_url_with_alt_sse() {
let query = BTreeMap::from([
("foo".to_string(), "bar".to_string()),
("key".to_string(), "blocked".to_string()),
]);
assert_eq!(
build_gemini_cli_v1internal_url(
"https://cloudcode-pa.googleapis.com/",
GeminiCliRequestUrlAction::StreamGenerateContent,
Some(&query),
)
.as_deref(),
Some("https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse&foo=bar")
);
}
#[test]
fn builds_gemini_cli_sync_url_without_stream_query() {
assert_eq!(
build_gemini_cli_v1internal_url(
"https://cloudcode-pa.googleapis.com",
GeminiCliRequestUrlAction::GenerateContent,
None,
)
.as_deref(),
Some("https://cloudcode-pa.googleapis.com/v1internal:generateContent")
);
assert_eq!(
GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH,
"/v1internal:retrieveUserQuota"
);
}
}
@@ -195,6 +195,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
@@ -914,6 +914,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: String::new(),
decrypted_auth_config: Some(auth_config.to_string()),
},
@@ -208,6 +208,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "upstream-key".to_string(),
decrypted_auth_config: None,
},
@@ -150,6 +150,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: Some(
r#"{
@@ -223,6 +223,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: Some(raw_auth_config.to_string()),
},
@@ -5,6 +5,7 @@ mod cache;
pub mod claude_code;
pub mod conversion;
mod diagnostics;
pub mod gemini_cli;
mod gemini_files;
mod generic_oauth;
pub mod grok;
@@ -39,6 +40,14 @@ pub use diagnostics::{
append_transport_diagnostics_to_value, build_request_trace_proxy_value,
build_transport_diagnostics,
};
pub use gemini_cli::{
build_gemini_cli_v1internal_request, build_gemini_cli_v1internal_url,
is_gemini_cli_provider_transport, resolve_gemini_cli_project_id,
GeminiCliRequestEnvelopeSupport, GeminiCliRequestEnvelopeUnsupportedReason,
GeminiCliRequestUrlAction, GEMINI_CLI_PROVIDER_TYPE, GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH,
GEMINI_CLI_USER_AGENT, GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
GEMINI_CLI_V1INTERNAL_PATH_TEMPLATE,
};
pub use gemini_files::{
build_gemini_files_headers, build_gemini_files_request_body, build_gemini_files_upstream_url,
gemini_files_transport_unsupported_reason, resolve_gemini_files_auth, GeminiFilesHeadersInput,
@@ -422,6 +422,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: Some(json!({"node_id":"proxy-node-1","kind":"manual"})),
fingerprint: Some(json!({"transport_profile":"chrome_136"})),
upstream_metadata: None,
decrypted_api_key: "sk-test".to_string(),
decrypted_auth_config: None,
},
@@ -685,6 +685,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: Some("{\"refresh_token\":\"rt-1\"}".to_string()),
},
@@ -157,6 +157,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
@@ -304,6 +304,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "sk-test".to_string(),
decrypted_auth_config: None,
},
@@ -292,6 +292,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
@@ -10,6 +10,9 @@ use crate::antigravity::{
AntigravityRequestUrlAction,
};
use crate::claude_code::build_claude_code_messages_url;
use crate::gemini_cli::{
build_gemini_cli_v1internal_url, is_gemini_cli_provider_transport, GeminiCliRequestUrlAction,
};
use crate::snapshot::GatewayProviderTransportSnapshot;
use crate::url::{
build_claude_messages_url, build_gemini_content_url, build_openai_chat_url,
@@ -284,7 +287,9 @@ fn build_transport_hook_url(
));
}
match aether_ai_formats::normalize_api_format_alias(params.provider_api_format).as_str() {
let normalized_provider_api_format =
aether_ai_formats::normalize_api_format_alias(params.provider_api_format);
match normalized_provider_api_format.as_str() {
"gemini:generate_content" => {
if let Some(auth) = resolve_local_vertex_api_key_query_auth(transport) {
return build_vertex_api_key_gemini_content_url(
@@ -339,6 +344,25 @@ fn build_transport_hook_url(
);
}
if is_gemini_cli_provider_transport(transport)
&& normalized_provider_api_format == "gemini:generate_content"
{
let query = params.request_query.map(|raw| {
form_urlencoded::parse(raw.as_bytes())
.into_owned()
.collect::<BTreeMap<String, String>>()
});
return build_gemini_cli_v1internal_url(
&transport.endpoint.base_url,
if params.upstream_is_stream {
GeminiCliRequestUrlAction::StreamGenerateContent
} else {
GeminiCliRequestUrlAction::GenerateContent
},
query.as_ref(),
);
}
None
}
@@ -581,6 +605,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "vertex-secret".to_string(),
decrypted_auth_config: None,
},
@@ -9,6 +9,7 @@ use crate::auth::{
};
use crate::claude_code::build_claude_code_passthrough_headers;
use crate::claude_code::local_claude_code_transport_unsupported_reason_with_network;
use crate::gemini_cli::is_gemini_cli_provider_transport;
use crate::grok::{is_grok_provider_transport, resolve_grok_session_auth};
use crate::kiro::{
build_kiro_provider_headers, build_kiro_provider_request_body, is_kiro_provider_transport,
@@ -48,6 +49,7 @@ pub struct SameFormatProviderRequestBehaviorParams<'a> {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SameFormatProviderRequestBehavior {
pub is_antigravity: bool,
pub is_gemini_cli: bool,
pub is_claude_code: bool,
pub is_vertex: bool,
pub is_kiro: bool,
@@ -103,6 +105,7 @@ pub fn classify_same_format_provider_request_behavior(
params: SameFormatProviderRequestBehaviorParams<'_>,
) -> SameFormatProviderRequestBehavior {
let is_antigravity = is_antigravity_provider_transport(transport);
let is_gemini_cli = is_gemini_cli_provider_transport(transport);
let is_claude_code = transport
.provider
.provider_type
@@ -137,6 +140,7 @@ pub fn classify_same_format_provider_request_behavior(
SameFormatProviderRequestBehavior {
is_antigravity,
is_gemini_cli,
is_claude_code,
is_vertex,
is_kiro,
@@ -518,6 +522,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
@@ -993,6 +998,7 @@ mod tests {
header_rules: None,
behavior: SameFormatProviderRequestBehavior {
is_antigravity: false,
is_gemini_cli: false,
is_claude_code: false,
is_vertex: false,
is_kiro: false,
@@ -70,6 +70,7 @@ pub struct GatewayProviderTransportKey {
pub expires_at_unix_secs: Option<u64>,
pub proxy: Option<serde_json::Value>,
pub fingerprint: Option<serde_json::Value>,
pub upstream_metadata: Option<serde_json::Value>,
pub decrypted_api_key: String,
pub decrypted_auth_config: Option<String>,
}
@@ -398,6 +399,7 @@ mod tests {
expires_at_unix_secs: Some(1_800_000_000),
proxy: Some(serde_json::json!({"node_id":"proxy-node-1"})),
fingerprint: Some(serde_json::json!({"transport_profile":"chrome_136"})),
upstream_metadata: None,
decrypted_api_key: "sk-live-openai".to_string(),
decrypted_auth_config: Some(
"{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}".to_string(),
@@ -108,6 +108,7 @@ pub(super) fn map_key(
expires_at_unix_secs: key.expires_at_unix_secs,
proxy: normalize_optional_json(key.proxy),
fingerprint: normalize_optional_json(key.fingerprint),
upstream_metadata: normalize_optional_json(key.upstream_metadata),
decrypted_api_key,
decrypted_auth_config,
})
@@ -349,6 +349,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
@@ -387,6 +387,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "vertex-secret".to_string(),
decrypted_auth_config: None,
},
@@ -159,6 +159,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "vertex-secret".to_string(),
decrypted_auth_config: None,
},
@@ -246,6 +246,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "vertex-secret".to_string(),
decrypted_auth_config: None,
},
@@ -315,6 +315,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},