diff --git a/Cargo.lock b/Cargo.lock index 7b5b75682..3fafa264e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -29,6 +29,7 @@ dependencies = [ "aether-data", "aether-data-contracts", "aether-provider-pool", + "aether-provider-transport", "axum", "base64 0.22.1", "chrono", diff --git a/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/payloads.rs b/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/payloads.rs index f82e03345..cffcd0d05 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/payloads.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/payloads.rs @@ -13,12 +13,14 @@ pub(super) fn key_api_formats_without_entry( } pub(super) fn endpoint_key_counts_by_format( + provider_type: &str, + endpoints: &[StoredProviderCatalogEndpoint], keys: &[StoredProviderCatalogKey], ) -> ( std::collections::BTreeMap, std::collections::BTreeMap, ) { - admin_provider_endpoints_pure::endpoint_key_counts_by_format(keys) + admin_provider_endpoints_pure::endpoint_key_counts_by_format(provider_type, endpoints, keys) } pub(super) fn build_admin_provider_endpoint_response( diff --git a/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/reads.rs b/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/reads.rs index c4b8e40c9..36e443015 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/reads.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/reads.rs @@ -38,7 +38,8 @@ pub(crate) async fn build_admin_provider_endpoints_payload( .await .ok() .unwrap_or_default(); - let (total_keys_by_format, active_keys_by_format) = endpoint_key_counts_by_format(&keys); + let (total_keys_by_format, active_keys_by_format) = + endpoint_key_counts_by_format(&provider.provider_type, &endpoints, &keys); let now_unix_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .ok() @@ -51,15 +52,17 @@ pub(crate) async fn build_admin_provider_endpoints_payload( .skip(skip) .take(limit) .map(|endpoint| { + let endpoint_api_format = + aether_ai_formats::normalize_api_format_alias(&endpoint.api_format); build_admin_provider_endpoint_response( &endpoint, &provider.name, total_keys_by_format - .get(endpoint.api_format.as_str()) + .get(endpoint_api_format.as_str()) .copied() .unwrap_or(0), active_keys_by_format - .get(endpoint.api_format.as_str()) + .get(endpoint_api_format.as_str()) .copied() .unwrap_or(0), now_unix_secs, @@ -92,22 +95,27 @@ pub(crate) async fn build_admin_endpoint_payload( .await .ok() .unwrap_or_default(); - let (total_keys_by_format, active_keys_by_format) = endpoint_key_counts_by_format(&keys); + let (total_keys_by_format, active_keys_by_format) = endpoint_key_counts_by_format( + &provider.provider_type, + std::slice::from_ref(&endpoint), + &keys, + ); let now_unix_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .ok() .map(|duration| duration.as_secs()) .unwrap_or(0); + let endpoint_api_format = aether_ai_formats::normalize_api_format_alias(&endpoint.api_format); Some(build_admin_provider_endpoint_response( &endpoint, &provider.name, total_keys_by_format - .get(endpoint.api_format.as_str()) + .get(endpoint_api_format.as_str()) .copied() .unwrap_or(0), active_keys_by_format - .get(endpoint.api_format.as_str()) + .get(endpoint_api_format.as_str()) .copied() .unwrap_or(0), now_unix_secs, diff --git a/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/update.rs b/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/update.rs index ac304eee0..67d78f28f 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/update.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/update.rs @@ -147,18 +147,23 @@ pub(super) async fn maybe_handle( .list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id)) .await .unwrap_or_default(); - let (total_keys_by_format, active_keys_by_format) = endpoint_key_counts_by_format(&keys); + let (total_keys_by_format, active_keys_by_format) = endpoint_key_counts_by_format( + &provider.provider_type, + std::slice::from_ref(&updated), + &keys, + ); + let updated_api_format = aether_ai_formats::normalize_api_format_alias(&updated.api_format); Ok(Some( Json(build_admin_provider_endpoint_response( &updated, &provider.name, total_keys_by_format - .get(updated.api_format.as_str()) + .get(updated_api_format.as_str()) .copied() .unwrap_or(0), active_keys_by_format - .get(updated.api_format.as_str()) + .get(updated_api_format.as_str()) .copied() .unwrap_or(0), now_unix_secs, diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs index eac259b79..70cf099dd 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs @@ -148,6 +148,91 @@ async fn gateway_handles_admin_provider_endpoints_locally_with_trusted_admin_pri upstream_handle.abort(); } +#[tokio::test] +async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpoint() { + let upstream_hits = Arc::new(Mutex::new(0usize)); + let upstream_hits_clone = Arc::clone(&upstream_hits); + let upstream = Router::new().route( + "/api/admin/endpoints/providers/provider-codex/endpoints", + any(move |_request: Request| { + let upstream_hits_inner = Arc::clone(&upstream_hits_clone); + async move { + *upstream_hits_inner.lock().expect("mutex should lock") += 1; + (StatusCode::OK, Body::from("unexpected upstream hit")) + } + }), + ); + + let mut inherited_key = sample_key( + "key-codex-oauth", + "provider-codex", + "openai:responses", + "codex-token", + ); + inherited_key.auth_type = "oauth".to_string(); + inherited_key.api_formats = None; + + let mut provider = sample_provider("provider-codex", "codex", 10); + provider.provider_type = "codex".to_string(); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![ + sample_endpoint( + "endpoint-codex-responses", + "provider-codex", + "openai:responses", + "https://chatgpt.com/backend-api/codex", + ) + .with_timestamps(Some(1_711_000_000), Some(1_711_000_100)), + sample_endpoint( + "endpoint-codex-image", + "provider-codex", + "openai:image", + "https://chatgpt.com/backend-api/codex", + ) + .with_timestamps(Some(1_710_000_000), Some(1_710_000_100)), + ], + vec![inherited_key], + )); + + let (upstream_url, upstream_handle) = start_server(upstream).await; + let gateway = build_router_with_state( + AppState::new() + .expect("gateway should build") + .with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests( + provider_catalog_repository, + )), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .get(format!( + "{gateway_url}/api/admin/endpoints/providers/provider-codex/endpoints?skip=0&limit=50" + )) + .header(crate::constants::GATEWAY_HEADER, "rust-phase3b") + .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123") + .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin") + .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123") + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + let payload: serde_json::Value = response.json().await.expect("json body should parse"); + let items = payload.as_array().expect("payload should be an array"); + assert_eq!(items.len(), 2); + assert_eq!(items[0]["api_format"], "openai:responses"); + assert_eq!(items[0]["total_keys"], 1); + assert_eq!(items[0]["active_keys"], 1); + assert_eq!(items[1]["api_format"], "openai:image"); + assert_eq!(items[1]["total_keys"], 1); + assert_eq!(items[1]["active_keys"], 1); + assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); + + gateway_handle.abort(); + upstream_handle.abort(); +} + #[tokio::test] async fn gateway_returns_service_unavailable_for_admin_provider_endpoint_create_when_catalog_writer_unavailable( ) { diff --git a/crates/aether-admin/Cargo.toml b/crates/aether-admin/Cargo.toml index 3af976a48..c6d95e611 100644 --- a/crates/aether-admin/Cargo.toml +++ b/crates/aether-admin/Cargo.toml @@ -13,6 +13,7 @@ aether-contracts.workspace = true aether-data.workspace = true aether-data-contracts.workspace = true aether-provider-pool.workspace = true +aether-provider-transport.workspace = true axum.workspace = true base64.workspace = true chrono.workspace = true diff --git a/crates/aether-admin/src/provider/endpoints.rs b/crates/aether-admin/src/provider/endpoints.rs index 084c2c16b..97b5f6405 100644 --- a/crates/aether-admin/src/provider/endpoints.rs +++ b/crates/aether-admin/src/provider/endpoints.rs @@ -1,6 +1,7 @@ use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, }; +use aether_provider_transport::provider_types::fixed_provider_key_inherits_api_formats; use chrono::{TimeZone, Utc}; use serde_json::{json, Value}; use std::collections::BTreeMap; @@ -41,23 +42,63 @@ pub fn key_api_formats_without_entry( ) } +fn active_endpoint_api_formats(endpoints: &[StoredProviderCatalogEndpoint]) -> Vec { + let mut formats = Vec::new(); + for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) { + let api_format = aether_ai_formats::normalize_api_format_alias(&endpoint.api_format); + if !formats.iter().any(|existing| existing == &api_format) { + formats.push(api_format); + } + } + formats +} + +fn configured_key_api_formats(key: &StoredProviderCatalogKey) -> Vec { + let Some(formats) = key + .api_formats + .as_ref() + .and_then(serde_json::Value::as_array) + else { + return Vec::new(); + }; + let mut normalized = Vec::new(); + for api_format in formats.iter().filter_map(serde_json::Value::as_str) { + let api_format = aether_ai_formats::normalize_api_format_alias(api_format); + if !normalized.iter().any(|existing| existing == &api_format) { + normalized.push(api_format); + } + } + normalized +} + pub fn endpoint_key_counts_by_format( + provider_type: &str, + endpoints: &[StoredProviderCatalogEndpoint], keys: &[StoredProviderCatalogKey], ) -> (BTreeMap, BTreeMap) { let mut total = BTreeMap::new(); let mut active = BTreeMap::new(); + let inherited_api_formats = active_endpoint_api_formats(endpoints); + for key in keys { - let Some(formats) = key - .api_formats - .as_ref() - .and_then(serde_json::Value::as_array) - else { + if fixed_provider_key_inherits_api_formats( + provider_type, + &key.auth_type, + key.encrypted_auth_config.as_deref(), + ) { + for api_format in &inherited_api_formats { + *total.entry(api_format.clone()).or_insert(0) += 1; + if key.is_active { + *active.entry(api_format.clone()).or_insert(0) += 1; + } + } continue; - }; - for api_format in formats.iter().filter_map(serde_json::Value::as_str) { - *total.entry(api_format.to_string()).or_insert(0) += 1; + } + + for api_format in configured_key_api_formats(key) { + *total.entry(api_format.clone()).or_insert(0) += 1; if key.is_active { - *active.entry(api_format.to_string()).or_insert(0) += 1; + *active.entry(api_format).or_insert(0) += 1; } } }