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

fix: hydrate Gemini CLI project metadata

Browse Source
This commit is contained in:
Mas0nShi
2026-05-21 17:10:17 +08:00
parent e53d5f07e8
commit 3e6ce6cf4a
7 changed files with 386 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs
+34 -22
View File
@@ -133,6 +133,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
)
.await;
let effective_headers = input.effective_headers(&parts.headers);
let mut transport = Arc::clone(&prepared.transport);
let Some(mut base_provider_request_body) =
super::super::request::build_same_format_provider_request_body(
@@ -199,7 +200,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
let antigravity_auth = if prepared.is_antigravity {
match classify_local_antigravity_request_support(
&prepared.transport,
&transport,
&base_provider_request_body,
AntigravityEnvelopeRequestType::Agent,
) {
@@ -222,20 +223,32 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
None
};
let gemini_cli_project_id = if prepared.behavior.is_gemini_cli {
match resolve_gemini_cli_project_id(&prepared.transport) {
match resolve_gemini_cli_project_id(&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;
match state
.hydrate_gemini_cli_project_metadata_for_transport(&transport)
.await
{
Some(hydrated) => {
let project_id = resolve_gemini_cli_project_id(&hydrated);
transport = Arc::new(hydrated);
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 {
@@ -308,14 +321,13 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
.provider_type
.trim()
.eq_ignore_ascii_case("grok");
let transport_profile =
crate::ai_serving::transport::resolve_transport_profile(&prepared.transport);
let transport_profile = crate::ai_serving::transport::resolve_transport_profile(&transport);
let upstream_url = if is_grok {
Some(build_grok_upstream_url(&prepared.transport, GROK_CHAT_PATH))
Some(build_grok_upstream_url(&transport, GROK_CHAT_PATH))
} else {
super::super::request::build_same_format_upstream_url(
parts,
&prepared.transport,
&transport,
&prepared.mapped_model,
prepared.provider_api_format.as_str(),
spec,
@@ -352,12 +364,12 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
}
let Some(provider_request_headers) = (if is_grok {
build_grok_browser_headers(GrokHeaderInput {
transport: &prepared.transport,
transport: &transport,
transport_profile: transport_profile.as_ref(),
request_headers: Some(effective_headers),
content_type: "application/json",
accept: "text/event-stream",
header_rules: prepared.transport.endpoint.header_rules.as_ref(),
header_rules: transport.endpoint.header_rules.as_ref(),
provider_request_body: &provider_request_body,
original_request_body: body_json,
})
@@ -366,12 +378,12 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
headers: effective_headers,
provider_request_body: &provider_request_body,
original_request_body: body_json,
header_rules: prepared.transport.endpoint.header_rules.as_ref(),
header_rules: transport.endpoint.header_rules.as_ref(),
behavior: prepared.behavior,
auth_header: prepared.auth_header.as_deref(),
auth_value: prepared.auth_value.as_deref(),
extra_headers: &extra_headers,
key_fingerprint: prepared.transport.key.fingerprint.as_ref(),
key_fingerprint: transport.key.fingerprint.as_ref(),
kiro_auth_config: prepared.kiro_auth.as_ref().map(|auth| &auth.auth_config),
kiro_machine_id: prepared
.kiro_auth
@@ -398,7 +410,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
};
Some(LocalSameFormatProviderCandidatePayloadParts {
transport: prepared.transport,
transport,
is_antigravity: prepared.is_antigravity,
is_gemini_cli: prepared.behavior.is_gemini_cli,
is_kiro: prepared.is_kiro,
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs
+73 -1
View File
@@ -29,6 +29,7 @@ pub(super) struct AdminProviderOAuthBatchImportEntry {
pub user_id: Option<String>,
pub email: Option<String>,
pub account_name: Option<String>,
pub project_id: Option<String>,
pub sso_rw_token: Option<String>,
pub cf_cookies: Option<String>,
pub cf_clearance: Option<String>,
@@ -74,6 +75,20 @@ fn coerce_admin_provider_oauth_import_str(value: Option<&serde_json::Value>) ->
.map(ToOwned::to_owned)
}
fn coerce_admin_provider_oauth_import_project_id(
value: Option<&serde_json::Value>,
) -> Option<String> {
match value {
Some(serde_json::Value::Object(object)) => coerce_admin_provider_oauth_import_str(
object
.get("id")
.or_else(|| object.get("project_id"))
.or_else(|| object.get("projectId")),
),
other => coerce_admin_provider_oauth_import_str(other),
}
}
fn grok_cookie_value(raw: &str, name: &str) -> Option<String> {
raw.trim()
.strip_prefix("Cookie:")
@@ -151,6 +166,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
user_id: grok_cookie_value(raw_token, "x-userid"),
email: None,
account_name: None,
project_id: None,
sso_rw_token: grok_cookie_value(raw_token, "sso-rw"),
cf_cookies: grok_cookie_profile(raw_token),
cf_clearance: grok_cookie_value(raw_token, "cf_clearance"),
@@ -250,6 +266,13 @@ fn extract_admin_provider_oauth_batch_import_entry(
.get("account_name")
.or_else(|| object.get("accountName")),
);
let project_id = coerce_admin_provider_oauth_import_project_id(
object
.get("project_id")
.or_else(|| object.get("projectId"))
.or_else(|| object.get("cloudaicompanionProject"))
.or_else(|| object.get("cloudAiCompanionProject")),
);
let sso_rw_token = coerce_admin_provider_oauth_import_str(
object
.get("sso_rw_token")
@@ -295,6 +318,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
user_id,
email,
account_name,
project_id,
sso_rw_token,
cf_cookies,
cf_clearance,
@@ -363,6 +387,19 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints(
auth_config: &mut serde_json::Map<String, serde_json::Value>,
) {
let provider_type = provider_type.trim().to_ascii_lowercase();
if provider_type == "gemini_cli" {
if let Some(project_id) = entry.project_id.as_ref() {
auth_config
.entry("project_id".to_string())
.or_insert_with(|| json!(project_id));
}
if let Some(plan_type) = entry.plan_type.as_ref() {
auth_config
.entry("plan_type".to_string())
.or_insert_with(|| json!(plan_type));
}
return;
}
if !matches!(provider_type.as_str(), "codex" | "chatgpt_web" | "grok") {
return;
}
@@ -513,7 +550,10 @@ pub(super) fn build_admin_provider_oauth_batch_task_state(
#[cfg(test)]
mod tests {
use super::parse_admin_provider_oauth_batch_import_entries;
use super::{
apply_admin_provider_oauth_batch_import_hints,
parse_admin_provider_oauth_batch_import_entries,
};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use serde_json::json;
@@ -633,4 +673,36 @@ mod tests {
assert_eq!(entries[0].user_id.as_deref(), Some("user-1"));
assert_eq!(entries[0].pool_tier.as_deref(), Some("heavy"));
}
#[test]
fn parses_gemini_cli_project_id_hint() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"gemini_cli",
r#"[{"refresh_token":"rt-1","projectId":"project-gemini-cli-1","planType":"free"}]"#,
);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].refresh_token.as_deref(), Some("rt-1"));
assert_eq!(
entries[0].project_id.as_deref(),
Some("project-gemini-cli-1")
);
assert_eq!(entries[0].plan_type.as_deref(), Some("free"));
}
#[test]
fn applies_gemini_cli_project_id_hint_to_auth_config() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"gemini_cli",
r#"{"refreshToken":"rt-1","cloudaicompanionProject":{"id":"project-gemini-cli-2"}}"#,
);
let mut auth_config = serde_json::Map::new();
apply_admin_provider_oauth_batch_import_hints("gemini_cli", &entries[0], &mut auth_config);
assert_eq!(
auth_config.get("project_id"),
Some(&json!("project-gemini-cli-2"))
);
}
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/gemini_cli.rs
+17 -4
View File
@@ -66,7 +66,7 @@ pub(crate) async fn refresh_gemini_cli_provider_quota_locally(
let mut auto_removed_count = 0usize;
for key in keys {
let transport = match state
let mut transport = match state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
.await?
{
@@ -102,14 +102,27 @@ pub(crate) async fn refresh_gemini_cli_provider_quota_locally(
}
};
let Some(project_id) = crate::provider_transport::resolve_gemini_cli_project_id(&transport)
else {
let project_id = match crate::provider_transport::resolve_gemini_cli_project_id(&transport)
{
Some(project_id) => Some(project_id),
None => state
.app()
.hydrate_gemini_cli_project_metadata_for_transport(&transport)
.await
.and_then(|hydrated| {
let project_id =
crate::provider_transport::resolve_gemini_cli_project_id(&hydrated);
transport = hydrated;
project_id
}),
};
let Some(project_id) = project_id 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",
"message": "缺少 Gemini CLI project_id，loadCodeAssist 未返回可用项目信息",
}));
continue;
};
apps/aether-gateway/src/handlers/admin/provider/query/models/mod.rs
+14 -2
View File
@@ -15,7 +15,7 @@ use crate::ai_serving::{
ANTIGRAVITY_V1INTERNAL_ENVELOPE_NAME, GEMINI_CHAT_SYNC_FINALIZE_REPORT_KIND,
OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
};
use crate::clock::current_unix_ms;
use crate::clock::{current_unix_ms, current_unix_secs};
use crate::execution_runtime;
use crate::handlers::admin::provider::shared::model_test_capabilities::{
admin_provider_model_supports_image_generation, admin_provider_model_test_capabilities_payload,
@@ -64,7 +64,7 @@ use aether_data_contracts::repository::provider_catalog::{
};
use aether_model_fetch::{
aggregate_models_for_cache, fetch_models_from_transports, json_string_list,
preset_models_for_provider, selected_models_fetch_endpoints,
merge_upstream_metadata, preset_models_for_provider, selected_models_fetch_endpoints,
};
use axum::{
body::{to_bytes, Body},
@@ -509,6 +509,18 @@ async fn provider_query_fetch_models_for_key(
)
.await;
}
if let Some(upstream_metadata) = outcome.upstream_metadata.as_ref() {
let merged_metadata =
merge_upstream_metadata(key.upstream_metadata.as_ref(), upstream_metadata);
state
.app()
.update_provider_catalog_key_upstream_metadata(
&key.id,
Some(&merged_metadata),
Some(current_unix_secs()),
)
.await?;
}
if unique_models.is_empty() && !all_errors.is_empty() {
if let Some(fallback) = provider_query_codex_preset_fallback(provider) {
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
+16 -3
View File
@@ -2639,7 +2639,7 @@ async fn provider_query_execute_standard_test_candidate(
route_path: &str,
trace_id: &str,
) -> Result<ProviderQueryExecutionOutcome, GatewayError> {
let Some(transport) = state
let Some(mut transport) = state
.read_provider_transport_snapshot(&provider.id, &candidate.endpoint.id, &candidate.key.id)
.await?
else {
@@ -2839,8 +2839,21 @@ async fn provider_query_execute_standard_test_candidate(
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 {
let project_id = match crate::provider_transport::resolve_gemini_cli_project_id(&transport)
{
Some(project_id) => Some(project_id),
None => state
.app()
.hydrate_gemini_cli_project_metadata_for_transport(&transport)
.await
.and_then(|hydrated| {
let project_id =
crate::provider_transport::resolve_gemini_cli_project_id(&hydrated);
transport = hydrated;
project_id
}),
};
let Some(project_id) = project_id else {
return Ok(provider_query_skipped_execution_outcome(
provider_request_body,
"Gemini CLI project_id is unavailable for v1internal request",
apps/aether-gateway/src/state/integrations.rs
+61 -3
View File
@@ -13,15 +13,16 @@ use aether_data_contracts::repository::provider_catalog::{
};
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_model_fetch::{
aggregate_models_for_cache, model_fetch_interval_minutes, ModelFetchAssociationStore,
ModelFetchTransportRuntime,
aggregate_models_for_cache, fetch_models_from_transports, merge_upstream_metadata,
model_fetch_interval_minutes, ModelFetchAssociationStore, ModelFetchTransportRuntime,
};
use aether_scheduler_core::SchedulerAffinityTarget;
use async_trait::async_trait;
use serde_json::Value;
use tracing::debug;
use tracing::{debug, warn};
use super::{AppState, GatewayError};
use crate::clock::current_unix_secs;
use crate::model_fetch::ModelFetchRuntimeState;
use crate::provider_transport::{GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth};
use crate::request_candidate_runtime::{
@@ -31,6 +32,63 @@ use crate::request_candidate_runtime::{
use crate::scheduler::state::SchedulerRuntimeState;
use crate::{execution_runtime, provider_transport};
impl AppState {
pub(crate) async fn hydrate_gemini_cli_project_metadata_for_transport(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<GatewayProviderTransportSnapshot> {
if !provider_transport::is_gemini_cli_provider_transport(transport) {
return None;
}
if provider_transport::resolve_gemini_cli_project_id(transport).is_some() {
return Some(transport.clone());
}
let outcome =
match fetch_models_from_transports(self, std::slice::from_ref(transport)).await {
Ok(outcome) => outcome,
Err(err) => {
warn!(
provider_id = %transport.provider.id,
endpoint_id = %transport.endpoint.id,
key_id = %transport.key.id,
error = %err,
"gemini_cli project metadata hydration failed"
);
return None;
}
};
let upstream_metadata = outcome.upstream_metadata.as_ref()?;
let merged_metadata =
merge_upstream_metadata(transport.key.upstream_metadata.as_ref(), upstream_metadata);
let mut hydrated = transport.clone();
hydrated.key.upstream_metadata = Some(merged_metadata.clone());
if provider_transport::resolve_gemini_cli_project_id(&hydrated).is_none() {
return None;
}
if let Err(err) = self
.update_provider_catalog_key_upstream_metadata(
&transport.key.id,
Some(&merged_metadata),
Some(current_unix_secs()),
)
.await
{
warn!(
provider_id = %transport.provider.id,
endpoint_id = %transport.endpoint.id,
key_id = %transport.key.id,
error = ?err,
"gemini_cli project metadata hydration could not persist metadata"
);
}
Some(hydrated)
}
}
#[async_trait]
impl provider_transport::TransportTunnelAffinityLookup for AppState {
async fn lookup_tunnel_attachment_owner(
apps/aether-gateway/src/tests/control/admin/provider_query.rs
+171 -1
View File
@@ -8,7 +8,9 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogEndpoint;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
};
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Json, Router};
@@ -4880,6 +4882,174 @@ async fn gateway_handles_gemini_cli_test_model_with_oauth_header_fallback() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_hydrates_gemini_cli_project_id_from_load_code_assist_for_test_model() {
let seen_urls = Arc::new(Mutex::new(Vec::<String>::new()));
let seen_urls_clone = Arc::clone(&seen_urls);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| {
let seen_urls_inner = Arc::clone(&seen_urls_clone);
async move {
seen_urls_inner
.lock()
.expect("mutex should lock")
.push(plan.url.clone());
if plan.url == "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" {
assert_eq!(plan.model_name.as_deref(), Some("loadCodeAssist"));
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer cached-gemini-cli-token")
);
assert_eq!(
plan.body.json_body.as_ref().and_then(|body| body
.get("metadata")
.and_then(|metadata| metadata.get("pluginType"))),
Some(&json!("GEMINI"))
);
return Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"cloudaicompanionProject": {
"id": "project-from-load-code-assist"
},
"currentTier": {
"id": "free"
}
}
}
}));
}
assert_eq!(
plan.url,
"https://cloudcode-pa.googleapis.com/v1internal:generateContent"
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["project"],
json!("project-from-load-code-assist")
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["model"],
json!("gemini-2.5-pro")
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-gemini-cli-test-model",
"choices": [{
"message": {
"role": "assistant",
"content": "Hello from hydrated Gemini CLI"
}
}]
}
},
"telemetry": {
"elapsed_ms": 19
}
}))
}
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-gemini", "Gemini", 10);
provider.provider_type = "gemini_cli".to_string();
let mut key = sample_key(
"key-gemini-cli",
"provider-gemini",
"gemini:generate_content",
"cached-gemini-cli-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"gemini_cli","refresh_token":"rt-gemini-cli-123"}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-gemini-cli",
"provider-gemini",
"gemini:generate_content",
"https://cloudcode-pa.googleapis.com",
)],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&provider_catalog_repository,
))
.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/provider-query/test-model"))
.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")
.json(&json!({
"provider_id": "provider-gemini",
"model": "gemini-2.5-pro",
"api_format": "gemini:generate_content"
}))
.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"], json!(true));
assert_eq!(
payload["data"]["response"]["choices"][0]["message"]["content"],
json!("Hello from hydrated Gemini CLI")
);
assert_eq!(
*seen_urls.lock().expect("mutex should lock"),
vec![
"https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist".to_string(),
"https://cloudcode-pa.googleapis.com/v1internal:generateContent".to_string(),
]
);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-gemini-cli".to_string()])
.await
.expect("key should reload");
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.get("gemini_cli"))
.and_then(|metadata| metadata.get("project_id")),
Some(&json!("project-from-load-code-assist"))
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_uses_compatible_gemini_cli_endpoint_when_api_format_is_omitted() {
let execution_runtime = Router::new().route(