From 135a862326475d1947d950a90049da6c6cddcc0e Mon Sep 17 00:00:00 2001 From: MMEXA Date: Wed, 3 Jun 2026 21:45:37 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20Antigravity=20OAuth=20?= =?UTF-8?q?=E9=85=8D=E9=A2=9D=E5=A4=8D=E6=A3=80=E7=BC=BA=20project?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../passthrough/provider/family/request.rs | 24 +- .../openai/responses/decision/request.rs | 118 ++++---- .../provider/oauth/dispatch/batch/parse.rs | 61 +++++ .../admin/provider/oauth/dispatch/import.rs | 103 ++++++- .../admin/provider/oauth/quota/antigravity.rs | 20 +- apps/aether-gateway/src/state/integrations.rs | 132 ++++++++- crates/aether-model-fetch/src/lib.rs | 9 +- crates/aether-model-fetch/src/strategy.rs | 251 ++++++++++++++++-- crates/aether-model-fetch/src/transport.rs | 116 +++++++- .../src/antigravity/auth.rs | 244 ++++++++++++++--- .../src/antigravity/mod.rs | 7 +- .../src/provider_types.rs | 22 +- 12 files changed, 977 insertions(+), 130 deletions(-) diff --git a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs index cd42341cb..8fc226798 100644 --- a/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request.rs @@ -14,6 +14,7 @@ use crate::ai_serving::transport::antigravity::{ build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers, classify_local_antigravity_request_support, AntigravityEnvelopeRequestType, AntigravityRequestEnvelopeSupport, AntigravityRequestSideSupport, + AntigravityRequestSideUnsupportedReason, }; use crate::ai_serving::transport::{ build_gemini_cli_v1internal_request, build_grok_browser_headers, build_grok_upstream_url, @@ -217,11 +218,32 @@ 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( + let mut antigravity_support = classify_local_antigravity_request_support( &transport, &base_provider_request_body, AntigravityEnvelopeRequestType::Agent, + ); + if matches!( + antigravity_support, + AntigravityRequestSideSupport::Unsupported( + AntigravityRequestSideUnsupportedReason::UnsupportedAuth( + crate::provider_transport::antigravity::AntigravityRequestAuthUnsupportedReason::MissingProjectId + ) + ) ) { + if let Some(hydrated) = state + .hydrate_antigravity_project_metadata_for_transport(&transport) + .await + { + transport = Arc::new(hydrated); + antigravity_support = classify_local_antigravity_request_support( + &transport, + &base_provider_request_body, + AntigravityEnvelopeRequestType::Agent, + ); + } + } + match antigravity_support { AntigravityRequestSideSupport::Supported(spec) => Some(spec.auth), AntigravityRequestSideSupport::Unsupported(_) => { mark_skipped_local_same_format_provider_candidate( diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs index 422ab7738..b046fdaa0 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs @@ -32,7 +32,7 @@ use crate::ai_serving::transport::antigravity::{ build_antigravity_safe_v1internal_request, build_antigravity_static_identity_headers, classify_local_antigravity_request_support, is_antigravity_provider_transport, AntigravityEnvelopeRequestType, AntigravityRequestEnvelopeSupport, - AntigravityRequestSideSupport, + AntigravityRequestSideSupport, AntigravityRequestSideUnsupportedReason, }; use crate::ai_serving::transport::auth::{ resolve_local_gemini_auth, resolve_local_openai_bearer_auth, resolve_local_standard_auth, @@ -118,11 +118,11 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( let provider_api_format = eligible.provider_api_format.as_str(); let normalized_provider_api_format = crate::ai_serving::normalize_api_format_alias(provider_api_format); - let transport = &eligible.transport; - let transport_profile = crate::ai_serving::transport::resolve_transport_profile(transport); - let is_antigravity = is_antigravity_provider_transport(transport); - let is_gemini_cli = is_gemini_cli_provider_transport(transport); - let is_kiro_claude_cli = is_kiro_claude_messages_transport(transport, provider_api_format); + let mut transport = Arc::clone(&eligible.transport); + let transport_profile = crate::ai_serving::transport::resolve_transport_profile(&transport); + let is_antigravity = is_antigravity_provider_transport(&transport); + let is_gemini_cli = is_gemini_cli_provider_transport(&transport); + let is_kiro_claude_cli = is_kiro_claude_messages_transport(&transport, provider_api_format); let is_grok = transport .provider .provider_type @@ -144,33 +144,34 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( .await); } let is_windsurf_cascade = - provider_api_format == "openai:chat" && is_windsurf_provider_transport(transport); + provider_api_format == "openai:chat" && is_windsurf_provider_transport(&transport); let same_format = api_format_alias_matches(provider_api_format, &client_api_format); let conversion_kind = request_conversion_kind(spec_metadata.api_format, provider_api_format); - let transport_unsupported_reason = - if is_grok && is_grok_text_provider_api_format(provider_api_format) { - None - } else if same_format && is_kiro_claude_cli { - local_kiro_request_transport_unsupported_reason_with_network(transport) - } else if same_format { - local_standard_transport_unsupported_reason_with_network(transport, provider_api_format) - } else if is_windsurf_cascade { - local_windsurf_request_transport_unsupported_reason_with_network(transport) - } else { - match conversion_kind { - Some(_) - if (is_antigravity || is_gemini_cli) - && normalized_provider_api_format == "gemini:generate_content" => - { - None - } - Some(kind) => crate::ai_serving::request_conversion_transport_unsupported_reason( - transport, kind, - ), - None => Some("transport_api_format_unsupported"), + let transport_unsupported_reason = if is_grok + && is_grok_text_provider_api_format(provider_api_format) + { + None + } else if same_format && is_kiro_claude_cli { + local_kiro_request_transport_unsupported_reason_with_network(&transport) + } else if same_format { + local_standard_transport_unsupported_reason_with_network(&transport, provider_api_format) + } else if is_windsurf_cascade { + local_windsurf_request_transport_unsupported_reason_with_network(&transport) + } else { + match conversion_kind { + Some(_) + if (is_antigravity || is_gemini_cli) + && normalized_provider_api_format == "gemini:generate_content" => + { + None } - }; + Some(kind) => { + crate::ai_serving::request_conversion_transport_unsupported_reason(&transport, kind) + } + None => Some("transport_api_format_unsupported"), + } + }; if let Some(skip_reason) = transport_unsupported_reason { mark_skipped_local_openai_responses_candidate( state, @@ -193,7 +194,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( let kiro_auth = if is_kiro_claude_cli { match crate::ai_serving::planner::candidate_preparation::resolve_candidate_oauth_auth( planner_state, - transport, + &transport, oauth_context, ) .await @@ -218,20 +219,20 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( }; let direct_auth = if is_grok && is_grok_text_provider_api_format(provider_api_format) { - crate::ai_serving::transport::resolve_grok_session_auth(transport) + crate::ai_serving::transport::resolve_grok_session_auth(&transport) } else if kiro_auth.is_some() { None } else if same_format { match crate::ai_serving::normalize_api_format_alias(provider_api_format).as_str() { - "gemini:generate_content" => resolve_local_gemini_auth(transport), - "claude:messages" => resolve_local_standard_auth(transport), + "gemini:generate_content" => resolve_local_gemini_auth(&transport), + "claude:messages" => resolve_local_standard_auth(&transport), "openai:responses" | "openai:responses:compact" => { - resolve_local_openai_bearer_auth(transport) + resolve_local_openai_bearer_auth(&transport) } _ => None, } } else { - conversion_kind.and_then(|kind| request_conversion_direct_auth(transport, kind)) + conversion_kind.and_then(|kind| request_conversion_direct_auth(&transport, kind)) }; let prepared_candidate = if let Some(kiro_auth) = kiro_auth.as_ref() { match prepare_header_authenticated_candidate_from_auth( @@ -257,7 +258,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( } else { match prepare_header_authenticated_candidate( planner_state, - transport, + &transport, candidate, direct_auth, oauth_context, @@ -409,11 +410,32 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( Some(body_json), ); let antigravity_auth = if is_antigravity { - match classify_local_antigravity_request_support( - transport, + let mut antigravity_support = classify_local_antigravity_request_support( + &transport, &base_provider_request_body, AntigravityEnvelopeRequestType::Agent, + ); + if matches!( + antigravity_support, + AntigravityRequestSideSupport::Unsupported( + AntigravityRequestSideUnsupportedReason::UnsupportedAuth( + crate::provider_transport::antigravity::AntigravityRequestAuthUnsupportedReason::MissingProjectId + ) + ) ) { + if let Some(hydrated) = state + .hydrate_antigravity_project_metadata_for_transport(&transport) + .await + { + transport = Arc::new(hydrated); + antigravity_support = classify_local_antigravity_request_support( + &transport, + &base_provider_request_body, + AntigravityEnvelopeRequestType::Agent, + ); + } + } + match antigravity_support { AntigravityRequestSideSupport::Supported(spec) => Some(spec.auth), AntigravityRequestSideSupport::Unsupported(_) => { mark_skipped_local_openai_responses_candidate( @@ -475,7 +497,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( candidate_index, candidate_id, spec_metadata.api_format, - transport, + &transport, provider_api_format, mapped_model, auth_header, @@ -499,7 +521,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( candidate_index, candidate_id, spec_metadata.api_format, - transport, + &transport, provider_api_format, mapped_model, auth_header, @@ -511,7 +533,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( .await); } if provider_api_format == "gemini:generate_content" - && is_gemini_cli_provider_transport(transport) + && is_gemini_cli_provider_transport(&transport) { return Ok(build_gemini_cli_openai_responses_payload_parts( state, @@ -523,7 +545,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( candidate_index, candidate_id, spec_metadata.api_format, - transport, + &transport, provider_api_format, mapped_model, auth_header, @@ -536,11 +558,11 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( } let Some(upstream_url) = (if is_grok && is_grok_text_provider_api_format(provider_api_format) { - Some(build_grok_upstream_url(transport, GROK_CHAT_PATH)) + Some(build_grok_upstream_url(&transport, GROK_CHAT_PATH)) } else if needs_bidirectional_conversion { build_cross_format_openai_responses_upstream_url( parts, - transport, + &transport, &mapped_model, spec_metadata.api_format, provider_api_format, @@ -549,7 +571,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( } else { build_local_openai_responses_upstream_url( parts, - transport, + &transport, api_format_alias_matches(provider_api_format, "openai:responses:compact"), ) }) else { @@ -576,7 +598,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( .unwrap_or_default(); let resolved_headers = if is_grok && is_grok_text_provider_api_format(provider_api_format) { let Some(headers) = build_grok_browser_headers(GrokHeaderInput { - transport, + transport: &transport, transport_profile: transport_profile.as_ref(), request_headers: Some(effective_headers), content_type: "application/json", @@ -610,7 +632,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( } else { let Some(resolved_headers) = build_standard_provider_request_headers(StandardProviderRequestHeadersInput { - transport, + transport: &transport, provider_api_format, same_format, headers: effective_headers, @@ -704,7 +726,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts( None }, upstream_is_stream, - transport: Arc::clone(transport), + transport: Arc::clone(&transport), transport_profile, image_request_summary: None, request_redacted: redaction.redacted, diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs index ced42afd6..1d00d0dfb 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs @@ -32,6 +32,8 @@ pub(super) struct AdminProviderOAuthBatchImportEntry { pub email: Option, pub account_name: Option, pub project_id: Option, + pub client_version: Option, + pub session_id: Option, pub sso_rw_token: Option, pub cf_cookies: Option, pub cf_clearance: Option, @@ -191,6 +193,8 @@ fn extract_admin_provider_oauth_batch_import_entry( email: None, account_name: None, project_id: None, + client_version: None, + session_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"), @@ -331,6 +335,19 @@ fn extract_admin_provider_oauth_batch_import_entry( .or_else(|| object.get("cloudaicompanionProject")) .or_else(|| object.get("cloudAiCompanionProject")), ); + let client_version = coerce_admin_provider_oauth_import_str( + object + .get("client_version") + .or_else(|| object.get("clientVersion")) + .or_else(|| object.get("antigravityClientVersion")), + ); + let session_id = coerce_admin_provider_oauth_import_str( + object + .get("session_id") + .or_else(|| object.get("sessionId")) + .or_else(|| object.get("vscode_session_id")) + .or_else(|| object.get("vscodeSessionId")), + ); let sso_rw_token = coerce_admin_provider_oauth_import_str( object .get("sso_rw_token") @@ -379,6 +396,8 @@ fn extract_admin_provider_oauth_batch_import_entry( email, account_name, project_id, + client_version, + session_id, sso_rw_token, cf_cookies, cf_clearance, @@ -470,6 +489,8 @@ fn parse_error_entry(error: String) -> AdminProviderOAuthBatchImportEntry { email: None, account_name: None, project_id: None, + client_version: None, + session_id: None, sso_rw_token: None, cf_cookies: None, cf_clearance: None, @@ -502,6 +523,29 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints( } return; } + if provider_type == "antigravity" { + 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(client_version) = entry.client_version.as_ref() { + auth_config + .entry("client_version".to_string()) + .or_insert_with(|| json!(client_version)); + } + if let Some(session_id) = entry.session_id.as_ref() { + auth_config + .entry("session_id".to_string()) + .or_insert_with(|| json!(session_id)); + } + if let Some(user_agent) = entry.user_agent.as_ref() { + auth_config + .entry("user_agent".to_string()) + .or_insert_with(|| json!(user_agent)); + } + return; + } if !matches!(provider_type.as_str(), "codex" | "chatgpt_web" | "grok") { return; } @@ -823,6 +867,23 @@ mod tests { ); } + #[test] + fn applies_antigravity_project_and_user_agent_hints_to_auth_config() { + let entries = parse_admin_provider_oauth_batch_import_entries( + "antigravity", + r#"{"refreshToken":"rt-1","cloudaicompanionProject":{"id":"project-antigravity-2"},"userAgent":"antigravity"}"#, + ); + let mut auth_config = serde_json::Map::new(); + + apply_admin_provider_oauth_batch_import_hints("antigravity", &entries[0], &mut auth_config); + + assert_eq!( + auth_config.get("project_id"), + Some(&json!("project-antigravity-2")) + ); + assert_eq!(auth_config.get("user_agent"), Some(&json!("antigravity"))); + } + #[test] fn parses_windsurf_json_credentials_for_native_import() { let entries = parse_admin_provider_oauth_batch_import_entries( diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs index 9551eb231..bf3a0dee4 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs @@ -106,6 +106,34 @@ fn import_payload_string_any( .map(ToOwned::to_owned) } +fn import_payload_project_id_any( + payload: &serde_json::Map, + keys: &[&str], +) -> Option { + keys.iter().find_map(|key| { + let value = payload.get(*key)?; + if let Some(string) = value + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + return Some(string.to_string()); + } + value + .as_object() + .and_then(|object| { + object + .get("id") + .or_else(|| object.get("project_id")) + .or_else(|| object.get("projectId")) + }) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + }) +} + fn import_payload_u64_any( payload: &serde_json::Map, keys: &[&str], @@ -129,6 +157,49 @@ fn apply_single_import_hints( auth_config: &mut serde_json::Map, ) { let provider_type = provider_type.trim().to_ascii_lowercase(); + if provider_type == "antigravity" { + if let Some(project_id) = import_payload_project_id_any( + payload, + &[ + "project_id", + "projectId", + "cloudaicompanionProject", + "cloudAiCompanionProject", + ], + ) { + auth_config + .entry("project_id".to_string()) + .or_insert_with(|| json!(project_id)); + } + for (target, keys) in [ + ( + "client_version", + &[ + "client_version", + "clientVersion", + "antigravityClientVersion", + ][..], + ), + ( + "session_id", + &[ + "session_id", + "sessionId", + "vscode_session_id", + "vscodeSessionId", + ][..], + ), + ("user_agent", &["user_agent", "userAgent"][..]), + ] { + let Some(value) = import_payload_string_any(payload, keys) else { + continue; + }; + auth_config + .entry(target.to_string()) + .or_insert_with(|| json!(value)); + } + return; + } if !matches!(provider_type.as_str(), "codex" | "chatgpt_web" | "grok") { return; } @@ -640,7 +711,8 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token( #[cfg(test)] mod tests { use super::{ - import_payload_string_any, import_payload_u64_any, sanitize_windsurf_import_error, + apply_single_import_hints, import_payload_string_any, import_payload_u64_any, + sanitize_windsurf_import_error, }; use aether_oauth::core::OAuthError; use serde_json::json; @@ -675,6 +747,35 @@ mod tests { ); } + #[test] + fn single_import_applies_antigravity_identity_hints() { + let payload = json!({ + "cloudaicompanionProject": { + "id": "project-antigravity-1" + }, + "clientVersion": "1.99.0", + "sessionId": "session-antigravity-1", + "userAgent": "antigravity" + }) + .as_object() + .cloned() + .expect("payload should be an object"); + let mut auth_config = serde_json::Map::new(); + + apply_single_import_hints("antigravity", &payload, &mut auth_config); + + assert_eq!( + auth_config.get("project_id"), + Some(&json!("project-antigravity-1")) + ); + assert_eq!(auth_config.get("client_version"), Some(&json!("1.99.0"))); + assert_eq!( + auth_config.get("session_id"), + Some(&json!("session-antigravity-1")) + ); + assert_eq!(auth_config.get("user_agent"), Some(&json!("antigravity"))); + } + #[test] fn windsurf_import_error_redacts_http_body() { let error = OAuthError::HttpStatus { diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs index 857e4b8f8..3e60995f1 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs @@ -68,7 +68,7 @@ pub(crate) async fn refresh_antigravity_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? { @@ -104,15 +104,25 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally( } }; - let Some((project_id, identity_headers)) = - state.resolve_local_antigravity_identity_headers(&transport) - else { + let identity = match state.resolve_local_antigravity_identity_headers(&transport) { + Some(identity) => Some(identity), + None => state + .app() + .hydrate_antigravity_project_metadata_for_transport(&transport) + .await + .and_then(|hydrated| { + let identity = state.resolve_local_antigravity_identity_headers(&hydrated); + transport = hydrated; + identity + }), + }; + let Some((project_id, identity_headers)) = identity else { failed_count += 1; results.push(json!({ "key_id": key.id, "key_name": key.name, "status": "error", - "message": "缺少 OAuth 认证信息，请先授权/刷新 Token", + "message": "缺少 Antigravity project_id，loadCodeAssist 未返回可用项目信息", })); continue; }; diff --git a/apps/aether-gateway/src/state/integrations.rs b/apps/aether-gateway/src/state/integrations.rs index dd6e6c289..195aeb88b 100644 --- a/apps/aether-gateway/src/state/integrations.rs +++ b/apps/aether-gateway/src/state/integrations.rs @@ -13,8 +13,9 @@ use aether_data_contracts::repository::provider_catalog::{ }; use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot; use aether_model_fetch::{ - aggregate_models_for_cache, fetch_models_from_transports, merge_upstream_metadata, - model_fetch_interval_minutes, ModelFetchAssociationStore, ModelFetchTransportRuntime, + aggregate_models_for_cache, build_antigravity_load_code_assist_plan, + fetch_models_from_transports, merge_upstream_metadata, model_fetch_interval_minutes, + ModelFetchAssociationStore, ModelFetchTransportRuntime, }; use aether_scheduler_core::SchedulerAffinityTarget; use async_trait::async_trait; @@ -33,6 +34,109 @@ use crate::scheduler::state::SchedulerRuntimeState; use crate::{execution_runtime, provider_transport}; impl AppState { + pub(crate) async fn hydrate_antigravity_project_metadata_for_transport( + &self, + transport: &GatewayProviderTransportSnapshot, + ) -> Option { + if !provider_transport::antigravity::is_antigravity_provider_transport(transport) { + return None; + } + if matches!( + provider_transport::antigravity::resolve_local_antigravity_request_auth(transport), + provider_transport::antigravity::AntigravityRequestAuthSupport::Supported(_) + ) { + return Some(transport.clone()); + } + + let plan = match build_antigravity_load_code_assist_plan(self, transport).await { + Ok(plan) => plan, + Err(err) => { + warn!( + provider_id = %transport.provider.id, + endpoint_id = %transport.endpoint.id, + key_id = %transport.key.id, + error = %err, + "antigravity project metadata hydration failed" + ); + return None; + } + }; + let result = + match execution_runtime::execute_execution_runtime_sync_plan(self, None, &plan).await { + Ok(result) => result, + Err(err) => { + warn!( + provider_id = %transport.provider.id, + endpoint_id = %transport.endpoint.id, + key_id = %transport.key.id, + error = ?err, + "antigravity project metadata hydration request failed" + ); + return None; + } + }; + if !(200..300).contains(&result.status_code) { + warn!( + provider_id = %transport.provider.id, + endpoint_id = %transport.endpoint.id, + key_id = %transport.key.id, + status_code = result.status_code, + "antigravity project metadata hydration returned non-success status" + ); + return None; + } + let Some(project_id) = result + .body + .as_ref() + .and_then(|body| body.json_body.as_ref()) + .and_then(extract_antigravity_load_code_assist_project_id) + else { + warn!( + provider_id = %transport.provider.id, + endpoint_id = %transport.endpoint.id, + key_id = %transport.key.id, + "antigravity project metadata hydration response missing project" + ); + return None; + }; + let upstream_metadata = serde_json::json!({ + "antigravity": { + "project_id": project_id, + "updated_at": current_unix_secs(), + } + }); + 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 !matches!( + provider_transport::antigravity::resolve_local_antigravity_request_auth(&hydrated), + provider_transport::antigravity::AntigravityRequestAuthSupport::Supported(_) + ) { + 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, + "antigravity project metadata hydration could not persist metadata" + ); + } + + Some(hydrated) + } + pub(crate) async fn hydrate_gemini_cli_project_metadata_for_transport( &self, transport: &GatewayProviderTransportSnapshot, @@ -89,6 +193,30 @@ impl AppState { } } +fn extract_antigravity_load_code_assist_project_id(value: &Value) -> Option { + let raw = value + .get("cloudaicompanionProject") + .or_else(|| value.get("cloudAiCompanionProject"))?; + if let Some(project_id) = raw + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + return Some(project_id.to_string()); + } + raw.as_object() + .and_then(|object| { + object + .get("id") + .or_else(|| object.get("project_id")) + .or_else(|| object.get("projectId")) + }) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + #[async_trait] impl provider_transport::TransportTunnelAffinityLookup for AppState { async fn lookup_tunnel_attachment_owner( diff --git a/crates/aether-model-fetch/src/lib.rs b/crates/aether-model-fetch/src/lib.rs index a83ba1d49..4f54ac9e7 100644 --- a/crates/aether-model-fetch/src/lib.rs +++ b/crates/aether-model-fetch/src/lib.rs @@ -23,8 +23,9 @@ pub use strategy::{ SelectedModelFetchStrategy, }; pub use transport::{ - build_antigravity_fetch_available_models_plan, build_gemini_cli_load_code_assist_plan, - build_kiro_list_available_models_plan, build_models_fetch_execution_plan, - build_standard_models_fetch_execution_plan, build_vertex_models_fetch_execution_plan, - build_windsurf_model_configs_execution_plan, ModelFetchTransportRuntime, + build_antigravity_fetch_available_models_plan, build_antigravity_load_code_assist_plan, + build_gemini_cli_load_code_assist_plan, build_kiro_list_available_models_plan, + build_models_fetch_execution_plan, build_standard_models_fetch_execution_plan, + build_vertex_models_fetch_execution_plan, build_windsurf_model_configs_execution_plan, + ModelFetchTransportRuntime, }; diff --git a/crates/aether-model-fetch/src/strategy.rs b/crates/aether-model-fetch/src/strategy.rs index 55b7b3b55..c98bfdc01 100644 --- a/crates/aether-model-fetch/src/strategy.rs +++ b/crates/aether-model-fetch/src/strategy.rs @@ -2,6 +2,9 @@ use std::collections::{BTreeMap, BTreeSet}; use std::time::{SystemTime, UNIX_EPOCH}; use aether_contracts::{ExecutionPlan, ExecutionResult, RequestBody}; +use aether_provider_transport::antigravity::{ + resolve_local_antigravity_request_auth, AntigravityRequestAuthSupport, +}; use aether_provider_transport::{ is_vertex_api_key_transport_context, resolve_transport_execution_timeouts, resolve_transport_profile, GatewayProviderTransportSnapshot, @@ -21,10 +24,10 @@ use crate::logic::{ parse_windsurf_model_configs_response, preset_models_for_provider, }; use crate::transport::{ - build_antigravity_fetch_available_models_plan, build_gemini_cli_load_code_assist_plan, - build_kiro_list_available_models_plan, build_standard_models_fetch_execution_plan, - build_vertex_models_fetch_execution_plan, build_windsurf_model_configs_execution_plan, - ModelFetchTransportRuntime, + build_antigravity_fetch_available_models_plan, build_antigravity_load_code_assist_plan, + build_gemini_cli_load_code_assist_plan, build_kiro_list_available_models_plan, + build_standard_models_fetch_execution_plan, build_vertex_models_fetch_execution_plan, + build_windsurf_model_configs_execution_plan, ModelFetchTransportRuntime, }; const ANTIGRAVITY_SANDBOX_BASE_URL: &str = "https://daily-cloudcode-pa.sandbox.googleapis.com"; @@ -261,25 +264,18 @@ async fn fetch_antigravity_models( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transport: &GatewayProviderTransportSnapshot, ) -> Result { - let auth_config = transport_auth_config(transport); - let project_id = auth_config - .as_ref() - .and_then(|value| value.get("project_id")) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .ok_or_else(|| "antigravity: missing auth_config.project_id (please re-auth)".to_string())? - .to_string(); + let (project_id, hydrated_transport, project_metadata) = + resolve_or_hydrate_antigravity_project(runtime, transport).await?; let mut errors = Vec::new(); for base_url in [ - ANTIGRAVITY_SANDBOX_BASE_URL, ANTIGRAVITY_DAILY_BASE_URL, ANTIGRAVITY_PROD_BASE_URL, + ANTIGRAVITY_SANDBOX_BASE_URL, ] { let plan = match build_antigravity_fetch_available_models_plan( runtime, - transport, + &hydrated_transport, base_url, &project_id, ) @@ -300,6 +296,9 @@ async fn fetch_antigravity_models( if (200..300).contains(&result.status_code) { let body_json = execution_result_json_body_allow_empty(&result)?; let (models, metadata) = parse_antigravity_models_response(&body_json)?; + let metadata = metadata + .map(|metadata| attach_antigravity_project_metadata(metadata, &project_id)) + .or(project_metadata.clone()); return Ok(build_success_outcome(models, metadata, true)); } @@ -320,6 +319,74 @@ async fn fetch_antigravity_models( }) } +async fn resolve_or_hydrate_antigravity_project( + runtime: &(impl ModelFetchTransportRuntime + ?Sized), + transport: &GatewayProviderTransportSnapshot, +) -> Result<(String, GatewayProviderTransportSnapshot, Option), String> { + if let Some(project_id) = resolve_antigravity_project_id_from_transport(transport) { + let metadata = Some(build_antigravity_project_metadata(&project_id)); + return Ok((project_id, transport.clone(), metadata)); + } + + let plan = build_antigravity_load_code_assist_plan(runtime, transport).await?; + let result = runtime.execute_model_fetch_execution_plan(&plan).await?; + if !(200..300).contains(&result.status_code) { + return Err(format!( + "antigravity: loadCodeAssist failed: {}", + execution_result_error_message(&result) + )); + } + let body_json = execution_result_json_body_allow_empty(&result)?; + let project_id = extract_cloud_ai_companion_project_id(&body_json) + .ok_or_else(|| "antigravity: loadCodeAssist response missing project_id".to_string())?; + let metadata = build_antigravity_project_metadata(&project_id); + let mut hydrated_transport = transport.clone(); + hydrated_transport.key.upstream_metadata = Some(metadata.clone()); + + Ok((project_id, hydrated_transport, Some(metadata))) +} + +fn resolve_antigravity_project_id_from_transport( + transport: &GatewayProviderTransportSnapshot, +) -> Option { + match resolve_local_antigravity_request_auth(transport) { + AntigravityRequestAuthSupport::Supported(auth) => Some(auth.project_id), + AntigravityRequestAuthSupport::Unsupported(_) => None, + } +} + +fn build_antigravity_project_metadata(project_id: &str) -> Value { + json!({ + "antigravity": { + "project_id": project_id, + "updated_at": now_unix_secs(), + } + }) +} + +fn attach_antigravity_project_metadata(mut metadata: Value, project_id: &str) -> Value { + let Value::Object(root) = &mut metadata else { + return build_antigravity_project_metadata(project_id); + }; + let antigravity = root + .entry("antigravity".to_string()) + .or_insert_with(|| json!({})); + let Some(object) = antigravity.as_object_mut() else { + *antigravity = json!({ + "project_id": project_id, + "updated_at": now_unix_secs(), + }); + return metadata; + }; + object + .entry("project_id".to_string()) + .or_insert_with(|| Value::String(project_id.to_string())); + object + .entry("updated_at".to_string()) + .or_insert_with(|| Value::from(now_unix_secs())); + metadata +} + async fn fetch_gemini_cli_models( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transport: &GatewayProviderTransportSnapshot, @@ -340,8 +407,8 @@ async fn fetch_gemini_cli_models( provider_meta.insert(key.to_string(), value); } } - if let Some(project_id) = - extract_gemini_cli_project_id(&body_json).or_else(|| { + if let Some(project_id) = extract_cloud_ai_companion_project_id(&body_json) + .or_else(|| { transport_auth_config(transport) .and_then(|value| value.get("project_id").cloned()) .and_then(|value| value.as_str().map(ToOwned::to_owned)) @@ -1300,8 +1367,10 @@ fn extract_gemini_cli_tier_metadata(body: &Value, key: &str) -> Option { (!out.is_empty()).then_some(Value::Object(out)) } -fn extract_gemini_cli_project_id(body: &Value) -> Option { - let raw = body.get("cloudaicompanionProject")?; +fn extract_cloud_ai_companion_project_id(body: &Value) -> Option { + let raw = body + .get("cloudaicompanionProject") + .or_else(|| body.get("cloudAiCompanionProject"))?; if let Some(value) = raw.as_str() { let value = value.trim(); if !value.is_empty() { @@ -1368,6 +1437,11 @@ mod tests { routes: Vec, } + struct OAuthRoutingTestRuntime { + executed_urls: Arc>>, + routes: Vec, + } + #[async_trait] impl ModelFetchTransportRuntime for TestRuntime { async fn resolve_local_oauth_request_auth( @@ -1459,6 +1533,62 @@ mod tests { } } + #[async_trait] + impl ModelFetchTransportRuntime for OAuthRoutingTestRuntime { + async fn resolve_local_oauth_request_auth( + &self, + _transport: &GatewayProviderTransportSnapshot, + ) -> Result, String> + { + Ok(Some( + aether_provider_transport::LocalResolvedOAuthRequestAuth::Header { + name: "authorization".to_string(), + value: "Bearer oauth-token".to_string(), + }, + )) + } + + async fn resolve_model_fetch_proxy( + &self, + _transport: &GatewayProviderTransportSnapshot, + ) -> Option { + None + } + + async fn execute_model_fetch_execution_plan( + &self, + plan: &aether_contracts::ExecutionPlan, + ) -> Result { + self.executed_urls + .lock() + .expect("executed_urls lock") + .push(plan.url.clone()); + let Some((_, route_result)) = self + .routes + .iter() + .find(|(url_part, _)| plan.url.contains(url_part)) + else { + return Err(format!("unexpected models fetch URL {}", plan.url)); + }; + let (status_code, response_body) = match route_result { + Ok((status_code, response_body)) => (*status_code, response_body.clone()), + Err(err) => return Err(err.clone()), + }; + Ok(ExecutionResult { + request_id: plan.request_id.clone(), + candidate_id: plan.candidate_id.clone(), + status_code, + headers: BTreeMap::new(), + body: Some(ResponseBody { + json_body: Some(response_body), + body_bytes_b64: None, + }), + telemetry: None, + error: None, + }) + } + } + fn sample_custom_aiplatform_transport() -> GatewayProviderTransportSnapshot { GatewayProviderTransportSnapshot { provider: GatewayProviderTransportProvider { @@ -1565,6 +1695,18 @@ mod tests { transport } + fn sample_antigravity_transport_without_project() -> GatewayProviderTransportSnapshot { + let mut transport = sample_custom_aiplatform_transport(); + transport.provider.provider_type = "antigravity".to_string(); + transport.provider.name = "Antigravity".to_string(); + transport.endpoint.base_url = "https://daily-cloudcode-pa.googleapis.com".to_string(); + transport.key.auth_type = "oauth".to_string(); + transport.key.decrypted_api_key = "__placeholder__".to_string(); + transport.key.decrypted_auth_config = + Some(r#"{"provider_type":"antigravity","refresh_token":"rt"}"#.to_string()); + transport + } + fn sample_windsurf_transport() -> GatewayProviderTransportSnapshot { let mut transport = sample_custom_aiplatform_transport(); transport.provider.provider_type = "windsurf".to_string(); @@ -1933,6 +2075,77 @@ mod tests { .is_none()); } + #[tokio::test] + async fn antigravity_model_fetch_hydrates_project_from_daily_load_code_assist() { + let executed_urls = Arc::new(Mutex::new(Vec::new())); + let runtime = OAuthRoutingTestRuntime { + executed_urls: Arc::clone(&executed_urls), + routes: vec![ + ( + "https://daily-cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" + .to_string(), + Ok(( + 200, + json!({ + "cloudaicompanionProject": { + "id": "project-from-antigravity-load" + } + }), + )), + ), + ( + "https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels" + .to_string(), + Ok(( + 200, + json!({ + "models": { + "chat_12345": { + "displayName": "Antigravity Chat", + "quotaInfo": { + "remainingFraction": 0.75 + } + } + } + }), + )), + ), + ], + }; + + let outcome = fetch_models_from_transports( + &runtime, + &[sample_antigravity_transport_without_project()], + ) + .await + .expect("antigravity models fetch should hydrate project and succeed"); + + let urls = executed_urls.lock().expect("executed_urls lock"); + assert_eq!( + urls.as_slice(), + &[ + "https://daily-cloudcode-pa.googleapis.com/v1internal:loadCodeAssist", + "https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels", + ] + ); + assert_eq!(outcome.fetched_model_ids, vec!["chat_12345"]); + assert_eq!( + outcome + .upstream_metadata + .as_ref() + .and_then(|value| value.pointer("/antigravity/project_id")), + Some(&json!("project-from-antigravity-load")) + ); + assert_eq!( + outcome + .upstream_metadata + .as_ref() + .and_then(|value| value + .pointer("/antigravity/quota_by_model/chat_12345/remaining_fraction")), + Some(&json!(0.75)) + ); + } + #[tokio::test] async fn kiro_transport_fetches_list_available_models() { let executed_urls = Arc::new(Mutex::new(Vec::new())); diff --git a/crates/aether-model-fetch/src/transport.rs b/crates/aether-model-fetch/src/transport.rs index cc7390470..9da7f1a5d 100644 --- a/crates/aether-model-fetch/src/transport.rs +++ b/crates/aether-model-fetch/src/transport.rs @@ -2,8 +2,9 @@ use std::collections::BTreeMap; use aether_contracts::{ExecutionPlan, ExecutionResult, ProxySnapshot, RequestBody}; use aether_provider_transport::antigravity::{ - build_antigravity_static_identity_headers, resolve_local_antigravity_request_auth, - AntigravityRequestAuthSupport, ANTIGRAVITY_REQUEST_USER_AGENT, + build_antigravity_static_client_headers, build_antigravity_static_identity_headers, + resolve_local_antigravity_request_auth, AntigravityRequestAuthSupport, + ANTIGRAVITY_REQUEST_USER_AGENT, }; use aether_provider_transport::auth::{ ensure_upstream_auth_header, resolve_local_gemini_auth, resolve_local_openai_bearer_auth, @@ -29,6 +30,7 @@ const CLAUDE_CLI_USER_AGENT: &str = "claude-code/1.0.1"; const GEMINI_CLI_USER_AGENT: &str = "GeminiCLI/0.1.5 (Windows; AMD64)"; const CLAUDE_VERSION_HEADER: &str = "2023-06-01"; const ANTIGRAVITY_FETCH_PROVIDER_API_FORMAT: &str = "antigravity:fetch_available_models"; +const ANTIGRAVITY_LOAD_CODE_ASSIST_PROVIDER_API_FORMAT: &str = "antigravity:load_code_assist"; const GEMINI_CLI_LOAD_CODE_ASSIST_PROVIDER_API_FORMAT: &str = "gemini_cli:load_code_assist"; const KIRO_LIST_AVAILABLE_MODELS_PROVIDER_API_FORMAT: &str = "kiro:list_available_models"; const WINDSURF_MODEL_CONFIGS_PROVIDER_API_FORMAT: &str = "windsurf:model_configs"; @@ -215,6 +217,49 @@ pub async fn build_antigravity_fetch_available_models_plan( .await } +pub async fn build_antigravity_load_code_assist_plan( + runtime: &(impl ModelFetchTransportRuntime + ?Sized), + transport: &GatewayProviderTransportSnapshot, +) -> Result { + let authorization = resolve_oauth_header_auth(runtime, transport) + .await? + .ok_or_else(|| { + "Antigravity loadCodeAssist requires OAuth authorization header".to_string() + })?; + + let mut headers = build_antigravity_static_client_headers(None, None); + headers.insert(authorization.0.clone(), authorization.1.clone()); + headers.insert("content-type".to_string(), "application/json".to_string()); + headers.insert("accept".to_string(), "application/json".to_string()); + headers + .entry("user-agent".to_string()) + .or_insert_with(|| ANTIGRAVITY_REQUEST_USER_AGENT.to_string()); + let protected_headers = vec![authorization.0]; + headers = apply_fetch_header_rules(transport, headers, &protected_headers)?; + + build_execution_plan( + runtime, + transport, + ModelFetchExecutionPlanRequest { + method: "POST".to_string(), + url: "https://daily-cloudcode-pa.googleapis.com/v1internal:loadCodeAssist".to_string(), + headers, + content_type: Some("application/json".to_string()), + body: RequestBody::from_json(json!({ + "metadata": { + "ideType": "ANTIGRAVITY", + "platform": "PLATFORM_UNSPECIFIED", + "pluginType": "GEMINI", + } + })), + client_api_format: "gemini:generate_content".to_string(), + provider_api_format: ANTIGRAVITY_LOAD_CODE_ASSIST_PROVIDER_API_FORMAT.to_string(), + model_name: Some("loadCodeAssist".to_string()), + }, + ) + .await +} + pub async fn build_gemini_cli_load_code_assist_plan( runtime: &(impl ModelFetchTransportRuntime + ?Sized), transport: &GatewayProviderTransportSnapshot, @@ -678,10 +723,10 @@ mod tests { use serde_json::json; use super::{ - build_antigravity_fetch_available_models_plan, build_gemini_cli_load_code_assist_plan, - build_kiro_list_available_models_plan, build_models_fetch_execution_plan, - build_standard_models_fetch_execution_plan, build_vertex_models_fetch_execution_plan, - ModelFetchTransportRuntime, + build_antigravity_fetch_available_models_plan, build_antigravity_load_code_assist_plan, + build_gemini_cli_load_code_assist_plan, build_kiro_list_available_models_plan, + build_models_fetch_execution_plan, build_standard_models_fetch_execution_plan, + build_vertex_models_fetch_execution_plan, ModelFetchTransportRuntime, }; struct TestRuntime { @@ -1006,6 +1051,65 @@ mod tests { .and_then(|value| value.get("project")), Some(&json!("project-1")) ); + assert_eq!( + plan.headers.get("user-agent").map(String::as_str), + Some("antigravity") + ); + assert_eq!( + plan.headers.get("x-client-name").map(String::as_str), + Some("antigravity") + ); + assert_eq!( + plan.headers.get("x-goog-api-client").map(String::as_str), + Some("gl-node/18.18.2 fire/0.8.6 grpc/1.10.x") + ); + assert_eq!( + plan.headers.get("x-client-version").map(String::as_str), + Some("1.2.3") + ); + assert_eq!( + plan.headers.get("x-vscode-sessionid").map(String::as_str), + Some("sess-1") + ); + } + + #[tokio::test] + async fn builds_antigravity_load_code_assist_plan_with_cli_headers() { + let runtime = TestRuntime { + oauth_auth: Some( + aether_provider_transport::LocalResolvedOAuthRequestAuth::Header { + name: "authorization".to_string(), + value: "Bearer oauth-token".to_string(), + }, + ), + proxy: None, + }; + let transport = sample_transport("antigravity", "gemini:generate_content", "oauth"); + let plan = build_antigravity_load_code_assist_plan(&runtime, &transport) + .await + .expect("plan"); + + assert_eq!(plan.method, "POST"); + assert_eq!( + plan.url, + "https://daily-cloudcode-pa.googleapis.com/v1internal:loadCodeAssist" + ); + assert_eq!( + plan.headers.get("authorization").map(String::as_str), + Some("Bearer oauth-token") + ); + assert_eq!( + plan.headers.get("user-agent").map(String::as_str), + Some("antigravity") + ); + assert_eq!( + plan.headers.get("x-client-name").map(String::as_str), + Some("antigravity") + ); + assert_eq!( + plan.headers.get("x-goog-api-client").map(String::as_str), + Some("gl-node/18.18.2 fire/0.8.6 grpc/1.10.x") + ); } #[tokio::test] diff --git a/crates/aether-provider-transport/src/antigravity/auth.rs b/crates/aether-provider-transport/src/antigravity/auth.rs index 919129c52..ce5dfc998 100644 --- a/crates/aether-provider-transport/src/antigravity/auth.rs +++ b/crates/aether-provider-transport/src/antigravity/auth.rs @@ -69,46 +69,26 @@ pub fn resolve_local_antigravity_request_auth( ); } - let Some(project_id) = find_string_by_paths( + let upstream_metadata = transport.key.upstream_metadata.as_ref(); + let Some(project_id) = find_antigravity_string( + upstream_metadata, &auth_config, - &[ - &["project_id"], - &["projectId"], - &["project", "id"], - &["project", "project_id"], - &["project", "projectId"], - &["antigravity", "project_id"], - &["antigravity", "projectId"], - &["metadata", "project_id"], - &["metadata", "projectId"], - ], + ANTIGRAVITY_PROJECT_ID_PATHS, ) else { return AntigravityRequestAuthSupport::Unsupported( AntigravityRequestAuthUnsupportedReason::MissingProjectId, ); }; - let client_version = find_string_by_paths( + let client_version = find_antigravity_string( + upstream_metadata, &auth_config, - &[ - &["client_version"], - &["clientVersion"], - &["antigravity", "client_version"], - &["antigravity", "clientVersion"], - &["metadata", "client_version"], - &["metadata", "clientVersion"], - ], + ANTIGRAVITY_CLIENT_VERSION_PATHS, ); - let session_id = find_string_by_paths( + let session_id = find_antigravity_string( + upstream_metadata, &auth_config, - &[ - &["session_id"], - &["sessionId"], - &["antigravity", "session_id"], - &["antigravity", "sessionId"], - &["metadata", "session_id"], - &["metadata", "sessionId"], - ], + ANTIGRAVITY_SESSION_ID_PATHS, ); AntigravityRequestAuthSupport::Supported(AntigravityRequestAuth { @@ -120,6 +100,16 @@ pub fn resolve_local_antigravity_request_auth( pub fn build_antigravity_static_identity_headers( auth: &AntigravityRequestAuth, +) -> BTreeMap { + build_antigravity_static_client_headers( + auth.client_version.as_deref(), + auth.session_id.as_deref(), + ) +} + +pub fn build_antigravity_static_client_headers( + client_version: Option<&str>, + session_id: Option<&str>, ) -> BTreeMap { let mut headers = BTreeMap::from([ ( @@ -132,26 +122,72 @@ pub fn build_antigravity_static_identity_headers( ), ]); - if let Some(client_version) = auth - .client_version - .as_deref() + if let Some(client_version) = client_version .map(str::trim) .filter(|value| !value.is_empty()) { headers.insert(String::from("x-client-version"), client_version.to_string()); } - if let Some(session_id) = auth - .session_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - { + if let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) { headers.insert(String::from("x-vscode-sessionid"), session_id.to_string()); } headers } +const ANTIGRAVITY_PROJECT_ID_PATHS: &[&[&str]] = &[ + &["project_id"], + &["projectId"], + &["project", "id"], + &["project", "project_id"], + &["project", "projectId"], + &["cloudaicompanionProject"], + &["cloudaicompanionProject", "id"], + &["cloudAiCompanionProject"], + &["cloudAiCompanionProject", "id"], + &["antigravity", "project_id"], + &["antigravity", "projectId"], + &["antigravity", "project", "id"], + &["antigravity", "cloudaicompanionProject"], + &["antigravity", "cloudaicompanionProject", "id"], + &["antigravity", "cloudAiCompanionProject"], + &["antigravity", "cloudAiCompanionProject", "id"], + &["metadata", "project_id"], + &["metadata", "projectId"], + &["metadata", "cloudaicompanionProject"], + &["metadata", "cloudaicompanionProject", "id"], + &["metadata", "cloudAiCompanionProject"], + &["metadata", "cloudAiCompanionProject", "id"], +]; + +const ANTIGRAVITY_CLIENT_VERSION_PATHS: &[&[&str]] = &[ + &["client_version"], + &["clientVersion"], + &["antigravity", "client_version"], + &["antigravity", "clientVersion"], + &["metadata", "client_version"], + &["metadata", "clientVersion"], +]; + +const ANTIGRAVITY_SESSION_ID_PATHS: &[&[&str]] = &[ + &["session_id"], + &["sessionId"], + &["antigravity", "session_id"], + &["antigravity", "sessionId"], + &["metadata", "session_id"], + &["metadata", "sessionId"], +]; + +fn find_antigravity_string( + upstream_metadata: Option<&Value>, + auth_config: &Value, + paths: &[&[&str]], +) -> Option { + upstream_metadata + .and_then(|metadata| find_string_by_paths(metadata, paths)) + .or_else(|| find_string_by_paths(auth_config, paths)) +} + fn find_string_by_paths(value: &Value, paths: &[&[&str]]) -> Option { for path in paths { let mut current = value; @@ -173,6 +209,20 @@ fn find_string_by_paths(value: &Value, paths: &[&[&str]]) -> Option { { return Some(string.to_string()); } + if let Some(string) = current + .as_object() + .and_then(|object| { + object + .get("id") + .or_else(|| object.get("project_id")) + .or_else(|| object.get("projectId")) + }) + .and_then(Value::as_str) + .map(str::trim) + .filter(|item| !item.is_empty()) + { + return Some(string.to_string()); + } } None @@ -211,3 +261,119 @@ fn is_blocked_auth_key(key: &str) -> bool { | "audience" ) } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{ + resolve_local_antigravity_request_auth, AntigravityRequestAuth, + AntigravityRequestAuthSupport, + }; + use crate::snapshot::{ + GatewayProviderTransportEndpoint, GatewayProviderTransportKey, + GatewayProviderTransportProvider, GatewayProviderTransportSnapshot, + }; + + fn sample_transport(auth_config: &str) -> GatewayProviderTransportSnapshot { + GatewayProviderTransportSnapshot { + provider: GatewayProviderTransportProvider { + id: "provider-1".to_string(), + name: "Antigravity".to_string(), + provider_type: "antigravity".to_string(), + website: None, + is_active: true, + keep_priority_on_conversion: false, + enable_format_conversion: true, + 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: Some("gemini".to_string()), + endpoint_kind: Some("generate_content".to_string()), + is_active: true, + base_url: "https://daily-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: Some(vec!["gemini:generate_content".to_string()]), + 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: None, + decrypted_api_key: "__placeholder__".to_string(), + decrypted_auth_config: Some(auth_config.to_string()), + }, + } + } + + #[test] + fn resolves_cloudaicompanion_project_object_from_auth_config() { + let transport = sample_transport( + r#"{ + "provider_type":"antigravity", + "refresh_token":"rt", + "cloudaicompanionProject":{"id":"project-from-auth-config"} + }"#, + ); + + assert_eq!( + resolve_local_antigravity_request_auth(&transport), + AntigravityRequestAuthSupport::Supported(AntigravityRequestAuth { + project_id: "project-from-auth-config".to_string(), + client_version: None, + session_id: None, + }) + ); + } + + #[test] + fn resolves_identity_from_antigravity_upstream_metadata() { + let mut transport = sample_transport( + r#"{ + "provider_type":"antigravity", + "refresh_token":"rt" + }"#, + ); + transport.key.upstream_metadata = Some(json!({ + "antigravity": { + "project_id": "project-from-metadata", + "client_version": "1.99.0", + "session_id": "session-from-metadata" + } + })); + + assert_eq!( + resolve_local_antigravity_request_auth(&transport), + AntigravityRequestAuthSupport::Supported(AntigravityRequestAuth { + project_id: "project-from-metadata".to_string(), + client_version: Some("1.99.0".to_string()), + session_id: Some("session-from-metadata".to_string()), + }) + ); + } +} diff --git a/crates/aether-provider-transport/src/antigravity/mod.rs b/crates/aether-provider-transport/src/antigravity/mod.rs index cdab73700..4d09239e2 100644 --- a/crates/aether-provider-transport/src/antigravity/mod.rs +++ b/crates/aether-provider-transport/src/antigravity/mod.rs @@ -4,9 +4,10 @@ mod request; mod url; pub use auth::{ - build_antigravity_static_identity_headers, resolve_local_antigravity_request_auth, - AntigravityRequestAuth, AntigravityRequestAuthSupport, AntigravityRequestAuthUnsupportedReason, - ANTIGRAVITY_PROVIDER_TYPE, ANTIGRAVITY_REQUEST_USER_AGENT, + build_antigravity_static_client_headers, build_antigravity_static_identity_headers, + resolve_local_antigravity_request_auth, AntigravityRequestAuth, AntigravityRequestAuthSupport, + AntigravityRequestAuthUnsupportedReason, ANTIGRAVITY_PROVIDER_TYPE, + ANTIGRAVITY_REQUEST_USER_AGENT, }; pub use policy::{ classify_local_antigravity_request_support, is_antigravity_provider_transport, diff --git a/crates/aether-provider-transport/src/provider_types.rs b/crates/aether-provider-transport/src/provider_types.rs index 42916ada7..8a3bc663b 100644 --- a/crates/aether-provider-transport/src/provider_types.rs +++ b/crates/aether-provider-transport/src/provider_types.rs @@ -384,8 +384,8 @@ const VERTEX_AI_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTe const ANTIGRAVITY_FIXED_PROVIDER_TEMPLATE: FixedProviderTemplate = FixedProviderTemplate { provider_type: "antigravity", - version: 1, - base_url: "https://cloudcode-pa.googleapis.com", + version: 2, + base_url: "https://daily-cloudcode-pa.googleapis.com", endpoints: &[FixedProviderEndpointTemplate { item_key: "gemini:generate_content", api_format: "gemini:generate_content", @@ -750,6 +750,24 @@ mod tests { ); } + #[test] + fn antigravity_fixed_provider_template_uses_daily_cloudcode_endpoint() { + let template = + fixed_provider_template("antigravity").expect("antigravity template should exist"); + assert_eq!( + template.base_url, + "https://daily-cloudcode-pa.googleapis.com" + ); + assert_eq!(template.version, 2); + + let endpoint = fixed_provider_endpoint_template_by_api_format( + "antigravity", + "gemini:generate_content", + ) + .expect("antigravity generateContent endpoint should exist"); + assert_eq!(endpoint.custom_path, None); + } + #[test] fn windsurf_fixed_provider_template_exposes_openai_chat() { let template = fixed_provider_template("windsurf").expect("windsurf template should exist");