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

refactor: 大规模模块拆分与代码精简，新增 ai-pipeline/data-contracts 独立 crate

Browse Source 
- 新增 aether-ai-pipeline 和 aether-data-contracts crate，将 pipeline 逻辑与数据契约从 gateway 中解耦
- 重构 admin handlers：拆分单体模块为 auth/billing/endpoint/features/model/observability/provider/system 等独立子模块
- 合并 chat/cli 重复代码路径：精简 conversion、finalize、planner 中的 sync/chat/cli 分支
- 重构 scheduler/executor/data 层，引入 facade 模式降低模块间耦合
- 移除冗余的 intent 模块，将 plan_fallback/policy/stream_path/sync_path 迁移至 executor
- 前端适配：调整 admin API 调用和 provider 模型测试对话框
This commit is contained in:
fawney19
2026-04-07 02:50:19 +08:00
parent 763ff03a7b
commit 5d96d6673b
732 changed files with 28593 additions and 20666 deletions
Download Patch File Download Diff File Expand all files Collapse all files

143
apps/aether-gateway/src/data/candidate_selection.rs Normal file
View File

@@ -0,0 +1,143 @@
use aether_data::DataLayerError;
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_scheduler_core::{
auth_constraints_allow_api_format, build_minimal_candidate_selection,
collect_global_model_names_for_required_capability, normalize_api_format,
resolve_requested_global_model_name, SchedulerAuthConstraints,
SchedulerMinimalCandidateSelectionCandidate,
};
use async_trait::async_trait;
use super::auth::GatewayAuthApiKeySnapshot;
#[async_trait]
pub(crate) trait MinimalCandidateSelectionRowSource {
async fn read_minimal_candidate_selection_rows_for_api_format_and_global_model(
&self,
api_format: &str,
global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>;
async fn read_minimal_candidate_selection_rows_for_api_format(
&self,
api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>;
}
pub(crate) async fn read_requested_model_rows(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
) -> Result<Option<(String, Vec<StoredMinimalCandidateSelectionRow>)>, DataLayerError> {
let exact_rows = state
.read_minimal_candidate_selection_rows_for_api_format_and_global_model(
api_format,
requested_model_name,
)
.await?;
if !exact_rows.is_empty() {
return Ok(Some((requested_model_name.to_string(), exact_rows)));
}
let rows = state
.read_minimal_candidate_selection_rows_for_api_format(api_format)
.await?;
let Some(resolved_global_model_name) =
resolve_requested_global_model_name(&rows, requested_model_name, api_format)
else {
return Ok(None);
};
Ok(Some((
resolved_global_model_name.clone(),
rows.into_iter()
.filter(|row| row.global_model_name == resolved_global_model_name)
.collect(),
)))
}
pub(crate) async fn read_minimal_candidate_selection(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format);
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !auth_constraints_allow_api_format(
auth_snapshot.map(auth_snapshot_constraints).as_ref(),
&normalized_api_format,
) {
return Ok(Vec::new());
}
let Some((resolved_global_model_name, rows)) =
read_requested_model_rows(state, &normalized_api_format, requested_model_name).await?
else {
return Ok(Vec::new());
};
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
let affinity_key = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
.filter(|value| !value.is_empty());
build_minimal_candidate_selection(
rows,
&normalized_api_format,
requested_model_name,
resolved_global_model_name.as_str(),
require_streaming,
auth_constraints.as_ref(),
affinity_key,
)
}
pub(crate) async fn read_global_model_names_for_required_capability(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
required_capability: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<String>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format);
let required_capability = required_capability.trim();
if normalized_api_format.is_empty() || required_capability.is_empty() {
return Ok(Vec::new());
}
if !auth_constraints_allow_api_format(
auth_snapshot.map(auth_snapshot_constraints).as_ref(),
&normalized_api_format,
) {
return Ok(Vec::new());
}
let rows = state
.read_minimal_candidate_selection_rows_for_api_format(&normalized_api_format)
.await?;
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
Ok(collect_global_model_names_for_required_capability(
rows,
&normalized_api_format,
required_capability,
require_streaming,
auth_constraints.as_ref(),
))
}
fn auth_snapshot_constraints(snapshot: &GatewayAuthApiKeySnapshot) -> SchedulerAuthConstraints {
SchedulerAuthConstraints {
allowed_providers: snapshot
.effective_allowed_providers()
.map(|items| items.to_vec()),
allowed_api_formats: snapshot
.effective_allowed_api_formats()
.map(|items| items.to_vec()),
allowed_models: snapshot
.effective_allowed_models()
.map(|items| items.to_vec()),
}
}

10
apps/aether-gateway/src/data/candidates.rs
View File

@@ -1,8 +1,8 @@
use aether_data::DataLayerError;
use aether_data_contracts::DataLayerError;
use super::state::GatewayDataState;
pub(crate) use aether_data::repository::candidates::{
pub(crate) use aether_data_contracts::repository::candidates::{
RequestCandidateFinalStatus, RequestCandidateTrace,
};
@@ -25,9 +25,9 @@ pub(crate) async fn read_request_candidate_trace(
mod tests {
use super::super::GatewayDataState;
use super::{read_request_candidate_trace, RequestCandidateFinalStatus};
use aether_data::repository::candidates::{
derive_request_candidate_final_status, InMemoryRequestCandidateRepository,
RequestCandidateStatus, StoredRequestCandidate,
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data_contracts::repository::candidates::{
derive_request_candidate_final_status, RequestCandidateStatus, StoredRequestCandidate,
};
use std::sync::Arc;

19
apps/aether-gateway/src/data/decision_trace.rs
View File

@@ -1,11 +1,13 @@
use std::collections::BTreeSet;
use aether_data::repository::candidates::build_decision_trace;
use aether_data::DataLayerError;
use aether_data_contracts::repository::candidates::build_decision_trace;
use aether_data_contracts::DataLayerError;
use super::state::GatewayDataState;
pub(crate) use aether_data::repository::candidates::{DecisionTrace, DecisionTraceCandidate};
pub(crate) use aether_data_contracts::repository::candidates::{
DecisionTrace, DecisionTraceCandidate,
};
pub(crate) async fn read_decision_trace(
state: &GatewayDataState,
@@ -63,12 +65,13 @@ fn unique_ids<'a>(items: impl Iterator<Item = &'a String>) -> Vec<String> {
mod tests {
use std::sync::Arc;
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogProvider,
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use super::{read_decision_trace, DecisionTrace, DecisionTraceCandidate};

1
apps/aether-gateway/src/data/mod.rs
View File

@@ -1,4 +1,5 @@
pub(crate) mod auth;
pub(crate) mod candidate_selection;
pub(crate) mod candidates;
mod config;
pub(crate) mod decision_trace;

294
apps/aether-gateway/src/data/state/auth.rs
View File

@@ -16,6 +16,27 @@ use aether_data::repository::auth::{
use sqlx::Row;
use uuid::Uuid;
fn postgres_error(error: sqlx::Error) -> DataLayerError {
DataLayerError::postgres(error)
}
trait SqlxResultExt<T> {
fn map_postgres_err(self) -> Result<T, DataLayerError>;
}
impl<T> SqlxResultExt<T> for Result<T, sqlx::Error> {
fn map_postgres_err(self) -> Result<T, DataLayerError> {
self.map_err(postgres_error)
}
}
fn row_get<T>(row: &sqlx::postgres::PgRow, column: &str) -> Result<T, DataLayerError>
where
for<'r> T: sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres>,
{
row.try_get(column).map_postgres_err()
}
const FIND_USER_SESSION_SQL: &str = r#"
SELECT
id,
@@ -500,15 +521,15 @@ SET
ELSE role
END,
allowed_providers = CASE
WHEN $4::BOOLEAN THEN $5
WHEN $4::BOOLEAN THEN $5::json
ELSE allowed_providers
END,
allowed_api_formats = CASE
WHEN $6::BOOLEAN THEN $7
WHEN $6::BOOLEAN THEN $7::json
ELSE allowed_api_formats
END,
allowed_models = CASE
WHEN $8::BOOLEAN THEN $9
WHEN $8::BOOLEAN THEN $9::json
ELSE allowed_models
END,
rate_limit = CASE
@@ -791,40 +812,40 @@ fn map_user_session_row(
row: &sqlx::postgres::PgRow,
) -> Result<StoredUserSessionRecord, DataLayerError> {
StoredUserSessionRecord::new(
row.try_get("id")?,
row.try_get("user_id")?,
row.try_get("client_device_id")?,
row.try_get("device_label")?,
row.try_get("refresh_token_hash")?,
row.try_get("prev_refresh_token_hash")?,
row.try_get("rotated_at")?,
row.try_get("last_seen_at")?,
row.try_get("expires_at")?,
row.try_get("revoked_at")?,
row.try_get("revoke_reason")?,
row.try_get("ip_address")?,
row.try_get("user_agent")?,
row.try_get("created_at")?,
row.try_get("updated_at")?,
row_get(row, "id")?,
row_get(row, "user_id")?,
row_get(row, "client_device_id")?,
row_get(row, "device_label")?,
row_get(row, "refresh_token_hash")?,
row_get(row, "prev_refresh_token_hash")?,
row_get(row, "rotated_at")?,
row_get(row, "last_seen_at")?,
row_get(row, "expires_at")?,
row_get(row, "revoked_at")?,
row_get(row, "revoke_reason")?,
row_get(row, "ip_address")?,
row_get(row, "user_agent")?,
row_get(row, "created_at")?,
row_get(row, "updated_at")?,
)
}
fn map_user_auth_row(row: &sqlx::postgres::PgRow) -> Result<StoredUserAuthRecord, DataLayerError> {
StoredUserAuthRecord::new(
row.try_get("id")?,
row.try_get("email")?,
row.try_get("email_verified")?,
row.try_get("username")?,
row.try_get("password_hash")?,
row.try_get("role")?,
row.try_get("auth_source")?,
row.try_get("allowed_providers")?,
row.try_get("allowed_api_formats")?,
row.try_get("allowed_models")?,
row.try_get("is_active")?,
row.try_get("is_deleted")?,
row.try_get("created_at")?,
row.try_get("last_login_at")?,
row_get(row, "id")?,
row_get(row, "email")?,
row_get(row, "email_verified")?,
row_get(row, "username")?,
row_get(row, "password_hash")?,
row_get(row, "role")?,
row_get(row, "auth_source")?,
row_get(row, "allowed_providers")?,
row_get(row, "allowed_api_formats")?,
row_get(row, "allowed_models")?,
row_get(row, "is_active")?,
row_get(row, "is_deleted")?,
row_get(row, "created_at")?,
row_get(row, "last_login_at")?,
)
}
@@ -835,7 +856,8 @@ async fn check_username_taken_in_tx(
let row = sqlx::query(CHECK_AUTH_USER_USERNAME_TAKEN_SQL)
.bind(username)
.fetch_optional(&mut **tx)
.await?;
.await
.map_postgres_err()?;
Ok(row.is_some())
}
@@ -849,7 +871,8 @@ async fn find_ldap_auth_user_for_update_in_tx(
let row = sqlx::query(FIND_LDAP_AUTH_USER_BY_DN_SQL)
.bind(ldap_dn)
.fetch_optional(&mut **tx)
.await?;
.await
.map_postgres_err()?;
if let Some(row) = row.as_ref() {
return map_user_auth_row(row).map(Some);
}
@@ -859,7 +882,8 @@ async fn find_ldap_auth_user_for_update_in_tx(
let row = sqlx::query(FIND_LDAP_AUTH_USER_BY_LDAP_USERNAME_SQL)
.bind(ldap_username)
.fetch_optional(&mut **tx)
.await?;
.await
.map_postgres_err()?;
if let Some(row) = row.as_ref() {
return map_user_auth_row(row).map(Some);
}
@@ -868,7 +892,8 @@ async fn find_ldap_auth_user_for_update_in_tx(
let row = sqlx::query(FIND_LDAP_AUTH_USER_BY_EMAIL_SQL)
.bind(email)
.fetch_optional(&mut **tx)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_user_auth_row).transpose()
}
@@ -876,26 +901,26 @@ fn map_wallet_snapshot_row(
row: &sqlx::postgres::PgRow,
) -> Result<StoredWalletSnapshot, DataLayerError> {
StoredWalletSnapshot::new(
row.try_get("id")?,
row.try_get("user_id")?,
row.try_get("api_key_id")?,
row.try_get("balance")?,
row.try_get("gift_balance")?,
row.try_get("limit_mode")?,
row.try_get("currency")?,
row.try_get("status")?,
row.try_get("total_recharged")?,
row.try_get("total_consumed")?,
row.try_get("total_refunded")?,
row.try_get("total_adjusted")?,
row.try_get("updated_at_unix_secs")?,
row_get(row, "id")?,
row_get(row, "user_id")?,
row_get(row, "api_key_id")?,
row_get(row, "balance")?,
row_get(row, "gift_balance")?,
row_get(row, "limit_mode")?,
row_get(row, "currency")?,
row_get(row, "status")?,
row_get(row, "total_recharged")?,
row_get(row, "total_consumed")?,
row_get(row, "total_refunded")?,
row_get(row, "total_adjusted")?,
row_get(row, "updated_at_unix_secs")?,
)
}
fn map_user_preference_row(
row: &sqlx::postgres::PgRow,
) -> Result<StoredUserPreferenceRecord, DataLayerError> {
let user_id = row.try_get::<String, _>("user_id")?;
let user_id = row_get::<String>(row, "user_id")?;
if user_id.trim().is_empty() {
return Err(DataLayerError::UnexpectedValue(
"user_preferences.user_id is empty".to_string(),
@@ -904,16 +929,16 @@ fn map_user_preference_row(
Ok(StoredUserPreferenceRecord {
user_id,
avatar_url: row.try_get("avatar_url")?,
bio: row.try_get("bio")?,
default_provider_id: row.try_get("default_provider_id")?,
default_provider_name: row.try_get("default_provider_name")?,
theme: row.try_get("theme")?,
language: row.try_get("language")?,
timezone: row.try_get("timezone")?,
email_notifications: row.try_get("email_notifications")?,
usage_alerts: row.try_get("usage_alerts")?,
announcement_notifications: row.try_get("announcement_notifications")?,
avatar_url: row_get(row, "avatar_url")?,
bio: row_get(row, "bio")?,
default_provider_id: row_get(row, "default_provider_id")?,
default_provider_name: row_get(row, "default_provider_name")?,
theme: row_get(row, "theme")?,
language: row_get(row, "language")?,
timezone: row_get(row, "timezone")?,
email_notifications: row_get(row, "email_notifications")?,
usage_alerts: row_get(row, "usage_alerts")?,
announcement_notifications: row_get(row, "announcement_notifications")?,
})
}
@@ -937,7 +962,8 @@ impl GatewayDataState {
.bind(email)
.bind(user_id)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(row.is_some())
}
@@ -953,7 +979,8 @@ impl GatewayDataState {
.bind(username)
.bind(user_id)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(row.is_some())
}
@@ -995,7 +1022,8 @@ impl GatewayDataState {
let row = sqlx::query(READ_USER_PREFERENCES_SQL)
.bind(user_id)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_user_preference_row).transpose()
}
@@ -1027,7 +1055,8 @@ impl GatewayDataState {
.bind(preferences.usage_alerts)
.bind(preferences.announcement_notifications)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_user_preference_row).transpose()
}
@@ -1049,7 +1078,8 @@ impl GatewayDataState {
let row = sqlx::query(FIND_ACTIVE_PROVIDER_NAME_SQL)
.bind(provider_id)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(row.and_then(|row| row.try_get("name").ok()))
}
@@ -1065,7 +1095,8 @@ impl GatewayDataState {
.bind(user_id)
.bind(session_id)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_user_session_row).transpose()
}
@@ -1079,7 +1110,8 @@ impl GatewayDataState {
let rows = sqlx::query(LIST_USER_SESSIONS_SQL)
.bind(user_id)
.fetch_all(&pool)
.await?;
.await
.map_postgres_err()?;
rows.iter().map(map_user_session_row).collect()
}
@@ -1100,7 +1132,8 @@ impl GatewayDataState {
.bind(&session.client_device_id)
.bind(now)
.execute(&pool)
.await?;
.await
.map_postgres_err()?;
let row = sqlx::query(CREATE_USER_SESSION_SQL)
.bind(&session.id)
.bind(&session.user_id)
@@ -1115,7 +1148,8 @@ impl GatewayDataState {
.bind(session.created_at.unwrap_or(now))
.bind(session.updated_at.unwrap_or(now))
.fetch_one(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(Some(map_user_session_row(&row)?))
}
@@ -1131,7 +1165,8 @@ impl GatewayDataState {
.bind(user_id)
.bind(settings)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(row
.as_ref()
.and_then(|row| row.try_get("model_capability_settings").ok())
@@ -1152,7 +1187,8 @@ impl GatewayDataState {
.bind(email)
.bind(username)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_user_auth_row).transpose()
}
@@ -1170,7 +1206,8 @@ impl GatewayDataState {
.bind(password_hash)
.bind(updated_at)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_user_auth_row).transpose()
}
@@ -1192,7 +1229,8 @@ impl GatewayDataState {
.bind(username)
.bind(password_hash)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_user_auth_row).transpose()
}
@@ -1224,7 +1262,8 @@ impl GatewayDataState {
.bind(allowed_models.map(serde_json::Value::from))
.bind(rate_limit)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_user_auth_row).transpose()
}
@@ -1260,7 +1299,8 @@ impl GatewayDataState {
.bind(is_active.is_some())
.bind(is_active)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_user_auth_row).transpose()
}
@@ -1276,7 +1316,8 @@ impl GatewayDataState {
.bind(user_id)
.bind(logged_in_at)
.execute(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(result.rows_affected() > 0)
}
@@ -1295,7 +1336,7 @@ impl GatewayDataState {
return Ok(None);
};
let mut tx = pool.begin().await?;
let mut tx = pool.begin().await.map_postgres_err()?;
let existing = find_ldap_auth_user_for_update_in_tx(
&mut tx,
ldap_dn.as_deref(),
@@ -1306,11 +1347,11 @@ impl GatewayDataState {
if let Some(existing) = existing {
if existing.is_deleted || !existing.is_active {
tx.commit().await?;
tx.commit().await.map_err(postgres_error)?;
return Ok(None);
}
if !existing.auth_source.eq_ignore_ascii_case("ldap") {
tx.commit().await?;
tx.commit().await.map_err(postgres_error)?;
return Ok(None);
}
@@ -1320,9 +1361,10 @@ impl GatewayDataState {
.bind(&email)
.bind(&existing.id)
.fetch_optional(&mut *tx)
.await?;
.await
.map_postgres_err()?;
if taken.is_some() {
tx.commit().await?;
tx.commit().await.map_err(postgres_error)?;
return Ok(None);
}
}
@@ -1334,8 +1376,9 @@ impl GatewayDataState {
.bind(ldap_username.as_deref())
.bind(logged_in_at)
.fetch_one(&mut *tx)
.await?;
tx.commit().await?;
.await
.map_postgres_err()?;
tx.commit().await.map_err(postgres_error)?;
return Ok(Some(map_user_auth_row(&row)?));
}
@@ -1366,7 +1409,8 @@ impl GatewayDataState {
.bind(ldap_username.as_deref())
.bind(logged_in_at)
.fetch_one(&mut *tx)
.await?;
.await
.map_postgres_err()?;
let user = map_user_auth_row(&user_row)?;
let gift_amount = if unlimited {
@@ -1381,7 +1425,8 @@ impl GatewayDataState {
.bind(if unlimited { "unlimited" } else { "finite" })
.bind(gift_amount)
.fetch_one(&mut *tx)
.await?;
.await
.map_postgres_err()?;
let wallet = map_wallet_snapshot_row(&wallet_row)?;
if gift_amount > 0.0 {
sqlx::query(CREATE_AUTH_USER_WALLET_GIFT_TX_SQL)
@@ -1390,14 +1435,15 @@ impl GatewayDataState {
.bind(gift_amount)
.bind(&user.id)
.execute(&mut *tx)
.await?;
.await
.map_postgres_err()?;
}
tx.commit().await?;
tx.commit().await.map_err(postgres_error)?;
return Ok(Some(user));
}
tx.commit().await?;
tx.commit().await.map_err(postgres_error)?;
Ok(None)
}
@@ -1411,7 +1457,7 @@ impl GatewayDataState {
let Some(pool) = self.postgres_pool() else {
return Ok(None);
};
let mut tx = pool.begin().await?;
let mut tx = pool.begin().await.map_postgres_err()?;
let gift_amount = if unlimited {
0.0
} else {
@@ -1424,7 +1470,8 @@ impl GatewayDataState {
.bind(if unlimited { "unlimited" } else { "finite" })
.bind(gift_amount)
.fetch_one(&mut *tx)
.await?;
.await
.map_postgres_err()?;
let wallet = map_wallet_snapshot_row(&row)?;
if gift_amount > 0.0 {
sqlx::query(CREATE_AUTH_USER_WALLET_GIFT_TX_SQL)
@@ -1433,9 +1480,10 @@ impl GatewayDataState {
.bind(gift_amount)
.bind(user_id)
.execute(&mut *tx)
.await?;
.await
.map_postgres_err()?;
}
tx.commit().await?;
tx.commit().await.map_err(postgres_error)?;
Ok(Some(wallet))
}
@@ -1451,7 +1499,8 @@ impl GatewayDataState {
.bind(user_id)
.bind(limit_mode)
.fetch_optional(&pool)
.await?;
.await
.map_postgres_err()?;
row.as_ref().map(map_wallet_snapshot_row).transpose()
}
@@ -1461,8 +1510,9 @@ impl GatewayDataState {
};
let row = sqlx::query(COUNT_ACTIVE_ADMIN_USERS_SQL)
.fetch_one(&pool)
.await?;
let total = row.try_get::<i64, _>("total")?.max(0) as u64;
.await
.map_postgres_err()?;
let total = row.try_get::<i64, _>("total").map_postgres_err()?.max(0) as u64;
Ok(total)
}
@@ -1474,8 +1524,9 @@ impl GatewayDataState {
};
let row = sqlx::query(COUNT_ACTIVE_LOCAL_ADMIN_USERS_WITH_VALID_PASSWORD_SQL)
.fetch_one(&pool)
.await?;
let total = row.try_get::<i64, _>("total")?.max(0) as u64;
.await
.map_postgres_err()?;
let total = row.try_get::<i64, _>("total").map_postgres_err()?.max(0) as u64;
Ok(total)
}
@@ -1495,8 +1546,9 @@ impl GatewayDataState {
.bind(user_id)
.bind(statuses)
.fetch_one(&pool)
.await?;
let total = row.try_get::<i64, _>("total")?.max(0) as u64;
.await
.map_postgres_err()?;
let total = row.try_get::<i64, _>("total").map_postgres_err()?.max(0) as u64;
Ok(total)
}
@@ -1512,8 +1564,9 @@ impl GatewayDataState {
.bind(user_id)
.bind(statuses)
.fetch_one(&pool)
.await?;
let total = row.try_get::<i64, _>("total")?.max(0) as u64;
.await
.map_postgres_err()?;
let total = row.try_get::<i64, _>("total").map_postgres_err()?.max(0) as u64;
Ok(total)
}
@@ -1527,7 +1580,8 @@ impl GatewayDataState {
let result = sqlx::query(DELETE_LOCAL_AUTH_USER_SQL)
.bind(user_id)
.execute(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(result.rows_affected() > 0)
}
@@ -1543,7 +1597,7 @@ impl GatewayDataState {
let Some(pool) = self.postgres_pool() else {
return Ok(None);
};
let mut tx = pool.begin().await?;
let mut tx = pool.begin().await.map_postgres_err()?;
let user_row = sqlx::query(CREATE_LOCAL_USER_SQL)
.bind(Uuid::new_v4().to_string())
.bind(email)
@@ -1551,7 +1605,8 @@ impl GatewayDataState {
.bind(username)
.bind(password_hash)
.fetch_one(&mut *tx)
.await?;
.await
.map_postgres_err()?;
let user = map_user_auth_row(&user_row)?;
let gift_amount = if unlimited {
0.0
@@ -1565,7 +1620,8 @@ impl GatewayDataState {
.bind(if unlimited { "unlimited" } else { "finite" })
.bind(gift_amount)
.fetch_one(&mut *tx)
.await?;
.await
.map_postgres_err()?;
let wallet = map_wallet_snapshot_row(&wallet_row)?;
if gift_amount > 0.0 {
sqlx::query(CREATE_AUTH_USER_WALLET_GIFT_TX_SQL)
@@ -1574,9 +1630,10 @@ impl GatewayDataState {
.bind(gift_amount)
.bind(&user.id)
.execute(&mut *tx)
.await?;
.await
.map_postgres_err()?;
}
tx.commit().await?;
tx.commit().await.map_err(postgres_error)?;
Ok(Some((user, wallet)))
}
@@ -1598,7 +1655,8 @@ impl GatewayDataState {
.bind(ip_address)
.bind(user_agent.map(|value| value.chars().take(1000).collect::<String>()))
.execute(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(result.rows_affected() > 0)
}
@@ -1618,7 +1676,8 @@ impl GatewayDataState {
.bind(device_label.chars().take(120).collect::<String>())
.bind(updated_at)
.execute(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(result.rows_affected() > 0)
}
@@ -1646,7 +1705,8 @@ impl GatewayDataState {
.bind(ip_address)
.bind(user_agent.map(|value| value.chars().take(1000).collect::<String>()))
.execute(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(result.rows_affected() > 0)
}
@@ -1666,7 +1726,8 @@ impl GatewayDataState {
.bind(revoked_at)
.bind(reason.chars().take(100).collect::<String>())
.execute(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(result.rows_affected() > 0)
}
@@ -1684,7 +1745,8 @@ impl GatewayDataState {
.bind(revoked_at)
.bind(reason.chars().take(100).collect::<String>())
.execute(&pool)
.await?;
.await
.map_postgres_err()?;
Ok(result.rows_affected())
}
@@ -2221,6 +2283,7 @@ mod tests {
StoredAuthApiKeySnapshot,
};
use super::UPDATE_LOCAL_AUTH_USER_ADMIN_FIELDS_SQL;
use crate::data::GatewayDataState;
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
@@ -2250,6 +2313,13 @@ mod tests {
.expect("snapshot should build")
}
#[test]
fn update_local_auth_user_admin_fields_sql_casts_json_case_values() {
assert!(UPDATE_LOCAL_AUTH_USER_ADMIN_FIELDS_SQL.contains("WHEN $4::BOOLEAN THEN $5::json"));
assert!(UPDATE_LOCAL_AUTH_USER_ADMIN_FIELDS_SQL.contains("WHEN $6::BOOLEAN THEN $7::json"));
assert!(UPDATE_LOCAL_AUTH_USER_ADMIN_FIELDS_SQL.contains("WHEN $8::BOOLEAN THEN $9::json"));
}
#[tokio::test]
async fn data_state_lists_auth_api_key_export_records() {
let repository = Arc::new(

23
apps/aether-gateway/src/data/state/integrations.rs
View File

@@ -5,15 +5,16 @@ use aether_data::repository::audit::RequestAuditReader;
use aether_data::repository::auth::{
AuthApiKeyLookupKey, ResolvedAuthApiKeySnapshotReader, StoredAuthApiKeySnapshot,
};
use aether_data::repository::billing::StoredBillingModelContext;
use aether_data::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_data::repository::provider_catalog::{
use aether_data::DataLayerError;
use aether_data_contracts::repository::billing::StoredBillingModelContext;
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_data_contracts::repository::candidates::DecisionTrace;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
use aether_data::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskLookupKey};
use aether_data::DataLayerError;
use aether_data_contracts::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::repository::video_tasks::{StoredVideoTask, VideoTaskLookupKey};
use aether_usage_runtime::{
UsageBillingEventEnricher, UsageEvent, UsageRecordWriter, UsageRuntimeAccess,
UsageSettlementWriter,
@@ -22,8 +23,8 @@ use aether_video_tasks_core::StoredVideoTaskReadSide;
use async_trait::async_trait;
use super::GatewayDataState;
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
use crate::provider_transport::ProviderTransportSnapshotSource;
use crate::scheduler::SchedulerCandidateSelectionRowSource;
#[async_trait]
impl RequestAuditReader for GatewayDataState {
@@ -38,7 +39,7 @@ impl RequestAuditReader for GatewayDataState {
&self,
request_id: &str,
attempted_only: bool,
) -> Result<Option<aether_data::repository::candidates::DecisionTrace>, DataLayerError> {
) -> Result<Option<DecisionTrace>, DataLayerError> {
GatewayDataState::read_decision_trace(self, request_id, attempted_only).await
}
@@ -102,7 +103,7 @@ impl ProviderTransportSnapshotSource for GatewayDataState {
}
#[async_trait]
impl SchedulerCandidateSelectionRowSource for GatewayDataState {
impl MinimalCandidateSelectionRowSource for GatewayDataState {
async fn read_minimal_candidate_selection_rows_for_api_format_and_global_model(
&self,
api_format: &str,
@@ -197,7 +198,7 @@ mod tests {
let state = GatewayDataState::with_billing_reader_for_tests(
std::sync::Arc::new(
aether_data::repository::billing::InMemoryBillingReadRepository::seed(vec![
aether_data::repository::billing::StoredBillingModelContext::new(
aether_data_contracts::repository::billing::StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-1".to_string()),

72
apps/aether-gateway/src/data/state/mod.rs
View File

@@ -25,27 +25,11 @@ use aether_data::repository::auth_modules::{
AuthModuleReadRepository, AuthModuleWriteRepository, StoredLdapModuleConfig,
StoredOAuthProviderModuleConfig,
};
use aether_data::repository::billing::{BillingReadRepository, StoredBillingModelContext};
use aether_data::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
};
use aether_data::repository::candidates::{
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
RequestCandidateWriteRepository, StoredRequestCandidate, UpsertRequestCandidateRecord,
};
use aether_data::repository::gemini_file_mappings::{
GeminiFileMappingListQuery, GeminiFileMappingReadRepository, GeminiFileMappingStats,
GeminiFileMappingWriteRepository, StoredGeminiFileMapping, StoredGeminiFileMappingListPage,
UpsertGeminiFileMappingRecord,
};
use aether_data::repository::global_models::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
GlobalModelReadRepository, GlobalModelWriteRepository, PublicCatalogModelListQuery,
PublicCatalogModelSearchQuery, PublicGlobalModelQuery, StoredAdminGlobalModel,
StoredAdminGlobalModelPage, StoredAdminProviderModel, StoredProviderActiveGlobalModel,
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
use aether_data::repository::management_tokens::{
CreateManagementTokenRecord, ManagementTokenListQuery, ManagementTokenReadRepository,
ManagementTokenWriteRepository, RegenerateManagementTokenSecret, StoredManagementToken,
@@ -55,40 +39,21 @@ use aether_data::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, StoredOAuthProviderConfig,
UpsertOAuthProviderConfigRecord,
};
use aether_data::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use aether_data::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeReadRepository, ProxyNodeTunnelStatusMutation,
ProxyNodeWriteRepository, StoredProxyNode, StoredProxyNodeEvent,
};
use aether_data::repository::quota::{
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
};
use aether_data::repository::settlement::{
SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
};
use aether_data::repository::shadow_results::{
merge_shadow_result_sample, RecordShadowResultSample, ShadowResultLookupKey,
ShadowResultReadRepository, ShadowResultWriteRepository, StoredShadowResult,
};
pub(crate) use aether_data::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
use aether_data::repository::usage::{
StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
UsageWriteRepository,
};
use aether_data::repository::users::{
StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, UserReadRepository,
};
pub(crate) use aether_data::repository::users::{
StoredUserPreferenceRecord, StoredUserSessionRecord,
};
use aether_data::repository::video_tasks::{
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatusCount, VideoTaskWriteRepository,
};
use aether_data::repository::wallet::{
AdjustWalletBalanceInput, AdminPaymentOrderListQuery, AdminWalletLedgerQuery,
AdminWalletListQuery, AdminWalletRefundRequestListQuery, CompleteAdminWalletRefundInput,
@@ -104,6 +69,43 @@ use aether_data::repository::wallet::{
WalletReadRepository, WalletWriteRepository,
};
use aether_data::{DataBackends, DataLayerError};
use aether_data_contracts::repository::billing::{
BillingReadRepository, StoredBillingModelContext,
};
use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
};
use aether_data_contracts::repository::candidates::{
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
RequestCandidateWriteRepository, StoredRequestCandidate, UpsertRequestCandidateRecord,
};
use aether_data_contracts::repository::global_models::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
GlobalModelReadRepository, GlobalModelWriteRepository, PublicCatalogModelListQuery,
PublicCatalogModelSearchQuery, PublicGlobalModelQuery, StoredAdminGlobalModel,
StoredAdminGlobalModelPage, StoredAdminProviderModel, StoredProviderActiveGlobalModel,
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::quota::{
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
};
use aether_data_contracts::repository::settlement::{
SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{
StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
UsageWriteRepository,
};
use aether_data_contracts::repository::video_tasks::{
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatusCount, VideoTaskWriteRepository,
};
#[derive(Clone, Default)]
pub(crate) struct GatewayDataState {

3
apps/aether-gateway/src/data/state/runtime.rs
View File

@@ -22,6 +22,7 @@ use super::{
VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskStatusCount, WalletLookupKey,
WalletMutationOutcome,
};
use aether_data_contracts::repository::usage::UsageAuditListQuery;
use aether_video_tasks_core::read_data_backed_video_task_response;
fn is_missing_shadow_results_relation_error(error: &DataLayerError) -> bool {
@@ -640,7 +641,7 @@ impl GatewayDataState {
pub(crate) async fn list_usage_audits(
&self,
query: &aether_data::repository::usage::UsageAuditListQuery,
query: &UsageAuditListQuery,
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.list_usage_audits(query).await,

46
apps/aether-gateway/src/data/state/testing/mod.rs
View File

@@ -1,6 +1,10 @@
use std::collections::BTreeMap;
use std::sync::{Arc, RwLock};
use aether_data_contracts::repository::candidates::RequestCandidateRepository;
use aether_data_contracts::repository::quota::ProviderQuotaRepository;
use aether_data_contracts::repository::usage::UsageRepository;
use super::{
AnnouncementReadRepository, AnnouncementWriteRepository, AuthApiKeyReadRepository,
AuthApiKeyWriteRepository, AuthModuleReadRepository, AuthModuleWriteRepository,
@@ -85,7 +89,7 @@ impl GatewayDataState {
#[cfg(test)]
pub(crate) fn with_request_candidate_repository_for_tests<T>(repository: Arc<T>) -> Self
where
T: aether_data::repository::candidates::RequestCandidateRepository + 'static,
T: RequestCandidateRepository + 'static,
{
let request_candidate_reader: Arc<dyn RequestCandidateReadRepository> = repository.clone();
let request_candidate_writer: Arc<dyn RequestCandidateWriteRepository> = repository;
@@ -385,7 +389,7 @@ impl GatewayDataState {
provider_quota_repository: Arc<T>,
) -> Self
where
T: aether_data::repository::quota::ProviderQuotaRepository + 'static,
T: ProviderQuotaRepository + 'static,
{
let provider_quota_reader: Arc<dyn ProviderQuotaReadRepository> =
provider_quota_repository.clone();
@@ -441,7 +445,7 @@ impl GatewayDataState {
provider_quota_repository: Arc<T>,
) -> Self
where
T: aether_data::repository::quota::ProviderQuotaRepository + 'static,
T: ProviderQuotaRepository + 'static,
U: ProviderCatalogReadRepository + ProviderCatalogWriteRepository + 'static,
V: GlobalModelReadRepository + GlobalModelWriteRepository + 'static,
{
@@ -600,8 +604,8 @@ impl GatewayDataState {
usage_repository: Arc<U>,
) -> Self
where
T: aether_data::repository::candidates::RequestCandidateRepository + 'static,
U: aether_data::repository::usage::UsageRepository + 'static,
T: RequestCandidateRepository + 'static,
U: UsageRepository + 'static,
{
let request_candidate_reader: Arc<dyn RequestCandidateReadRepository> =
request_candidate_repository.clone();
@@ -659,7 +663,7 @@ impl GatewayDataState {
gemini_file_mapping_repository: Arc<U>,
) -> Self
where
T: aether_data::repository::candidates::RequestCandidateRepository + 'static,
T: RequestCandidateRepository + 'static,
U: aether_data::repository::gemini_file_mappings::GeminiFileMappingRepository + 'static,
{
let request_candidate_reader: Arc<dyn RequestCandidateReadRepository> =
@@ -879,7 +883,7 @@ impl GatewayDataState {
usage_repository: Arc<TUsage>,
) -> Self
where
TUsage: aether_data::repository::usage::UsageRepository + 'static,
TUsage: UsageRepository + 'static,
TWallet: aether_data::repository::wallet::WalletRepository + 'static,
{
let wallet_reader: Arc<dyn WalletReadRepository> = wallet_repository.clone();
@@ -1842,7 +1846,7 @@ impl GatewayDataState {
provider_quota_repository: Arc<T>,
) -> Self
where
T: aether_data::repository::quota::ProviderQuotaRepository + 'static,
T: ProviderQuotaRepository + 'static,
{
let provider_quota_reader: Arc<dyn ProviderQuotaReadRepository> =
provider_quota_repository.clone();
@@ -1899,7 +1903,7 @@ impl GatewayDataState {
request_candidate_repository: Arc<dyn RequestCandidateReadRepository>,
) -> Self
where
T: aether_data::repository::quota::ProviderQuotaRepository + 'static,
T: ProviderQuotaRepository + 'static,
{
let provider_quota_reader: Arc<dyn ProviderQuotaReadRepository> =
provider_quota_repository.clone();
@@ -1959,7 +1963,7 @@ impl GatewayDataState {
request_candidate_repository: Arc<dyn RequestCandidateReadRepository>,
) -> Self
where
T: aether_data::repository::quota::ProviderQuotaRepository + 'static,
T: ProviderQuotaRepository + 'static,
{
let provider_quota_reader: Arc<dyn ProviderQuotaReadRepository> =
provider_quota_repository.clone();
@@ -2021,7 +2025,7 @@ impl GatewayDataState {
encryption_key: impl Into<String>,
) -> Self
where
T: aether_data::repository::candidates::RequestCandidateRepository + 'static,
T: RequestCandidateRepository + 'static,
U: ProviderCatalogReadRepository + ProviderCatalogWriteRepository + 'static,
{
let request_candidate_reader: Arc<dyn RequestCandidateReadRepository> =
@@ -2089,7 +2093,7 @@ impl GatewayDataState {
encryption_key: impl Into<String>,
) -> Self
where
T: aether_data::repository::candidates::RequestCandidateRepository + 'static,
T: RequestCandidateRepository + 'static,
U: ProviderCatalogReadRepository + ProviderCatalogWriteRepository + 'static,
V: aether_data::repository::shadow_results::ShadowResultRepository + 'static,
{
@@ -2162,9 +2166,9 @@ impl GatewayDataState {
encryption_key: impl Into<String>,
) -> Self
where
T: aether_data::repository::candidates::RequestCandidateRepository + 'static,
T: RequestCandidateRepository + 'static,
U: ProviderCatalogReadRepository + ProviderCatalogWriteRepository + 'static,
V: aether_data::repository::usage::UsageRepository + 'static,
V: UsageRepository + 'static,
{
let request_candidate_reader: Arc<dyn RequestCandidateReadRepository> =
request_candidate_repository.clone();
@@ -2236,9 +2240,9 @@ impl GatewayDataState {
encryption_key: impl Into<String>,
) -> Self
where
T: aether_data::repository::candidates::RequestCandidateRepository + 'static,
T: RequestCandidateRepository + 'static,
U: ProviderCatalogReadRepository + ProviderCatalogWriteRepository + 'static,
V: aether_data::repository::usage::UsageRepository + 'static,
V: UsageRepository + 'static,
W: aether_data::repository::wallet::WalletRepository + 'static,
{
let request_candidate_reader: Arc<dyn RequestCandidateReadRepository> =
@@ -2307,7 +2311,7 @@ impl GatewayDataState {
encryption_key: impl Into<String>,
) -> Self
where
T: aether_data::repository::quota::ProviderQuotaRepository + 'static,
T: ProviderQuotaRepository + 'static,
{
let provider_quota_reader: Arc<dyn ProviderQuotaReadRepository> =
provider_quota_repository.clone();
@@ -2360,7 +2364,7 @@ impl GatewayDataState {
#[cfg(test)]
pub(crate) fn with_usage_repository_for_tests<T>(repository: Arc<T>) -> Self
where
T: aether_data::repository::usage::UsageRepository + 'static,
T: UsageRepository + 'static,
{
let usage_reader: Arc<dyn UsageReadRepository> = repository.clone();
let usage_writer: Arc<dyn UsageWriteRepository> = repository;
@@ -2516,7 +2520,7 @@ impl GatewayDataState {
) -> Self
where
TWallet: aether_data::repository::wallet::WalletRepository + 'static,
TUsage: aether_data::repository::usage::UsageRepository + 'static,
TUsage: UsageRepository + 'static,
{
let wallet_reader: Arc<dyn WalletReadRepository> = wallet_repository.clone();
let wallet_writer: Arc<dyn WalletWriteRepository> = wallet_repository;
@@ -2572,7 +2576,7 @@ impl GatewayDataState {
wallet_repository: Arc<TWallet>,
) -> Self
where
TUsage: aether_data::repository::usage::UsageRepository + 'static,
TUsage: UsageRepository + 'static,
TWallet: aether_data::repository::wallet::WalletRepository + 'static,
{
let usage_reader: Arc<dyn UsageReadRepository> = usage_repository.clone();
@@ -2626,7 +2630,7 @@ impl GatewayDataState {
#[cfg(test)]
pub(crate) fn with_provider_quota_repository_for_tests<T>(repository: Arc<T>) -> Self
where
T: aether_data::repository::quota::ProviderQuotaRepository + 'static,
T: ProviderQuotaRepository + 'static,
{
let provider_quota_reader: Arc<dyn ProviderQuotaReadRepository> = repository.clone();
let provider_quota_writer: Arc<dyn ProviderQuotaWriteRepository> = repository;

17
apps/aether-gateway/src/data/state/testing/video_tasks.rs
View File

@@ -1,5 +1,8 @@
use std::sync::Arc;
use aether_data_contracts::repository::candidates::RequestCandidateRepository;
use aether_data_contracts::repository::video_tasks::VideoTaskRepository;
use super::{
AuthApiKeyReadRepository, GatewayDataConfig, GatewayDataState, ProviderCatalogReadRepository,
RequestCandidateReadRepository, RequestCandidateWriteRepository, VideoTaskReadRepository,
@@ -59,7 +62,7 @@ impl GatewayDataState {
repository: Arc<T>,
) -> Self
where
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
T: VideoTaskRepository + 'static,
{
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;
@@ -110,7 +113,7 @@ impl GatewayDataState {
#[cfg(test)]
pub(crate) fn with_video_task_repository_for_tests<T>(repository: Arc<T>) -> Self
where
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
T: VideoTaskRepository + 'static,
{
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;
@@ -165,7 +168,7 @@ impl GatewayDataState {
encryption_key: impl Into<String>,
) -> Self
where
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
T: VideoTaskRepository + 'static,
{
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;
@@ -219,8 +222,8 @@ impl GatewayDataState {
request_candidate_repository: Arc<U>,
) -> Self
where
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
U: aether_data::repository::candidates::RequestCandidateRepository + 'static,
T: VideoTaskRepository + 'static,
U: RequestCandidateRepository + 'static,
{
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;
@@ -284,9 +287,9 @@ impl GatewayDataState {
encryption_key: impl Into<String>,
) -> Self
where
T: aether_data::repository::video_tasks::VideoTaskRepository + 'static,
T: VideoTaskRepository + 'static,
U: ProviderCatalogReadRepository + 'static,
V: aether_data::repository::candidates::RequestCandidateRepository + 'static,
V: RequestCandidateRepository + 'static,
{
let video_task_reader: Arc<dyn VideoTaskReadRepository> = repository.clone();
let video_task_writer: Arc<dyn VideoTaskWriteRepository> = repository;

60
apps/aether-gateway/src/data/tests.rs
View File

@@ -4,28 +4,31 @@ use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidate_selection::{
InMemoryMinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredProviderModelMapping,
};
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::shadow_results::{
InMemoryShadowResultRepository, RecordShadowResultSample, ShadowResultLookupKey,
ShadowResultMatchStatus, ShadowResultReadRepository, ShadowResultSampleOrigin,
ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
};
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredRequestUsageAudit};
use aether_data::repository::video_tasks::{
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus,
VideoTaskWriteRepository,
};
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data::repository::video_tasks::InMemoryVideoTaskRepository;
use aether_data::DataLayerError;
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use aether_data_contracts::repository::video_tasks::{
UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus, VideoTaskWriteRepository,
};
use aether_scheduler_core::{build_minimal_candidate_selection, SchedulerAuthConstraints};
use async_trait::async_trait;
use serde_json::json;
@@ -643,14 +646,31 @@ async fn data_state_reads_minimal_candidate_selection_with_auth_filters() {
.expect("auth snapshot should read")
.expect("auth snapshot should exist");
let selection = crate::scheduler::read_minimal_candidate_selection(
&state,
let rows = state
.list_minimal_candidate_selection_rows("openai:chat", "gpt-4.1")
.await
.expect("minimal candidate selection rows should read");
let auth_constraints = SchedulerAuthConstraints {
allowed_providers: auth_snapshot
.effective_allowed_providers()
.map(|items| items.to_vec()),
allowed_api_formats: auth_snapshot
.effective_allowed_api_formats()
.map(|items| items.to_vec()),
allowed_models: auth_snapshot
.effective_allowed_models()
.map(|items| items.to_vec()),
};
let selection = build_minimal_candidate_selection(
rows,
"openai:chat",
"gpt-4.1",
"gpt-4.1",
false,
Some(&auth_snapshot),
Some(&auth_constraints),
Some(auth_snapshot.api_key_id.as_str()),
)
.await
.expect("selection should read");
assert_eq!(selection.len(), 2);