i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-12 22:20:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add Gemini CLI v1internal quota support

Browse Source
This commit is contained in:
Mas0nShi
2026-05-22 16:58:04 +08:00
parent 9533bd7043
commit ce02f1ae8c
17 changed files with 985 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-admin/src/provider/quota.rs
+63 -10
View File
@@ -263,9 +263,9 @@ pub fn parse_gemini_cli_retrieve_user_quota_response(
let mut quota_by_model = serde_json::Map::new();
for bucket in buckets {
let Some(bucket_object) = bucket.as_object() else {
if !bucket.is_object() {
continue;
};
}
let model_id = first_json_string_by_paths(
bucket,
&[
@@ -354,6 +354,54 @@ pub fn parse_gemini_cli_retrieve_user_quota_response(
],
)
.or_else(|| remaining_fraction.map(|value| value <= 1e-9));
let remaining_amount = first_json_f64_by_paths(
bucket,
&[
&["remainingAmount"],
&["remaining_amount"],
&["remaining"],
&["remaining_value"],
&["quotaInfo", "remainingAmount"],
&["quotaInfo", "remaining_amount"],
&["quotaInfo", "remaining"],
&["quotaInfo", "remaining_value"],
&["quota", "remainingAmount"],
&["quota", "remaining_amount"],
&["quota", "remaining"],
&["quota", "remaining_value"],
],
);
let explicit_total = first_json_f64_by_paths(
bucket,
&[
&["limit"],
&["limitAmount"],
&["limit_amount"],
&["total"],
&["totalAmount"],
&["total_amount"],
&["quotaInfo", "limit"],
&["quotaInfo", "limitAmount"],
&["quotaInfo", "limit_amount"],
&["quotaInfo", "total"],
&["quotaInfo", "totalAmount"],
&["quotaInfo", "total_amount"],
&["quota", "limit"],
&["quota", "limitAmount"],
&["quota", "limit_amount"],
&["quota", "total"],
&["quota", "totalAmount"],
&["quota", "total_amount"],
],
)
.filter(|value| *value > 0.0);
let total_amount = explicit_total.or_else(|| {
remaining_amount
.zip(remaining_fraction)
.and_then(|(remaining, fraction)| {
(fraction > 0.0).then_some((remaining / fraction).round())
})
});
let mut payload = serde_json::Map::new();
payload.insert("display_name".to_string(), json!(display_name));
@@ -379,15 +427,11 @@ pub fn parse_gemini_cli_retrieve_user_quota_response(
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 let Some(value) = total_amount {
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));
}
if let Some(value) = remaining_amount {
payload.insert("remaining".to_string(), json!(value));
}
quota_by_model.insert(quota_key, serde_json::Value::Object(payload));
@@ -1740,6 +1784,7 @@ mod tests {
"tokenType": "model",
"displayName": "Gemini 2.5 Pro",
"remainingFraction": 0.25,
"remainingAmount": "25",
"resetTime": "2030-01-01T00:00:00Z",
"isExhausted": false
},
@@ -1764,6 +1809,14 @@ mod tests {
parsed["quota_by_model"]["gemini-2.5-pro"]["remaining_fraction"],
json!(0.25)
);
assert_eq!(
parsed["quota_by_model"]["gemini-2.5-pro"]["remaining"],
json!(25.0)
);
assert_eq!(
parsed["quota_by_model"]["gemini-2.5-pro"]["total"],
json!(100.0)
);
assert_eq!(
parsed["quota_by_model"]["gemini-2.5-pro"]["reset_at"],
json!(1_893_456_000u64)
crates/aether-model-fetch/src/strategy.rs
+113 -1
View File
@@ -327,6 +327,11 @@ async fn fetch_gemini_cli_models(
if let Some(plan_type) = extract_gemini_cli_plan_type(&body_json) {
provider_meta.insert("plan_type".to_string(), Value::String(plan_type));
}
for key in ["paidTier", "currentTier"] {
if let Some(value) = extract_gemini_cli_tier_metadata(&body_json, key) {
provider_meta.insert(key.to_string(), value);
}
}
if let Some(project_id) =
extract_gemini_cli_project_id(&body_json).or_else(|| {
transport_auth_config(transport)
@@ -1192,7 +1197,9 @@ fn normalize_api_format(value: &str) -> String {
fn extract_gemini_cli_plan_type(body: &Value) -> Option<String> {
for key in ["paidTier", "currentTier"] {
let tier = body.get(key)?;
let Some(tier) = body.get(key) else {
continue;
};
let raw = if let Some(value) = tier.as_str() {
value.trim().to_string()
} else if let Some(value) = tier
@@ -1218,6 +1225,40 @@ fn extract_gemini_cli_plan_type(body: &Value) -> Option<String> {
None
}
fn extract_gemini_cli_tier_metadata(body: &Value, key: &str) -> Option<Value> {
let tier = body.get(key)?;
if let Some(text) = tier
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
{
return Some(Value::String(text.to_string()));
}
let object = tier.as_object()?;
let mut out = serde_json::Map::new();
for field in [
"id",
"tierType",
"name",
"displayName",
"availableCredits",
"remainingCredits",
"consumedCredits",
"totalCredits",
"unlimited",
"hasCredits",
] {
let Some(value) = object.get(field) else {
continue;
};
if value.is_string() || value.is_number() || value.is_boolean() || value.is_null() {
out.insert(field.to_string(), value.clone());
}
}
(!out.is_empty()).then_some(Value::Object(out))
}
fn extract_gemini_cli_project_id(body: &Value) -> Option<String> {
let raw = body.get("cloudaicompanionProject")?;
if let Some(value) = raw.as_str() {
@@ -1414,6 +1455,16 @@ mod tests {
transport
}
fn sample_gemini_cli_transport() -> GatewayProviderTransportSnapshot {
let mut transport = sample_custom_aiplatform_transport();
transport.provider.provider_type = "gemini_cli".to_string();
transport.provider.name = "Gemini CLI".to_string();
transport.endpoint.base_url = "https://cloudcode-pa.googleapis.com".to_string();
transport.key.auth_type = "bearer".to_string();
transport.key.decrypted_api_key = "gemini-cli-access-token".to_string();
transport
}
#[test]
fn strategy_selection_keeps_codex_on_standard_transport_fetch() {
let strategy = select_model_fetch_strategy(&[sample_codex_transport()])
@@ -1566,6 +1617,67 @@ mod tests {
assert_eq!(outcome.cached_models.len(), 1);
}
#[tokio::test]
async fn gemini_cli_load_code_assist_preserves_paid_tier_credits() {
let executed_urls = Arc::new(Mutex::new(Vec::new()));
let runtime = TestRuntime {
executed_urls: Arc::clone(&executed_urls),
response_body: json!({
"cloudaicompanionProject": {
"id": "project-from-load-code-assist"
},
"currentTier": {
"id": "free-tier"
},
"paidTier": {
"id": "g1-pro-tier",
"availableCredits": 123.5,
"consumedCredits": 7,
"totalCredits": 200,
"privateField": {
"ignored": true
}
}
}),
status_code: 200,
};
let outcome = fetch_models_from_transports(&runtime, &[sample_gemini_cli_transport()])
.await
.expect("models fetch should succeed");
let urls = executed_urls.lock().expect("executed_urls lock");
assert_eq!(
urls.as_slice(),
&["https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"]
);
assert_eq!(
outcome
.upstream_metadata
.as_ref()
.and_then(|value| value.pointer("/gemini_cli/project_id")),
Some(&json!("project-from-load-code-assist"))
);
assert_eq!(
outcome
.upstream_metadata
.as_ref()
.and_then(|value| value.pointer("/gemini_cli/plan_type")),
Some(&json!("g1-pro-tier"))
);
assert_eq!(
outcome
.upstream_metadata
.as_ref()
.and_then(|value| value.pointer("/gemini_cli/paidTier/availableCredits")),
Some(&json!(123.5))
);
assert!(outcome
.upstream_metadata
.as_ref()
.and_then(|value| value.pointer("/gemini_cli/paidTier/privateField"))
.is_none());
}
#[tokio::test]
async fn kiro_transport_fetches_list_available_models() {
let executed_urls = Arc::new(Mutex::new(Vec::new()));
crates/aether-provider-transport/src/provider_types.rs
+15 -2
View File
@@ -321,12 +321,12 @@ const KIRO_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTemplat
const GEMINI_CLI_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTemplate {
provider_type: "gemini_cli",
version: 1,
version: 2,
base_url: "https://cloudcode-pa.googleapis.com",
endpoints: &[FixedProviderEndpointTemplate {
item_key: "gemini:generate_content",
api_format: "gemini:generate_content",
custom_path: None,
custom_path: Some("/v1internal:{action}"),
config_defaults: EMPTY_ENDPOINT_CONFIG_DEFAULTS,
}],
runtime_policy: GEMINI_CLI_RUNTIME_POLICY,
@@ -675,6 +675,19 @@ mod tests {
assert!(!template.runtime_policy.supports_local_same_format_transport);
}
#[test]
fn gemini_cli_fixed_provider_template_uses_v1internal_endpoint_path() {
let template =
fixed_provider_template("gemini_cli").expect("gemini_cli template should exist");
assert_eq!(template.base_url, "https://cloudcode-pa.googleapis.com");
assert_eq!(template.version, 2);
let endpoint =
fixed_provider_endpoint_template_by_api_format("gemini_cli", "gemini:generate_content")
.expect("gemini_cli generateContent endpoint should exist");
assert_eq!(endpoint.custom_path, Some("/v1internal:{action}"));
}
#[test]
fn fixed_provider_key_inheritance_keeps_oauth_and_kiro_configured_bearer_keys_open() {
assert!(fixed_provider_key_inherits_api_formats(