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
@@ -28,6 +28,12 @@ pub(crate) fn maybe_normalize_provider_private_sync_report_payload(
let mut normalized = payload.clone();
normalized.report_context = normalize_provider_private_report_context(Some(report_context));
if let (Some(body_json), Some(context)) = (
payload.body_json.as_ref(),
normalized.report_context.as_mut(),
) {
maybe_attach_gemini_cli_v1internal_credits_context(report_context, body_json, context);
}
if let Some(body_json) = payload.body_json.clone() {
normalized.body_json = normalize_provider_private_response_value(body_json, report_context);
@@ -55,6 +61,44 @@ pub(crate) fn maybe_normalize_provider_private_sync_report_payload(
Ok(Some(normalized))
}
fn maybe_attach_gemini_cli_v1internal_credits_context(
original_report_context: &Value,
body_json: &Value,
normalized_report_context: &mut Value,
) {
if !original_report_context
.get("envelope_name")
.and_then(Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case("gemini_cli:v1internal"))
{
return;
}
let mut credits = serde_json::Map::new();
for (source, target) in [
("remainingCredits", "remainingCredits"),
("consumedCredits", "consumedCredits"),
("traceId", "traceId"),
] {
if let Some(value) = body_json
.get(source)
.cloned()
.filter(|value| !value.is_null())
{
credits.insert(target.to_string(), value);
}
}
if credits.is_empty() {
return;
}
if let Some(object) = normalized_report_context.as_object_mut() {
object.insert(
"gemini_cli_v1internal_credits".to_string(),
Value::Object(credits),
);
}
}
fn normalize_provider_private_stream_bytes(
report_context: &Value,
body: &[u8],
@@ -1859,6 +1859,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,
},
@@ -198,6 +198,7 @@ mod tests {
}
}
})),
upstream_metadata: None,
decrypted_api_key: "sk-test".to_string(),
decrypted_auth_config: None,
},
@@ -254,6 +255,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,
},
@@ -144,6 +144,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: String::new(),
decrypted_auth_config: None,
},
@@ -591,6 +591,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,
},
@@ -90,6 +90,11 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
"envelope_name".to_string(),
json!(super::super::ANTIGRAVITY_ENVELOPE_NAME),
);
} else if resolved.is_gemini_cli {
extra_fields.insert(
"envelope_name".to_string(),
json!(crate::ai_serving::transport::GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME),
);
}
let provider_api_format = resolved.provider_api_format.clone();
let effective_headers = input.effective_headers(&parts.headers);
@@ -133,7 +138,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
.and_then(serde_json::Value::as_bool)
.unwrap_or(false),
upstream_is_stream: resolved.upstream_is_stream,
has_envelope: resolved.is_kiro || resolved.is_antigravity,
has_envelope: resolved.is_kiro || resolved.is_antigravity || resolved.is_gemini_cli,
needs_conversion: false,
extra_fields,
}),
@@ -147,6 +152,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
let super::request::LocalSameFormatProviderCandidatePayloadParts {
transport,
is_antigravity: _,
is_gemini_cli: _,
is_kiro: _,
auth_header,
auth_value,
@@ -12,9 +12,12 @@ use crate::ai_serving::transport::antigravity::{
classify_local_antigravity_request_support, AntigravityEnvelopeRequestType,
AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport,
};
use crate::ai_serving::transport::gemini_cli::resolve_gemini_cli_project_id;
use crate::ai_serving::transport::{
build_grok_browser_headers, build_grok_upstream_url, build_same_format_provider_headers,
GrokHeaderInput, SameFormatProviderHeadersInput, GROK_CHAT_PATH,
build_gemini_cli_v1internal_request, build_grok_browser_headers, build_grok_upstream_url,
build_same_format_provider_headers, GeminiCliRequestEnvelopeSupport, GrokHeaderInput,
SameFormatProviderHeadersInput, GEMINI_CLI_USER_AGENT, GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
GROK_CHAT_PATH,
};
use crate::ai_serving::{CandidateFailureDiagnostic, GatewayProviderTransportSnapshot};
use crate::AppState;
@@ -88,6 +91,7 @@ pub(crate) fn resolve_same_format_provider_transport_unsupported_reason_for_trac
pub(crate) struct LocalSameFormatProviderCandidatePayloadParts {
pub(super) transport: Arc<GatewayProviderTransportSnapshot>,
pub(super) is_antigravity: bool,
pub(super) is_gemini_cli: bool,
pub(super) is_kiro: bool,
pub(super) auth_header: Option<String>,
pub(super) auth_value: Option<String>,
@@ -217,6 +221,26 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
} else {
None
};
let gemini_cli_project_id = if prepared.behavior.is_gemini_cli {
match resolve_gemini_cli_project_id(&prepared.transport) {
Some(project_id) => Some(project_id),
None => {
mark_skipped_local_same_format_provider_candidate(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"transport_auth_unavailable",
)
.await;
return None;
}
}
} else {
None
};
let provider_request_body = if let Some(antigravity_auth) = antigravity_auth.as_ref() {
match build_antigravity_safe_v1internal_request(
antigravity_auth,
@@ -246,6 +270,34 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
return None;
}
}
} else if let Some(project_id) = gemini_cli_project_id.as_deref() {
match build_gemini_cli_v1internal_request(
project_id,
trace_id,
&prepared.mapped_model,
&base_provider_request_body,
) {
GeminiCliRequestEnvelopeSupport::Supported(envelope) => envelope,
GeminiCliRequestEnvelopeSupport::Unsupported(_) => {
mark_skipped_local_same_format_provider_candidate_with_extra_data(
state,
input,
trace_id,
candidate,
attempt.candidate_index,
&attempt.candidate_id,
"provider_request_body_missing",
same_format_provider_request_body_failure_extra_data(
body_json,
attempt.eligible.provider_api_format.as_str(),
prepared.transport.endpoint.body_rules.as_ref(),
"gemini_cli_v1internal_envelope",
),
)
.await;
return None;
}
}
} else {
base_provider_request_body
};
@@ -291,10 +343,13 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
return None;
};
let extra_headers = antigravity_auth
let mut extra_headers = antigravity_auth
.as_ref()
.map(build_antigravity_static_identity_headers)
.unwrap_or_default();
if prepared.behavior.is_gemini_cli {
extra_headers.insert("user-agent".to_string(), GEMINI_CLI_USER_AGENT.to_string());
}
let Some(provider_request_headers) = (if is_grok {
build_grok_browser_headers(GrokHeaderInput {
transport: &prepared.transport,
@@ -345,6 +400,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
Some(LocalSameFormatProviderCandidatePayloadParts {
transport: prepared.transport,
is_antigravity: prepared.is_antigravity,
is_gemini_cli: prepared.behavior.is_gemini_cli,
is_kiro: prepared.is_kiro,
auth_header: prepared.auth_header,
auth_value: prepared.auth_value,
@@ -431,6 +431,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "sk-upstream".to_string(),
decrypted_auth_config: None,
},
@@ -68,6 +68,7 @@ fn sample_transport(base_url: &str, api_format: &str) -> GatewayProviderTranspor
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: None,
},
@@ -103,6 +103,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,
},
+25 -18
View File
@@ -18,6 +18,10 @@ pub(crate) mod grok {
pub(crate) use aether_provider_transport::grok::*;
}
pub(crate) mod gemini_cli {
pub(crate) use aether_provider_transport::gemini_cli::*;
}
pub(crate) mod oauth_refresh {
pub(crate) use aether_provider_transport::oauth_refresh::*;
}
@@ -58,21 +62,22 @@ pub(crate) use aether_provider_transport::{
apply_transport_request_body_semantics, body_rules_are_locally_supported,
body_rules_handle_path, body_rules_have_enabled_rules,
build_cross_format_openai_chat_upstream_url, build_cross_format_openai_responses_upstream_url,
build_gemini_files_headers, build_gemini_files_request_body, build_gemini_files_upstream_url,
build_grok_app_chat_body, build_grok_browser_headers, build_grok_upstream_url,
build_kiro_cross_format_upstream_url, build_local_openai_chat_upstream_url,
build_local_openai_responses_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_passthrough_headers, build_request_trace_proxy_value,
build_same_format_provider_headers, build_same_format_provider_request_body,
build_same_format_provider_upstream_url, build_standard_plan_fallback_headers,
build_standard_plan_fallback_openai_chat_url,
build_gemini_cli_v1internal_request, build_gemini_files_headers,
build_gemini_files_request_body, build_gemini_files_upstream_url, build_grok_app_chat_body,
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
build_local_openai_chat_upstream_url, build_local_openai_responses_upstream_url,
build_openai_image_headers, build_openai_image_upstream_url, build_passthrough_headers,
build_request_trace_proxy_value, build_same_format_provider_headers,
build_same_format_provider_request_body, build_same_format_provider_upstream_url,
build_standard_plan_fallback_headers, build_standard_plan_fallback_openai_chat_url,
build_standard_plan_fallback_openai_responses_url, build_standard_provider_request_headers,
build_transport_request_url, build_transport_request_url_for_request_body,
build_video_create_headers, build_video_create_request_body, build_video_create_upstream_url,
candidate_common_transport_skip_reason, candidate_transport_pair_skip_reason,
classify_same_format_provider_request_behavior, ensure_upstream_auth_header,
gemini_files_transport_unsupported_reason, header_rules_are_locally_supported,
header_rules_have_enabled_rules, local_gemini_transport_unsupported_reason_with_network,
header_rules_have_enabled_rules, is_gemini_cli_provider_transport,
local_gemini_transport_unsupported_reason_with_network,
local_openai_chat_transport_unsupported_reason,
local_standard_transport_unsupported_reason_with_network,
openai_image_transport_unsupported_reason, request_conversion_direct_auth,
@@ -88,13 +93,15 @@ pub(crate) use aether_provider_transport::{
supports_local_generic_oauth_request_auth_resolution,
supports_local_oauth_request_auth_resolution, transport_proxy_is_locally_supported,
video_create_transport_unsupported_reason, CandidateTransportPolicyFacts,
GatewayProviderTransportSnapshot, GeminiFilesHeadersInput, GeminiFilesRequestBodyError,
GeminiFilesRequestBodyParts, GrokHeaderInput, LocalResolvedOAuthRequestAuth,
ProviderOpenAiImageHeadersInput, ProviderVideoCreateFamily, ProviderVideoCreateHeadersInput,
SameFormatProviderFamily, SameFormatProviderHeadersInput, SameFormatProviderRequestBehavior,
SameFormatProviderRequestBehaviorParams, SameFormatProviderRequestBodyInput,
SameFormatProviderUpstreamUrlParams, StandardPlanFallbackAcceptPolicy,
StandardPlanFallbackHeadersInput, StandardProviderRequestHeaders,
StandardProviderRequestHeadersInput, TransportRequestBodySemanticsError,
TransportRequestUrlParams, GROK_CHAT_PATH, GROK_INTERNAL_HEADER, GROK_RATE_LIMITS_PATH,
GatewayProviderTransportSnapshot, GeminiCliRequestEnvelopeSupport, GeminiFilesHeadersInput,
GeminiFilesRequestBodyError, GeminiFilesRequestBodyParts, GrokHeaderInput,
LocalResolvedOAuthRequestAuth, ProviderOpenAiImageHeadersInput, ProviderVideoCreateFamily,
ProviderVideoCreateHeadersInput, SameFormatProviderFamily, SameFormatProviderHeadersInput,
SameFormatProviderRequestBehavior, SameFormatProviderRequestBehaviorParams,
SameFormatProviderRequestBodyInput, SameFormatProviderUpstreamUrlParams,
StandardPlanFallbackAcceptPolicy, StandardPlanFallbackHeadersInput,
StandardProviderRequestHeaders, StandardProviderRequestHeadersInput,
TransportRequestBodySemanticsError, TransportRequestUrlParams, GEMINI_CLI_USER_AGENT,
GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME, GROK_CHAT_PATH, GROK_INTERNAL_HEADER,
GROK_RATE_LIMITS_PATH,
};
@@ -3431,6 +3431,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,
},
+1
View File
@@ -196,6 +196,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,
},
@@ -375,6 +375,17 @@ fn invalid_gemini_provider_success_message(
{
return None;
}
let normalized_body_json = report_context
.filter(|context| {
context
.get("has_envelope")
.and_then(Value::as_bool)
.unwrap_or(false)
})
.and_then(|context| {
crate::ai_serving::normalize_provider_private_response_value(body_json.clone(), context)
});
let body_json = normalized_body_json.as_ref().unwrap_or(body_json);
if crate::ai_serving::gemini_generate_content_response_has_visible_output(body_json) {
return None;
}
@@ -2477,6 +2488,39 @@ mod tests {
assert!(message.contains("visible model output"));
}
#[test]
fn invalid_gemini_provider_success_unwraps_gemini_cli_v1internal_envelope() {
let plan = test_gemini_chat_plan();
let report_context = json!({
"has_envelope": true,
"envelope_name": "gemini_cli:v1internal",
"provider_api_format": "gemini:generate_content",
});
let body = json!({
"response": {
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Hello from Gemini CLI"}]
},
"finishReason": "STOP"
}]
},
"remainingCredits": 41,
"consumedCredits": 1,
"traceId": "trace-upstream-sync-1"
});
let message = invalid_gemini_provider_success_message(
&plan,
Some(&report_context),
StatusCode::OK.as_u16(),
Some(&body),
);
assert!(message.is_none());
}
#[tokio::test]
async fn sync_attempt_terminal_guard_marks_dropped_pending_attempt_cancelled() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
@@ -4,6 +4,7 @@ use std::pin::Pin;
use super::antigravity::refresh_antigravity_provider_quota_locally;
use super::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
use super::codex::refresh_codex_provider_quota_locally;
use super::gemini_cli::refresh_gemini_cli_provider_quota_locally;
use super::grok::refresh_grok_provider_quota_locally;
use super::kiro::refresh_kiro_provider_quota_locally;
use crate::handlers::admin::request::AdminAppState;
@@ -34,6 +35,10 @@ const PROVIDER_QUOTA_REFRESH_HANDLERS: &[(&str, ProviderQuotaRefreshHandler)] =
refresh_chatgpt_web_provider_quota_locally_boxed,
),
("codex", refresh_codex_provider_quota_locally_boxed),
(
"gemini_cli",
refresh_gemini_cli_provider_quota_locally_boxed,
),
("grok", refresh_grok_provider_quota_locally_boxed),
("kiro", refresh_kiro_provider_quota_locally_boxed),
];
@@ -104,6 +109,22 @@ fn refresh_codex_provider_quota_locally_boxed<'a>(
))
}
fn refresh_gemini_cli_provider_quota_locally_boxed<'a>(
state: &'a AdminAppState<'a>,
provider: &'a StoredProviderCatalogProvider,
endpoint: &'a StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> ProviderQuotaRefreshFuture<'a> {
Box::pin(refresh_gemini_cli_provider_quota_locally(
state,
provider,
endpoint,
keys,
proxy_override,
))
}
fn refresh_kiro_provider_quota_locally_boxed<'a>(
state: &'a AdminAppState<'a>,
provider: &'a StoredProviderCatalogProvider,
@@ -0,0 +1,258 @@
use super::shared::{
build_provider_quota_execution_plan, build_quota_snapshot_payload,
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
extract_execution_error_message, oauth_refresh_auto_removed_result,
persist_provider_quota_refresh_state, quota_key_auto_removed,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_gemini_cli_retrieve_user_quota_response;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_provider_pool::build_gemini_cli_pool_quota_request;
use serde_json::json;
use std::time::{SystemTime, UNIX_EPOCH};
async fn execute_gemini_cli_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
authorization: (String, String),
project_id: &str,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = match proxy_override {
Some(proxy) => Some(proxy.clone()),
None => {
state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let spec = build_gemini_cli_pool_quota_request(
&transport.key.id,
&transport.endpoint.base_url,
authorization,
project_id,
);
let plan = build_provider_quota_execution_plan(
transport,
spec,
proxy,
state.resolve_transport_profile(transport),
timeouts,
);
execute_provider_quota_plan(state, transport, plan, "gemini_cli").await
}
pub(crate) async fn refresh_gemini_cli_provider_quota_locally(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let mut results = Vec::new();
let mut success_count = 0usize;
let mut failed_count = 0usize;
let mut auto_removed_count = 0usize;
for key in keys {
let transport = match state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
.await?
{
Some(transport) => transport,
None => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "Provider transport snapshot unavailable",
}));
continue;
}
};
let authorization = match state.resolve_local_oauth_header_auth(&transport).await? {
Some(auth) => auth,
_ => {
if quota_key_auto_removed(state, &key.id).await? {
auto_removed_count += 1;
results.push(oauth_refresh_auto_removed_result(&key));
continue;
}
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "缺少 OAuth 认证信息,请先授权/刷新 Token",
}));
continue;
}
};
let Some(project_id) = crate::provider_transport::resolve_gemini_cli_project_id(&transport)
else {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "缺少 Gemini CLI project_id,请先刷新模型或在 auth_config/upstream_metadata 中写入 project_id",
}));
continue;
};
let result = match execute_gemini_cli_quota_plan(
state,
&transport,
authorization,
&project_id,
proxy_override.as_ref(),
)
.await?
{
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("retrieveUserQuota 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
}
};
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
.unwrap_or(0);
let mut metadata_update = None::<serde_json::Value>;
let (mut oauth_invalid_at_unix_secs, mut oauth_invalid_reason) =
quota_refresh_success_invalid_state(&key);
let mut status = "error".to_string();
let mut message = None::<String>;
if result.status_code == 200 {
if let Some(body_json) = result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
{
metadata_update =
parse_gemini_cli_retrieve_user_quota_response(body_json, now_unix_secs)
.map(|metadata| json!({ "gemini_cli": metadata }));
if metadata_update.is_some() {
status = "success".to_string();
} else {
status = "no_metadata".to_string();
message = Some("响应中未包含配额 buckets".to_string());
}
} else {
status = "no_metadata".to_string();
message = Some("响应中未包含配额信息".to_string());
}
} else {
let err_msg = extract_execution_error_message(&result);
message = Some(match err_msg.as_deref() {
Some(detail) if !detail.is_empty() => {
format!(
"retrieveUserQuota 返回状态码 {}: {}",
result.status_code, detail
)
}
_ => format!("retrieveUserQuota 返回状态码 {}", result.status_code),
});
if result.status_code == 403 {
let reason = err_msg
.clone()
.filter(|value| !value.trim().is_empty())
.unwrap_or_else(|| "账户访问被禁止".to_string());
oauth_invalid_at_unix_secs = Some(now_unix_secs);
oauth_invalid_reason = Some(format!("账户访问被禁止: {reason}"));
metadata_update = Some(json!({
"gemini_cli": {
"is_forbidden": true,
"forbidden_reason": reason,
"forbidden_at": now_unix_secs,
"updated_at": now_unix_secs,
}
}));
status = "forbidden".to_string();
}
}
if !persist_provider_quota_refresh_state(
state,
&key.id,
metadata_update.as_ref(),
oauth_invalid_at_unix_secs,
oauth_invalid_reason,
None,
)
.await?
{
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "Key 状态写入失败",
}));
continue;
}
if status == "success" {
success_count += 1;
} else {
failed_count += 1;
}
let mut payload = serde_json::Map::new();
payload.insert("key_id".to_string(), json!(key.id));
payload.insert("key_name".to_string(), json!(key.name));
payload.insert("status".to_string(), json!(status));
if let Some(message) = message {
payload.insert("message".to_string(), json!(message));
}
if let Some(metadata) = metadata_update
.as_ref()
.and_then(|value| value.get("gemini_cli"))
.cloned()
{
payload.insert("metadata".to_string(), metadata);
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"gemini_cli",
key.status_snapshot.as_ref(),
metadata_update.as_ref(),
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
results.push(serde_json::Value::Object(payload));
}
Ok(Some(json!({
"success": success_count,
"failed": failed_count,
"total": results.len(),
"results": results,
"message": format!("已处理 {} 个 Key", results.len()),
"auto_removed": auto_removed_count,
})))
}
@@ -2,6 +2,7 @@ pub(crate) mod antigravity;
pub(crate) mod chatgpt_web;
pub(crate) mod codex;
pub(crate) mod dispatch;
pub(crate) mod gemini_cli;
pub(crate) mod grok;
pub(crate) mod kiro;
pub(crate) mod shared;
@@ -781,6 +781,18 @@ fn admin_pool_build_grok_account_quota_from_snapshot(
fn admin_pool_build_gemini_cli_account_quota_from_snapshot(
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
if let Some(credits) = quota_snapshot
.get("credits")
.and_then(serde_json::Value::as_object)
{
if let Some(remaining) = admin_pool_json_to_f64(credits.get("remaining")) {
return Some(format!(
"AI Credits 剩余 {}",
admin_pool_format_quota_value(remaining)
));
}
}
let now = chrono::Utc::now().timestamp();
let mut active = admin_pool_quota_windows(quota_snapshot)
.into_iter()
@@ -2836,6 +2836,33 @@ async fn provider_query_execute_standard_test_candidate(
upstream_is_stream,
require_body_stream_field,
);
if crate::provider_transport::is_gemini_cli_provider_transport(&transport)
&& normalized_provider_api_format == "gemini:generate_content"
{
let Some(project_id) = crate::provider_transport::resolve_gemini_cli_project_id(&transport)
else {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
"Gemini CLI project_id is unavailable for v1internal request",
));
};
provider_request_body = match crate::provider_transport::build_gemini_cli_v1internal_request(
project_id.as_str(),
trace_id,
request_model,
&provider_request_body,
) {
crate::provider_transport::GeminiCliRequestEnvelopeSupport::Supported(envelope) => {
envelope
}
crate::provider_transport::GeminiCliRequestEnvelopeSupport::Unsupported(_) => {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
"Gemini CLI v1internal envelope could not be built",
));
}
};
}
let uses_vertex_query_auth =
crate::provider_transport::uses_vertex_api_key_query_auth(&transport, provider_api_format);
@@ -2956,6 +2983,13 @@ async fn provider_query_execute_standard_test_candidate(
request_headers
.entry("content-type".to_string())
.or_insert_with(|| "application/json".to_string());
if crate::provider_transport::is_gemini_cli_provider_transport(&transport)
&& normalized_provider_api_format == "gemini:generate_content"
{
request_headers
.entry("user-agent".to_string())
.or_insert_with(|| crate::provider_transport::GEMINI_CLI_USER_AGENT.to_string());
}
let protected_headers = if uses_vertex_query_auth {
vec!["content-type"]
} else {
@@ -51,6 +51,7 @@ fn sample_openai_image_transport(provider_type: &str) -> AdminGatewayProviderTra
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: String::new(),
decrypted_auth_config: Some(
json!({
@@ -1318,6 +1318,7 @@ fn build_gemini_cli_quota_status_snapshot(
"reset_at": reset_at,
"reset_seconds": reset_seconds,
"plan_type": serde_json::Value::Null,
"credits": metadata.get("credits").cloned().unwrap_or(serde_json::Value::Null),
"windows": windows,
}))
}
@@ -760,6 +760,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: decrypted_auth_config.map(ToOwned::to_owned),
},
@@ -251,6 +251,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,
},
@@ -314,6 +314,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,
},
@@ -9,6 +9,7 @@ use aether_usage_runtime::{
report_request_id, GatewayStreamReportRequest, GatewaySyncReportRequest,
GEMINI_FILE_MAPPING_TTL_SECONDS,
};
use base64::Engine as _;
use regex::Regex;
use serde_json::{json, Value};
use tracing::warn;
@@ -263,6 +264,121 @@ fn grok_upstream_response_body(report_context: Option<&Value>) -> Option<&Value>
.and_then(|response| response.get("body"))
}
fn gemini_cli_credits_from_report_context(
report_context: Option<&Value>,
now_unix_secs: u64,
) -> Option<Value> {
report_context
.and_then(|context| context.get("gemini_cli_v1internal_credits"))
.and_then(|value| {
admin_provider_quota_pure::parse_gemini_cli_v1internal_credits_response(
value,
now_unix_secs,
)
})
}
fn gemini_cli_credits_from_stream_payload(
payload: &GatewayStreamReportRequest,
now_unix_secs: u64,
) -> Option<Value> {
let body_base64 = payload.provider_body_base64.as_deref()?;
let body = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.ok()?;
let text = std::str::from_utf8(&body).ok()?;
let mut latest = None::<Value>;
for raw_line in text.lines() {
let line = raw_line.trim_matches('\r').trim();
let data = line.strip_prefix("data:").map(str::trim).unwrap_or(line);
if data.is_empty() || data == "[DONE]" || data.starts_with(':') {
continue;
}
let Ok(value) = serde_json::from_str::<Value>(data) else {
continue;
};
if let Some(credits) =
admin_provider_quota_pure::parse_gemini_cli_v1internal_credits_response(
&value,
now_unix_secs,
)
{
latest = Some(credits);
}
}
latest
}
async fn sync_gemini_cli_credits_from_report(
state: &AppState,
report_context: Option<&Value>,
credits: Option<Value>,
) -> Result<bool, GatewayError> {
let Some(credits) = credits else {
return Ok(false);
};
let key_id = match report_context_key_id(report_context) {
Some(value) => value,
None => return Ok(false),
};
let Some(key) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
else {
return Ok(false);
};
let Some(provider) = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
.await?
.into_iter()
.next()
else {
return Ok(false);
};
if !provider
.provider_type
.trim()
.eq_ignore_ascii_case("gemini_cli")
{
return Ok(false);
}
let now_unix_secs = current_unix_secs();
let mut gemini_cli_bucket = key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("gemini_cli"))
.and_then(Value::as_object)
.cloned()
.unwrap_or_else(serde_json::Map::new);
gemini_cli_bucket.insert("credits".to_string(), credits);
gemini_cli_bucket.insert("updated_at".to_string(), json!(now_unix_secs));
let updated_upstream_metadata = merge_metadata_object(
key.upstream_metadata.as_ref(),
"gemini_cli",
Value::Object(gemini_cli_bucket),
);
let updated_status_snapshot = sync_provider_key_quota_status_snapshot(
key.status_snapshot.as_ref(),
provider.provider_type.as_str(),
updated_upstream_metadata.as_ref(),
"report_effect",
);
let mut updated_key = key;
updated_key.upstream_metadata = updated_upstream_metadata;
updated_key.status_snapshot = updated_status_snapshot;
updated_key.updated_at_unix_secs = Some(now_unix_secs);
Ok(state
.update_provider_catalog_key(&updated_key)
.await?
.is_some())
}
fn grok_quota_reset_after_seconds(
body_json: Option<&Value>,
report_context: Option<&Value>,
@@ -464,6 +580,23 @@ async fn apply_local_sync_report_effect(state: &AppState, payload: &GatewaySyncR
"gateway failed to persist grok realtime quota from sync response"
);
}
let now_unix_secs = current_unix_secs();
if let Err(err) = sync_gemini_cli_credits_from_report(
state,
payload.report_context.as_ref(),
gemini_cli_credits_from_report_context(payload.report_context.as_ref(), now_unix_secs),
)
.await
{
warn!(
event_name = "gemini_cli_realtime_credits_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist gemini cli realtime credits from sync response"
);
}
}
async fn apply_local_stream_report_effect(state: &AppState, payload: &GatewayStreamReportRequest) {
@@ -500,6 +633,22 @@ async fn apply_local_stream_report_effect(state: &AppState, payload: &GatewayStr
"gateway failed to persist grok realtime quota from stream response"
);
}
let now_unix_secs = current_unix_secs();
let credits =
gemini_cli_credits_from_report_context(payload.report_context.as_ref(), now_unix_secs)
.or_else(|| gemini_cli_credits_from_stream_payload(payload, now_unix_secs));
if let Err(err) =
sync_gemini_cli_credits_from_report(state, payload.report_context.as_ref(), credits).await
{
warn!(
event_name = "gemini_cli_realtime_credits_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist gemini cli realtime credits from stream response"
);
}
}
async fn apply_local_gemini_file_mapping_report_effect(
+1
View File
@@ -1817,6 +1817,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("{\"project_id\":\"demo\"}".to_string()),
},
@@ -453,6 +453,9 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
trace_id: String,
url: String,
has_model_field: bool,
project: String,
user_prompt_id: String,
envelope_model: String,
accept: String,
authorization: String,
exact_temperature: f64,
@@ -574,7 +577,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com".to_string(),
"https://cloudcode-pa.googleapis.com".to_string(),
Some(serde_json::json!([
{"action":"set","key":"x-endpoint-tag","value":"gemini-cli-oauth-local"}
])),
@@ -595,7 +598,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"gemini_cli","refresh_token":"rt-gemini-cli-stream-local-123"}"#,
r#"{"provider_type":"gemini_cli","refresh_token":"rt-gemini-cli-stream-local-123","project_id":"gemini-cli-project-1"}"#,
)
.expect("auth config should encrypt");
StoredProviderCatalogKey::new(
@@ -735,6 +738,27 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("model"))
.is_some(),
project: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("project"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
user_prompt_id: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("user_prompt_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
envelope_model: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("model"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
accept: payload
.get("headers")
.and_then(|value| value.get("accept"))
@@ -750,6 +774,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
exact_temperature: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("generationConfig"))
.and_then(|value| value.get("temperature"))
.and_then(|value| value.as_f64())
@@ -763,6 +788,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
metadata_mode: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("metadata"))
.and_then(|value| value.get("mode"))
.and_then(|value| value.as_str())
@@ -771,6 +797,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
metadata_source: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("metadata"))
.and_then(|value| value.get("source"))
.and_then(|value| value.as_str())
@@ -779,6 +806,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
tool_config_present: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("toolConfig"))
.is_some(),
proxy_node_id: payload
@@ -795,7 +823,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
});
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"candidates\\\":[]}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"response\\\":{\\\"candidates\\\":[]},\\\"remainingCredits\\\":42,\\\"consumedCredits\\\":1,\\\"traceId\\\":\\\"trace-upstream-1\\\"}\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":34,\"upstream_bytes\":26}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
@@ -909,9 +937,21 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
);
assert_eq!(
seen_execution_runtime_request.url,
"https://generativelanguage.googleapis.com/custom/v1beta/models/gemini-cli-upstream:streamGenerateContent?alt=sse"
"https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert!(seen_execution_runtime_request.has_model_field);
assert_eq!(
seen_execution_runtime_request.project,
"gemini-cli-project-1"
);
assert_eq!(
seen_execution_runtime_request.user_prompt_id,
"trace-gemini-cli-oauth-local-stream-123"
);
assert_eq!(
seen_execution_runtime_request.envelope_model,
"gemini-cli-upstream"
);
assert!(!seen_execution_runtime_request.has_model_field);
assert_eq!(seen_execution_runtime_request.accept, "text/event-stream");
assert_eq!(
seen_execution_runtime_request.authorization,
@@ -780,6 +780,9 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
trace_id: String,
url: String,
has_model_field: bool,
project: String,
user_prompt_id: String,
envelope_model: String,
authorization: String,
exact_temperature: f64,
endpoint_tag: String,
@@ -900,7 +903,7 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com".to_string(),
"https://cloudcode-pa.googleapis.com".to_string(),
Some(serde_json::json!([
{"action":"set","key":"x-endpoint-tag","value":"gemini-cli-oauth-local"}
])),
@@ -921,7 +924,7 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"gemini_cli","refresh_token":"rt-gemini-cli-local-123"}"#,
r#"{"provider_type":"gemini_cli","refresh_token":"rt-gemini-cli-local-123","project_id":"gemini-cli-project-1"}"#,
)
.expect("auth config should encrypt");
StoredProviderCatalogKey::new(
@@ -1062,6 +1065,27 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("model"))
.is_some(),
project: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("project"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
user_prompt_id: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("user_prompt_id"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
envelope_model: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("model"))
.and_then(|value| value.as_str())
.unwrap_or_default()
.to_string(),
authorization: payload
.get("headers")
.and_then(|value| value.get("authorization"))
@@ -1071,6 +1095,7 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
exact_temperature: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("generationConfig"))
.and_then(|value| value.get("temperature"))
.and_then(|value| value.as_f64())
@@ -1084,6 +1109,7 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
metadata_mode: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("metadata"))
.and_then(|value| value.get("mode"))
.and_then(|value| value.as_str())
@@ -1092,6 +1118,7 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
metadata_source: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("metadata"))
.and_then(|value| value.get("source"))
.and_then(|value| value.as_str())
@@ -1100,6 +1127,7 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
tool_config_present: payload
.get("body")
.and_then(|value| value.get("json_body"))
.and_then(|value| value.get("request"))
.and_then(|value| value.get("toolConfig"))
.is_some(),
proxy_node_id: payload
@@ -1123,18 +1151,23 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
},
"body": {
"json_body": {
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Hello from Gemini CLI"}]
},
"finishReason": "STOP"
}],
"usageMetadata": {
"promptTokenCount": 1,
"candidatesTokenCount": 2,
"totalTokenCount": 3
"response": {
"candidates": [{
"content": {
"role": "model",
"parts": [{"text": "Hello from Gemini CLI"}]
},
"finishReason": "STOP"
}],
"usageMetadata": {
"promptTokenCount": 1,
"candidatesTokenCount": 2,
"totalTokenCount": 3
}
}
,"remainingCredits": 41,
"consumedCredits": 1,
"traceId": "trace-upstream-sync-1"
}
},
"telemetry": {
@@ -1203,7 +1236,9 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let response_status = response.status();
let response_body = response.text().await.expect("response body should read");
assert_eq!(response_status, StatusCode::OK, "body={response_body}");
let seen_refresh_request = seen_refresh
.lock()
@@ -1238,9 +1273,21 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
);
assert_eq!(
seen_execution_runtime_request.url,
"https://generativelanguage.googleapis.com/custom/v1beta/models/gemini-cli-upstream:generateContent"
"https://cloudcode-pa.googleapis.com/v1internal:generateContent"
);
assert!(seen_execution_runtime_request.has_model_field);
assert_eq!(
seen_execution_runtime_request.project,
"gemini-cli-project-1"
);
assert_eq!(
seen_execution_runtime_request.user_prompt_id,
"trace-gemini-cli-oauth-local-sync-123"
);
assert_eq!(
seen_execution_runtime_request.envelope_model,
"gemini-cli-upstream"
);
assert!(!seen_execution_runtime_request.has_model_field);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer refreshed-gemini-cli-access-token"
@@ -1782,6 +1782,7 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
"refresh_codex_provider_quota_locally",
"refresh_kiro_provider_quota_locally",
"refresh_antigravity_provider_quota_locally",
"refresh_gemini_cli_provider_quota_locally",
"refresh_chatgpt_web_provider_quota_locally",
] {
assert!(
@@ -1358,6 +1358,210 @@ async fn gateway_refresh_kiro_quota_reconciles_missing_fixed_endpoint_before_ref
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_refreshes_admin_provider_quota_locally_for_gemini_cli_with_trusted_admin_principal(
) {
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeRequest {
url: String,
authorization: String,
provider_api_format: String,
request_body: Option<serde_json::Value>,
}
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-gemini-cli/refresh-quota",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
url: plan.url.clone(),
authorization: plan
.headers
.get("authorization")
.cloned()
.unwrap_or_default(),
provider_api_format: plan.provider_api_format.clone(),
request_body: plan.body.json_body.clone(),
});
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"buckets": [
{
"modelId": "gemini-2.5-pro",
"tokenType": "model",
"displayName": "Gemini 2.5 Pro",
"remainingFraction": 0.25,
"resetTime": "2030-01-01T00:00:00Z",
"isExhausted": false
},
{
"modelId": "gemini-2.5-flash",
"tokenType": "model",
"displayName": "Gemini 2.5 Flash",
"quotaInfo": {
"remainingFraction": 0.0,
"resetTime": "2030-01-01T01:00:00Z",
"isExhausted": true
}
}
]
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let mut key = sample_key(
"key-gemini-cli-quota",
"provider-gemini-cli",
"gemini:generate_content",
"cached-gemini-cli-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"gemini_cli","project_id":"gemini-cli-project-1"}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![StoredProviderCatalogProvider::new(
"provider-gemini-cli".to_string(),
"gemini_cli".to_string(),
Some("https://example.com".to_string()),
"gemini_cli".to_string(),
)
.expect("provider should build")],
vec![sample_endpoint(
"endpoint-gemini-cli-quota",
"provider-gemini-cli",
"gemini:generate_content",
"https://cloudcode-pa.googleapis.com",
)],
vec![key],
));
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-gemini-cli/refresh-quota"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], 1);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(
payload["results"][0]["quota_snapshot"]["provider_type"],
"gemini_cli"
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["windows"][0]["model"],
"gemini-2.5-pro"
);
let seen_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
assert_eq!(
seen_request.url,
"https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota"
);
assert_eq!(seen_request.authorization, "Bearer cached-gemini-cli-token");
assert_eq!(
seen_request.provider_api_format,
"gemini_cli:retrieve_user_quota"
);
assert_eq!(
seen_request.request_body,
Some(json!({
"project": "gemini-cli-project-1",
"userAgent": "GeminiCLI/0.1.5 (Windows; AMD64)"
}))
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-gemini-cli-quota".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
let upstream_metadata = reloaded[0]
.upstream_metadata
.as_ref()
.expect("upstream metadata should persist");
assert_eq!(
upstream_metadata["gemini_cli"]["quota_by_model"]["gemini-2.5-pro"]["remaining_fraction"],
json!(0.25)
);
assert_eq!(
upstream_metadata["gemini_cli"]["quota_by_model"]["gemini-2.5-flash"]["is_exhausted"],
json!(true)
);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_refresh_quota_reconciles_unsupported_fixed_provider_endpoints_before_clear_message(
) {
@@ -1370,14 +1574,6 @@ async fn gateway_refresh_quota_reconciles_unsupported_fixed_provider_endpoints_b
"https://api.anthropic.com",
"Claude Code 暂不支持自动刷新额度",
),
(
"provider-gemini-cli-reconcile",
"gemini_cli",
1usize,
"gemini:generate_content",
"https://cloudcode-pa.googleapis.com",
"Gemini CLI 暂不支持自动刷新额度",
),
(
"provider-vertex-ai-reconcile",
"vertex_ai",
@@ -4773,8 +4773,19 @@ async fn gateway_handles_gemini_cli_test_model_with_oauth_header_fallback() {
assert_eq!(plan.provider_api_format, "gemini:generate_content");
assert_eq!(
plan.url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent"
"https://cloudcode-pa.googleapis.com/v1internal:generateContent"
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["project"],
json!("project-1")
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["model"],
json!("gemini-2.5-pro")
);
assert!(plan.body.json_body.as_ref().unwrap()["request"]
.get("contents")
.is_some());
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer cached-gemini-cli-token")
@@ -4818,7 +4829,7 @@ async fn gateway_handles_gemini_cli_test_model_with_oauth_header_fallback() {
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"gemini_cli"}"#,
r#"{"provider_type":"gemini_cli","project_id":"project-1"}"#,
)
.expect("auth config should encrypt"),
);
@@ -4828,7 +4839,7 @@ async fn gateway_handles_gemini_cli_test_model_with_oauth_header_fallback() {
"endpoint-gemini-cli",
"provider-gemini",
"gemini:generate_content",
"https://generativelanguage.googleapis.com",
"https://cloudcode-pa.googleapis.com",
)],
vec![key],
));
+298 -2
View File
@@ -255,6 +255,231 @@ pub fn parse_antigravity_usage_response(
}))
}
pub fn parse_gemini_cli_retrieve_user_quota_response(
value: &serde_json::Value,
updated_at_unix_secs: u64,
) -> Option<serde_json::Value> {
let buckets = value.get("buckets")?.as_array()?;
let mut quota_by_model = serde_json::Map::new();
for bucket in buckets {
let Some(bucket_object) = bucket.as_object() else {
continue;
};
let model_id = first_json_string_by_paths(
bucket,
&[
&["modelId"],
&["model_id"],
&["model"],
&["modelName"],
&["metadata", "modelId"],
&["metadata", "model_id"],
&["labels", "modelId"],
&["labels", "model_id"],
],
);
let token_type = first_json_string_by_paths(
bucket,
&[
&["tokenType"],
&["token_type"],
&["metadata", "tokenType"],
&["metadata", "token_type"],
&["labels", "tokenType"],
&["labels", "token_type"],
],
);
let Some(quota_key) = model_id
.as_deref()
.or(token_type.as_deref())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
else {
continue;
};
let display_name = first_json_string_by_paths(
bucket,
&[
&["displayName"],
&["display_name"],
&["metadata", "displayName"],
&["metadata", "display_name"],
],
)
.or_else(|| model_id.clone())
.or_else(|| token_type.clone())
.unwrap_or_else(|| quota_key.clone());
let remaining_fraction = first_json_f64_by_paths(
bucket,
&[
&["remainingFraction"],
&["remaining_fraction"],
&["quotaInfo", "remainingFraction"],
&["quotaInfo", "remaining_fraction"],
&["quota", "remainingFraction"],
&["quota", "remaining_fraction"],
],
)
.map(|value| value.clamp(0.0, 1.0));
let reset_time = first_json_value_by_paths(
bucket,
&[
&["resetTime"],
&["reset_time"],
&["nextResetTime"],
&["next_reset_time"],
&["quotaInfo", "resetTime"],
&["quotaInfo", "reset_time"],
&["quota", "resetTime"],
&["quota", "reset_time"],
],
)
.cloned()
.filter(|value| !value.is_null());
let reset_at = reset_time
.as_ref()
.and_then(parse_gemini_cli_reset_timestamp);
let is_exhausted = first_json_bool_by_paths(
bucket,
&[
&["isExhausted"],
&["is_exhausted"],
&["exhausted"],
&["quotaInfo", "isExhausted"],
&["quotaInfo", "is_exhausted"],
&["quota", "isExhausted"],
&["quota", "is_exhausted"],
],
)
.or_else(|| remaining_fraction.map(|value| value <= 1e-9));
let mut payload = serde_json::Map::new();
payload.insert("display_name".to_string(), json!(display_name));
if let Some(model_id) = model_id {
payload.insert("model_id".to_string(), json!(model_id));
}
if let Some(token_type) = token_type {
payload.insert("token_type".to_string(), json!(token_type));
}
if let Some(remaining_fraction) = remaining_fraction {
payload.insert("remaining_fraction".to_string(), json!(remaining_fraction));
payload.insert(
"used_percent".to_string(),
json!(((1.0 - remaining_fraction) * 100.0).clamp(0.0, 100.0)),
);
}
if let Some(reset_time) = reset_time {
payload.insert("reset_time".to_string(), reset_time);
}
if let Some(reset_at) = reset_at {
payload.insert("reset_at".to_string(), json!(reset_at));
}
if let Some(is_exhausted) = is_exhausted {
payload.insert("is_exhausted".to_string(), json!(is_exhausted));
}
if bucket_object.contains_key("limit") {
if let Some(value) = bucket_object.get("limit").and_then(coerce_json_f64) {
payload.insert("total".to_string(), json!(value));
}
}
if bucket_object.contains_key("remaining") {
if let Some(value) = bucket_object.get("remaining").and_then(coerce_json_f64) {
payload.insert("remaining".to_string(), json!(value));
}
}
quota_by_model.insert(quota_key, serde_json::Value::Object(payload));
}
if quota_by_model.is_empty() {
return None;
}
Some(json!({
"updated_at": updated_at_unix_secs,
"quota_by_model": quota_by_model,
}))
}
pub fn parse_gemini_cli_v1internal_credits_response(
value: &serde_json::Value,
updated_at_unix_secs: u64,
) -> Option<serde_json::Value> {
let mut credits = serde_json::Map::new();
if let Some(value) = value.get("remainingCredits").and_then(coerce_json_f64) {
credits.insert("remaining".to_string(), json!(value));
}
if let Some(value) = value.get("consumedCredits").and_then(coerce_json_f64) {
credits.insert("consumed".to_string(), json!(value));
}
if let Some(value) = coerce_json_string(value.get("traceId")) {
credits.insert("trace_id".to_string(), json!(value));
}
if credits.is_empty() {
return None;
}
credits.insert("updated_at".to_string(), json!(updated_at_unix_secs));
Some(serde_json::Value::Object(credits))
}
fn first_json_value_by_paths<'a>(
value: &'a serde_json::Value,
paths: &[&[&str]],
) -> Option<&'a serde_json::Value> {
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 {
return Some(current);
}
}
None
}
fn first_json_string_by_paths(value: &serde_json::Value, paths: &[&[&str]]) -> Option<String> {
paths
.iter()
.find_map(|path| coerce_json_string(first_json_value_by_paths(value, &[*path])))
}
fn first_json_f64_by_paths(value: &serde_json::Value, paths: &[&[&str]]) -> Option<f64> {
paths
.iter()
.find_map(|path| first_json_value_by_paths(value, &[*path]).and_then(coerce_json_f64))
}
fn first_json_bool_by_paths(value: &serde_json::Value, paths: &[&[&str]]) -> Option<bool> {
paths
.iter()
.find_map(|path| first_json_value_by_paths(value, &[*path]).and_then(coerce_json_bool))
}
fn parse_gemini_cli_reset_timestamp(value: &serde_json::Value) -> Option<u64> {
if let Some(value) = coerce_json_u64(value) {
return Some(if value > 1_000_000_000_000 {
value / 1000
} else {
value
});
}
let raw = value.as_str()?.trim();
if raw.is_empty() {
return None;
}
chrono::DateTime::parse_from_rfc3339(raw)
.ok()
.and_then(|timestamp| u64::try_from(timestamp.timestamp()).ok())
}
pub fn normalize_codex_plan_type(value: Option<&str>) -> Option<String> {
value
.map(str::trim)
@@ -1141,8 +1366,9 @@ mod tests {
use super::{
codex_build_invalid_state, codex_runtime_invalid_reason,
parse_chatgpt_web_conversation_init_response, parse_codex_backend_me_response,
parse_codex_wham_usage_response, OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
parse_codex_wham_usage_response, parse_gemini_cli_retrieve_user_quota_response,
parse_gemini_cli_v1internal_credits_response, OAUTH_ACCOUNT_BLOCK_PREFIX,
OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::json;
@@ -1504,6 +1730,76 @@ mod tests {
assert!(parsed.get("secondary_used_percent").is_none());
}
#[test]
fn parses_gemini_cli_retrieve_user_quota_buckets() {
let parsed = parse_gemini_cli_retrieve_user_quota_response(
&json!({
"buckets": [
{
"modelId": "gemini-2.5-pro",
"tokenType": "model",
"displayName": "Gemini 2.5 Pro",
"remainingFraction": 0.25,
"resetTime": "2030-01-01T00:00:00Z",
"isExhausted": false
},
{
"modelId": "gemini-2.5-flash",
"tokenType": "model",
"displayName": "Gemini 2.5 Flash",
"quotaInfo": {
"remainingFraction": 0.0,
"resetTime": 1_893_459_600_000u64,
"isExhausted": true
}
}
]
}),
1_777_000_000,
)
.expect("gemini cli quota should parse");
assert_eq!(parsed.get("updated_at"), Some(&json!(1_777_000_000u64)));
assert_eq!(
parsed["quota_by_model"]["gemini-2.5-pro"]["remaining_fraction"],
json!(0.25)
);
assert_eq!(
parsed["quota_by_model"]["gemini-2.5-pro"]["reset_at"],
json!(1_893_456_000u64)
);
assert_eq!(
parsed["quota_by_model"]["gemini-2.5-flash"]["is_exhausted"],
json!(true)
);
assert_eq!(
parsed["quota_by_model"]["gemini-2.5-flash"]["used_percent"],
json!(100.0)
);
}
#[test]
fn parses_gemini_cli_v1internal_credits() {
let parsed = parse_gemini_cli_v1internal_credits_response(
&json!({
"response": {"candidates": []},
"remainingCredits": "41.5",
"consumedCredits": 1,
"traceId": "trace-upstream-sync-1"
}),
1_777_000_123,
)
.expect("gemini cli credits should parse");
assert_eq!(parsed.get("remaining"), Some(&json!(41.5)));
assert_eq!(parsed.get("consumed"), Some(&json!(1.0)));
assert_eq!(
parsed.get("trace_id"),
Some(&json!("trace-upstream-sync-1"))
);
assert_eq!(parsed.get("updated_at"), Some(&json!(1_777_000_123u64)));
}
#[test]
fn parses_chatgpt_web_image_quota_from_conversation_init() {
let parsed = parse_chatgpt_web_conversation_init_response(
@@ -1368,6 +1368,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,
},
@@ -687,6 +687,7 @@ mod tests {
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
upstream_metadata: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: Some(
r#"{"project_id":"project-1","client_version":"1.2.3","session_id":"sess-1"}"#
+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))
@@ -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,
},