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

修复 Antigravity OAuth 配额复检缺 project

Browse Source
This commit is contained in:
MMEXA
2026-06-19 22:21:22 +08:00
parent 16584067d7
commit 6c4e730e60
12 changed files with 977 additions and 130 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
+23 -1
View File
@@ -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,
@@ -230,11 +231,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(
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs
+70 -48
View File
@@ -33,7 +33,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,
@@ -119,11 +119,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
@@ -145,33 +145,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,
@@ -194,7 +195,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
@@ -219,20 +220,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(
@@ -258,7 +259,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,
@@ -414,11 +415,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(
@@ -480,7 +502,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,
@@ -504,7 +526,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,
@@ -516,7 +538,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,
@@ -528,7 +550,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,
@@ -541,11 +563,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,
@@ -554,7 +576,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 {
@@ -581,7 +603,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",
@@ -615,7 +637,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,
@@ -709,7 +731,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,
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs
+61
View File
@@ -32,6 +32,8 @@ pub(super) struct AdminProviderOAuthBatchImportEntry {
pub email: Option<String>,
pub account_name: Option<String>,
pub project_id: Option<String>,
pub client_version: Option<String>,
pub session_id: Option<String>,
pub sso_rw_token: Option<String>,
pub cf_cookies: Option<String>,
pub cf_clearance: Option<String>,
@@ -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(
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
+102 -1
View File
@@ -106,6 +106,34 @@ fn import_payload_string_any(
.map(ToOwned::to_owned)
}
fn import_payload_project_id_any(
payload: &serde_json::Map<String, serde_json::Value>,
keys: &[&str],
) -> Option<String> {
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<String, serde_json::Value>,
keys: &[&str],
@@ -129,6 +157,49 @@ fn apply_single_import_hints(
auth_config: &mut serde_json::Map<String, serde_json::Value>,
) {
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 {
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
+15 -5
View File
@@ -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;
};
apps/aether-gateway/src/state/integrations.rs
+130 -2
View File
@@ -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<GatewayProviderTransportSnapshot> {
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<String> {
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(
crates/aether-model-fetch/src/lib.rs
+5 -4
View File
@@ -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,
};
crates/aether-model-fetch/src/strategy.rs
+232 -19
View File
@@ -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<ModelsFetchOutcome, String> {
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<Value>), 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<String> {
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<Value> {
(!out.is_empty()).then_some(Value::Object(out))
}
fn extract_gemini_cli_project_id(body: &Value) -> Option<String> {
let raw = body.get("cloudaicompanionProject")?;
fn extract_cloud_ai_companion_project_id(body: &Value) -> Option<String> {
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<ModelFetchRoute>,
}
struct OAuthRoutingTestRuntime {
executed_urls: Arc<Mutex<Vec<String>>>,
routes: Vec<ModelFetchRoute>,
}
#[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<Option<aether_provider_transport::LocalResolvedOAuthRequestAuth>, 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<aether_contracts::ProxySnapshot> {
None
}
async fn execute_model_fetch_execution_plan(
&self,
plan: &aether_contracts::ExecutionPlan,
) -> Result<ExecutionResult, String> {
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()));
crates/aether-model-fetch/src/transport.rs
+110 -6
View File
@@ -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<ExecutionPlan, String> {
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]
crates/aether-provider-transport/src/antigravity/auth.rs
+205 -39
View File
@@ -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<String, String> {
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<String, String> {
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<String> {
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<String> {
for path in paths {
let mut current = value;
@@ -173,6 +209,20 @@ fn find_string_by_paths(value: &Value, paths: &[&[&str]]) -> Option<String> {
{
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()),
})
);
}
}
crates/aether-provider-transport/src/antigravity/mod.rs
+4 -3
View File
@@ -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,
crates/aether-provider-transport/src/provider_types.rs
+20 -2
View File
@@ -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");