From c5ae9c2c7732d83414432405a678533c88864e7f Mon Sep 17 00:00:00 2001 From: ZheFox <77232781+zhefox@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:06:39 +0800 Subject: [PATCH] fix(antigravity): sync discovered models into catalog --- .../admin/provider/oauth/quota/antigravity.rs | 62 +++++++++++++++++++ .../tests/control/admin/endpoints/quota.rs | 29 +++++++++ crates/aether-model-fetch/src/lib.rs | 5 +- crates/aether-model-fetch/src/strategy.rs | 10 ++- 4 files changed, 103 insertions(+), 3 deletions(-) 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 d3e26fc79..2cbd0848a 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 @@ -5,6 +5,7 @@ use super::shared::{ quota_key_auto_removed, quota_refresh_success_invalid_state, resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome, }; +use crate::handlers::admin::provider::shared::payloads::AdminImportProviderModelsRequest; use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot}; use crate::GatewayError; use aether_admin::provider::quota::{ @@ -22,6 +23,63 @@ use std::collections::BTreeMap; use std::time::{SystemTime, UNIX_EPOCH}; use tracing::warn; +fn antigravity_discovered_model_ids(metadata_update: Option<&serde_json::Value>) -> Vec { + metadata_update + .and_then(|value| value.pointer("/antigravity/quota_by_model")) + .and_then(serde_json::Value::as_object) + .into_iter() + .flat_map(|models| models.keys()) + .map(String::as_str) + .filter(|model_id| aether_model_fetch::antigravity_model_id_is_routable(model_id)) + .map(ToOwned::to_owned) + .collect() +} + +async fn sync_antigravity_discovered_models( + state: &AdminAppState<'_>, + provider_id: &str, + metadata_update: Option<&serde_json::Value>, +) { + if !state.has_global_model_data_reader() || !state.has_global_model_data_writer() { + return; + } + let model_ids = antigravity_discovered_model_ids(metadata_update); + if model_ids.is_empty() { + return; + } + + let result = state + .build_admin_import_provider_models_payload( + provider_id, + AdminImportProviderModelsRequest { + model_ids, + tiered_pricing: None, + price_per_request: None, + }, + ) + .await; + match result { + Ok(payload) => { + let errors = payload + .get("errors") + .and_then(serde_json::Value::as_array) + .map(Vec::len) + .unwrap_or(0); + if errors > 0 { + warn!( + provider_id, + errors, "Antigravity discovered-model catalog sync completed with item errors" + ); + } + } + Err(error) => warn!( + provider_id, + error = %error, + "Antigravity discovered-model catalog sync failed" + ), + } +} + async fn execute_antigravity_quota_plan( state: &AdminAppState<'_>, transport: &AdminGatewayProviderTransportSnapshot, @@ -330,6 +388,10 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally( continue; } + if status == "success" { + sync_antigravity_discovered_models(state, &provider.id, metadata_update.as_ref()).await; + } + if status == "success" { success_count += 1; } else { diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs index 444d550f2..7120cf871 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs @@ -4,8 +4,12 @@ use std::sync::{Arc, Mutex}; use aether_crypto::{ decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY, }; +use aether_data::repository::global_models::InMemoryGlobalModelReadRepository; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository; +use aether_data_contracts::repository::global_models::{ + AdminProviderModelListQuery, GlobalModelReadRepository, +}; use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; @@ -2330,6 +2334,12 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru }, "gemini-2.5-pro": { "displayName": "Gemini 2.5 Pro" + }, + "gemini-3.7-flash-tiered": { + "displayName": "Gemini 3.7 Flash" + }, + "chat_23310": { + "displayName": "Internal Chat" } } }), @@ -2431,6 +2441,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru )], vec![key], )); + let global_model_repository = Arc::new(InMemoryGlobalModelReadRepository::default()); let (upstream_url, upstream_handle) = start_server(upstream).await; let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; @@ -2440,6 +2451,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru GatewayDataState::with_provider_catalog_repository_for_tests( provider_catalog_repository.clone(), ) + .with_global_model_repository_for_tests(global_model_repository.clone()) .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), ), ); @@ -2541,6 +2553,23 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru .and_then(|value| value.get("remaining_fraction")), Some(&json!(0.25)) ); + let imported_provider_models = global_model_repository + .list_admin_provider_models(&AdminProviderModelListQuery { + provider_id: "provider-antigravity".to_string(), + is_active: None, + offset: 0, + limit: 100, + }) + .await + .expect("imported Antigravity provider models should read"); + let imported_model_names = imported_provider_models + .iter() + .map(|model| model.provider_model_name.as_str()) + .collect::>(); + assert!(imported_model_names.contains("claude-sonnet-4")); + assert!(imported_model_names.contains("gemini-2.5-pro")); + assert!(imported_model_names.contains("gemini-3.7-flash-tiered")); + assert!(!imported_model_names.contains("chat_23310")); assert_eq!( reloaded[0] .upstream_metadata diff --git a/crates/aether-model-fetch/src/lib.rs b/crates/aether-model-fetch/src/lib.rs index bb3432805..8ef4b5855 100644 --- a/crates/aether-model-fetch/src/lib.rs +++ b/crates/aether-model-fetch/src/lib.rs @@ -21,8 +21,9 @@ pub use logic::{ upstream_metadata_namespace_updates, ModelFetchRunSummary, ModelsFetchPage, ModelsFetchSuccess, }; pub use strategy::{ - fetch_models_from_transports, fetch_models_from_transports_for_client_version, - ModelFetchStrategy, ModelFetchStrategyKind, ModelsFetchOutcome, SelectedModelFetchStrategy, + antigravity_model_id_is_routable, fetch_models_from_transports, + fetch_models_from_transports_for_client_version, ModelFetchStrategy, ModelFetchStrategyKind, + ModelsFetchOutcome, SelectedModelFetchStrategy, }; pub use transport::{ build_antigravity_fetch_available_models_plan, build_antigravity_load_code_assist_plan, diff --git a/crates/aether-model-fetch/src/strategy.rs b/crates/aether-model-fetch/src/strategy.rs index 64951adb8..249bec490 100644 --- a/crates/aether-model-fetch/src/strategy.rs +++ b/crates/aether-model-fetch/src/strategy.rs @@ -1018,7 +1018,7 @@ fn parse_antigravity_models_response(body: &Value) -> Result<(Vec, Option let mut quota_by_model = serde_json::Map::new(); for (model_id, model_data) in models_object { let model_id = model_id.trim(); - if model_id.is_empty() || ANTIGRAVITY_BLOCKED_MODELS.contains(&model_id) { + if !antigravity_model_id_is_routable(model_id) { continue; } let model_object = model_data.as_object().cloned().unwrap_or_default(); @@ -1054,6 +1054,14 @@ fn parse_antigravity_models_response(body: &Value) -> Result<(Vec, Option Ok((models, upstream_metadata)) } +pub fn antigravity_model_id_is_routable(model_id: &str) -> bool { + let model_id = model_id.trim(); + !model_id.is_empty() + && !ANTIGRAVITY_BLOCKED_MODELS + .iter() + .any(|blocked| blocked.eq_ignore_ascii_case(model_id)) +} + fn parse_kiro_available_models_response( body: &Value, ) -> Result<(Vec, Option), String> {