i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-07 09:57:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: support Gemini CLI v1internal quota

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs
+45 -5
View File
@@ -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,
apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs
+63 -16
View File
@@ -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"
apps/aether-gateway/src/tests/architecture/admin_provider.rs
+1
View File
@@ -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!(
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+204 -8
View File
@@ -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",
apps/aether-gateway/src/tests/control/admin/provider_query.rs
+14 -3
View File
@@ -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],
));