Merge pull request #552 from RWDai/issue-549-fix

Bound models route data reads

apps/aether-gateway/src/handlers/public/support/models/route.rs

@@ -1,7 +1,11 @@
-use std::time::{SystemTime, UNIX_EPOCH};
+use std::fmt::Debug;
+use std::future::Future;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
 use axum::{body::Body, response::Response};
+use tokio::time::timeout;
+use tracing::warn;
 
 use super::models_responses::{
     build_claude_model_detail_response, build_claude_models_list_response,
@@ -15,6 +19,59 @@ use super::models_shared::{
 };
 use super::{query_param_value, AppState, GatewayPublicRequestContext};
 
+#[cfg(not(test))]
+const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_secs(5);
+#[cfg(test)]
+const MODELS_ROUTE_READ_TIMEOUT: Duration = Duration::from_millis(50);
+
+async fn await_models_route_read<T, E, Fut>(operation: &'static str, future: Fut) -> Option<T>
+where
+    E: Debug,
+    Fut: Future<Output = Result<T, E>>,
+{
+    match timeout(MODELS_ROUTE_READ_TIMEOUT, future).await {
+        Ok(Ok(value)) => Some(value),
+        Ok(Err(error)) => {
+            warn!(
+                event_name = "models_route_read_error",
+                log_type = "ops",
+                operation,
+                error = ?error,
+                "gateway local models route read failed"
+            );
+            None
+        }
+        Err(_) => {
+            warn!(
+                event_name = "models_route_read_timeout",
+                log_type = "ops",
+                operation,
+                timeout_ms = MODELS_ROUTE_READ_TIMEOUT.as_millis() as u64,
+                "gateway local models route read timed out"
+            );
+            None
+        }
+    }
+}
+
+fn build_models_read_fallback_response(
+    request_context: &GatewayPublicRequestContext,
+    api_format: &str,
+) -> Response<Body> {
+    let route_kind = request_context
+        .control_decision
+        .as_ref()
+        .and_then(|decision| decision.route_kind.as_deref());
+    match route_kind {
+        Some("detail") => {
+            let model_id = models_detail_id(&request_context.request_path)
+                .unwrap_or_else(|| "unknown".to_string());
+            build_models_not_found_response(&model_id, api_format)
+        }
+        _ => build_empty_models_list_response(api_format),
+    }
+}
+
 fn sort_and_dedup_model_rows(
     mut rows: Vec<StoredMinimalCandidateSelectionRow>,
 ) -> Vec<StoredMinimalCandidateSelectionRow> {
@@ -47,10 +104,11 @@ async fn list_model_rows_for_client_format(
 ) -> Option<Vec<StoredMinimalCandidateSelectionRow>> {
     let mut collected = Vec::new();
     for query_format in models_query_api_formats(api_format) {
-        let rows = state
+        let rows = await_models_route_read(
-            .list_minimal_candidate_selection_rows_for_api_format(query_format)
+            "candidate_selection_by_api_format",
-            .await
+            state.list_minimal_candidate_selection_rows_for_api_format(query_format),
-            .ok()?;
+        )
+        .await?;
         let mut filtered = filter_rows_for_models(rows, auth_snapshot, query_format);
         collected.append(&mut filtered);
     }
@@ -65,13 +123,14 @@ async fn list_model_rows_for_client_format_and_global_model(
 ) -> Option<Vec<StoredMinimalCandidateSelectionRow>> {
     let mut collected = Vec::new();
     for query_format in models_query_api_formats(api_format) {
-        let rows = state
+        let rows = await_models_route_read(
-            .list_minimal_candidate_selection_rows_for_api_format_and_global_model(
+            "candidate_selection_by_global_model",
+            state.list_minimal_candidate_selection_rows_for_api_format_and_global_model(
                 query_format,
                 global_model_name,
-            )
+            ),
-            .await
+        )
-            .ok()?;
+        .await?;
         let mut filtered = filter_rows_for_models(rows, auth_snapshot, query_format);
         collected.append(&mut filtered);
     }
@@ -96,21 +155,38 @@ pub(super) async fn maybe_build_local_models_route_response(
         .duration_since(UNIX_EPOCH)
         .unwrap_or_default()
         .as_secs();
-    let auth_snapshot = state
+    let auth_snapshot = match await_models_route_read(
-        .data
+        "auth_api_key_snapshot",
-        .read_auth_api_key_snapshot(
+        state.data.read_auth_api_key_snapshot(
             &auth_context.user_id,
             &auth_context.api_key_id,
             now_unix_secs,
-        )
+        ),
-        .await
+    )
-        .ok()
+    .await
-        .flatten();
+    {
+        Some(snapshot) => snapshot,
+        None => {
+            return Some(build_models_read_fallback_response(
+                request_context,
+                api_format,
+            ))
+        }
+    };
     let auth_snapshot = auth_snapshot.as_ref();
 
     match decision.route_kind.as_deref() {
         Some("list") => {
-            let rows = list_model_rows_for_client_format(state, api_format, auth_snapshot).await?;
+            let rows =
+                match list_model_rows_for_client_format(state, api_format, auth_snapshot).await {
+                    Some(rows) => rows,
+                    None => {
+                        return Some(build_models_read_fallback_response(
+                            request_context,
+                            api_format,
+                        ))
+                    }
+                };
             if rows.is_empty() {
                 return Some(build_empty_models_list_response(api_format));
             }
@@ -156,13 +232,22 @@ pub(super) async fn maybe_build_local_models_route_response(
         }
         Some("detail") => {
             let model_id = models_detail_id(&request_context.request_path)?;
-            let rows = list_model_rows_for_client_format_and_global_model(
+            let rows = match list_model_rows_for_client_format_and_global_model(
                 state,
                 api_format,
                 &model_id,
                 auth_snapshot,
             )
-            .await?;
+            .await
+            {
+                Some(rows) => rows,
+                None => {
+                    return Some(build_models_read_fallback_response(
+                        request_context,
+                        api_format,
+                    ))
+                }
+            };
             let Some(row) = rows.first() else {
                 return Some(build_models_not_found_response(&model_id, api_format));
             };

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

@@ -10,7 +10,15 @@ 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::DataLayerError;
+use aether_data_contracts::repository::candidate_selection::{
+    MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
+    StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
+    StoredRequestedModelCandidateRowsQuery,
+};
+use async_trait::async_trait;
 use axum::response::IntoResponse;
+use std::future::pending;
 
 fn gemini_operation_status_label(status: VideoTaskStatus) -> &'static str {
     match status {
@@ -118,6 +126,63 @@ fn sample_gemini_video_task(
     }
 }
 
+struct PendingMinimalCandidateSelectionReadRepository;
+
+impl PendingMinimalCandidateSelectionReadRepository {
+    async fn pending_rows(
+        &self,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        pending::<Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>>().await
+    }
+}
+
+#[async_trait]
+impl MinimalCandidateSelectionReadRepository for PendingMinimalCandidateSelectionReadRepository {
+    async fn list_for_exact_api_format(
+        &self,
+        _api_format: &str,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        self.pending_rows().await
+    }
+
+    async fn list_for_exact_api_format_and_global_model(
+        &self,
+        _api_format: &str,
+        _global_model_name: &str,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        self.pending_rows().await
+    }
+
+    async fn list_for_exact_api_format_and_requested_model(
+        &self,
+        _api_format: &str,
+        _requested_model_name: &str,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        self.pending_rows().await
+    }
+
+    async fn list_for_exact_api_format_and_requested_model_page(
+        &self,
+        _query: &StoredRequestedModelCandidateRowsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        self.pending_rows().await
+    }
+
+    async fn list_pool_key_rows_for_group(
+        &self,
+        _query: &StoredPoolKeyCandidateRowsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        self.pending_rows().await
+    }
+
+    async fn list_pool_key_rows_for_group_key_ids(
+        &self,
+        _query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
+    ) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
+        self.pending_rows().await
+    }
+}
+
 #[tokio::test]
 async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
     let fallback_probe_hits = Arc::new(Mutex::new(0usize));
@@ -175,6 +240,119 @@ async fn gateway_handles_public_openai_models_without_hitting_fallback_probe() {
     fallback_probe_handle.abort();
 }
 
+#[tokio::test]
+async fn gateway_returns_empty_openai_models_when_candidate_rows_stall() {
+    let fallback_probe_hits = Arc::new(Mutex::new(0usize));
+    let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
+    let fallback_probe = Router::new().route(
+        "/{*path}",
+        any(move |_request: Request| {
+            let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
+            async move {
+                *fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
+                (StatusCode::OK, Body::from("proxied"))
+            }
+        }),
+    );
+
+    let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
+        Some(hash_api_key("sk-openai-models-stalled")),
+        unrestricted_models_snapshot("key-stalled", "user-stalled"),
+    )]));
+    let candidate_repository = Arc::new(PendingMinimalCandidateSelectionReadRepository);
+
+    let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
+    let gateway = build_router_with_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,
+                    auth_repository,
+                ),
+            ),
+    );
+    let (gateway_url, gateway_handle) = start_server(gateway).await;
+
+    let response = reqwest::Client::builder()
+        .timeout(std::time::Duration::from_millis(500))
+        .build()
+        .expect("client should build")
+        .get(format!("{gateway_url}/v1/models"))
+        .header("authorization", "Bearer sk-openai-models-stalled")
+        .send()
+        .await
+        .expect("request should return before client timeout");
+
+    assert_eq!(response.status(), StatusCode::OK);
+    let payload: serde_json::Value = response.json().await.expect("json body should parse");
+    assert_eq!(payload["object"], "list");
+    assert_eq!(
+        payload["data"]
+            .as_array()
+            .expect("data should be an array")
+            .len(),
+        0
+    );
+    assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
+
+    gateway_handle.abort();
+    fallback_probe_handle.abort();
+}
+
+#[tokio::test]
+async fn gateway_returns_not_found_for_openai_model_detail_when_candidate_rows_stall() {
+    let fallback_probe_hits = Arc::new(Mutex::new(0usize));
+    let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits);
+    let fallback_probe = Router::new().route(
+        "/{*path}",
+        any(move |_request: Request| {
+            let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone);
+            async move {
+                *fallback_probe_hits_inner.lock().expect("mutex should lock") += 1;
+                (StatusCode::OK, Body::from("proxied"))
+            }
+        }),
+    );
+
+    let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
+        Some(hash_api_key("sk-openai-model-detail-stalled")),
+        unrestricted_models_snapshot("key-detail-stalled", "user-detail-stalled"),
+    )]));
+    let candidate_repository = Arc::new(PendingMinimalCandidateSelectionReadRepository);
+
+    let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await;
+    let gateway = build_router_with_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,
+                    auth_repository,
+                ),
+            ),
+    );
+    let (gateway_url, gateway_handle) = start_server(gateway).await;
+
+    let response = reqwest::Client::builder()
+        .timeout(std::time::Duration::from_millis(500))
+        .build()
+        .expect("client should build")
+        .get(format!("{gateway_url}/v1/models/gpt-stalled"))
+        .header("authorization", "Bearer sk-openai-model-detail-stalled")
+        .send()
+        .await
+        .expect("request should return before client timeout");
+
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
+    let payload: serde_json::Value = response.json().await.expect("json body should parse");
+    assert_eq!(payload["error"]["code"], "model_not_found");
+    assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0);
+
+    gateway_handle.abort();
+    fallback_probe_handle.abort();
+}
+
 #[tokio::test]
 async fn gateway_handles_public_openai_models_with_cross_format_candidates_without_hitting_fallback_probe(
 ) {