From ccfc4cbddc3ee319d88259b50858125a9927f2d4 Mon Sep 17 00:00:00 2001 From: fawney19 Date: Wed, 27 May 2026 13:56:39 +0800 Subject: [PATCH 1/6] Cache provider catalog lookups --- apps/aether-gateway/src/data/state/catalog.rs | 84 +++- apps/aether-gateway/src/data/state/core.rs | 13 +- apps/aether-gateway/src/data/state/mod.rs | 1 + .../src/data/state/provider_catalog_cache.rs | 360 ++++++++++++++++++ .../maintenance/runtime/pool_quota_probe.rs | 26 +- .../src/orchestration/effects.rs | 5 + apps/aether-gateway/src/state/core.rs | 2 + .../src/repository/provider_catalog/types.rs | 2 + ...00_add_provider_key_name_hotspot_index.sql | 15 + ...00_add_provider_key_name_hotspot_index.sql | 2 + ...00_add_provider_key_name_hotspot_index.sql | 2 + .../schema/bootstrap/postgres/004_indexes.sql | 8 + .../src/lifecycle/bootstrap/postgres.rs | 2 +- .../src/lifecycle/migrate/tests.rs | 5 + .../src/repository/pool_scores/postgres.rs | 15 +- .../src/repository/provider_catalog/mysql.rs | 115 ++++-- 16 files changed, 582 insertions(+), 75 deletions(-) create mode 100644 apps/aether-gateway/src/data/state/provider_catalog_cache.rs create mode 100644 crates/aether-data/migrations/mysql/20260527000000_add_provider_key_name_hotspot_index.sql create mode 100644 crates/aether-data/migrations/postgres/20260527000000_add_provider_key_name_hotspot_index.sql create mode 100644 crates/aether-data/migrations/sqlite/20260527000000_add_provider_key_name_hotspot_index.sql diff --git a/apps/aether-gateway/src/data/state/catalog.rs b/apps/aether-gateway/src/data/state/catalog.rs index 0a9c0de45..6a52aa8ab 100644 --- a/apps/aether-gateway/src/data/state/catalog.rs +++ b/apps/aether-gateway/src/data/state/catalog.rs @@ -334,7 +334,7 @@ impl GatewayDataState { encrypted_auth_config: Option<&str>, expires_at_unix_secs: Option, ) -> Result { - match &self.provider_catalog_writer { + let updated = match &self.provider_catalog_writer { Some(repository) => { repository .update_key_oauth_credentials( @@ -346,17 +346,25 @@ impl GatewayDataState { .await } None => Ok(false), + }?; + if updated { + self.clear_provider_catalog_cache(); } + Ok(updated) } pub(crate) async fn create_provider_catalog_key( &self, key: &StoredProviderCatalogKey, ) -> Result, DataLayerError> { - match &self.provider_catalog_writer { + let created = match &self.provider_catalog_writer { Some(repository) => repository.create_key(key).await.map(Some), None => Ok(None), + }?; + if created.is_some() { + self.clear_provider_catalog_cache(); } + Ok(created) } pub(crate) async fn create_provider_catalog_provider( @@ -364,33 +372,45 @@ impl GatewayDataState { provider: &StoredProviderCatalogProvider, shift_existing_priorities_from: Option, ) -> Result, DataLayerError> { - match &self.provider_catalog_writer { + let created = match &self.provider_catalog_writer { Some(repository) => repository .create_provider(provider, shift_existing_priorities_from) .await .map(Some), None => Ok(None), + }?; + if created.is_some() { + self.clear_provider_catalog_cache(); } + Ok(created) } pub(crate) async fn update_provider_catalog_provider( &self, provider: &StoredProviderCatalogProvider, ) -> Result, DataLayerError> { - match &self.provider_catalog_writer { + let updated = match &self.provider_catalog_writer { Some(repository) => repository.update_provider(provider).await.map(Some), None => Ok(None), + }?; + if updated.is_some() { + self.clear_provider_catalog_cache(); } + Ok(updated) } pub(crate) async fn delete_provider_catalog_provider( &self, provider_id: &str, ) -> Result { - match &self.provider_catalog_writer { + let deleted = match &self.provider_catalog_writer { Some(repository) => repository.delete_provider(provider_id).await, None => Ok(false), + }?; + if deleted { + self.clear_provider_catalog_cache(); } + Ok(deleted) } pub(crate) async fn cleanup_deleted_provider_catalog_refs( @@ -399,54 +419,74 @@ impl GatewayDataState { endpoint_ids: &[String], key_ids: &[String], ) -> Result<(), DataLayerError> { - match &self.provider_catalog_writer { + let cleaned = match &self.provider_catalog_writer { Some(repository) => { repository .cleanup_deleted_provider_refs(provider_id, endpoint_ids, key_ids) .await } None => Ok(()), + }; + if !endpoint_ids.is_empty() || !key_ids.is_empty() { + self.clear_provider_catalog_cache(); } + cleaned } pub(crate) async fn create_provider_catalog_endpoint( &self, endpoint: &StoredProviderCatalogEndpoint, ) -> Result, DataLayerError> { - match &self.provider_catalog_writer { + let created = match &self.provider_catalog_writer { Some(repository) => repository.create_endpoint(endpoint).await.map(Some), None => Ok(None), + }?; + if created.is_some() { + self.clear_provider_catalog_cache(); } + Ok(created) } pub(crate) async fn update_provider_catalog_endpoint( &self, endpoint: &StoredProviderCatalogEndpoint, ) -> Result, DataLayerError> { - match &self.provider_catalog_writer { + let updated = match &self.provider_catalog_writer { Some(repository) => repository.update_endpoint(endpoint).await.map(Some), None => Ok(None), + }?; + if updated.is_some() { + self.clear_provider_catalog_cache(); } + Ok(updated) } pub(crate) async fn delete_provider_catalog_endpoint( &self, endpoint_id: &str, ) -> Result { - match &self.provider_catalog_writer { + let deleted = match &self.provider_catalog_writer { Some(repository) => repository.delete_endpoint(endpoint_id).await, None => Ok(false), + }?; + if deleted { + self.clear_provider_catalog_cache(); } + Ok(deleted) } pub(crate) async fn update_provider_catalog_key( &self, key: &StoredProviderCatalogKey, ) -> Result, DataLayerError> { - match &self.provider_catalog_writer { + let updated = match &self.provider_catalog_writer { Some(repository) => repository.update_key(key).await.map(Some), None => Ok(None), + }?; + if updated.is_some() { + self.clear_provider_catalog_cache(); } + Ok(updated) } pub(crate) async fn update_provider_catalog_key_upstream_metadata( @@ -455,34 +495,46 @@ impl GatewayDataState { upstream_metadata: Option<&serde_json::Value>, updated_at_unix_secs: Option, ) -> Result { - match &self.provider_catalog_writer { + let updated = match &self.provider_catalog_writer { Some(repository) => { repository .update_key_upstream_metadata(key_id, upstream_metadata, updated_at_unix_secs) .await } None => Ok(false), + }?; + if updated { + self.clear_provider_catalog_cache(); } + Ok(updated) } pub(crate) async fn delete_provider_catalog_key( &self, key_id: &str, ) -> Result { - match &self.provider_catalog_writer { + let deleted = match &self.provider_catalog_writer { Some(repository) => repository.delete_key(key_id).await, None => Ok(false), + }?; + if deleted { + self.clear_provider_catalog_cache(); } + Ok(deleted) } pub(crate) async fn clear_provider_catalog_key_oauth_invalid_marker( &self, key_id: &str, ) -> Result { - match &self.provider_catalog_writer { + let updated = match &self.provider_catalog_writer { Some(repository) => repository.clear_key_oauth_invalid_marker(key_id).await, None => Ok(false), + }?; + if updated { + self.clear_provider_catalog_cache(); } + Ok(updated) } pub(crate) async fn update_provider_catalog_key_health_state( @@ -492,7 +544,7 @@ impl GatewayDataState { health_by_format: Option<&serde_json::Value>, circuit_breaker_by_format: Option<&serde_json::Value>, ) -> Result { - match &self.provider_catalog_writer { + let updated = match &self.provider_catalog_writer { Some(repository) => { repository .update_key_health_state( @@ -504,6 +556,10 @@ impl GatewayDataState { .await } None => Ok(false), + }?; + if updated { + self.clear_provider_catalog_cache(); } + Ok(updated) } } diff --git a/apps/aether-gateway/src/data/state/core.rs b/apps/aether-gateway/src/data/state/core.rs index 3247d7acc..1ae4a56dd 100644 --- a/apps/aether-gateway/src/data/state/core.rs +++ b/apps/aether-gateway/src/data/state/core.rs @@ -1,5 +1,6 @@ use aether_data::{DataBackends, DataLayerError, DatabaseDriver}; use aether_data_contracts::repository::candidate_selection::MinimalCandidateSelectionReadRepository; +use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository; use aether_runtime_state::RuntimeQueueStore; use std::sync::Arc; @@ -99,7 +100,11 @@ impl GatewayDataState { let request_candidate_reader = backends.read().request_candidates(); let request_candidate_writer = backends.write().request_candidates(); let gemini_file_mapping_writer = backends.write().gemini_file_mappings(); - let provider_catalog_reader = backends.read().provider_catalog(); + let provider_catalog_reader = backends.read().provider_catalog().map(|repository| { + Arc::new( + super::provider_catalog_cache::CachedProviderCatalogReadRepository::new(repository), + ) as Arc + }); let provider_catalog_writer = backends.write().provider_catalog(); let pool_score_reader = backends.read().pool_scores(); let pool_score_writer = backends.write().pool_scores(); @@ -276,6 +281,12 @@ impl GatewayDataState { } } + pub(crate) fn clear_provider_catalog_cache(&self) { + if let Some(repository) = &self.provider_catalog_reader { + repository.clear_local_cache(); + } + } + pub(crate) fn has_request_candidate_reader(&self) -> bool { self.request_candidate_reader.is_some() } diff --git a/apps/aether-gateway/src/data/state/mod.rs b/apps/aether-gateway/src/data/state/mod.rs index 40da5c651..e4cfa864b 100644 --- a/apps/aether-gateway/src/data/state/mod.rs +++ b/apps/aether-gateway/src/data/state/mod.rs @@ -324,6 +324,7 @@ mod core; mod integrations; mod models; mod pool_scores; +mod provider_catalog_cache; mod referrals; mod routing_profiles; mod runtime; diff --git a/apps/aether-gateway/src/data/state/provider_catalog_cache.rs b/apps/aether-gateway/src/data/state/provider_catalog_cache.rs new file mode 100644 index 000000000..13c5d3423 --- /dev/null +++ b/apps/aether-gateway/src/data/state/provider_catalog_cache.rs @@ -0,0 +1,360 @@ +use std::collections::HashMap; +use std::future::Future; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use aether_cache::ExpiringMap; +use aether_data::DataLayerError; +use aether_data_contracts::repository::provider_catalog::{ + ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, + StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, +}; +use async_trait::async_trait; +use tokio::sync::Notify; + +const PROVIDER_CATALOG_CACHE_TTL: Duration = Duration::from_secs(5); +const PROVIDER_CATALOG_CACHE_MAX_ENTRIES: usize = 1024; + +pub(super) struct CachedProviderCatalogReadRepository { + inner: Arc, + entries: ExpiringMap, + inflight: Mutex>, + inflight_notify: Notify, + next_inflight_token: AtomicU64, + epoch: AtomicU64, +} + +impl CachedProviderCatalogReadRepository { + pub(super) fn new(inner: Arc) -> Self { + Self { + inner, + entries: ExpiringMap::new(), + inflight: Mutex::new(HashMap::new()), + inflight_notify: Notify::new(), + next_inflight_token: AtomicU64::new(1), + epoch: AtomicU64::new(0), + } + } + + async fn get_or_load( + &self, + key: ProviderCatalogCacheKey, + load: F, + ) -> Result + where + F: Fn() -> Fut, + Fut: Future>, + { + if let Some(value) = self.entries.get_fresh(&key, PROVIDER_CATALOG_CACHE_TTL) { + return Ok(value); + } + + loop { + let notified = self.inflight_notify.notified(); + match self.register_inflight(&key) { + InflightRegistration::Bypass => return load().await, + InflightRegistration::Follower => { + notified.await; + if let Some(value) = self.entries.get_fresh(&key, PROVIDER_CATALOG_CACHE_TTL) { + return Ok(value); + } + } + InflightRegistration::Leader(token) => { + let mut guard = InflightGuard::new(self, key.clone(), token); + let load_epoch = self.epoch.load(Ordering::Acquire); + let result = load().await; + if let Ok(value) = &result { + if load_epoch == self.epoch.load(Ordering::Acquire) { + self.entries.insert( + key.clone(), + value.clone(), + PROVIDER_CATALOG_CACHE_TTL, + PROVIDER_CATALOG_CACHE_MAX_ENTRIES, + ); + } + } + guard.finish(); + return result; + } + } + } + } + + fn register_inflight(&self, key: &ProviderCatalogCacheKey) -> InflightRegistration { + match self.inflight.lock() { + Ok(mut inflight) => { + if inflight.contains_key(key) { + return InflightRegistration::Follower; + } + let token = self.next_inflight_token.fetch_add(1, Ordering::AcqRel); + inflight.insert(key.clone(), token); + InflightRegistration::Leader(token) + } + Err(_) => InflightRegistration::Bypass, + } + } + + fn finish_inflight(&self, key: &ProviderCatalogCacheKey, token: u64) { + let mut removed = false; + if let Ok(mut inflight) = self.inflight.lock() { + if inflight.get(key).copied() == Some(token) { + inflight.remove(key); + removed = true; + } + } + if removed { + self.inflight_notify.notify_waiters(); + } + } + + fn clear(&self) { + self.epoch.fetch_add(1, Ordering::AcqRel); + self.entries.clear(); + let mut cleared_inflight = false; + if let Ok(mut inflight) = self.inflight.lock() { + cleared_inflight = !inflight.is_empty(); + inflight.clear(); + } + if cleared_inflight { + self.inflight_notify.notify_waiters(); + } + } +} + +#[async_trait] +impl ProviderCatalogReadRepository for CachedProviderCatalogReadRepository { + fn clear_local_cache(&self) { + self.clear(); + self.inner.clear_local_cache(); + } + + async fn list_providers( + &self, + active_only: bool, + ) -> Result, DataLayerError> { + match self + .get_or_load( + ProviderCatalogCacheKey::Providers { active_only }, + || async move { + self.inner + .list_providers(active_only) + .await + .map(ProviderCatalogCacheValue::Providers) + }, + ) + .await? + { + ProviderCatalogCacheValue::Providers(items) => Ok(items), + _ => Ok(Vec::new()), + } + } + + async fn list_providers_by_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + let key = ProviderCatalogCacheKey::ProvidersByIds(normalize_ids(provider_ids)); + match self + .get_or_load(key, || async move { + self.inner + .list_providers_by_ids(provider_ids) + .await + .map(ProviderCatalogCacheValue::Providers) + }) + .await? + { + ProviderCatalogCacheValue::Providers(items) => Ok(items), + _ => Ok(Vec::new()), + } + } + + async fn list_endpoints_by_ids( + &self, + endpoint_ids: &[String], + ) -> Result, DataLayerError> { + self.inner.list_endpoints_by_ids(endpoint_ids).await + } + + async fn list_endpoints_by_provider_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + let key = ProviderCatalogCacheKey::EndpointsByProviderIds(normalize_ids(provider_ids)); + match self + .get_or_load(key, || async move { + self.inner + .list_endpoints_by_provider_ids(provider_ids) + .await + .map(ProviderCatalogCacheValue::Endpoints) + }) + .await? + { + ProviderCatalogCacheValue::Endpoints(items) => Ok(items), + _ => Ok(Vec::new()), + } + } + + async fn list_keys_by_ids( + &self, + key_ids: &[String], + ) -> Result, DataLayerError> { + self.inner.list_keys_by_ids(key_ids).await + } + + async fn list_keys_by_provider_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + let key = ProviderCatalogCacheKey::KeysByProviderIds(normalize_ids(provider_ids)); + match self + .get_or_load(key, || async move { + self.inner + .list_keys_by_provider_ids(provider_ids) + .await + .map(ProviderCatalogCacheValue::Keys) + }) + .await? + { + ProviderCatalogCacheValue::Keys(items) => Ok(items), + _ => Ok(Vec::new()), + } + } + + async fn list_key_summaries_by_provider_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + let key = ProviderCatalogCacheKey::KeySummariesByProviderIds(normalize_ids(provider_ids)); + match self + .get_or_load(key, || async move { + self.inner + .list_key_summaries_by_provider_ids(provider_ids) + .await + .map(ProviderCatalogCacheValue::Keys) + }) + .await? + { + ProviderCatalogCacheValue::Keys(items) => Ok(items), + _ => Ok(Vec::new()), + } + } + + async fn list_key_maintenance_summaries_by_provider_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + let key = ProviderCatalogCacheKey::KeyMaintenanceSummariesByProviderIds(normalize_ids( + provider_ids, + )); + match self + .get_or_load(key, || async move { + self.inner + .list_key_maintenance_summaries_by_provider_ids(provider_ids) + .await + .map(ProviderCatalogCacheValue::KeyMaintenanceSummaries) + }) + .await? + { + ProviderCatalogCacheValue::KeyMaintenanceSummaries(items) => Ok(items), + _ => Ok(Vec::new()), + } + } + + async fn list_keys_page( + &self, + query: &ProviderCatalogKeyListQuery, + ) -> Result { + self.inner.list_keys_page(query).await + } + + async fn list_key_stats_by_provider_ids( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + let key = ProviderCatalogCacheKey::KeyStatsByProviderIds(normalize_ids(provider_ids)); + match self + .get_or_load(key, || async move { + self.inner + .list_key_stats_by_provider_ids(provider_ids) + .await + .map(ProviderCatalogCacheValue::KeyStats) + }) + .await? + { + ProviderCatalogCacheValue::KeyStats(items) => Ok(items), + _ => Ok(Vec::new()), + } + } +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +enum ProviderCatalogCacheKey { + Providers { active_only: bool }, + ProvidersByIds(Vec), + EndpointsByProviderIds(Vec), + KeysByProviderIds(Vec), + KeySummariesByProviderIds(Vec), + KeyMaintenanceSummariesByProviderIds(Vec), + KeyStatsByProviderIds(Vec), +} + +#[derive(Clone)] +enum ProviderCatalogCacheValue { + Providers(Vec), + Endpoints(Vec), + Keys(Vec), + KeyMaintenanceSummaries(Vec), + KeyStats(Vec), +} + +enum InflightRegistration { + Leader(u64), + Follower, + Bypass, +} + +struct InflightGuard<'a> { + cache: &'a CachedProviderCatalogReadRepository, + key: Option, + token: u64, +} + +impl<'a> InflightGuard<'a> { + fn new( + cache: &'a CachedProviderCatalogReadRepository, + key: ProviderCatalogCacheKey, + token: u64, + ) -> Self { + Self { + cache, + key: Some(key), + token, + } + } + + fn finish(&mut self) { + if let Some(key) = self.key.take() { + self.cache.finish_inflight(&key, self.token); + } + } +} + +impl Drop for InflightGuard<'_> { + fn drop(&mut self) { + self.finish(); + } +} + +fn normalize_ids(ids: &[String]) -> Vec { + let mut normalized = ids + .iter() + .map(|id| id.trim()) + .filter(|id| !id.is_empty()) + .map(ToOwned::to_owned) + .collect::>(); + normalized.sort(); + normalized.dedup(); + normalized +} diff --git a/apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs b/apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs index ea9c82a75..694e55b99 100644 --- a/apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs +++ b/apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs @@ -2,7 +2,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use aether_data_contracts::repository::pool_scores::{ - GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, PoolMemberHardState, + ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus, StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL, }; @@ -22,7 +22,6 @@ use crate::admin_api::{ }; use crate::{AppState, GatewayError}; -use crate::ai_serving::provider_key_pool_score_id; use crate::ai_serving::provider_key_pool_score_scope; use crate::handlers::shared::provider_pool::{ admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig, @@ -38,6 +37,7 @@ const POOL_QUOTA_PROBE_REDIS_PREFIX: &str = "ap:quota_probe:last"; const POOL_QUOTA_PROBE_DEFAULT_SCAN_INTERVAL_SECONDS: u64 = 60; const POOL_QUOTA_PROBE_MIN_SCAN_INTERVAL_SECONDS: u64 = 15; const POOL_QUOTA_PROBE_DEFAULT_MAX_KEYS_PER_PROVIDER: usize = 50; +const POOL_QUOTA_PROBE_PROVIDER_SCORE_READ_LIMIT: usize = 100_000; const POOL_QUOTA_PROBE_DEFAULT_GLOBAL_CONCURRENCY: usize = 16; const POOL_QUOTA_PROBE_PROVIDER_LOCK_TTL_MS: u64 = 30_000; const POOL_QUOTA_PROBE_BURST_TRIGGER_LOCK_TTL_MS: u64 = 30_000; @@ -386,21 +386,25 @@ async fn load_provider_key_account_scores( return BTreeMap::new(); } let scope = provider_key_pool_score_scope(); - let score_ids = key_ids - .iter() - .map(|key_id| { - let identity = - PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.clone()); - provider_key_pool_score_id(&identity, &scope) - }) - .collect::>(); + let key_ids = key_ids.iter().map(String::as_str).collect::>(); match state .data - .get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery { ids: score_ids }) + .list_pool_member_scores(&ListPoolMemberScoresQuery { + pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(), + pool_id: provider_id.to_string(), + capability: Some(scope.capability), + scope_kind: Some(scope.scope_kind), + scope_id: scope.scope_id, + hard_states: Vec::new(), + probe_statuses: None, + offset: 0, + limit: POOL_QUOTA_PROBE_PROVIDER_SCORE_READ_LIMIT, + }) .await { Ok(scores) => scores .into_iter() + .filter(|score| key_ids.contains(score.member_id.as_str())) .map(|score| (score.member_id.clone(), score)) .collect(), Err(err) => { diff --git a/apps/aether-gateway/src/orchestration/effects.rs b/apps/aether-gateway/src/orchestration/effects.rs index d52f11616..a7598d51c 100644 --- a/apps/aether-gateway/src/orchestration/effects.rs +++ b/apps/aether-gateway/src/orchestration/effects.rs @@ -645,6 +645,11 @@ async fn record_health_success_effect( .circuit_breaker_by_format .as_ref() .and_then(|current| project_local_key_circuit_closed(Some(current), api_format)); + if current_key.health_by_format.as_ref() == Some(&health_by_format) + && circuit_breaker_by_format.as_ref() == current_key.circuit_breaker_by_format.as_ref() + { + return; + } let circuit_breaker_update = circuit_breaker_by_format .as_ref() .or(current_key.circuit_breaker_by_format.as_ref()); diff --git a/apps/aether-gateway/src/state/core.rs b/apps/aether-gateway/src/state/core.rs index d682b31ee..c69dc9fd2 100644 --- a/apps/aether-gateway/src/state/core.rs +++ b/apps/aether-gateway/src/state/core.rs @@ -573,12 +573,14 @@ impl AppState { pub(crate) fn invalidate_provider_routing_caches(&self) { self.data.clear_minimal_candidate_selection_cache(); + self.data.clear_provider_catalog_cache(); self.clear_provider_transport_snapshot_cache(); self.invalidate_scheduler_affinity_cache(); } pub(crate) fn invalidate_provider_health_routing_caches(&self) { self.data.clear_minimal_candidate_selection_cache(); + self.data.clear_provider_catalog_cache(); self.clear_provider_transport_snapshot_cache(); } diff --git a/crates/aether-data-contracts/src/repository/provider_catalog/types.rs b/crates/aether-data-contracts/src/repository/provider_catalog/types.rs index 1dac60764..035a916f9 100644 --- a/crates/aether-data-contracts/src/repository/provider_catalog/types.rs +++ b/crates/aether-data-contracts/src/repository/provider_catalog/types.rs @@ -580,6 +580,8 @@ impl StoredProviderCatalogKeyStats { #[async_trait] pub trait ProviderCatalogReadRepository: Send + Sync { + fn clear_local_cache(&self) {} + async fn list_providers( &self, active_only: bool, diff --git a/crates/aether-data/migrations/mysql/20260527000000_add_provider_key_name_hotspot_index.sql b/crates/aether-data/migrations/mysql/20260527000000_add_provider_key_name_hotspot_index.sql new file mode 100644 index 000000000..b60243da1 --- /dev/null +++ b/crates/aether-data/migrations/mysql/20260527000000_add_provider_key_name_hotspot_index.sql @@ -0,0 +1,15 @@ +SET @aether_provider_key_name_index_sql := IF( + ( + SELECT COUNT(*) + FROM information_schema.statistics + WHERE table_schema = DATABASE() + AND table_name = 'provider_api_keys' + AND index_name = 'idx_provider_api_keys_provider_name_id' + ) = 0, + 'CREATE INDEX idx_provider_api_keys_provider_name_id ON provider_api_keys (provider_id, name, id)', + 'DO 0' +); + +PREPARE aether_provider_key_name_index_stmt FROM @aether_provider_key_name_index_sql; +EXECUTE aether_provider_key_name_index_stmt; +DEALLOCATE PREPARE aether_provider_key_name_index_stmt; diff --git a/crates/aether-data/migrations/postgres/20260527000000_add_provider_key_name_hotspot_index.sql b/crates/aether-data/migrations/postgres/20260527000000_add_provider_key_name_hotspot_index.sql new file mode 100644 index 000000000..5a2fcab99 --- /dev/null +++ b/crates/aether-data/migrations/postgres/20260527000000_add_provider_key_name_hotspot_index.sql @@ -0,0 +1,2 @@ +CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id + ON public.provider_api_keys USING btree (provider_id, name, id); diff --git a/crates/aether-data/migrations/sqlite/20260527000000_add_provider_key_name_hotspot_index.sql b/crates/aether-data/migrations/sqlite/20260527000000_add_provider_key_name_hotspot_index.sql new file mode 100644 index 000000000..49f4d606d --- /dev/null +++ b/crates/aether-data/migrations/sqlite/20260527000000_add_provider_key_name_hotspot_index.sql @@ -0,0 +1,2 @@ +CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id + ON provider_api_keys (provider_id, name, id); diff --git a/crates/aether-data/schema/bootstrap/postgres/004_indexes.sql b/crates/aether-data/schema/bootstrap/postgres/004_indexes.sql index 955f73552..89fab1fed 100644 --- a/crates/aether-data/schema/bootstrap/postgres/004_indexes.sql +++ b/crates/aether-data/schema/bootstrap/postgres/004_indexes.sql @@ -197,6 +197,14 @@ CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_default_sort ON public +-- +-- Name: idx_provider_api_keys_provider_name_id; Type: INDEX; Schema: public; Owner: - +-- + +CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_name_id ON public.provider_api_keys USING btree (provider_id, name, id); + + + -- -- Name: idx_provider_api_keys_provider_active_priority_id; Type: INDEX; Schema: public; Owner: - -- diff --git a/crates/aether-data/src/lifecycle/bootstrap/postgres.rs b/crates/aether-data/src/lifecycle/bootstrap/postgres.rs index 9e7891c07..f0d8101af 100644 --- a/crates/aether-data/src/lifecycle/bootstrap/postgres.rs +++ b/crates/aether-data/src/lifecycle/bootstrap/postgres.rs @@ -7,7 +7,7 @@ use tracing::info; // Generated by build.rs from schema/bootstrap/postgres. pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str = include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql")); -pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260524000000; +pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260527000000; const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#" SELECT COUNT(*)::BIGINT diff --git a/crates/aether-data/src/lifecycle/migrate/tests.rs b/crates/aether-data/src/lifecycle/migrate/tests.rs index 3ea3b22ef..5f20372ff 100644 --- a/crates/aether-data/src/lifecycle/migrate/tests.rs +++ b/crates/aether-data/src/lifecycle/migrate/tests.rs @@ -314,6 +314,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() { 20260520010000, 20260522000000, 20260524000000, + 20260527000000, ] ); } @@ -390,6 +391,7 @@ fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role() assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("ix_usage_counter_deltas_unprocessed")); assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_entitlement_usage_entitlement_date")); assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_default_sort")); + assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_name_id")); assert!( EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_provider_api_keys_provider_active_priority_id") ); @@ -668,6 +670,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() { 20260520000000, 20260520010000, 20260524000000, + 20260527000000, ] ); assert_eq!( @@ -692,6 +695,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() { 20260520000000, 20260520010000, 20260524000000, + 20260527000000, ] ); } @@ -1217,6 +1221,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() { 20260520010000, 20260522000000, 20260524000000, + 20260527000000, ] ); } diff --git a/crates/aether-data/src/repository/pool_scores/postgres.rs b/crates/aether-data/src/repository/pool_scores/postgres.rs index 8ab176bea..076819ddd 100644 --- a/crates/aether-data/src/repository/pool_scores/postgres.rs +++ b/crates/aether-data/src/repository/pool_scores/postgres.rs @@ -341,14 +341,13 @@ impl PoolScoreReadRepository for PostgresPoolMemberScoreRepository { if query.ids.is_empty() { return Ok(Vec::new()); } - let mut builder = QueryBuilder::::new(SCORE_COLUMNS); - let mut where_clause = WhereClause::new(); - push_in(&mut builder, &mut where_clause, "id", &query.ids); - let rows = builder - .build() - .fetch_all(&self.pool) - .await - .map_postgres_err()?; + let rows = sqlx::query(&format!( + "{SCORE_COLUMNS} WHERE id = ANY($1) ORDER BY id ASC" + )) + .bind(&query.ids) + .fetch_all(&self.pool) + .await + .map_postgres_err()?; rows.iter().map(map_score_row).collect() } } diff --git a/crates/aether-data/src/repository/provider_catalog/mysql.rs b/crates/aether-data/src/repository/provider_catalog/mysql.rs index e372841bb..5e05d36d3 100644 --- a/crates/aether-data/src/repository/provider_catalog/mysql.rs +++ b/crates/aether-data/src/repository/provider_catalog/mysql.rs @@ -1,5 +1,5 @@ use async_trait::async_trait; -use sqlx::{mysql::MySqlRow, Row}; +use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row}; use super::{ InMemoryProviderCatalogReadRepository, ProviderCatalogKeyListQuery, @@ -16,6 +16,33 @@ pub struct MysqlProviderCatalogReadRepository { pool: MysqlPool, } +const KEY_SELECT_SQL: &str = r#" +SELECT + id, provider_id, name, auth_type, capabilities, is_active, api_formats, + auth_type_by_format, allow_auth_channel_mismatch_formats, + COALESCE(api_key, encrypted_key) AS api_key, + auth_config, note, internal_priority, rate_multipliers, + global_priority_by_format, allowed_models, + expires_at AS expires_at_unix_secs, + cache_ttl_minutes, max_probe_interval_minutes, proxy, fingerprint, + rpm_limit, concurrent_limit, learned_rpm_limit, concurrent_429_count, + rpm_429_count, last_429_at AS last_429_at_unix_secs, last_429_type, + adjustment_history, utilization_samples, + last_probe_increase_at AS last_probe_increase_at_unix_secs, + last_rpm_peak, request_count, total_tokens, total_cost_usd, + success_count, error_count, total_response_time_ms, + last_used_at AS last_used_at_unix_secs, auto_fetch_models, + last_models_fetch_at AS last_models_fetch_at_unix_secs, + last_models_fetch_error, locked_models, model_include_patterns, + model_exclude_patterns, upstream_metadata, + oauth_invalid_at AS oauth_invalid_at_unix_secs, + oauth_invalid_reason, status_snapshot, + created_at AS created_at_unix_ms, + updated_at AS updated_at_unix_secs, + health_by_format, circuit_breaker_by_format +FROM provider_api_keys +"#; + impl MysqlProviderCatalogReadRepository { pub fn new(pool: MysqlPool) -> Self { Self { pool } @@ -71,37 +98,26 @@ WHERE api_format IS NOT NULL } async fn load_keys(&self) -> Result, DataLayerError> { - let rows = sqlx::query( - r#" -SELECT - id, provider_id, name, auth_type, capabilities, is_active, api_formats, - auth_type_by_format, allow_auth_channel_mismatch_formats, - COALESCE(api_key, encrypted_key) AS api_key, - auth_config, note, internal_priority, rate_multipliers, - global_priority_by_format, allowed_models, - expires_at AS expires_at_unix_secs, - cache_ttl_minutes, max_probe_interval_minutes, proxy, fingerprint, - rpm_limit, concurrent_limit, learned_rpm_limit, concurrent_429_count, - rpm_429_count, last_429_at AS last_429_at_unix_secs, last_429_type, - adjustment_history, utilization_samples, - last_probe_increase_at AS last_probe_increase_at_unix_secs, - last_rpm_peak, request_count, total_tokens, total_cost_usd, - success_count, error_count, total_response_time_ms, - last_used_at AS last_used_at_unix_secs, auto_fetch_models, - last_models_fetch_at AS last_models_fetch_at_unix_secs, - last_models_fetch_error, locked_models, model_include_patterns, - model_exclude_patterns, upstream_metadata, - oauth_invalid_at AS oauth_invalid_at_unix_secs, - oauth_invalid_reason, status_snapshot, - created_at AS created_at_unix_ms, - updated_at AS updated_at_unix_secs, - health_by_format, circuit_breaker_by_format -FROM provider_api_keys -"#, - ) - .fetch_all(&self.pool) - .await - .map_sql_err()?; + let rows = sqlx::query(KEY_SELECT_SQL) + .fetch_all(&self.pool) + .await + .map_sql_err()?; + rows.iter().map(map_key_row).collect() + } + + async fn list_keys_by_provider_ids_direct( + &self, + provider_ids: &[String], + ) -> Result, DataLayerError> { + if provider_ids.is_empty() { + return Ok(Vec::new()); + } + + let rows = build_list_keys_by_provider_ids_query(provider_ids) + .build() + .fetch_all(&self.pool) + .await + .map_sql_err()?; rows.iter().map(map_key_row).collect() } @@ -912,20 +928,14 @@ impl ProviderCatalogReadRepository for MysqlProviderCatalogReadRepository { &self, provider_ids: &[String], ) -> Result, DataLayerError> { - self.load_memory() - .await? - .list_keys_by_provider_ids(provider_ids) - .await + self.list_keys_by_provider_ids_direct(provider_ids).await } async fn list_key_summaries_by_provider_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { - self.load_memory() - .await? - .list_key_summaries_by_provider_ids(provider_ids) - .await + self.list_keys_by_provider_ids_direct(provider_ids).await } async fn list_key_maintenance_summaries_by_provider_ids( @@ -1162,6 +1172,19 @@ fn optional_json_to_string( optional_json_ref_to_string(value.as_ref(), field_name) } +fn build_list_keys_by_provider_ids_query(provider_ids: &[String]) -> QueryBuilder<'_, MySql> { + let mut builder = QueryBuilder::::new(KEY_SELECT_SQL); + builder.push("WHERE provider_id IN ("); + { + let mut separated = builder.separated(", "); + for provider_id in provider_ids { + separated.push_bind(provider_id.clone()); + } + } + builder.push(") ORDER BY provider_id ASC, name ASC, id ASC"); + builder +} + fn key_insert_sql() -> &'static str { r#" INSERT INTO provider_api_keys ( @@ -1568,6 +1591,7 @@ mod tests { StoredProviderCatalogProvider, }; use serde_json::json; + use sqlx::Execute; #[tokio::test] async fn repository_builds_from_lazy_pool() { @@ -1580,6 +1604,17 @@ mod tests { let _repository = MysqlProviderCatalogReadRepository::new(pool); } + #[test] + fn list_keys_by_provider_ids_query_targets_index_aligned_ordering() { + let provider_ids = vec!["provider-a".to_string(), "provider-b".to_string()]; + let mut builder = super::build_list_keys_by_provider_ids_query(&provider_ids); + let query = builder.build(); + let sql = query.sql(); + + assert!(sql.contains("WHERE provider_id IN (")); + assert!(sql.contains("ORDER BY provider_id ASC, name ASC, id ASC")); + } + #[tokio::test] async fn mysql_provider_catalog_repository_round_trips_when_url_is_set() { let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL") From b9ac97ebc3d395405f538f4fef79d8f2789455c2 Mon Sep 17 00:00:00 2001 From: fawney19 Date: Wed, 27 May 2026 14:24:57 +0800 Subject: [PATCH 2/6] Simplify request detail cost overview --- .../usage/components/RequestDetailDrawer.vue | 118 ++++-------------- 1 file changed, 23 insertions(+), 95 deletions(-) diff --git a/frontend/src/features/usage/components/RequestDetailDrawer.vue b/frontend/src/features/usage/components/RequestDetailDrawer.vue index 938671b36..2d5048785 100644 --- a/frontend/src/features/usage/components/RequestDetailDrawer.vue +++ b/frontend/src/features/usage/components/RequestDetailDrawer.vue @@ -203,83 +203,33 @@
- -
-
- 总费用 - +
+ {{ priceSourceLabel }} + | + + 总费用 + ${{ ((typeof detail.cost === 'object' ? detail.cost?.total : detail.cost) || detail.total_cost || 0).toFixed(6) }} -
-
- 响应时间 - {{ detail.response_time_ms ? formatResponseTime(detail.response_time_ms).value : 'N/A' }} - {{ detail.response_time_ms ? formatResponseTime(detail.response_time_ms).unit : '' }} -
- -
- -
-
-
- 首字时间 -
-
- {{ formatDurationMs(detail.first_byte_time_ms) }} -
-
-
-
- 生成耗时 -
-
- {{ formatDurationMs(detailGenerationTimeMs) }} -
-
-
-
- 输出 Tokens -
-
- {{ formatNumber(detailOutputTokens) }} -
-
+ + | + + 耗时 + + {{ formatDurationMs(detail.first_byte_time_ms) }} / {{ formatDurationMs(detail.response_time_ms) }} + + + | + + 输出速度 + {{ formatOutputRateValue(detailOutputRate) }}tps +
- -
- {{ priceSourceLabel }} - | - - 总费用 = Token费用 ${{ tokenCostTotal.toFixed(6) }} - - - - -
- - +
Token 计费 + ${{ tokenCostTotal.toFixed(6) }} (输入 {{ formatNumber(displayInputTokens) }} + 缓存创建 {{ cacheCreationSummaryText }} + 缓存读取 {{ formatNumber(detail.cache_read_input_tokens || 0) }}) -
+
按次计费 + ${{ perRequestCost.toFixed(6) }}
{ return detail.value.tokens?.output ?? detail.value.output_tokens ?? 0 }) -const detailGenerationTimeMs = computed(() => { - if (!detail.value) return null - return getGenerationTimeMs({ - response_time_ms: detail.value.response_time_ms, - first_byte_time_ms: detail.value.first_byte_time_ms, - is_stream: detail.value.is_stream, - upstream_is_stream: detail.value.upstream_is_stream, - }) -}) - const detailOutputRate = computed(() => { if (!detail.value) return null return getDisplayOutputRate({ @@ -1129,11 +1070,6 @@ const detailOutputRate = computed(() => { }) }) -const hasDetailPerformanceBreakdown = computed(() => { - if (!detail.value) return false - return detail.value.first_byte_time_ms != null || detailGenerationTimeMs.value != null || detailOutputTokens.value > 0 -}) - // 监听标签页切换 watch(activeTab, (newTab) => { if (!['request-headers', 'response-headers'].includes(newTab) && viewMode.value === 'compare') { @@ -2378,14 +2314,6 @@ function formatPixels(value: number): string { return `${formatCompactNumber(value)} px` } -// 格式化响应时间，自动选择合适的单位 -function formatResponseTime(ms: number): { value: string; unit: string } { - if (ms >= 1_000) { - return { value: (ms / 1_000).toFixed(2), unit: 's' } - } - return { value: ms.toString(), unit: 'ms' } -} - // 格式化价格，修复浮点数精度问题 function formatPrice(price: number): string { // 处理浮点数精度问题，最多保留4位小数，去掉尾部的0 From d488b1a6802b69390b6d1560ed4b27c4eee4fb1d Mon Sep 17 00:00:00 2001 From: fawney19 Date: Wed, 27 May 2026 15:06:51 +0800 Subject: [PATCH 3/6] fix auth refresh request body --- .../src/api/__tests__/auth-turnstile.spec.ts | 14 +++++++- frontend/src/api/__tests__/client.spec.ts | 34 +++++++++++++++++++ frontend/src/api/auth.ts | 2 +- frontend/src/api/client.ts | 2 +- 4 files changed, 49 insertions(+), 3 deletions(-) diff --git a/frontend/src/api/__tests__/auth-turnstile.spec.ts b/frontend/src/api/__tests__/auth-turnstile.spec.ts index 7744e71c6..1f78005dd 100644 --- a/frontend/src/api/__tests__/auth-turnstile.spec.ts +++ b/frontend/src/api/__tests__/auth-turnstile.spec.ts @@ -1,12 +1,14 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' -const { postMock } = vi.hoisted(() => ({ +const { postMock, setTokenMock } = vi.hoisted(() => ({ postMock: vi.fn(), + setTokenMock: vi.fn(), })) vi.mock('@/api/client', () => ({ default: { post: postMock, + setToken: setTokenMock, }, })) @@ -15,6 +17,7 @@ import { authApi } from '@/api/auth' describe('authApi turnstile payloads', () => { beforeEach(() => { postMock.mockReset() + setTokenMock.mockReset() postMock.mockResolvedValue({ data: {} }) }) @@ -42,4 +45,13 @@ describe('authApi turnstile payloads', () => { turnstile_token: 'turnstile-token', }) }) + + it('refreshes auth token without a request body', async () => { + postMock.mockResolvedValue({ data: { access_token: 'new-access-token' } }) + + await authApi.refreshToken() + + expect(postMock).toHaveBeenCalledWith('/api/auth/refresh') + expect(setTokenMock).toHaveBeenCalledWith('new-access-token') + }) }) diff --git a/frontend/src/api/__tests__/client.spec.ts b/frontend/src/api/__tests__/client.spec.ts index 93e33508c..c9c8657ac 100644 --- a/frontend/src/api/__tests__/client.spec.ts +++ b/frontend/src/api/__tests__/client.spec.ts @@ -1,7 +1,12 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import type { AxiosAdapter, AxiosInstance, InternalAxiosRequestConfig } from 'axios' import apiClient, { AUTH_STATE_CHANGE_EVENT } from '@/api/client' +type TestableApiClient = typeof apiClient & { + client: AxiosInstance +} + describe('apiClient auth state change event', () => { beforeEach(() => { localStorage.clear() @@ -28,4 +33,33 @@ describe('apiClient auth state change event', () => { window.removeEventListener(AUTH_STATE_CHANGE_EVENT, handler as EventListener) }) + + it('sends auth refresh without a request body', async () => { + const rawClient = apiClient as TestableApiClient + const previousAdapter = rawClient.client.defaults.adapter + const requests: InternalAxiosRequestConfig[] = [] + + rawClient.client.defaults.adapter = (async (config: InternalAxiosRequestConfig) => { + requests.push(config) + return { + data: { access_token: 'new-access-token' }, + status: 200, + statusText: 'OK', + headers: {}, + config, + } + }) as AxiosAdapter + + try { + const response = await apiClient.refreshToken() + + expect(response.data.access_token).toBe('new-access-token') + expect(requests).toHaveLength(1) + expect(requests[0].url).toBe('/api/auth/refresh') + expect(requests[0].method).toBe('post') + expect(requests[0].data).toBeUndefined() + } finally { + rawClient.client.defaults.adapter = previousAdapter + } + }) }) diff --git a/frontend/src/api/auth.ts b/frontend/src/api/auth.ts index 3eb111a5a..b12e3e771 100644 --- a/frontend/src/api/auth.ts +++ b/frontend/src/api/auth.ts @@ -164,7 +164,7 @@ export const authApi = { }, async refreshToken(): Promise { - const response = await apiClient.post('/api/auth/refresh', {}) + const response = await apiClient.post('/api/auth/refresh') apiClient.setToken(response.data.access_token) return response.data }, diff --git a/frontend/src/api/client.ts b/frontend/src/api/client.ts index 58a269bb7..2eb5df726 100644 --- a/frontend/src/api/client.ts +++ b/frontend/src/api/client.ts @@ -323,7 +323,7 @@ class ApiClient { } async refreshToken(): Promise { - return this.client.post('/api/auth/refresh', {}) + return this.client.post('/api/auth/refresh') } // 以下方法直接委托给 axios client，Demo 模式由 adapter 统一处理 From 6673ab6d4af470318cb46bed514ffd3c5fc90490 Mon Sep 17 00:00:00 2001 From: fawney19 Date: Wed, 27 May 2026 15:08:23 +0800 Subject: [PATCH 4/6] Add fast model directive service tier --- .../handlers/admin/system/shared/modules.rs | 2 +- apps/aether-gateway/src/system_features.rs | 199 ++++++++++++++++-- crates/aether-admin/src/system.rs | 9 +- crates/aether-ai-formats/src/api.rs | 2 +- .../src/formats/shared/model_directives.rs | 198 ++++++++++++++++- crates/aether-ai-formats/src/lib.rs | 2 +- crates/aether-scheduler-core/src/model.rs | 18 ++ 7 files changed, 396 insertions(+), 34 deletions(-) diff --git a/apps/aether-gateway/src/handlers/admin/system/shared/modules.rs b/apps/aether-gateway/src/handlers/admin/system/shared/modules.rs index 132f6e557..a925a1765 100644 --- a/apps/aether-gateway/src/handlers/admin/system/shared/modules.rs +++ b/apps/aether-gateway/src/handlers/admin/system/shared/modules.rs @@ -113,7 +113,7 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[ AdminModuleDefinition { name: "model_directives", display_name: "模型后缀参数", - description: "允许通过模型名后缀覆盖推理参数", + description: "允许通过模型名后缀覆盖推理参数或服务层级", category: "integration", env_key: "MODEL_DIRECTIVES_AVAILABLE", default_available: true, diff --git a/apps/aether-gateway/src/system_features.rs b/apps/aether-gateway/src/system_features.rs index bd7e3f5a1..1fd746d3a 100644 --- a/apps/aether-gateway/src/system_features.rs +++ b/apps/aether-gateway/src/system_features.rs @@ -88,15 +88,16 @@ pub(crate) async fn reasoning_model_directive_enabled_for_api_format_and_model( return false; } - let Some(suffix) = requested_model.and_then(reasoning_suffix_from_model) else { + let Some(suffixes) = requested_model.and_then(model_directive_suffixes_from_model) else { return false; }; - settings + let mappings = settings .as_ref() - .and_then(|settings| settings.api_format_mappings(&api_format)) - .map(|mappings| mappings.contains_key(&suffix)) - .unwrap_or_else(|| DEFAULT_REASONING_SUFFIXES.contains(&suffix.as_str())) + .and_then(|settings| settings.api_format_mappings(&api_format)); + suffixes + .iter() + .all(|suffix| suffix_supported_for_api_format(&api_format, suffix, mappings.as_ref())) } pub(crate) async fn reasoning_model_directive_mapping_for_api_format_and_model( @@ -113,17 +114,17 @@ pub(crate) async fn reasoning_model_directive_mapping_for_api_format_and_model( { return None; } - let suffix = requested_model.and_then(reasoning_suffix_from_model)?; + let suffixes = requested_model.and_then(model_directive_suffixes_from_model)?; let api_format = crate::ai_serving::normalize_api_format_alias(api_format); let settings = read_reasoning_model_directive_settings(state).await; - settings + let mappings = settings .as_ref() - .and_then(|settings| settings.api_format_mappings(&api_format)) - .and_then(|mappings| mappings.get(&suffix).cloned()) - .or_else(|| default_reasoning_mapping(&api_format, &suffix)) + .and_then(|settings| settings.api_format_mappings(&api_format)); + model_directive_mapping_for_suffixes(&api_format, &suffixes, mappings.as_ref()) } -const DEFAULT_REASONING_SUFFIXES: &[&str] = &["low", "medium", "high", "xhigh", "max"]; +const DEFAULT_MODEL_DIRECTIVE_SUFFIXES: &[&str] = + &["low", "medium", "high", "xhigh", "max", "fast"]; #[derive(Debug, Clone, Default)] struct ReasoningModelDirectiveSettings { @@ -183,22 +184,66 @@ impl ReasoningModelDirectiveSettings { } } -fn reasoning_suffix_from_model(model: &str) -> Option { - let model = model.trim(); - let (base_model, suffix) = model.rsplit_once('-')?; - if base_model.trim().is_empty() { +fn model_directive_suffixes_from_model(model: &str) -> Option> { + let mut base_model = model.trim(); + let mut suffixes = Vec::new(); + let mut has_reasoning_effort = false; + let mut has_service_tier = false; + while let Some((candidate_base, suffix)) = base_model.rsplit_once('-') { + let Some(suffix) = normalize_reasoning_suffix(suffix) else { + break; + }; + match model_directive_suffix_kind(&suffix)? { + ModelDirectiveSuffixKind::ReasoningEffort => { + if has_reasoning_effort { + return None; + } + has_reasoning_effort = true; + } + ModelDirectiveSuffixKind::ServiceTier => { + if has_service_tier { + return None; + } + has_service_tier = true; + } + } + suffixes.push(suffix); + base_model = candidate_base.trim(); + } + if base_model.is_empty() || suffixes.is_empty() { return None; } - normalize_reasoning_suffix(suffix) + suffixes.sort_by_key(|suffix| { + DEFAULT_MODEL_DIRECTIVE_SUFFIXES + .iter() + .position(|value| value == suffix) + .unwrap_or(usize::MAX) + }); + Some(suffixes) } fn normalize_reasoning_suffix(suffix: &str) -> Option { let normalized = suffix.trim().to_ascii_lowercase(); - DEFAULT_REASONING_SUFFIXES + DEFAULT_MODEL_DIRECTIVE_SUFFIXES .contains(&normalized.as_str()) .then_some(normalized) } +enum ModelDirectiveSuffixKind { + ReasoningEffort, + ServiceTier, +} + +fn model_directive_suffix_kind(suffix: &str) -> Option { + match suffix { + "low" | "medium" | "high" | "xhigh" | "max" => { + Some(ModelDirectiveSuffixKind::ReasoningEffort) + } + "fast" => Some(ModelDirectiveSuffixKind::ServiceTier), + _ => None, + } +} + fn normalize_reasoning_mappings( mappings: &serde_json::Map, ) -> serde_json::Map { @@ -212,10 +257,16 @@ fn normalize_reasoning_mappings( fn default_reasoning_mapping(api_format: &str, suffix: &str) -> Option { match api_format { + "openai:chat" if suffix == "fast" => { + Some(serde_json::json!({ "service_tier": "priority" })) + } "openai:chat" => { let effort = openai_reasoning_effort_value(suffix)?; Some(serde_json::json!({ "reasoning_effort": effort })) } + "openai:responses" | "openai:responses:compact" if suffix == "fast" => { + Some(serde_json::json!({ "service_tier": "priority" })) + } "openai:responses" | "openai:responses:compact" => { let effort = openai_reasoning_effort_value(suffix)?; Some(serde_json::json!({ "reasoning": { "effort": effort } })) @@ -251,6 +302,54 @@ fn default_reasoning_mapping(api_format: &str, suffix: &str) -> Option>, +) -> bool { + mappings + .map(|mappings| mappings.contains_key(suffix)) + .unwrap_or_else(|| default_reasoning_mapping(api_format, suffix).is_some()) +} + +fn model_directive_mapping_for_suffixes( + api_format: &str, + suffixes: &[String], + mappings: Option<&serde_json::Map>, +) -> Option { + let mut combined = serde_json::json!({}); + for suffix in suffixes { + let mapping = mappings + .and_then(|mappings| mappings.get(suffix).cloned()) + .or_else(|| { + mappings + .is_none() + .then(|| default_reasoning_mapping(api_format, suffix)) + .flatten() + })?; + deep_merge_json(&mut combined, &mapping); + } + Some(combined) +} + +fn deep_merge_json(target: &mut serde_json::Value, patch: &serde_json::Value) { + match (target, patch) { + (serde_json::Value::Object(target_object), serde_json::Value::Object(patch_object)) => { + for (key, patch_value) in patch_object { + match target_object.get_mut(key) { + Some(target_value) => deep_merge_json(target_value, patch_value), + None => { + target_object.insert(key.clone(), patch_value.clone()); + } + } + } + } + (target, patch) => { + *target = patch.clone(); + } + } +} + fn openai_reasoning_effort_value(suffix: &str) -> Option<&'static str> { match suffix { "low" => Some("low"), @@ -297,7 +396,11 @@ fn parse_reasoning_model_directive_settings( #[cfg(test)] mod tests { - use super::parse_reasoning_model_directive_settings; + use super::{ + default_reasoning_mapping, model_directive_mapping_for_suffixes, + model_directive_suffixes_from_model, parse_reasoning_model_directive_settings, + suffix_supported_for_api_format, + }; use serde_json::json; #[test] @@ -332,4 +435,64 @@ mod tests { ); assert_eq!(settings.api_format_enabled("gemini:generate_content"), None); } + + #[test] + fn default_fast_suffix_maps_to_openai_priority_service_tier() { + assert_eq!( + default_reasoning_mapping("openai:chat", "fast"), + Some(json!({ "service_tier": "priority" })) + ); + assert_eq!( + default_reasoning_mapping("openai:responses", "fast"), + Some(json!({ "service_tier": "priority" })) + ); + assert_eq!( + default_reasoning_mapping("openai:responses:compact", "fast"), + Some(json!({ "service_tier": "priority" })) + ); + assert_eq!(default_reasoning_mapping("claude:messages", "fast"), None); + assert_eq!( + default_reasoning_mapping("gemini:generate_content", "fast"), + None + ); + } + + #[test] + fn combined_suffixes_are_order_insensitive() { + let expected = Some(vec!["xhigh".to_string(), "fast".to_string()]); + assert_eq!( + model_directive_suffixes_from_model("gpt-5.4-fast-xhigh"), + expected + ); + assert_eq!( + model_directive_suffixes_from_model("gpt-5.4-xhigh-fast"), + expected + ); + assert_eq!( + model_directive_mapping_for_suffixes( + "openai:chat", + expected.as_ref().expect("suffixes should parse"), + None, + ), + Some(json!({ + "reasoning_effort": "xhigh", + "service_tier": "priority" + })) + ); + } + + #[test] + fn combined_suffix_requires_each_suffix_to_support_api_format() { + let suffixes = model_directive_suffixes_from_model("gpt-5.4-xhigh-fast").expect("suffixes"); + assert!(suffixes + .iter() + .all(|suffix| { suffix_supported_for_api_format("openai:responses", suffix, None) })); + assert!(!suffixes + .iter() + .all(|suffix| { suffix_supported_for_api_format("claude:messages", suffix, None) })); + assert_eq!( + model_directive_mapping_for_suffixes("claude:messages", &suffixes, None), + None + ); + } } diff --git a/crates/aether-admin/src/system.rs b/crates/aether-admin/src/system.rs index 694ca7747..2c0a09843 100644 --- a/crates/aether-admin/src/system.rs +++ b/crates/aether-admin/src/system.rs @@ -1729,7 +1729,8 @@ pub fn admin_system_config_default_value(key: &str) -> Option "medium": { "reasoning_effort": "medium" }, "high": { "reasoning_effort": "high" }, "xhigh": { "reasoning_effort": "xhigh" }, - "max": { "reasoning_effort": "xhigh" } + "max": { "reasoning_effort": "xhigh" }, + "fast": { "service_tier": "priority" } } }, "openai:responses": { @@ -1739,7 +1740,8 @@ pub fn admin_system_config_default_value(key: &str) -> Option "medium": { "reasoning": { "effort": "medium" } }, "high": { "reasoning": { "effort": "high" } }, "xhigh": { "reasoning": { "effort": "xhigh" } }, - "max": { "reasoning": { "effort": "xhigh" } } + "max": { "reasoning": { "effort": "xhigh" } }, + "fast": { "service_tier": "priority" } } }, "openai:responses:compact": { @@ -1749,7 +1751,8 @@ pub fn admin_system_config_default_value(key: &str) -> Option "medium": { "reasoning": { "effort": "medium" } }, "high": { "reasoning": { "effort": "high" } }, "xhigh": { "reasoning": { "effort": "xhigh" } }, - "max": { "reasoning": { "effort": "xhigh" } } + "max": { "reasoning": { "effort": "xhigh" } }, + "fast": { "service_tier": "priority" } } }, "claude:messages": { diff --git a/crates/aether-ai-formats/src/api.rs b/crates/aether-ai-formats/src/api.rs index 16f8f6aba..56cb07682 100644 --- a/crates/aether-ai-formats/src/api.rs +++ b/crates/aether-ai-formats/src/api.rs @@ -77,7 +77,7 @@ pub use crate::formats::shared::model_directives::{ apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort, extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model, normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride, - ReasoningEffort, + ReasoningEffort, ServiceTier, }; pub use crate::formats::shared::passthrough::{ resolve_stream_spec as resolve_local_same_format_stream_spec, diff --git a/crates/aether-ai-formats/src/formats/shared/model_directives.rs b/crates/aether-ai-formats/src/formats/shared/model_directives.rs index 3592a0adf..345cbd354 100644 --- a/crates/aether-ai-formats/src/formats/shared/model_directives.rs +++ b/crates/aether-ai-formats/src/formats/shared/model_directives.rs @@ -9,6 +9,7 @@ pub struct ModelDirective { #[derive(Debug, Clone, PartialEq, Eq)] pub enum ModelOverride { ReasoningEffort(ReasoningEffort), + ServiceTier(ServiceTier), } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -79,19 +80,90 @@ impl ReasoningEffort { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ServiceTier { + Priority, +} + +impl ServiceTier { + pub fn parse(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "fast" => Some(Self::Priority), + _ => None, + } + } + + pub fn as_openai_value(self) -> &'static str { + match self { + Self::Priority => "priority", + } + } +} + pub fn parse_model_directive(model: &str) -> Option { - let model = model.trim(); - let (base_model, suffix) = model.rsplit_once('-')?; - let base_model = base_model.trim(); + let (base_model, overrides) = parse_model_directive_parts(model)?; + Some(ModelDirective { + base_model, + overrides, + }) +} + +fn parse_model_directive_parts(model: &str) -> Option<(String, Vec)> { + let mut base_model = model.trim(); + let mut overrides = ModelOverrideAccumulator::default(); + while let Some((candidate_base, suffix)) = base_model.rsplit_once('-') { + let Some(override_item) = parse_model_override(suffix) else { + break; + }; + overrides.insert(override_item)?; + base_model = candidate_base.trim(); + } if base_model.is_empty() { return None; } + let overrides = overrides.into_overrides()?; + Some((base_model.to_string(), overrides)) +} - let reasoning_effort = ReasoningEffort::parse(suffix)?; - Some(ModelDirective { - base_model: base_model.to_string(), - overrides: vec![ModelOverride::ReasoningEffort(reasoning_effort)], - }) +fn parse_model_override(suffix: &str) -> Option { + ReasoningEffort::parse(suffix) + .map(ModelOverride::ReasoningEffort) + .or_else(|| ServiceTier::parse(suffix).map(ModelOverride::ServiceTier)) +} + +#[derive(Default)] +struct ModelOverrideAccumulator { + reasoning_effort: Option, + service_tier: Option, +} + +impl ModelOverrideAccumulator { + fn insert(&mut self, override_item: ModelOverride) -> Option<()> { + match override_item { + ModelOverride::ReasoningEffort(value) => { + if self.reasoning_effort.replace(value).is_some() { + return None; + } + } + ModelOverride::ServiceTier(value) => { + if self.service_tier.replace(value).is_some() { + return None; + } + } + } + Some(()) + } + + fn into_overrides(self) -> Option> { + let mut overrides = Vec::new(); + if let Some(reasoning_effort) = self.reasoning_effort { + overrides.push(ModelOverride::ReasoningEffort(reasoning_effort)); + } + if let Some(service_tier) = self.service_tier { + overrides.push(ModelOverride::ServiceTier(service_tier)); + } + (!overrides.is_empty()).then_some(overrides) + } } pub fn model_directive_base_model(model: &str) -> Option { @@ -149,18 +221,23 @@ pub fn apply_model_directive_overrides_from_model( source_model: &str, ) -> Option { let directive = parse_model_directive(source_model)?; + let mut patched_body = provider_request_body.clone(); for override_item in &directive.overrides { match override_item { ModelOverride::ReasoningEffort(effort) => { apply_reasoning_effort_override( - provider_request_body, + &mut patched_body, provider_api_format, provider_model, *effort, )?; } + ModelOverride::ServiceTier(tier) => { + apply_service_tier_override(&mut patched_body, provider_api_format, *tier)?; + } } } + *provider_request_body = patched_body; Some(directive) } @@ -215,6 +292,21 @@ fn apply_reasoning_effort_override( } } +fn apply_service_tier_override( + provider_request_body: &mut Value, + provider_api_format: &str, + tier: ServiceTier, +) -> Option<()> { + match crate::normalize_api_format_alias(provider_api_format).as_str() { + "openai:chat" | "openai:responses" | "openai:responses:compact" => set_object_string( + provider_request_body, + "service_tier", + tier.as_openai_value(), + ), + _ => None, + } +} + fn set_object_string(body: &mut Value, key: &str, value: &str) -> Option<()> { body.as_object_mut()? .insert(key.to_string(), Value::String(value.to_string())); @@ -368,7 +460,7 @@ mod tests { use super::{ apply_model_directive_overrides_from_model, parse_model_directive, ModelDirective, - ModelOverride, ReasoningEffort, + ModelOverride, ReasoningEffort, ServiceTier, }; #[test] @@ -389,12 +481,37 @@ mod tests { ); } + #[test] + fn parses_supported_service_tier_suffixes() { + assert_eq!( + parse_model_directive("gpt-5.4-fast"), + Some(ModelDirective { + base_model: "gpt-5.4".to_string(), + overrides: vec![ModelOverride::ServiceTier(ServiceTier::Priority)], + }) + ); + } + + #[test] + fn parses_combined_suffixes_in_canonical_order() { + let expected = Some(ModelDirective { + base_model: "gpt-5.4".to_string(), + overrides: vec![ + ModelOverride::ReasoningEffort(ReasoningEffort::XHigh), + ModelOverride::ServiceTier(ServiceTier::Priority), + ], + }); + assert_eq!(parse_model_directive("gpt-5.4-fast-xhigh"), expected); + assert_eq!(parse_model_directive("gpt-5.4-xhigh-fast"), expected); + } + #[test] fn ignores_unknown_or_incomplete_suffixes() { assert_eq!(parse_model_directive("gpt-5.4-ultra"), None); assert_eq!(parse_model_directive("gpt-5.4"), None); assert_eq!(parse_model_directive("-high"), None); assert_eq!(parse_model_directive("gpt-5.4-high-json"), None); + assert_eq!(parse_model_directive("gpt-5.4-low-high"), None); } #[test] @@ -446,4 +563,65 @@ mod tests { 2048 ); } + + #[test] + fn applies_fast_suffix_to_openai_service_tier() { + let mut openai_chat = json!({"model": "gpt-5-upstream"}); + apply_model_directive_overrides_from_model( + &mut openai_chat, + "openai:chat", + "gpt-5-upstream", + "gpt-5.4-fast", + ) + .expect("directive should apply"); + assert_eq!(openai_chat["service_tier"], "priority"); + + let mut responses = json!({"model": "gpt-5-upstream"}); + apply_model_directive_overrides_from_model( + &mut responses, + "openai:responses", + "gpt-5-upstream", + "gpt-5.4-fast", + ) + .expect("directive should apply"); + assert_eq!(responses["service_tier"], "priority"); + } + + #[test] + fn applies_combined_suffixes_to_openai_body() { + let mut openai_chat = json!({"model": "gpt-5-upstream", "reasoning_effort": "low"}); + apply_model_directive_overrides_from_model( + &mut openai_chat, + "openai:chat", + "gpt-5-upstream", + "gpt-5.4-fast-xhigh", + ) + .expect("directive should apply"); + assert_eq!(openai_chat["reasoning_effort"], "xhigh"); + assert_eq!(openai_chat["service_tier"], "priority"); + + let mut reversed = json!({"model": "gpt-5-upstream", "reasoning_effort": "low"}); + apply_model_directive_overrides_from_model( + &mut reversed, + "openai:chat", + "gpt-5-upstream", + "gpt-5.4-xhigh-fast", + ) + .expect("directive should apply"); + assert_eq!(reversed, openai_chat); + } + + #[test] + fn unsupported_combined_suffix_leaves_body_unchanged() { + let mut claude = json!({"model": "claude-sonnet-4-5"}); + let original = claude.clone(); + assert!(apply_model_directive_overrides_from_model( + &mut claude, + "claude:messages", + "claude-sonnet-4-5", + "gpt-5.4-fast-xhigh", + ) + .is_none()); + assert_eq!(claude, original); + } } diff --git a/crates/aether-ai-formats/src/lib.rs b/crates/aether-ai-formats/src/lib.rs index 182d7c1d0..67160b7a5 100644 --- a/crates/aether-ai-formats/src/lib.rs +++ b/crates/aether-ai-formats/src/lib.rs @@ -25,7 +25,7 @@ pub use formats::shared::model_directives::{ apply_model_directive_overrides_from_request, claude_model_uses_adaptive_effort, extract_gemini_model_from_path, gemini_model_uses_thinking_level, model_directive_base_model, normalize_model_directive_model, parse_model_directive, ModelDirective, ModelOverride, - ReasoningEffort, + ReasoningEffort, ServiceTier, }; pub use formats::shared::request::{ endpoint_config_forces_upstream_stream_policy, enforce_request_body_stream_field, diff --git a/crates/aether-scheduler-core/src/model.rs b/crates/aether-scheduler-core/src/model.rs index 9c8ee3128..044b15a90 100644 --- a/crates/aether-scheduler-core/src/model.rs +++ b/crates/aether-scheduler-core/src/model.rs @@ -548,6 +548,24 @@ mod tests { .as_deref(), Some("gpt-5.4") ); + + let row = sample_row("gpt-5.4", "gpt-5.4-upstream"); + assert!(row_supports_requested_model_with_model_directives( + &row, + "gpt-5.4-fast-xhigh", + "openai:chat", + true + )); + assert_eq!( + resolve_requested_global_model_name_with_model_directives( + &[row], + "gpt-5.4-xhigh-fast", + "openai:chat", + true + ) + .as_deref(), + Some("gpt-5.4") + ); } #[test] From 21e52722e6a3eaaec77e1b2c272c3969c95f34b5 Mon Sep 17 00:00:00 2001 From: fawney19 Date: Wed, 27 May 2026 15:39:39 +0800 Subject: [PATCH 5/6] Prefer provider request body for usage badges --- .../src/repository/usage/types.rs | 43 ++++++++++--- .../list_recent_usage_audits_prefix.sql | 40 ++++++------ .../queries/list_usage_audits_prefix.sql | 40 ++++++------ .../src/request_metadata.rs | 61 +++++++++++++++++-- 4 files changed, 135 insertions(+), 49 deletions(-) diff --git a/crates/aether-data-contracts/src/repository/usage/types.rs b/crates/aether-data-contracts/src/repository/usage/types.rs index 2a0ac9daf..1a11fd251 100644 --- a/crates/aether-data-contracts/src/repository/usage/types.rs +++ b/crates/aether-data-contracts/src/repository/usage/types.rs @@ -418,19 +418,33 @@ impl StoredRequestUsageAudit { } pub fn provider_reasoning_effort(&self) -> Option { + if self + .provider_request_body + .as_ref() + .and_then(Value::as_object) + .is_some() + { + return extract_provider_reasoning_effort_from_body( + self.provider_request_body.as_ref(), + ); + } + self.request_metadata_string(PROVIDER_REASONING_EFFORT_METADATA_KEY) .and_then(normalize_provider_reasoning_effort) - .or_else(|| { - extract_provider_reasoning_effort_from_body(self.provider_request_body.as_ref()) - }) } pub fn provider_service_tier(&self) -> Option { + if self + .provider_request_body + .as_ref() + .and_then(Value::as_object) + .is_some() + { + return extract_provider_service_tier_from_body(self.provider_request_body.as_ref()); + } + self.request_metadata_string(PROVIDER_SERVICE_TIER_METADATA_KEY) .and_then(normalize_provider_service_tier) - .or_else(|| { - extract_provider_service_tier_from_body(self.provider_request_body.as_ref()) - }) } pub fn body_ref(&self, field: UsageBodyField) -> Option<&str> { @@ -2398,7 +2412,7 @@ mod tests { } #[test] - fn provider_reasoning_effort_reads_metadata_and_provider_request_body() { + fn provider_reasoning_effort_prefers_provider_request_body_over_metadata() { let mut usage = sample_usage(); usage.provider_request_body = Some(json!({ "reasoning": { "effort": "XHigh" }, @@ -2413,6 +2427,10 @@ mod tests { "provider_service_tier": "standard" })); + assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("xhigh")); + assert_eq!(usage.provider_service_tier().as_deref(), Some("priority")); + + usage.provider_request_body = None; assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("max")); assert_eq!(usage.provider_service_tier().as_deref(), Some("standard")); @@ -2424,6 +2442,17 @@ mod tests { assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("high")); assert_eq!(usage.provider_service_tier().as_deref(), Some("priority")); + + usage.request_metadata = Some(json!({ + "provider_reasoning_effort": "max", + "provider_service_tier": "standard" + })); + usage.provider_request_body = Some(json!({ + "model": "gpt-5" + })); + + assert_eq!(usage.provider_reasoning_effort(), None); + assert_eq!(usage.provider_service_tier(), None); } #[test] diff --git a/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql b/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql index 1ccf6d9b6..7a6200cd8 100644 --- a/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql +++ b/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql @@ -114,15 +114,17 @@ SELECT OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL - OR COALESCE( - NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''), - NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), - NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') - ) IS NOT NULL - OR COALESCE( - NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''), - NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') - ) IS NOT NULL + OR CASE + WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN COALESCE( + NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), + NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') + ) + ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') + END IS NOT NULL + OR CASE + WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') + ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') + END IS NOT NULL OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false') THEN jsonb_strip_nulls(jsonb_build_object( @@ -135,16 +137,18 @@ SELECT 'request_path_and_query', NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''), 'provider_reasoning_effort', - COALESCE( - NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''), - NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), - NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') - ), + CASE + WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN COALESCE( + NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), + NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') + ) + ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') + END, 'provider_service_tier', - COALESCE( - NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''), - NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') - ), + CASE + WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') + ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') + END, 'client_requested_stream', CASE WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') diff --git a/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql b/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql index 1ccf6d9b6..7a6200cd8 100644 --- a/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql +++ b/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql @@ -114,15 +114,17 @@ SELECT OR NULLIF(BTRIM("usage".request_metadata->>'user_agent'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL - OR COALESCE( - NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''), - NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), - NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') - ) IS NOT NULL - OR COALESCE( - NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''), - NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') - ) IS NOT NULL + OR CASE + WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN COALESCE( + NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), + NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') + ) + ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') + END IS NOT NULL + OR CASE + WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') + ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') + END IS NOT NULL OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false') THEN jsonb_strip_nulls(jsonb_build_object( @@ -135,16 +137,18 @@ SELECT 'request_path_and_query', NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''), 'provider_reasoning_effort', - COALESCE( - NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), ''), - NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), - NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') - ), + CASE + WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN COALESCE( + NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), + NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') + ) + ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') + END, 'provider_service_tier', - COALESCE( - NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''), - NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') - ), + CASE + WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') + ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') + END, 'client_requested_stream', CASE WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') diff --git a/crates/aether-usage-runtime/src/request_metadata.rs b/crates/aether-usage-runtime/src/request_metadata.rs index 7816b908f..eb60f906d 100644 --- a/crates/aether-usage-runtime/src/request_metadata.rs +++ b/crates/aether-usage-runtime/src/request_metadata.rs @@ -77,15 +77,20 @@ pub(crate) fn attach_provider_request_body_metadata( metadata: Option, provider_request_body: Option<&Value>, ) -> Option { + let provider_body_is_object = provider_request_body.and_then(Value::as_object).is_some(); let reasoning_effort = extract_provider_reasoning_effort_from_body(provider_request_body); let service_tier = extract_provider_service_tier_from_body(provider_request_body); - if reasoning_effort.is_none() && service_tier.is_none() { + if !provider_body_is_object && reasoning_effort.is_none() && service_tier.is_none() { return metadata; } let mut object = match metadata { Some(Value::Object(object)) => object, _ => Map::new(), }; + if provider_body_is_object { + object.remove(PROVIDER_REASONING_EFFORT_METADATA_KEY); + object.remove(PROVIDER_SERVICE_TIER_METADATA_KEY); + } if let Some(reasoning_effort) = reasoning_effort { object.insert( PROVIDER_REASONING_EFFORT_METADATA_KEY.to_string(), @@ -98,7 +103,7 @@ pub(crate) fn attach_provider_request_body_metadata( Value::String(service_tier), ); } - Some(Value::Object(object)) + (!object.is_empty()).then_some(Value::Object(object)) } fn copy_allowed_metadata_fields(source: &Map, target: &mut Map) { @@ -438,10 +443,11 @@ mod tests { use std::collections::BTreeMap; use super::{ - build_usage_request_metadata_seed, merge_usage_request_metadata, - merge_usage_request_metadata_owned, sanitize_usage_request_metadata, - sanitize_usage_request_metadata_ref, MAX_USAGE_REQUEST_METADATA_BYTES, - MAX_USAGE_REQUEST_METADATA_DEPTH, MAX_USAGE_REQUEST_METADATA_NODES, + attach_provider_request_body_metadata, build_usage_request_metadata_seed, + merge_usage_request_metadata, merge_usage_request_metadata_owned, + sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref, + MAX_USAGE_REQUEST_METADATA_BYTES, MAX_USAGE_REQUEST_METADATA_DEPTH, + MAX_USAGE_REQUEST_METADATA_NODES, }; fn sample_plan() -> ExecutionPlan { @@ -678,6 +684,49 @@ mod tests { assert_eq!(metadata, None); } + #[test] + fn provider_request_body_metadata_uses_final_provider_body_as_source_of_truth() { + let metadata = Some(json!({ + "trace_id": "trace-1", + "provider_reasoning_effort": "high", + "provider_service_tier": "priority" + })); + + let updated = attach_provider_request_body_metadata( + metadata.clone(), + Some(&json!({ + "model": "gpt-5", + "reasoning": { "effort": "low" }, + "service_tier": "standard" + })), + ) + .expect("metadata should remain"); + + assert_eq!( + updated, + json!({ + "trace_id": "trace-1", + "provider_reasoning_effort": "low", + "provider_service_tier": "standard" + }) + ); + + let cleared = attach_provider_request_body_metadata( + metadata, + Some(&json!({ + "model": "gpt-5" + })), + ) + .expect("metadata should retain unrelated fields"); + + assert_eq!( + cleared, + json!({ + "trace_id": "trace-1" + }) + ); + } + #[test] fn owned_merge_matches_filtered_merge_for_trusted_objects() { let base = Some(json!({ From d77a572dc7651372392982d24091c5505d2e3935 Mon Sep 17 00:00:00 2001 From: fawney19 Date: Wed, 27 May 2026 15:48:12 +0800 Subject: [PATCH 6/6] fix: cast usage provider body before jsonb type checks --- .../postgres/queries/list_recent_usage_audits_prefix.sql | 8 ++++---- .../usage/postgres/queries/list_usage_audits_prefix.sql | 8 ++++---- crates/aether-data/src/repository/usage/postgres/tests.rs | 2 ++ 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql b/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql index 7a6200cd8..d3f146505 100644 --- a/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql +++ b/crates/aether-data/src/repository/usage/postgres/queries/list_recent_usage_audits_prefix.sql @@ -115,14 +115,14 @@ SELECT OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL OR CASE - WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN COALESCE( + WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE( NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') ) ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') END IS NOT NULL OR CASE - WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') + WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') END IS NOT NULL OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') @@ -138,7 +138,7 @@ SELECT NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''), 'provider_reasoning_effort', CASE - WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN COALESCE( + WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE( NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') ) @@ -146,7 +146,7 @@ SELECT END, 'provider_service_tier', CASE - WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') + WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') END, 'client_requested_stream', diff --git a/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql b/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql index 7a6200cd8..d3f146505 100644 --- a/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql +++ b/crates/aether-data/src/repository/usage/postgres/queries/list_usage_audits_prefix.sql @@ -115,14 +115,14 @@ SELECT OR NULLIF(BTRIM("usage".request_metadata->>'request_path'), '') IS NOT NULL OR NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), '') IS NOT NULL OR CASE - WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN COALESCE( + WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE( NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') ) ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') END IS NOT NULL OR CASE - WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') + WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') END IS NOT NULL OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false') @@ -138,7 +138,7 @@ SELECT NULLIF(BTRIM("usage".request_metadata->>'request_path_and_query'), ''), 'provider_reasoning_effort', CASE - WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN COALESCE( + WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN COALESCE( NULLIF(BTRIM("usage".provider_request_body->>'reasoning_effort'), ''), NULLIF(BTRIM("usage".provider_request_body->'reasoning'->>'effort'), '') ) @@ -146,7 +146,7 @@ SELECT END, 'provider_service_tier', CASE - WHEN jsonb_typeof("usage".provider_request_body) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') + WHEN jsonb_typeof("usage".provider_request_body::jsonb) = 'object' THEN NULLIF(BTRIM("usage".provider_request_body->>'service_tier'), '') ELSE NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') END, 'client_requested_stream', diff --git a/crates/aether-data/src/repository/usage/postgres/tests.rs b/crates/aether-data/src/repository/usage/postgres/tests.rs index 44bdd68b1..ed6e9dfc8 100644 --- a/crates/aether-data/src/repository/usage/postgres/tests.rs +++ b/crates/aether-data/src/repository/usage/postgres/tests.rs @@ -707,6 +707,8 @@ fn usage_sql_uses_json_null_placeholders_for_usage_payload_columns() { super::LIST_RECENT_USAGE_AUDITS_PREFIX, ] { assert!(sql.contains("jsonb_strip_nulls(jsonb_build_object(")); + assert!(sql.contains("jsonb_typeof(\"usage\".provider_request_body::jsonb)")); + assert!(!sql.contains("jsonb_typeof(\"usage\".provider_request_body)")); assert!(sql.contains("'client_ip'")); assert!(sql.contains("request_metadata->>'client_ip'")); assert!(sql.contains("'user_agent'"));