i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-07 18:07:47 +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
apps/aether-gateway/src/control/route/ai.rs
+16 -5
View File
@@ -95,7 +95,7 @@ pub(super) fn classify_ai_public_route(
"openai:video",
true,
))
} else if is_gemini_models_route(normalized_path) {
} else if method == http::Method::POST && is_gemini_models_route(normalized_path) {
if normalized_path.ends_with(":predictLongRunning") {
Some(classified(
"ai_public",
@@ -134,7 +134,9 @@ pub(super) fn classify_ai_public_route(
true,
))
}
} else if is_gemini_operation_route(normalized_path) {
} else if is_gemini_operation_method(method, normalized_path)
&& is_gemini_operation_route(normalized_path)
{
Some(classified(
"ai_public",
"gemini",
@@ -142,9 +144,7 @@ pub(super) fn classify_ai_public_route(
"gemini:video",
true,
))
} else if (method == http::Method::POST && normalized_path == "/upload/v1beta/files")
|| normalized_path.starts_with("/v1beta/files")
{
} else if is_gemini_files_method(method, normalized_path) {
Some(classified(
"ai_public",
"gemini",
@@ -156,3 +156,14 @@ pub(super) fn classify_ai_public_route(
None
}
}
fn is_gemini_operation_method(method: &http::Method, normalized_path: &str) -> bool {
method == http::Method::GET
|| (method == http::Method::POST && normalized_path.ends_with(":cancel"))
}
fn is_gemini_files_method(method: &http::Method, normalized_path: &str) -> bool {
(method == http::Method::POST && normalized_path == "/upload/v1beta/files")
|| ((method == http::Method::GET || method == http::Method::DELETE)
&& normalized_path.starts_with("/v1beta/files"))
}
apps/aether-gateway/src/control/tests/ai.rs
+17
View File
@@ -219,6 +219,23 @@ fn classifies_gemini_generate_content_api_key_without_cli_marker() {
assert!(decision.is_execution_runtime_candidate());
}
#[test]
fn does_not_classify_gemini_options_preflight_as_ai_execution_route() {
let headers = headers(&[
("origin", "http://localhost:3000"),
("access-control-request-method", "POST"),
(
"access-control-request-headers",
"authorization,content-type",
),
]);
let uri: Uri = "/v1beta/models/gemini-3-flash-preview:streamGenerateContent?alt=sse"
.parse()
.expect("uri should parse");
assert!(classify_control_route(&http::Method::OPTIONS, &uri, &headers).is_none());
}
#[test]
fn classifies_gemini_embed_content_as_embedding_route() {
let headers = headers(&[("x-goog-api-key", "gemini-key")]);
apps/aether-gateway/src/handlers/admin/request/provider/builders.rs
+19 -1
View File
@@ -295,6 +295,25 @@ impl<'a> AdminAppState<'a> {
use crate::handlers::public::{admin_requested_force_stream, normalize_admin_base_url};
use aether_admin::provider::endpoints as admin_provider_endpoints_pure;
let (fields, payload) = patch.into_parts();
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
if provider_type == "gemini_cli"
&& [
"base_url",
"custom_path",
"header_rules",
"body_rules",
"max_retries",
"is_active",
"config",
"proxy",
"format_acceptance_config",
]
.iter()
.any(|field| fields.contains(field))
{
return Err("Gemini CLI Endpoint 由系统固定管理，不允许修改".to_string());
}
if self.provider_type_is_fixed(&provider.provider_type)
&& (fields.contains("base_url") || fields.contains("custom_path"))
@@ -326,7 +345,6 @@ impl<'a> AdminAppState<'a> {
&update_fields,
)?;
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
if provider_type == "codex"
&& crate::ai_serving::is_openai_responses_format(&existing_endpoint.api_format)
{
apps/aether-gateway/src/handlers/shared/catalog.rs
+235 -4
View File
@@ -553,6 +553,52 @@ fn provider_quota_metadata_string(
})
}
fn provider_quota_metadata_value_by_path<'a>(
metadata: &'a Map<String, Value>,
path: &[&str],
) -> Option<&'a Value> {
let (first, rest) = path.split_first()?;
let mut current = metadata.get(*first)?;
for segment in rest {
current = current.as_object()?.get(*segment)?;
}
Some(current)
}
fn provider_quota_metadata_number_by_paths(
metadata: &Map<String, Value>,
paths: &[&[&str]],
) -> Option<f64> {
paths.iter().find_map(|path| {
provider_quota_metadata_value_by_path(metadata, path)
.and_then(admin_provider_quota_pure::coerce_json_f64)
.filter(|value| value.is_finite())
})
}
fn provider_quota_metadata_bool_by_paths(
metadata: &Map<String, Value>,
paths: &[&[&str]],
) -> Option<bool> {
paths.iter().find_map(|path| {
provider_quota_metadata_value_by_path(metadata, path)
.and_then(admin_provider_quota_pure::coerce_json_bool)
})
}
fn provider_quota_metadata_string_by_paths(
metadata: &Map<String, Value>,
paths: &[&[&str]],
) -> Option<String> {
paths.iter().find_map(|path| {
provider_quota_metadata_value_by_path(metadata, path)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn quota_windows_usage_ratio(windows: &[Value]) -> Option<f64> {
windows
.iter()
@@ -1232,12 +1278,144 @@ fn build_grok_quota_status_snapshot(
}))
}
fn gemini_cli_plan_type(metadata: &Map<String, Value>) -> Option<String> {
provider_quota_metadata_string(metadata, &["plan_type", "tier", "plan"]).or_else(|| {
provider_quota_metadata_string_by_paths(
metadata,
&[
&["paidTier", "id"],
&["paidTier", "tierType"],
&["paidTier", "name"],
&["currentTier", "id"],
&["currentTier", "tierType"],
&["currentTier", "name"],
],
)
})
}
fn gemini_cli_credits_status_snapshot(metadata: &Map<String, Value>) -> Option<Value> {
let remaining = provider_quota_metadata_number_by_paths(
metadata,
&[
&["credits", "remaining"],
&["credits", "remainingCredits"],
&["credits", "available"],
&["credits", "availableCredits"],
&["credits", "balance"],
&["remainingCredits"],
&["availableCredits"],
&["paidTier", "remainingCredits"],
&["paidTier", "availableCredits"],
&["currentTier", "remainingCredits"],
&["currentTier", "availableCredits"],
],
);
let balance = provider_quota_metadata_number_by_paths(
metadata,
&[
&["credits", "balance"],
&["credits", "remaining"],
&["credits", "available"],
&["paidTier", "availableCredits"],
&["currentTier", "availableCredits"],
&["availableCredits"],
&["remainingCredits"],
],
)
.or(remaining);
let consumed = provider_quota_metadata_number_by_paths(
metadata,
&[
&["credits", "consumed"],
&["credits", "consumedCredits"],
&["consumedCredits"],
&["paidTier", "consumedCredits"],
&["currentTier", "consumedCredits"],
],
);
let total = provider_quota_metadata_number_by_paths(
metadata,
&[
&["credits", "total"],
&["credits", "totalCredits"],
&["totalCredits"],
&["paidTier", "totalCredits"],
&["currentTier", "totalCredits"],
],
);
let unlimited = provider_quota_metadata_bool_by_paths(
metadata,
&[
&["credits", "unlimited"],
&["credits_unlimited"],
&["paidTier", "unlimited"],
&["currentTier", "unlimited"],
],
);
let explicit_has_credits = provider_quota_metadata_bool_by_paths(
metadata,
&[
&["credits", "has_credits"],
&["has_credits"],
&["paidTier", "hasCredits"],
&["currentTier", "hasCredits"],
],
);
let trace_id = provider_quota_metadata_string_by_paths(
metadata,
&[
&["credits", "trace_id"],
&["credits", "traceId"],
&["trace_id"],
&["traceId"],
],
);
let updated_at = provider_quota_metadata_value_by_path(metadata, &["credits", "updated_at"])
.or_else(|| provider_quota_metadata_value_by_path(metadata, &["credits", "updatedAt"]))
.or_else(|| metadata.get("updated_at"))
.and_then(|value| provider_quota_timestamp_unix_secs(Some(value)));
if remaining.is_none()
&& balance.is_none()
&& consumed.is_none()
&& total.is_none()
&& unlimited.is_none()
&& explicit_has_credits.is_none()
&& trace_id.is_none()
{
return None;
}
let has_credits = explicit_has_credits
.or_else(|| unlimited.filter(|value| *value))
.or_else(|| remaining.or(balance).map(|value| value > 0.0));
let mut credits = Map::new();
credits.insert("has_credits".to_string(), json!(has_credits));
credits.insert("balance".to_string(), json!(balance));
credits.insert("remaining".to_string(), json!(remaining.or(balance)));
credits.insert("consumed".to_string(), json!(consumed));
credits.insert("total".to_string(), json!(total));
credits.insert("unlimited".to_string(), json!(unlimited));
credits.insert("trace_id".to_string(), json!(trace_id));
credits.insert("updated_at".to_string(), json!(updated_at));
Some(Value::Object(credits))
}
fn build_gemini_cli_quota_status_snapshot(
upstream_metadata: Option<&Value>,
source: &str,
) -> Option<Value> {
let metadata = provider_quota_metadata_bucket(upstream_metadata, "gemini_cli")?;
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
let credits = gemini_cli_credits_status_snapshot(metadata);
let plan_type = gemini_cli_plan_type(metadata);
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"))
.or_else(|| {
credits
.as_ref()
.and_then(|value| value.get("updated_at"))
.and_then(|value| provider_quota_timestamp_unix_secs(Some(value)))
});
let windows = provider_quota_model_bucket(metadata)
.map(|models| {
models
@@ -1253,7 +1431,11 @@ fn build_gemini_cli_quota_status_snapshot(
})
.unwrap_or_default();
if windows.is_empty() && observed_at_unix_secs.is_none() {
if windows.is_empty()
&& observed_at_unix_secs.is_none()
&& credits.is_none()
&& plan_type.is_none()
{
return None;
}
@@ -1317,8 +1499,8 @@ fn build_gemini_cli_quota_status_snapshot(
"updated_at": observed_at_unix_secs,
"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),
"plan_type": plan_type,
"credits": credits,
"windows": windows,
}))
}
@@ -2273,6 +2455,55 @@ mod tests {
assert_eq!(auto.get("used_value"), Some(&json!(90.0)));
}
#[test]
fn provider_key_status_snapshot_payload_backfills_gemini_cli_account_credits() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"gemini_cli": {
"updated_at": 1_778_067_246u64,
"plan_type": "g1-pro-tier",
"paidTier": {
"availableCredits": 123.5,
"consumedCredits": 7.0,
"totalCredits": 200.0
},
"quota_by_model": {
"gemini-2.5-pro": {
"display_name": "Gemini 2.5 Pro",
"remaining_fraction": 0.75,
"is_exhausted": false
}
}
}
}));
let payload = provider_key_status_snapshot_payload(&key, "gemini_cli");
let quota = payload
.get("quota")
.and_then(Value::as_object)
.expect("quota snapshot should be object");
let credits = quota
.get("credits")
.and_then(Value::as_object)
.expect("credits snapshot should exist");
let windows = quota
.get("windows")
.and_then(Value::as_array)
.expect("Gemini CLI model windows should exist");
assert_eq!(quota.get("provider_type"), Some(&json!("gemini_cli")));
assert_eq!(quota.get("code"), Some(&json!("ok")));
assert_eq!(quota.get("plan_type"), Some(&json!("g1-pro-tier")));
assert_eq!(quota.get("updated_at"), Some(&json!(1_778_067_246u64)));
assert_eq!(credits.get("remaining"), Some(&json!(123.5)));
assert_eq!(credits.get("balance"), Some(&json!(123.5)));
assert_eq!(credits.get("consumed"), Some(&json!(7.0)));
assert_eq!(credits.get("total"), Some(&json!(200.0)));
assert_eq!(credits.get("has_credits"), Some(&json!(true)));
assert_eq!(windows.len(), 1);
assert_eq!(windows[0].get("remaining_ratio"), Some(&json!(0.75)));
}
#[test]
fn provider_key_status_snapshot_payload_preserves_existing_materialized_quota_snapshot() {
let mut key = sample_catalog_key();
apps/aether-gateway/src/tests/control/admin/pool.rs
+93
View File
@@ -2085,6 +2085,99 @@ async fn gateway_renders_gemini_cli_account_quota_from_status_snapshot() {
assert_eq!(keys[0]["account_quota"], json!("Gemini 2.5 Pro 冷却中"));
}
#[tokio::test]
async fn gateway_prefers_gemini_cli_account_credits_over_model_quota_text() {
let mut provider = sample_provider("provider-gemini-cli", "gemini_cli", 10)
.with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true,
"skip_exhausted_accounts": true
}
})),
);
provider.provider_type = "gemini_cli".to_string();
let mut key = sample_key(
"key-gemini-cli-credits",
"provider-gemini-cli",
"gemini:generate_content",
"oauth-placeholder",
);
key.name = "gemini cli credits".to_string();
key.auth_type = "oauth".to_string();
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "gemini_cli",
"code": "ok",
"freshness": "fresh",
"source": "report_effect",
"observed_at": 1_775_553_285u64,
"exhausted": false,
"usage_ratio": 0.25,
"updated_at": 1_775_553_285u64,
"plan_type": "g1-pro-tier",
"credits": {
"remaining": 123.5,
"consumed": 7.0,
"has_credits": true
},
"windows": [
{
"code": "model:gemini-2.5-pro",
"label": "Gemini 2.5 Pro",
"scope": "model",
"unit": "percent",
"model": "gemini-2.5-pro",
"used_ratio": 0.25,
"remaining_ratio": 0.75,
"is_exhausted": false
}
]
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-gemini-cli/keys?page=1&page_size=50&status=all",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys[0]["scheduling_status"], json!("available"));
assert_eq!(keys[0]["account_quota"], json!("AI Credits 剩余 123.5"));
}
#[tokio::test]
async fn gateway_formats_codex_quota_countdown_from_reset_after_seconds() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(