i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-05 00:47:48 +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));
};