Merge branch 'review-pr-596'

apps/aether-gateway/src/state/catalog.rs

@@ -284,20 +284,30 @@ impl AppState {
         &self,
         record: &global_models::UpsertAdminProviderModelRecord,
     ) -> Result<Option<global_models::StoredAdminProviderModel>, GatewayError> {
-        self.data
+        let created = self
+            .data
             .create_admin_provider_model(record)
             .await
-            .map_err(|err| GatewayError::Internal(err.to_string()))
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        if created.is_some() {
+            self.invalidate_provider_routing_caches();
+        }
+        Ok(created)
     }
 
     pub(crate) async fn update_admin_provider_model(
         &self,
         record: &global_models::UpsertAdminProviderModelRecord,
     ) -> Result<Option<global_models::StoredAdminProviderModel>, GatewayError> {
-        self.data
+        let updated = self
+            .data
             .update_admin_provider_model(record)
             .await
-            .map_err(|err| GatewayError::Internal(err.to_string()))
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        if updated.is_some() {
+            self.invalidate_provider_routing_caches();
+        }
+        Ok(updated)
     }
 
     pub(crate) async fn delete_admin_provider_model(
@@ -305,40 +315,60 @@ impl AppState {
         provider_id: &str,
         model_id: &str,
     ) -> Result<bool, GatewayError> {
-        self.data
+        let deleted = self
+            .data
             .delete_admin_provider_model(provider_id, model_id)
             .await
-            .map_err(|err| GatewayError::Internal(err.to_string()))
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        if deleted {
+            self.invalidate_provider_routing_caches();
+        }
+        Ok(deleted)
     }
 
     pub(crate) async fn create_admin_global_model(
         &self,
         record: &global_models::CreateAdminGlobalModelRecord,
     ) -> Result<Option<global_models::StoredAdminGlobalModel>, GatewayError> {
-        self.data
+        let created = self
+            .data
             .create_admin_global_model(record)
             .await
-            .map_err(|err| GatewayError::Internal(err.to_string()))
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        if created.is_some() {
+            self.invalidate_provider_routing_caches();
+        }
+        Ok(created)
     }
 
     pub(crate) async fn update_admin_global_model(
         &self,
         record: &global_models::UpdateAdminGlobalModelRecord,
     ) -> Result<Option<global_models::StoredAdminGlobalModel>, GatewayError> {
-        self.data
+        let updated = self
+            .data
             .update_admin_global_model(record)
             .await
-            .map_err(|err| GatewayError::Internal(err.to_string()))
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        if updated.is_some() {
+            self.invalidate_provider_routing_caches();
+        }
+        Ok(updated)
     }
 
     pub(crate) async fn delete_admin_global_model(
         &self,
         global_model_id: &str,
     ) -> Result<bool, GatewayError> {
-        self.data
+        let deleted = self
+            .data
             .delete_admin_global_model(global_model_id)
             .await
-            .map_err(|err| GatewayError::Internal(err.to_string()))
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        if deleted {
+            self.invalidate_provider_routing_caches();
+        }
+        Ok(deleted)
     }
 
     pub(crate) async fn list_provider_model_stats(
@@ -840,13 +870,30 @@ impl AppState {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::Arc;
+    use std::sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    };
     use std::time::Duration;
 
-    use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
+    use aether_data::repository::{
+        global_models::InMemoryGlobalModelReadRepository,
+        provider_catalog::InMemoryProviderCatalogReadRepository,
+    };
+    use aether_data::DataLayerError;
+    use aether_data_contracts::repository::candidate_selection::{
+        MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
+        StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
+        StoredRequestedModelCandidateRowsQuery,
+    };
+    use aether_data_contracts::repository::global_models::{
+        CreateAdminGlobalModelRecord, StoredAdminGlobalModel, UpdateAdminGlobalModelRecord,
+        UpsertAdminProviderModelRecord,
+    };
     use aether_data_contracts::repository::provider_catalog::{
         StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
     };
+    use async_trait::async_trait;
 
     use crate::cache::SchedulerAffinityTarget;
     use crate::data::GatewayDataState;
@@ -897,6 +944,195 @@ mod tests {
         .expect("key should build")
     }
 
+    fn sample_admin_global_model() -> StoredAdminGlobalModel {
+        StoredAdminGlobalModel::new(
+            "global-1".to_string(),
+            "gpt-5".to_string(),
+            "GPT 5".to_string(),
+            true,
+            None,
+            None,
+            None,
+            None,
+            0,
+            0,
+            0,
+            Some(1_711_000_000),
+            Some(1_711_000_000),
+        )
+        .expect("global model should build")
+    }
+
+    fn sample_provider_model_record(
+        id: &str,
+        global_model_id: &str,
+        is_active: bool,
+    ) -> UpsertAdminProviderModelRecord {
+        UpsertAdminProviderModelRecord::new(
+            id.to_string(),
+            "provider-1".to_string(),
+            global_model_id.to_string(),
+            "gpt-5-upstream".to_string(),
+            None,
+            None,
+            None,
+            None,
+            None,
+            Some(true),
+            None,
+            None,
+            is_active,
+            true,
+            None,
+        )
+        .expect("provider model record should build")
+    }
+
+    #[derive(Debug, Default)]
+    struct ClearCountingCandidateSelectionReadRepository {
+        clear_count: AtomicUsize,
+    }
+
+    impl ClearCountingCandidateSelectionReadRepository {
+        fn clear_count(&self) -> usize {
+            self.clear_count.load(Ordering::SeqCst)
+        }
+    }
+
+    #[async_trait]
+    impl MinimalCandidateSelectionReadRepository for ClearCountingCandidateSelectionReadRepository {
+        fn clear_local_cache(&self) {
+            self.clear_count.fetch_add(1, Ordering::SeqCst);
+        }
+
+        async fn list_for_exact_api_format(
+            &self,
+            _api_format: &str,
+        ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+            Ok(Vec::new())
+        }
+
+        async fn list_for_exact_api_format_and_global_model(
+            &self,
+            _api_format: &str,
+            _global_model_name: &str,
+        ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+            Ok(Vec::new())
+        }
+
+        async fn list_for_exact_api_format_and_requested_model(
+            &self,
+            _api_format: &str,
+            _requested_model_name: &str,
+        ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+            Ok(Vec::new())
+        }
+
+        async fn list_for_exact_api_format_and_requested_model_page(
+            &self,
+            _query: &StoredRequestedModelCandidateRowsQuery,
+        ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+            Ok(Vec::new())
+        }
+
+        async fn list_pool_key_rows_for_group(
+            &self,
+            _query: &StoredPoolKeyCandidateRowsQuery,
+        ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+            Ok(Vec::new())
+        }
+
+        async fn list_pool_key_rows_for_group_key_ids(
+            &self,
+            _query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
+        ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+            Ok(Vec::new())
+        }
+    }
+
+    #[tokio::test]
+    async fn admin_model_writes_invalidate_candidate_selection_cache() {
+        let candidate_repository =
+            Arc::new(ClearCountingCandidateSelectionReadRepository::default());
+        let global_model_repository = Arc::new(
+            InMemoryGlobalModelReadRepository::seed(Vec::new())
+                .with_admin_global_models(vec![sample_admin_global_model()]),
+        );
+        let state = AppState::new()
+            .expect("app state should build")
+            .with_data_state_for_tests(
+                GatewayDataState::with_minimal_candidate_selection_reader_for_tests(
+                    candidate_repository.clone(),
+                )
+                .with_global_model_repository_for_tests(global_model_repository),
+            );
+
+        assert_eq!(candidate_repository.clear_count(), 0);
+
+        let provider_model = sample_provider_model_record("model-1", "global-1", true);
+        state
+            .create_admin_provider_model(&provider_model)
+            .await
+            .expect("provider model create should succeed")
+            .expect("provider model should create");
+        assert_eq!(candidate_repository.clear_count(), 1);
+
+        let disabled_provider_model = sample_provider_model_record("model-1", "global-1", false);
+        state
+            .update_admin_provider_model(&disabled_provider_model)
+            .await
+            .expect("provider model update should succeed")
+            .expect("provider model should update");
+        assert_eq!(candidate_repository.clear_count(), 2);
+
+        assert!(state
+            .delete_admin_provider_model("provider-1", "model-1")
+            .await
+            .expect("provider model delete should succeed"));
+        assert_eq!(candidate_repository.clear_count(), 3);
+
+        let created_global_model = CreateAdminGlobalModelRecord::new(
+            "global-2".to_string(),
+            "gpt-4.1".to_string(),
+            "GPT 4.1".to_string(),
+            true,
+            None,
+            None,
+            None,
+            None,
+        )
+        .expect("global model create record should build");
+        state
+            .create_admin_global_model(&created_global_model)
+            .await
+            .expect("global model create should succeed")
+            .expect("global model should create");
+        assert_eq!(candidate_repository.clear_count(), 4);
+
+        let disabled_global_model = UpdateAdminGlobalModelRecord::new(
+            "global-1".to_string(),
+            "GPT 5".to_string(),
+            false,
+            None,
+            None,
+            None,
+            None,
+        )
+        .expect("global model update record should build");
+        state
+            .update_admin_global_model(&disabled_global_model)
+            .await
+            .expect("global model update should succeed")
+            .expect("global model should update");
+        assert_eq!(candidate_repository.clear_count(), 5);
+
+        assert!(state
+            .delete_admin_global_model("global-2")
+            .await
+            .expect("global model delete should succeed"));
+        assert_eq!(candidate_repository.clear_count(), 6);
+    }
+
     #[tokio::test]
     async fn provider_catalog_update_invalidates_scheduler_affinity_and_transport_snapshot_cache() {
         let provider = sample_provider();

apps/aether-gateway/src/tests/frontdoor/ai.rs

@@ -10,15 +10,21 @@ use crate::tests::{
     to_bytes, AppState, Arc, Body, Json, Mutex, Request, Router, StatusCode, EXECUTION_PATH_HEADER,
     EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
 };
+use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
 use aether_data::DataLayerError;
 use aether_data_contracts::repository::candidate_selection::{
     MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
     StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
     StoredRequestedModelCandidateRowsQuery,
 };
+use aether_data_contracts::repository::global_models::{
+    StoredAdminGlobalModel, UpdateAdminGlobalModelRecord,
+};
 use async_trait::async_trait;
 use axum::response::IntoResponse;
+use std::collections::HashMap;
 use std::future::pending;
+use std::sync::atomic::{AtomicBool, Ordering};
 
 fn gemini_operation_status_label(status: VideoTaskStatus) -> &'static str {
     match status {
@@ -183,6 +189,120 @@ impl MinimalCandidateSelectionReadRepository for PendingMinimalCandidateSelectio
     }
 }
 
+struct CachedToggleMinimalCandidateSelectionReadRepository {
+    row: StoredMinimalCandidateSelectionRow,
+    active: AtomicBool,
+    cached_rows_by_api_format: Mutex<HashMap<String, Vec<StoredMinimalCandidateSelectionRow>>>,
+}
+
+impl CachedToggleMinimalCandidateSelectionReadRepository {
+    fn new(row: StoredMinimalCandidateSelectionRow) -> Self {
+        Self {
+            row,
+            active: AtomicBool::new(true),
+            cached_rows_by_api_format: Mutex::new(HashMap::new()),
+        }
+    }
+
+    fn set_active(&self, active: bool) {
+        self.active.store(active, Ordering::SeqCst);
+    }
+
+    fn rows_for_api_format(&self, api_format: &str) -> Vec<StoredMinimalCandidateSelectionRow> {
+        let api_format = api_format.trim().to_string();
+        let mut cached = self
+            .cached_rows_by_api_format
+            .lock()
+            .expect("candidate row cache lock");
+        if let Some(rows) = cached.get(&api_format) {
+            return rows.clone();
+        }
+
+        let rows = if self.active.load(Ordering::SeqCst)
+            && self
+                .row
+                .endpoint_api_format
+                .eq_ignore_ascii_case(&api_format)
+        {
+            vec![self.row.clone()]
+        } else {
+            Vec::new()
+        };
+        cached.insert(api_format, rows.clone());
+        rows
+    }
+}
+
+#[async_trait]
+impl MinimalCandidateSelectionReadRepository
+    for CachedToggleMinimalCandidateSelectionReadRepository
+{
+    fn clear_local_cache(&self) {
+        self.cached_rows_by_api_format
+            .lock()
+            .expect("candidate row cache lock")
+            .clear();
+    }
+
+    async fn list_for_exact_api_format(
+        &self,
+        api_format: &str,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        Ok(self.rows_for_api_format(api_format))
+    }
+
+    async fn list_for_exact_api_format_and_global_model(
+        &self,
+        api_format: &str,
+        global_model_name: &str,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        Ok(self
+            .rows_for_api_format(api_format)
+            .into_iter()
+            .filter(|row| row.global_model_name == global_model_name)
+            .collect())
+    }
+
+    async fn list_for_exact_api_format_and_requested_model(
+        &self,
+        api_format: &str,
+        requested_model_name: &str,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        Ok(self
+            .rows_for_api_format(api_format)
+            .into_iter()
+            .filter(|row| row.global_model_name == requested_model_name)
+            .collect())
+    }
+
+    async fn list_for_exact_api_format_and_requested_model_page(
+        &self,
+        query: &StoredRequestedModelCandidateRowsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        Ok(self
+            .rows_for_api_format(&query.api_format)
+            .into_iter()
+            .filter(|row| row.global_model_name == query.requested_model_name)
+            .skip(query.offset as usize)
+            .take(query.limit as usize)
+            .collect())
+    }
+
+    async fn list_pool_key_rows_for_group(
+        &self,
+        _query: &StoredPoolKeyCandidateRowsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        Ok(Vec::new())
+    }
+
+    async fn list_pool_key_rows_for_group_key_ids(
+        &self,
+        _query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        Ok(Vec::new())
+    }
+}
+
 #[tokio::test]
 async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
     let fallback_probe_hits = Arc::new(Mutex::new(0usize));
@@ -240,6 +360,102 @@ async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
     fallback_probe_handle.abort();
 }
 
+#[tokio::test]
+async fn gateway_openai_models_list_drops_disabled_global_model_after_cache_invalidation() {
+    let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
+        Some(hash_api_key("sk-openai-models-cache")),
+        unrestricted_models_snapshot("key-models-cache", "user-models-cache"),
+    )]));
+    let row = sample_models_candidate_row(
+        "provider-openai-cache",
+        "openai",
+        "openai:chat",
+        "gpt-5",
+        10,
+    );
+    let global_model_id = row.global_model_id.clone();
+    let candidate_repository = Arc::new(CachedToggleMinimalCandidateSelectionReadRepository::new(
+        row.clone(),
+    ));
+    let global_model_repository = Arc::new(
+        InMemoryGlobalModelReadRepository::seed(Vec::new()).with_admin_global_models(vec![
+            StoredAdminGlobalModel::new(
+                global_model_id.clone(),
+                row.global_model_name.clone(),
+                "GPT 5".to_string(),
+                true,
+                None,
+                None,
+                None,
+                None,
+                0,
+                1,
+                0,
+                Some(1_711_000_000),
+                Some(1_711_000_000),
+            )
+            .expect("global model should build"),
+        ]),
+    );
+    let state = AppState::new()
+        .expect("gateway should build")
+        .with_data_state_for_tests(
+            crate::data::GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
+                candidate_repository.clone(),
+                auth_repository,
+            )
+            .with_global_model_repository_for_tests(global_model_repository),
+        );
+    let gateway = build_router_with_state(state.clone());
+    let (gateway_url, gateway_handle) = start_server(gateway).await;
+    let client = reqwest::Client::new();
+
+    let response = client
+        .get(format!("{gateway_url}/v1/models"))
+        .header("authorization", "Bearer sk-openai-models-cache")
+        .send()
+        .await
+        .expect("initial models request should succeed");
+    assert_eq!(response.status(), StatusCode::OK);
+    let payload: serde_json::Value = response.json().await.expect("json body should parse");
+    assert_eq!(payload["data"][0]["id"], "gpt-5");
+
+    candidate_repository.set_active(false);
+    let disabled_global_model = UpdateAdminGlobalModelRecord::new(
+        global_model_id,
+        "GPT 5".to_string(),
+        false,
+        None,
+        None,
+        None,
+        None,
+    )
+    .expect("global model update record should build");
+    state
+        .update_admin_global_model(&disabled_global_model)
+        .await
+        .expect("global model update should succeed")
+        .expect("global model should update");
+
+    let response = client
+        .get(format!("{gateway_url}/v1/models"))
+        .header("authorization", "Bearer sk-openai-models-cache")
+        .send()
+        .await
+        .expect("models request after disable should succeed");
+    assert_eq!(response.status(), StatusCode::OK);
+    let payload: serde_json::Value = response.json().await.expect("json body should parse");
+    assert_eq!(
+        payload["data"]
+            .as_array()
+            .expect("data should be an array")
+            .len(),
+        0
+    );
+
+    gateway_handle.abort();
+}
+
 #[tokio::test]
 async fn gateway_returns_empty_openai_models_when_candidate_rows_stall() {
     let fallback_probe_hits = Arc::new(Mutex::new(0usize));