i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-10 21:20:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

Bound models route data reads

Browse Source
This commit is contained in:
RWDai
2026-05-23 19:36:20 +08:00
parent 6447fda852
commit 0655e868a2
2 changed files with 283 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/public/support/models/route.rs
+105 -20
View File
@@ -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
.list_minimal_candidate_selection_rows_for_api_format(query_format)
.await
.ok()?;
let rows = await_models_route_read(
"candidate_selection_by_api_format",
state.list_minimal_candidate_selection_rows_for_api_format(query_format),
)
.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
.list_minimal_candidate_selection_rows_for_api_format_and_global_model(
let rows = await_models_route_read(
"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
.data
.read_auth_api_key_snapshot(
let auth_snapshot = match await_models_route_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()
.flatten();
),
)
.await
{
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
+178
View File
@@ -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(
) {