Merge branch 'fawney19:main' into main

apps/aether-gateway/src/data/candidate_selection.rs

@@ -1,7 +1,7 @@
 use aether_data::DataLayerError;
 use aether_data_contracts::repository::candidate_selection::{
-    StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsQuery,
-    StoredRequestedModelCandidateRowsQuery,
+    StoredApiFormatCandidateRowsQuery, StoredMinimalCandidateSelectionRow,
+    StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
 };
 use aether_scheduler_core::{
     auth_constraints_allow_api_format, collect_global_model_names_for_required_capability,
@@ -39,6 +39,19 @@ pub(crate) trait MinimalCandidateSelectionRowSource {
         api_format: &str,
     ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>;
 
+    async fn read_minimal_candidate_selection_rows_for_api_format_page(
+        &self,
+        query: &StoredApiFormatCandidateRowsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        Ok(self
+            .read_minimal_candidate_selection_rows_for_api_format(&query.api_format)
+            .await?
+            .into_iter()
+            .skip(query.offset as usize)
+            .take(query.limit as usize)
+            .collect())
+    }
+
     async fn read_pool_key_candidate_rows_for_group(
         &self,
         query: &StoredPoolKeyCandidateRowsQuery,
@@ -231,6 +244,30 @@ pub(crate) async fn read_requested_model_rows_fast_path_page(
     })
 }
 
+pub(crate) async fn read_api_format_rows_fallback_page(
+    state: &(impl MinimalCandidateSelectionRowSource + Sync),
+    api_format: &str,
+    offset: u32,
+    limit: u32,
+) -> Result<RequestedModelCandidateRowsPage, DataLayerError> {
+    let limit = limit.max(1);
+    let rows = state
+        .read_minimal_candidate_selection_rows_for_api_format_page(
+            &StoredApiFormatCandidateRowsQuery {
+                api_format: api_format.to_string(),
+                offset,
+                limit,
+            },
+        )
+        .await?;
+    let scanned_rows = rows.len() as u32;
+    Ok(RequestedModelCandidateRowsPage {
+        rows,
+        scanned_rows,
+        end_of_requested_name: scanned_rows < limit,
+    })
+}
+
 pub(crate) async fn enumerate_minimal_candidate_selection_with_required_capabilities(
     state: &(impl MinimalCandidateSelectionRowSource + Sync),
     api_format: &str,

apps/aether-gateway/src/data/state/candidate_cache.rs

@@ -7,9 +7,10 @@ use std::time::Duration;
 use aether_cache::ExpiringMap;
 use aether_data::DataLayerError;
 use aether_data_contracts::repository::candidate_selection::{
-    MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
-    StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
-    StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
+    MinimalCandidateSelectionReadRepository, StoredApiFormatCandidateRowsQuery,
+    StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
+    StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
+    StoredRequestedModelCandidateRowsQuery,
 };
 use async_trait::async_trait;
 use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore};
@@ -389,6 +390,19 @@ impl MinimalCandidateSelectionReadRepository for CachedMinimalCandidateSelection
             .await
     }
 
+    async fn list_for_exact_api_format_page(
+        &self,
+        query: &StoredApiFormatCandidateRowsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        let key = CandidateSelectionCacheKey::ApiFormatPage {
+            api_format: normalize_api_format_key(&query.api_format),
+            offset: query.offset,
+            limit: query.limit,
+        };
+        self.get_or_load(key, || self.inner.list_for_exact_api_format_page(query))
+            .await
+    }
+
     async fn list_for_exact_api_format_and_global_model(
         &self,
         api_format: &str,
@@ -480,6 +494,11 @@ enum CandidateSelectionCacheKey {
     ApiFormat {
         api_format: String,
     },
+    ApiFormatPage {
+        api_format: String,
+        offset: u32,
+        limit: u32,
+    },
     ApiFormatAndGlobalModel {
         api_format: String,
         global_model_name: String,
@@ -1044,6 +1063,36 @@ mod tests {
         assert!(cache.inflight.lock().unwrap().is_empty());
     }
 
+    #[tokio::test]
+    async fn candidate_selection_cache_keys_api_format_pages_by_offset() {
+        let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
+        let cache = CachedMinimalCandidateSelectionReadRepository::new(inner.clone());
+        let first = StoredApiFormatCandidateRowsQuery {
+            api_format: "openai:chat".to_string(),
+            offset: 0,
+            limit: 256,
+        };
+        let second = StoredApiFormatCandidateRowsQuery {
+            offset: 256,
+            ..first.clone()
+        };
+
+        cache
+            .list_for_exact_api_format_page(&first)
+            .await
+            .expect("first page should load");
+        cache
+            .list_for_exact_api_format_page(&first)
+            .await
+            .expect("first page should be cached");
+        cache
+            .list_for_exact_api_format_page(&second)
+            .await
+            .expect("second page should load independently");
+
+        assert_eq!(inner.calls(), 2);
+    }
+
     #[tokio::test]
     async fn candidate_selection_cache_releases_inflight_when_leader_is_cancelled() {
         let inner = Arc::new(FirstLoadPendingThenFastRepository::new());

apps/aether-gateway/src/data/state/integrations.rs

@@ -176,6 +176,14 @@ impl MinimalCandidateSelectionRowSource for GatewayDataState {
             .await
     }
 
+    async fn read_minimal_candidate_selection_rows_for_api_format_page(
+        &self,
+        query: &aether_data_contracts::repository::candidate_selection::StoredApiFormatCandidateRowsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        self.list_minimal_candidate_selection_rows_for_api_format_page(query)
+            .await
+    }
+
     async fn read_pool_key_candidate_rows_for_group(
         &self,
         query: &aether_data_contracts::repository::candidate_selection::StoredPoolKeyCandidateRowsQuery,

apps/aether-gateway/src/data/state/mod.rs

@@ -97,9 +97,9 @@ use aether_data_contracts::repository::billing::{
     UserPlanEntitlementRecord,
 };
 use aether_data_contracts::repository::candidate_selection::{
-    MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
-    StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
-    StoredRequestedModelCandidateRowsQuery,
+    MinimalCandidateSelectionReadRepository, StoredApiFormatCandidateRowsQuery,
+    StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsByKeyIdsQuery,
+    StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
 };
 use aether_data_contracts::repository::candidates::{
     PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,

apps/aether-gateway/src/data/state/models.rs

@@ -2,11 +2,12 @@ use super::{
     AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
     DataLayerError, GatewayDataState, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
     PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage,
-    StoredAdminProviderModel, StoredMinimalCandidateSelectionRow,
-    StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
-    StoredProviderActiveGlobalModel, StoredProviderModelStats, StoredPublicCatalogModel,
-    StoredPublicGlobalModel, StoredPublicGlobalModelPage, StoredRequestedModelCandidateRowsQuery,
-    UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
+    StoredAdminProviderModel, StoredApiFormatCandidateRowsQuery,
+    StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsByKeyIdsQuery,
+    StoredPoolKeyCandidateRowsQuery, StoredProviderActiveGlobalModel, StoredProviderModelStats,
+    StoredPublicCatalogModel, StoredPublicGlobalModel, StoredPublicGlobalModelPage,
+    StoredRequestedModelCandidateRowsQuery, UpdateAdminGlobalModelRecord,
+    UpsertAdminProviderModelRecord,
 };
 
 impl GatewayDataState {
@@ -98,6 +99,23 @@ impl GatewayDataState {
         .await
     }
 
+    pub(crate) async fn list_minimal_candidate_selection_rows_for_api_format_page(
+        &self,
+        query: &StoredApiFormatCandidateRowsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        crate::request_diagnostics::observe_db_operation(
+            "candidate_selection",
+            self.database_pool_summary(),
+            async {
+                match &self.minimal_candidate_selection_reader {
+                    Some(repository) => repository.list_for_exact_api_format_page(query).await,
+                    None => Ok(Vec::new()),
+                }
+            },
+        )
+        .await
+    }
+
     pub(crate) async fn list_pool_key_candidate_rows_for_group(
         &self,
         query: &StoredPoolKeyCandidateRowsQuery,