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

feat: 扩展 Rust gateway 全功能模块，新增 billing/crypto/wallet crate 及完整数据层

Browse Source 
- 新增 aether-billing、aether-crypto、aether-wallet 独立 crate
- aether-data 扩展 repository 层：announcements、auth_modules、billing、
  candidate_selection、gemini_file_mappings、global_models、management_tokens、
  oauth_providers、proxy_nodes、quota、users、wallet 等模块
- aether-gateway 新增 api/auth/billing/control/middleware/scheduler/usage/
  video_tasks/hooks/maintenance/model_fetch/provider_transport 等功能模块
- 重构 executor decision 和 gateway state 为模块目录结构
- 新增 gateway router、frontdoor 路由层及对应测试
- Python 侧 API 路由重构，新增 compat/support 模块
- 前端 Logo 组件更新及 Provider 管理页面调整
This commit is contained in:
fawney19
2026-03-31 19:19:04 +08:00
parent b5a0070023
commit ddf18fed9a
690 changed files with 235087 additions and 16301 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
crates/aether-gateway/src/data/auth.rs
View File

@@ -12,6 +12,7 @@ pub(crate) struct StoredGatewayAuthApiKeySnapshot {
pub(crate) user_auth_source: String,
pub(crate) user_is_active: bool,
pub(crate) user_is_deleted: bool,
pub(crate) user_rate_limit: Option<i32>,
pub(crate) user_allowed_providers: Option<Vec<String>>,
pub(crate) user_allowed_api_formats: Option<Vec<String>>,
pub(crate) user_allowed_models: Option<Vec<String>>,
@@ -40,6 +41,7 @@ impl StoredGatewayAuthApiKeySnapshot {
user_auth_source: snapshot.user_auth_source,
user_is_active: snapshot.user_is_active,
user_is_deleted: snapshot.user_is_deleted,
user_rate_limit: snapshot.user_rate_limit,
user_allowed_providers: snapshot.user_allowed_providers,
user_allowed_api_formats: snapshot.user_allowed_api_formats,
user_allowed_models: snapshot.user_allowed_models,
@@ -57,6 +59,24 @@ impl StoredGatewayAuthApiKeySnapshot {
currently_usable,
}
}
pub(crate) fn effective_allowed_providers(&self) -> Option<&[String]> {
self.api_key_allowed_providers
.as_deref()
.or(self.user_allowed_providers.as_deref())
}
pub(crate) fn effective_allowed_api_formats(&self) -> Option<&[String]> {
self.api_key_allowed_api_formats
.as_deref()
.or(self.user_allowed_api_formats.as_deref())
}
pub(crate) fn effective_allowed_models(&self) -> Option<&[String]> {
self.api_key_allowed_models
.as_deref()
.or(self.user_allowed_models.as_deref())
}
}
pub(crate) async fn read_auth_api_key_snapshot(
@@ -75,10 +95,25 @@ pub(crate) async fn read_auth_api_key_snapshot(
.map(|snapshot| StoredGatewayAuthApiKeySnapshot::from_stored(snapshot, now_unix_secs)))
}
pub(crate) async fn read_auth_api_key_snapshot_by_key_hash(
state: &GatewayDataState,
key_hash: &str,
now_unix_secs: u64,
) -> Result<Option<StoredGatewayAuthApiKeySnapshot>, DataLayerError> {
let snapshot = state
.find_auth_api_key_snapshot(AuthApiKeyLookupKey::KeyHash(key_hash))
.await?;
Ok(snapshot
.map(|snapshot| StoredGatewayAuthApiKeySnapshot::from_stored(snapshot, now_unix_secs)))
}
#[cfg(test)]
mod tests {
use super::super::GatewayDataState;
use super::{read_auth_api_key_snapshot, StoredGatewayAuthApiKeySnapshot};
use super::{
read_auth_api_key_snapshot, read_auth_api_key_snapshot_by_key_hash,
StoredGatewayAuthApiKeySnapshot,
};
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
@@ -134,6 +169,7 @@ mod tests {
user_auth_source: "local".to_string(),
user_is_active: true,
user_is_deleted: false,
user_rate_limit: None,
user_allowed_providers: Some(vec!["openai".to_string()]),
user_allowed_api_formats: Some(vec!["openai:chat".to_string()]),
user_allowed_models: Some(vec!["gpt-4.1".to_string()]),
@@ -152,4 +188,33 @@ mod tests {
}
);
}
#[tokio::test]
async fn reads_auth_snapshot_by_key_hash() {
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-lookup".to_string()),
sample_snapshot("key-1", "user-1"),
)]));
let state = GatewayDataState::with_auth_api_key_reader_for_tests(repository);
let snapshot = read_auth_api_key_snapshot_by_key_hash(&state, "hash-lookup", 150)
.await
.expect("read should succeed")
.expect("snapshot should exist");
assert_eq!(snapshot.user_id, "user-1");
assert_eq!(snapshot.api_key_id, "key-1");
assert_eq!(
snapshot.effective_allowed_providers(),
Some(&["openai".to_string()][..])
);
assert_eq!(
snapshot.effective_allowed_api_formats(),
Some(&["openai:chat".to_string()][..])
);
assert_eq!(
snapshot.effective_allowed_models(),
Some(&["gpt-4.1".to_string()][..])
);
}
}

57
crates/aether-gateway/src/data/config.rs
View File

@@ -1,9 +1,23 @@
use aether_data::postgres::PostgresPoolConfig;
use aether_data::redis::RedisClientConfig;
use aether_data::DataLayerConfig;
use std::fmt;
#[derive(Debug, Clone, Default)]
#[derive(Clone, Default)]
pub struct GatewayDataConfig {
postgres: Option<PostgresPoolConfig>,
redis: Option<RedisClientConfig>,
encryption_key: Option<String>,
}
impl fmt::Debug for GatewayDataConfig {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("GatewayDataConfig")
.field("postgres", &self.postgres)
.field("redis", &self.redis)
.field("has_encryption_key", &self.encryption_key.is_some())
.finish()
}
}
impl GatewayDataConfig {
@@ -14,6 +28,8 @@ impl GatewayDataConfig {
pub fn from_postgres_config(postgres: PostgresPoolConfig) -> Self {
Self {
postgres: Some(postgres),
redis: None,
encryption_key: None,
}
}
@@ -28,14 +44,49 @@ impl GatewayDataConfig {
self.postgres.as_ref()
}
pub fn redis(&self) -> Option<&RedisClientConfig> {
self.redis.as_ref()
}
pub fn with_redis_config(mut self, redis: RedisClientConfig) -> Self {
self.redis = Some(redis);
self
}
pub fn with_redis_url(
self,
url: impl Into<String>,
key_prefix: Option<impl Into<String>>,
) -> Self {
self.with_redis_config(RedisClientConfig {
url: url.into(),
key_prefix: key_prefix.map(Into::into),
})
}
pub fn with_encryption_key(mut self, encryption_key: impl Into<String>) -> Self {
let encryption_key = encryption_key.into();
let encryption_key = encryption_key.trim();
self.encryption_key = if encryption_key.is_empty() {
None
} else {
Some(encryption_key.to_string())
};
self
}
pub fn encryption_key(&self) -> Option<&str> {
self.encryption_key.as_deref()
}
pub fn is_enabled(&self) -> bool {
self.postgres.is_some()
self.postgres.is_some() || self.redis.is_some()
}
pub fn to_data_layer_config(&self) -> DataLayerConfig {
DataLayerConfig {
postgres: self.postgres.clone(),
redis: None,
redis: self.redis.clone(),
}
}
}

11
crates/aether-gateway/src/data/mod.rs
View File

@@ -4,18 +4,17 @@ mod config;
mod decision_trace;
mod gemini;
mod openai;
mod request_audit;
mod state;
mod usage;
mod video_tasks;
#[cfg(test)]
mod tests;
pub(crate) use crate::gateway::provider_transport::GatewayProviderTransportSnapshot;
pub(crate) use auth::StoredGatewayAuthApiKeySnapshot;
pub(crate) use candidates::RequestCandidateTrace;
pub(crate) use candidates::{RequestCandidateFinalStatus, RequestCandidateTrace};
pub use config::GatewayDataConfig;
pub(crate) use decision_trace::DecisionTrace;
pub(crate) use request_audit::RequestAuditBundle;
pub(crate) use state::GatewayDataState;
pub(crate) use usage::RequestUsageAudit;
pub(crate) use state::{
GatewayDataState, StoredSystemConfigEntry, StoredUserPreferenceRecord, StoredUserSessionRecord,
};

56
crates/aether-gateway/src/data/request_audit.rs
View File

@@ -1,56 +0,0 @@
use aether_data::DataLayerError;
use super::auth::StoredGatewayAuthApiKeySnapshot;
use super::decision_trace::DecisionTrace;
use super::state::GatewayDataState;
use super::usage::RequestUsageAudit;
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub(crate) struct RequestAuditBundle {
pub(crate) request_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) usage: Option<RequestUsageAudit>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) decision_trace: Option<DecisionTrace>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) auth_snapshot: Option<StoredGatewayAuthApiKeySnapshot>,
}
pub(crate) async fn read_request_audit_bundle(
state: &GatewayDataState,
request_id: &str,
attempted_only: bool,
now_unix_secs: u64,
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
let usage = state.read_request_usage_audit(request_id).await?;
let decision_trace = state
.read_decision_trace(request_id, attempted_only)
.await?;
let auth_snapshot = if let Some(usage) = usage.as_ref() {
match (
usage.usage.user_id.as_deref(),
usage.usage.api_key_id.as_deref(),
) {
(Some(user_id), Some(api_key_id)) => {
state
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
.await?
}
_ => None,
}
} else {
None
};
if usage.is_none() && decision_trace.is_none() && auth_snapshot.is_none() {
return Ok(None);
}
Ok(Some(RequestAuditBundle {
request_id: request_id.to_string(),
usage,
decision_trace,
auth_snapshot,
}))
}

454
crates/aether-gateway/src/data/state.rs
View File

@@ -1,454 +0,0 @@
use std::fmt;
use std::sync::Arc;
use aether_data::repository::auth::{
AuthApiKeyLookupKey, AuthApiKeyReadRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidates::{RequestCandidateReadRepository, StoredRequestCandidate};
use aether_data::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use aether_data::repository::shadow_results::{
merge_shadow_result_sample, RecordShadowResultSample, ShadowResultLookupKey,
ShadowResultReadRepository, ShadowResultWriteRepository, StoredShadowResult,
};
use aether_data::repository::usage::{StoredRequestUsageAudit, UsageReadRepository};
use aether_data::repository::video_tasks::{
StoredVideoTask, VideoTaskLookupKey, VideoTaskReadRepository,
};
use aether_data::{DataBackends, DataLayerError};
use super::auth::{read_auth_api_key_snapshot, StoredGatewayAuthApiKeySnapshot};
use super::candidates::{read_request_candidate_trace, RequestCandidateTrace};
use super::config::GatewayDataConfig;
use super::decision_trace::{read_decision_trace, DecisionTrace};
use super::request_audit::{read_request_audit_bundle, RequestAuditBundle};
use super::usage::{read_request_usage_audit, RequestUsageAudit};
use super::video_tasks::read_video_task_response;
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
#[derive(Clone, Default)]
pub(crate) struct GatewayDataState {
config: GatewayDataConfig,
backends: Option<DataBackends>,
auth_api_key_reader: Option<Arc<dyn AuthApiKeyReadRepository>>,
request_candidate_reader: Option<Arc<dyn RequestCandidateReadRepository>>,
provider_catalog_reader: Option<Arc<dyn ProviderCatalogReadRepository>>,
usage_reader: Option<Arc<dyn UsageReadRepository>>,
video_task_reader: Option<Arc<dyn VideoTaskReadRepository>>,
shadow_result_reader: Option<Arc<dyn ShadowResultReadRepository>>,
shadow_result_writer: Option<Arc<dyn ShadowResultWriteRepository>>,
}
impl fmt::Debug for GatewayDataState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("GatewayDataState")
.field("config", &self.config)
.field("has_backends", &self.backends.is_some())
.field(
"has_auth_api_key_reader",
&self.auth_api_key_reader.is_some(),
)
.field(
"has_request_candidate_reader",
&self.request_candidate_reader.is_some(),
)
.field(
"has_provider_catalog_reader",
&self.provider_catalog_reader.is_some(),
)
.field("has_usage_reader", &self.usage_reader.is_some())
.field("has_video_task_reader", &self.video_task_reader.is_some())
.field(
"has_shadow_result_reader",
&self.shadow_result_reader.is_some(),
)
.field(
"has_shadow_result_writer",
&self.shadow_result_writer.is_some(),
)
.finish()
}
}
impl GatewayDataState {
pub(crate) fn disabled() -> Self {
Self::default()
}
pub(crate) fn from_config(config: GatewayDataConfig) -> Result<Self, DataLayerError> {
if !config.is_enabled() {
return Ok(Self {
config,
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
});
}
let backends = DataBackends::from_config(config.to_data_layer_config())?;
let auth_api_key_reader = backends.read().auth_api_keys();
let request_candidate_reader = backends.read().request_candidates();
let provider_catalog_reader = backends.read().provider_catalog();
let usage_reader = backends.read().usage();
let video_task_reader = backends.read().video_tasks();
let shadow_result_reader = backends.read().shadow_results();
let shadow_result_writer = backends.write().shadow_results();
Ok(Self {
config,
backends: Some(backends),
auth_api_key_reader,
request_candidate_reader,
provider_catalog_reader,
usage_reader,
video_task_reader,
shadow_result_reader,
shadow_result_writer,
})
}
pub(crate) fn has_backends(&self) -> bool {
self.backends.is_some()
}
pub(crate) fn has_auth_api_key_reader(&self) -> bool {
self.auth_api_key_reader.is_some()
}
pub(crate) fn has_request_candidate_reader(&self) -> bool {
self.request_candidate_reader.is_some()
}
pub(crate) fn has_provider_catalog_reader(&self) -> bool {
self.provider_catalog_reader.is_some()
}
pub(crate) fn has_usage_reader(&self) -> bool {
self.usage_reader.is_some()
}
pub(crate) fn has_video_task_reader(&self) -> bool {
self.video_task_reader.is_some()
}
pub(crate) fn has_shadow_result_writer(&self) -> bool {
self.shadow_result_writer.is_some()
}
pub(crate) fn has_shadow_result_reader(&self) -> bool {
self.shadow_result_reader.is_some()
}
pub(super) async fn find_video_task(
&self,
key: VideoTaskLookupKey<'_>,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
match &self.video_task_reader {
Some(repository) => repository.find(key).await,
None => Ok(None),
}
}
pub(super) async fn find_auth_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
match &self.auth_api_key_reader {
Some(repository) => repository.find_api_key_snapshot(key).await,
None => Ok(None),
}
}
pub(super) async fn list_request_candidates_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_by_request_id(request_id).await,
None => Ok(Vec::new()),
}
}
pub(super) async fn list_provider_catalog_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_providers_by_ids(provider_ids).await,
None => Ok(Vec::new()),
}
}
pub(super) async fn list_provider_catalog_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_endpoints_by_ids(endpoint_ids).await,
None => Ok(Vec::new()),
}
}
pub(super) async fn list_provider_catalog_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_keys_by_ids(key_ids).await,
None => Ok(Vec::new()),
}
}
pub(super) async fn find_request_usage_by_request_id(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.find_by_request_id(request_id).await,
None => Ok(None),
}
}
pub(crate) async fn read_request_candidate_trace(
&self,
request_id: &str,
attempted_only: bool,
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
read_request_candidate_trace(self, request_id, attempted_only).await
}
pub(crate) async fn read_decision_trace(
&self,
request_id: &str,
attempted_only: bool,
) -> Result<Option<DecisionTrace>, DataLayerError> {
read_decision_trace(self, request_id, attempted_only).await
}
pub(crate) async fn read_request_usage_audit(
&self,
request_id: &str,
) -> Result<Option<RequestUsageAudit>, DataLayerError> {
read_request_usage_audit(self, request_id).await
}
pub(crate) async fn read_request_audit_bundle(
&self,
request_id: &str,
attempted_only: bool,
now_unix_secs: u64,
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
read_request_audit_bundle(self, request_id, attempted_only, now_unix_secs).await
}
pub(crate) async fn read_auth_api_key_snapshot(
&self,
user_id: &str,
api_key_id: &str,
now_unix_secs: u64,
) -> Result<Option<StoredGatewayAuthApiKeySnapshot>, DataLayerError> {
read_auth_api_key_snapshot(self, user_id, api_key_id, now_unix_secs).await
}
pub(crate) async fn read_video_task_response(
&self,
route_family: Option<&str>,
request_path: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
read_video_task_response(self, route_family, request_path).await
}
#[cfg(test)]
pub(crate) async fn write_shadow_result(
&self,
result: aether_data::repository::shadow_results::UpsertShadowResult,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
match &self.shadow_result_writer {
Some(repository) => repository.upsert(result).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn record_shadow_result_sample(
&self,
sample: RecordShadowResultSample,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
let Some(writer) = &self.shadow_result_writer else {
return Ok(None);
};
let existing = match &self.shadow_result_reader {
Some(reader) => {
reader
.find(ShadowResultLookupKey::TraceFingerprint {
trace_id: &sample.trace_id,
request_fingerprint: &sample.request_fingerprint,
})
.await?
}
None => None,
};
let merged = merge_shadow_result_sample(existing.as_ref(), sample);
writer.upsert(merged).await.map(Some)
}
pub(crate) async fn list_recent_shadow_results(
&self,
limit: usize,
) -> Result<Vec<StoredShadowResult>, DataLayerError> {
match &self.shadow_result_reader {
Some(repository) => repository.list_recent(limit).await,
None => Ok(Vec::new()),
}
}
#[cfg(test)]
pub(crate) fn with_video_task_reader_for_tests(
repository: Arc<dyn VideoTaskReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: Some(repository),
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_request_candidate_reader_for_tests(
repository: Arc<dyn RequestCandidateReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: Some(repository),
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_usage_reader_for_tests(repository: Arc<dyn UsageReadRepository>) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: Some(repository),
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_auth_api_key_reader_for_tests(
repository: Arc<dyn AuthApiKeyReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: Some(repository),
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_decision_trace_readers_for_tests(
request_candidate_repository: Arc<dyn RequestCandidateReadRepository>,
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: Some(request_candidate_repository),
provider_catalog_reader: Some(provider_catalog_repository),
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_request_audit_readers_for_tests(
auth_api_key_repository: Arc<dyn AuthApiKeyReadRepository>,
request_candidate_repository: Arc<dyn RequestCandidateReadRepository>,
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
usage_repository: Arc<dyn UsageReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: Some(auth_api_key_repository),
request_candidate_reader: Some(request_candidate_repository),
provider_catalog_reader: Some(provider_catalog_repository),
usage_reader: Some(usage_repository),
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: None,
}
}
#[cfg(test)]
pub(crate) fn with_shadow_result_writer_for_tests(
repository: Arc<dyn ShadowResultWriteRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: None,
shadow_result_writer: Some(repository),
}
}
#[cfg(test)]
pub(crate) fn with_shadow_result_repository_for_tests<T>(repository: Arc<T>) -> Self
where
T: aether_data::repository::shadow_results::ShadowResultRepository + 'static,
{
let shadow_result_reader: Arc<dyn ShadowResultReadRepository> = repository.clone();
let shadow_result_writer: Arc<dyn ShadowResultWriteRepository> = repository;
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
request_candidate_reader: None,
provider_catalog_reader: None,
usage_reader: None,
video_task_reader: None,
shadow_result_reader: Some(shadow_result_reader),
shadow_result_writer: Some(shadow_result_writer),
}
}
}

2354
crates/aether-gateway/src/data/state/auth.rs Normal file
View File

File diff suppressed because it is too large Load Diff

466
crates/aether-gateway/src/data/state/catalog.rs Normal file
View File

@@ -0,0 +1,466 @@
use super::*;
impl GatewayDataState {
pub(crate) async fn list_request_candidates_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_by_request_id(request_id).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_request_candidates_by_provider_id(
&self,
provider_id: &str,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_by_provider_id(provider_id, limit).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_recent_request_candidates(
&self,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_recent(limit).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_finalized_request_candidates_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => {
repository
.list_finalized_by_endpoint_ids_since(endpoint_ids, since_unix_secs, limit)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn count_finalized_request_candidate_statuses_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
) -> Result<Vec<PublicHealthStatusCount>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => {
repository
.count_finalized_statuses_by_endpoint_ids_since(endpoint_ids, since_unix_secs)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn aggregate_finalized_request_candidate_timeline_by_endpoint_ids_since(
&self,
endpoint_ids: &[String],
since_unix_secs: u64,
until_unix_secs: u64,
segments: u32,
) -> Result<Vec<PublicHealthTimelineBucket>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => {
repository
.aggregate_finalized_timeline_by_endpoint_ids_since(
endpoint_ids,
since_unix_secs,
until_unix_secs,
segments,
)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn upsert_request_candidate(
&self,
candidate: UpsertRequestCandidateRecord,
) -> Result<Option<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_writer {
Some(repository) => repository.upsert(candidate).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn delete_request_candidates_created_before(
&self,
created_before_unix_secs: u64,
limit: usize,
) -> Result<usize, DataLayerError> {
match &self.request_candidate_writer {
Some(repository) => {
repository
.delete_created_before(created_before_unix_secs, limit)
.await
}
None => Ok(0),
}
}
pub(crate) async fn touch_auth_api_key_last_used(
&self,
api_key_id: &str,
) -> Result<bool, DataLayerError> {
match &self.auth_api_key_writer {
Some(repository) => repository.touch_last_used_at(api_key_id).await,
None => Ok(false),
}
}
pub(crate) async fn upsert_gemini_file_mapping(
&self,
record: UpsertGeminiFileMappingRecord,
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
match &self.gemini_file_mapping_writer {
Some(repository) => repository.upsert(record).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn list_gemini_file_mappings(
&self,
query: &GeminiFileMappingListQuery,
) -> Result<StoredGeminiFileMappingListPage, DataLayerError> {
match &self.gemini_file_mapping_reader {
Some(repository) => repository.list_mappings(query).await,
None => Ok(StoredGeminiFileMappingListPage {
items: Vec::new(),
total: 0,
}),
}
}
pub(crate) async fn summarize_gemini_file_mappings(
&self,
now_unix_secs: u64,
) -> Result<GeminiFileMappingStats, DataLayerError> {
match &self.gemini_file_mapping_reader {
Some(repository) => repository.summarize_mappings(now_unix_secs).await,
None => Ok(GeminiFileMappingStats {
total_mappings: 0,
active_mappings: 0,
expired_mappings: 0,
by_mime_type: Vec::new(),
}),
}
}
pub(crate) async fn delete_gemini_file_mapping_by_file_name(
&self,
file_name: &str,
) -> Result<bool, DataLayerError> {
match &self.gemini_file_mapping_writer {
Some(repository) => repository.delete_by_file_name(file_name).await,
None => Ok(false),
}
}
pub(crate) async fn delete_gemini_file_mapping_by_id(
&self,
mapping_id: &str,
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
match &self.gemini_file_mapping_writer {
Some(repository) => repository.delete_by_id(mapping_id).await,
None => Ok(None),
}
}
pub(crate) async fn delete_expired_gemini_file_mappings(
&self,
now_unix_secs: u64,
) -> Result<usize, DataLayerError> {
match &self.gemini_file_mapping_writer {
Some(repository) => repository.delete_expired_before(now_unix_secs).await,
None => Ok(0),
}
}
pub(crate) async fn cache_set_string_with_ttl(
&self,
key: &str,
value: &str,
ttl_seconds: u64,
) -> Result<(), DataLayerError> {
let Some(runner) = self.kv_runner() else {
return Ok(());
};
runner.setex(key, value, Some(ttl_seconds)).await?;
Ok(())
}
pub(crate) async fn cache_delete_key(&self, key: &str) -> Result<(), DataLayerError> {
let Some(runner) = self.kv_runner() else {
return Ok(());
};
let _deleted = runner.del(key).await?;
Ok(())
}
pub(in crate::gateway) async fn list_provider_catalog_providers_by_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_providers_by_ids(provider_ids).await,
None => Ok(Vec::new()),
}
}
pub(in crate::gateway) async fn list_provider_catalog_providers(
&self,
active_only: bool,
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_providers(active_only).await,
None => Ok(Vec::new()),
}
}
pub(in crate::gateway) async fn list_provider_catalog_endpoints_by_ids(
&self,
endpoint_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_endpoints_by_ids(endpoint_ids).await,
None => Ok(Vec::new()),
}
}
pub(in crate::gateway) async fn list_provider_catalog_endpoints_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => {
repository
.list_endpoints_by_provider_ids(provider_ids)
.await
}
None => Ok(Vec::new()),
}
}
pub(in crate::gateway) async fn list_provider_catalog_keys_by_ids(
&self,
key_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_keys_by_ids(key_ids).await,
None => Ok(Vec::new()),
}
}
pub(in crate::gateway) async fn list_provider_catalog_keys_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_keys_by_provider_ids(provider_ids).await,
None => Ok(Vec::new()),
}
}
pub(in crate::gateway) async fn list_provider_catalog_key_page(
&self,
query: &ProviderCatalogKeyListQuery,
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => repository.list_keys_page(query).await,
None => Ok(StoredProviderCatalogKeyPage {
items: Vec::new(),
total: 0,
}),
}
}
pub(in crate::gateway) async fn list_provider_catalog_key_stats_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
match &self.provider_catalog_reader {
Some(repository) => {
repository
.list_key_stats_by_provider_ids(provider_ids)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn update_provider_catalog_key_oauth_credentials(
&self,
key_id: &str,
encrypted_api_key: &str,
encrypted_auth_config: Option<&str>,
expires_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => {
repository
.update_key_oauth_credentials(
key_id,
encrypted_api_key,
encrypted_auth_config,
expires_at_unix_secs,
)
.await
}
None => Ok(false),
}
}
pub(crate) async fn create_provider_catalog_key(
&self,
key: &StoredProviderCatalogKey,
) -> Result<Option<StoredProviderCatalogKey>, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository.create_key(key).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn create_provider_catalog_provider(
&self,
provider: &StoredProviderCatalogProvider,
shift_existing_priorities_from: Option<i32>,
) -> Result<Option<StoredProviderCatalogProvider>, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository
.create_provider(provider, shift_existing_priorities_from)
.await
.map(Some),
None => Ok(None),
}
}
pub(crate) async fn update_provider_catalog_provider(
&self,
provider: &StoredProviderCatalogProvider,
) -> Result<Option<StoredProviderCatalogProvider>, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository.update_provider(provider).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn delete_provider_catalog_provider(
&self,
provider_id: &str,
) -> Result<bool, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository.delete_provider(provider_id).await,
None => Ok(false),
}
}
pub(crate) async fn cleanup_deleted_provider_catalog_refs(
&self,
provider_id: &str,
endpoint_ids: &[String],
key_ids: &[String],
) -> Result<(), DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => {
repository
.cleanup_deleted_provider_refs(provider_id, endpoint_ids, key_ids)
.await
}
None => Ok(()),
}
}
pub(crate) async fn create_provider_catalog_endpoint(
&self,
endpoint: &StoredProviderCatalogEndpoint,
) -> Result<Option<StoredProviderCatalogEndpoint>, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository.create_endpoint(endpoint).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn update_provider_catalog_endpoint(
&self,
endpoint: &StoredProviderCatalogEndpoint,
) -> Result<Option<StoredProviderCatalogEndpoint>, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository.update_endpoint(endpoint).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn delete_provider_catalog_endpoint(
&self,
endpoint_id: &str,
) -> Result<bool, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository.delete_endpoint(endpoint_id).await,
None => Ok(false),
}
}
pub(crate) async fn update_provider_catalog_key(
&self,
key: &StoredProviderCatalogKey,
) -> Result<Option<StoredProviderCatalogKey>, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository.update_key(key).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn delete_provider_catalog_key(
&self,
key_id: &str,
) -> Result<bool, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository.delete_key(key_id).await,
None => Ok(false),
}
}
pub(crate) async fn clear_provider_catalog_key_oauth_invalid_marker(
&self,
key_id: &str,
) -> Result<bool, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => repository.clear_key_oauth_invalid_marker(key_id).await,
None => Ok(false),
}
}
pub(crate) async fn update_provider_catalog_key_health_state(
&self,
key_id: &str,
is_active: bool,
health_by_format: Option<&serde_json::Value>,
circuit_breaker_by_format: Option<&serde_json::Value>,
) -> Result<bool, DataLayerError> {
match &self.provider_catalog_writer {
Some(repository) => {
repository
.update_key_health_state(
key_id,
is_active,
health_by_format,
circuit_breaker_by_format,
)
.await
}
None => Ok(false),
}
}
}

437
crates/aether-gateway/src/data/state/core.rs Normal file
View File

@@ -0,0 +1,437 @@
use super::*;
fn current_system_config_updated_at_unix_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
impl GatewayDataState {
pub(crate) fn disabled() -> Self {
Self::default()
}
pub(crate) fn from_config(config: GatewayDataConfig) -> Result<Self, DataLayerError> {
if !config.is_enabled() {
return Ok(Self {
config,
backends: None,
auth_api_key_reader: None,
auth_api_key_writer: None,
auth_module_reader: None,
auth_module_writer: None,
announcement_reader: None,
announcement_writer: None,
management_token_reader: None,
management_token_writer: None,
oauth_provider_reader: None,
oauth_provider_writer: None,
proxy_node_reader: None,
proxy_node_writer: None,
billing_reader: None,
gemini_file_mapping_reader: None,
gemini_file_mapping_writer: None,
global_model_reader: None,
global_model_writer: None,
minimal_candidate_selection_reader: None,
request_candidate_reader: None,
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
usage_writer: None,
user_reader: None,
user_preferences: None,
usage_worker_runner: None,
video_task_reader: None,
video_task_writer: None,
wallet_reader: None,
wallet_writer: None,
shadow_result_reader: None,
shadow_result_writer: None,
system_config_values: None,
});
}
let backends = DataBackends::from_config(config.to_data_layer_config())?;
let auth_api_key_reader = backends.read().auth_api_keys();
let auth_api_key_writer = backends.write().auth_api_keys();
let auth_module_reader = backends.read().auth_modules();
let auth_module_writer = backends.write().auth_modules();
let announcement_reader = backends.read().announcements();
let announcement_writer = backends.write().announcements();
let management_token_reader = backends.read().management_tokens();
let management_token_writer = backends.write().management_tokens();
let oauth_provider_reader = backends.read().oauth_providers();
let oauth_provider_writer = backends.write().oauth_providers();
let proxy_node_reader = backends.read().proxy_nodes();
let proxy_node_writer = backends.write().proxy_nodes();
let billing_reader = backends.read().billing();
let gemini_file_mapping_reader = backends.read().gemini_file_mappings();
let global_model_reader = backends.read().global_models();
let global_model_writer = backends.write().global_models();
let minimal_candidate_selection_reader = backends.read().minimal_candidate_selection();
let request_candidate_reader = backends.read().request_candidates();
let request_candidate_writer = backends.write().request_candidates();
let gemini_file_mapping_writer = backends.write().gemini_file_mappings();
let provider_catalog_reader = backends.read().provider_catalog();
let provider_catalog_writer = backends.write().provider_catalog();
let provider_quota_reader = backends.read().provider_quotas();
let provider_quota_writer = backends.write().provider_quotas();
let usage_reader = backends.read().usage();
let usage_writer = backends.write().usage();
let user_reader = backends.read().users();
let usage_worker_runner = backends.workers().redis();
let video_task_reader = backends.read().video_tasks();
let video_task_writer = backends.write().video_tasks();
let wallet_reader = backends.read().wallets();
let wallet_writer = backends.write().wallets();
let shadow_result_reader = backends.read().shadow_results();
let shadow_result_writer = backends.write().shadow_results();
Ok(Self {
config,
backends: Some(backends),
auth_api_key_reader,
auth_api_key_writer,
auth_module_reader,
auth_module_writer,
announcement_reader,
announcement_writer,
management_token_reader,
management_token_writer,
oauth_provider_reader,
oauth_provider_writer,
proxy_node_reader,
proxy_node_writer,
billing_reader,
gemini_file_mapping_reader,
gemini_file_mapping_writer,
global_model_reader,
global_model_writer,
minimal_candidate_selection_reader,
request_candidate_reader,
request_candidate_writer,
provider_catalog_reader,
provider_catalog_writer,
provider_quota_reader,
provider_quota_writer,
usage_reader,
usage_writer,
user_reader,
user_preferences: None,
usage_worker_runner,
video_task_reader,
video_task_writer,
wallet_reader,
wallet_writer,
shadow_result_reader,
shadow_result_writer,
system_config_values: None,
})
}
pub(crate) fn has_backends(&self) -> bool {
self.backends.is_some()
}
pub(crate) fn has_auth_api_key_reader(&self) -> bool {
self.auth_api_key_reader.is_some()
}
pub(crate) fn has_auth_api_key_writer(&self) -> bool {
self.auth_api_key_writer.is_some()
}
pub(crate) fn has_auth_module_writer(&self) -> bool {
self.auth_module_writer.is_some()
}
pub(crate) fn has_announcement_reader(&self) -> bool {
self.announcement_reader.is_some()
}
pub(crate) fn has_announcement_writer(&self) -> bool {
self.announcement_writer.is_some()
}
pub(crate) fn has_management_token_reader(&self) -> bool {
self.management_token_reader.is_some()
}
pub(crate) fn has_management_token_writer(&self) -> bool {
self.management_token_writer.is_some()
}
pub(crate) fn has_gemini_file_mapping_reader(&self) -> bool {
self.gemini_file_mapping_reader.is_some()
}
pub(crate) fn has_gemini_file_mapping_writer(&self) -> bool {
self.gemini_file_mapping_writer.is_some()
}
pub(crate) fn has_global_model_reader(&self) -> bool {
self.global_model_reader.is_some()
}
pub(crate) fn has_global_model_writer(&self) -> bool {
self.global_model_writer.is_some()
}
pub(crate) fn has_redis_backend(&self) -> bool {
self.backends
.as_ref()
.and_then(|backends| backends.redis())
.is_some()
}
#[allow(dead_code)]
pub(crate) fn has_minimal_candidate_selection_reader(&self) -> bool {
self.minimal_candidate_selection_reader.is_some()
}
pub(crate) fn has_request_candidate_reader(&self) -> bool {
self.request_candidate_reader.is_some()
}
pub(crate) fn has_request_candidate_writer(&self) -> bool {
self.request_candidate_writer.is_some()
}
pub(crate) fn has_provider_catalog_reader(&self) -> bool {
self.provider_catalog_reader.is_some()
}
pub(crate) fn has_provider_catalog_writer(&self) -> bool {
self.provider_catalog_writer.is_some()
}
pub(crate) fn has_proxy_node_reader(&self) -> bool {
self.proxy_node_reader.is_some()
}
pub(crate) fn has_proxy_node_writer(&self) -> bool {
self.proxy_node_writer.is_some()
}
pub(crate) fn oauth_refresh_lock_runner(&self) -> Option<RedisLockRunner> {
self.backends
.as_ref()
.and_then(|backends| backends.locks().redis())
}
pub(crate) fn kv_runner(&self) -> Option<RedisKvRunner> {
self.backends
.as_ref()
.and_then(|backends| backends.redis())
.and_then(|backend| backend.kv_runner(RedisKvRunnerConfig::default()).ok())
}
pub(crate) fn postgres_pool(&self) -> Option<aether_data::postgres::PostgresPool> {
self.backends
.as_ref()
.and_then(|backends| backends.postgres())
.map(|backend| backend.pool_clone())
}
pub(crate) fn postgres_max_connections(&self) -> Option<u32> {
self.config.postgres().map(|config| config.max_connections)
}
pub(crate) fn has_provider_quota_writer(&self) -> bool {
self.provider_quota_writer.is_some()
}
pub(crate) fn has_usage_reader(&self) -> bool {
self.usage_reader.is_some()
}
pub(crate) fn has_user_reader(&self) -> bool {
self.user_reader.is_some()
}
pub(crate) fn has_usage_writer(&self) -> bool {
self.usage_writer.is_some()
}
pub(crate) fn has_usage_worker_runner(&self) -> bool {
self.usage_worker_runner.is_some()
}
pub(crate) fn has_video_task_reader(&self) -> bool {
self.video_task_reader.is_some()
}
pub(crate) fn has_video_task_writer(&self) -> bool {
self.video_task_writer.is_some()
}
pub(crate) fn has_wallet_reader(&self) -> bool {
self.wallet_reader.is_some()
}
pub(crate) fn has_wallet_writer(&self) -> bool {
self.wallet_writer.is_some()
}
pub(crate) fn has_shadow_result_writer(&self) -> bool {
self.shadow_result_writer.is_some()
}
pub(crate) fn has_shadow_result_reader(&self) -> bool {
self.shadow_result_reader.is_some()
}
#[allow(dead_code)]
pub(in crate::gateway) fn encryption_key(&self) -> Option<&str> {
self.config.encryption_key()
}
pub(crate) async fn find_system_config_value(
&self,
key: &str,
) -> Result<Option<serde_json::Value>, DataLayerError> {
if let Some(values) = &self.system_config_values {
return Ok(values
.read()
.expect("system config values lock")
.get(key)
.map(|entry| entry.value.clone()));
}
match self
.backends
.as_ref()
.and_then(|backends| backends.postgres())
{
Some(backend) => backend.find_system_config_value(key).await,
None => Ok(None),
}
}
pub(crate) async fn upsert_system_config_value(
&self,
key: &str,
value: &serde_json::Value,
description: Option<&str>,
) -> Result<serde_json::Value, DataLayerError> {
Ok(self
.upsert_system_config_entry(key, value, description)
.await?
.value)
}
pub(crate) async fn list_system_config_entries(
&self,
) -> Result<Vec<StoredSystemConfigEntry>, DataLayerError> {
if let Some(values) = &self.system_config_values {
return Ok(values
.read()
.expect("system config values lock")
.values()
.cloned()
.collect());
}
match self
.backends
.as_ref()
.and_then(|backends| backends.postgres())
{
Some(backend) => backend.list_system_config_entries().await.map(|rows| {
rows.into_iter()
.map(|(key, value, description, updated_at_unix_secs)| {
StoredSystemConfigEntry {
key,
value,
description,
updated_at_unix_secs,
}
})
.collect()
}),
None => Ok(Vec::new()),
}
}
pub(crate) async fn upsert_system_config_entry(
&self,
key: &str,
value: &serde_json::Value,
description: Option<&str>,
) -> Result<StoredSystemConfigEntry, DataLayerError> {
if let Some(values) = &self.system_config_values {
let mut values = values.write().expect("system config values lock");
let description = description
.map(ToOwned::to_owned)
.or_else(|| values.get(key).and_then(|entry| entry.description.clone()));
let entry = StoredSystemConfigEntry {
key: key.to_string(),
value: value.clone(),
description,
updated_at_unix_secs: Some(current_system_config_updated_at_unix_secs()),
};
values.insert(key.to_string(), entry.clone());
return Ok(entry);
}
match self
.backends
.as_ref()
.and_then(|backends| backends.postgres())
{
Some(backend) => backend
.upsert_system_config_entry(key, value, description)
.await
.map(
|(key, value, description, updated_at_unix_secs)| StoredSystemConfigEntry {
key,
value,
description,
updated_at_unix_secs,
},
),
None => Ok(StoredSystemConfigEntry {
key: key.to_string(),
value: value.clone(),
description: description.map(ToOwned::to_owned),
updated_at_unix_secs: Some(current_system_config_updated_at_unix_secs()),
}),
}
}
pub(crate) async fn delete_system_config_value(
&self,
key: &str,
) -> Result<bool, DataLayerError> {
if let Some(values) = &self.system_config_values {
return Ok(values
.write()
.expect("system config values lock")
.remove(key)
.is_some());
}
match self
.backends
.as_ref()
.and_then(|backends| backends.postgres())
{
Some(backend) => backend.delete_system_config_value(key).await,
None => Ok(false),
}
}
pub(crate) async fn read_admin_system_stats(
&self,
) -> Result<(u64, u64, u64, u64), DataLayerError> {
match self
.backends
.as_ref()
.and_then(|backends| backends.postgres())
{
Some(backend) => backend.read_admin_system_stats().await,
None => Ok((0, 0, 0, 0)),
}
}
}

430
crates/aether-gateway/src/data/state/mod.rs Normal file
View File

@@ -0,0 +1,430 @@
use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use std::sync::RwLock;
use aether_data::redis::{RedisKvRunner, RedisKvRunnerConfig, RedisLockRunner, RedisStreamRunner};
use aether_data::repository::announcements::{
AnnouncementListQuery, AnnouncementReadRepository, AnnouncementWriteRepository,
CreateAnnouncementRecord, StoredAnnouncement, StoredAnnouncementPage, UpdateAnnouncementRecord,
};
use aether_data::repository::auth::{
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthApiKeyWriteRepository,
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
};
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,
StoredManagementTokenListPage, StoredManagementTokenWithUser, UpdateManagementTokenRecord,
};
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::shadow_results::{
merge_shadow_result_sample, RecordShadowResultSample, ShadowResultLookupKey,
ShadowResultReadRepository, ShadowResultWriteRepository, StoredShadowResult,
};
use aether_data::repository::usage::{
StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
UsageWriteRepository,
};
use aether_data::repository::users::{
StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, UserReadRepository,
};
use aether_data::repository::video_tasks::{
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
VideoTaskQueryFilter, VideoTaskReadRepository, VideoTaskStatusCount, VideoTaskWriteRepository,
};
use aether_data::repository::wallet::{
StoredUsageSettlement, StoredWalletSnapshot, UsageSettlementInput, WalletLookupKey,
WalletReadRepository, WalletWriteRepository,
};
use aether_data::{DataBackends, DataLayerError};
use chrono::{DateTime, Utc};
use super::auth::{
read_auth_api_key_snapshot, read_auth_api_key_snapshot_by_key_hash,
StoredGatewayAuthApiKeySnapshot,
};
use super::candidates::{read_request_candidate_trace, RequestCandidateTrace};
use super::config::GatewayDataConfig;
use super::decision_trace::{read_decision_trace, DecisionTrace};
use super::video_tasks::read_video_task_response;
use crate::gateway::provider_transport::{
read_provider_transport_snapshot, GatewayProviderTransportSnapshot,
};
use crate::gateway::scheduler::{
read_minimal_candidate_selection, GatewayMinimalCandidateSelectionCandidate,
};
use crate::gateway::usage::{
read_request_audit_bundle, read_request_usage_audit, RequestAuditBundle, RequestUsageAudit,
};
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
#[derive(Debug, Clone)]
pub(crate) struct StoredSystemConfigEntry {
pub(crate) key: String,
pub(crate) value: serde_json::Value,
pub(crate) description: Option<String>,
pub(crate) updated_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub(crate) struct StoredUserSessionRecord {
pub(crate) id: String,
pub(crate) user_id: String,
pub(crate) client_device_id: String,
pub(crate) device_label: Option<String>,
pub(crate) refresh_token_hash: String,
pub(crate) prev_refresh_token_hash: Option<String>,
pub(crate) rotated_at: Option<DateTime<Utc>>,
pub(crate) last_seen_at: Option<DateTime<Utc>>,
pub(crate) expires_at: Option<DateTime<Utc>>,
pub(crate) revoked_at: Option<DateTime<Utc>>,
pub(crate) revoke_reason: Option<String>,
pub(crate) ip_address: Option<String>,
pub(crate) user_agent: Option<String>,
pub(crate) created_at: Option<DateTime<Utc>>,
pub(crate) updated_at: Option<DateTime<Utc>>,
}
impl StoredUserSessionRecord {
pub(crate) const REFRESH_GRACE_SECONDS: i64 = 10;
pub(crate) const TOUCH_INTERVAL_SECONDS: i64 = 300;
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
id: String,
user_id: String,
client_device_id: String,
device_label: Option<String>,
refresh_token_hash: String,
prev_refresh_token_hash: Option<String>,
rotated_at: Option<DateTime<Utc>>,
last_seen_at: Option<DateTime<Utc>>,
expires_at: Option<DateTime<Utc>>,
revoked_at: Option<DateTime<Utc>>,
revoke_reason: Option<String>,
ip_address: Option<String>,
user_agent: Option<String>,
created_at: Option<DateTime<Utc>>,
updated_at: Option<DateTime<Utc>>,
) -> Result<Self, DataLayerError> {
if id.trim().is_empty() {
return Err(DataLayerError::UnexpectedValue(
"user_sessions.id is empty".to_string(),
));
}
if user_id.trim().is_empty() {
return Err(DataLayerError::UnexpectedValue(
"user_sessions.user_id is empty".to_string(),
));
}
if client_device_id.trim().is_empty() {
return Err(DataLayerError::UnexpectedValue(
"user_sessions.client_device_id is empty".to_string(),
));
}
if refresh_token_hash.trim().is_empty() {
return Err(DataLayerError::UnexpectedValue(
"user_sessions.refresh_token_hash is empty".to_string(),
));
}
Ok(Self {
id,
user_id,
client_device_id,
device_label,
refresh_token_hash,
prev_refresh_token_hash,
rotated_at,
last_seen_at,
expires_at,
revoked_at,
revoke_reason,
ip_address,
user_agent,
created_at,
updated_at,
})
}
pub(crate) fn hash_refresh_token(token: &str) -> String {
use sha2::Digest;
let mut hasher = sha2::Sha256::new();
hasher.update(token.as_bytes());
format!("{:x}", hasher.finalize())
}
pub(crate) fn verify_refresh_token(&self, token: &str, now: DateTime<Utc>) -> (bool, bool) {
let token_hash = Self::hash_refresh_token(token);
if self.refresh_token_hash == token_hash {
return (true, false);
}
let Some(prev_hash) = self.prev_refresh_token_hash.as_ref() else {
return (false, false);
};
let Some(rotated_at) = self.rotated_at else {
return (false, false);
};
if prev_hash == &token_hash
&& now.signed_duration_since(rotated_at).num_seconds() <= Self::REFRESH_GRACE_SECONDS
{
return (true, true);
}
(false, false)
}
pub(crate) fn is_revoked(&self) -> bool {
self.revoked_at.is_some()
}
pub(crate) fn is_expired(&self, now: DateTime<Utc>) -> bool {
self.expires_at.is_none_or(|expires_at| expires_at <= now)
}
pub(crate) fn should_touch(&self, now: DateTime<Utc>) -> bool {
self.last_seen_at
.map(|last_seen_at| {
now.signed_duration_since(last_seen_at).num_seconds()
>= Self::TOUCH_INTERVAL_SECONDS
})
.unwrap_or(true)
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub(crate) struct StoredUserPreferenceRecord {
pub(crate) user_id: String,
pub(crate) avatar_url: Option<String>,
pub(crate) bio: Option<String>,
pub(crate) default_provider_id: Option<String>,
pub(crate) default_provider_name: Option<String>,
pub(crate) theme: String,
pub(crate) language: String,
pub(crate) timezone: String,
pub(crate) email_notifications: bool,
pub(crate) usage_alerts: bool,
pub(crate) announcement_notifications: bool,
}
impl StoredUserPreferenceRecord {
pub(crate) fn default_for_user(user_id: impl Into<String>) -> Self {
Self {
user_id: user_id.into(),
avatar_url: None,
bio: None,
default_provider_id: None,
default_provider_name: None,
theme: "light".to_string(),
language: "zh-CN".to_string(),
timezone: "Asia/Shanghai".to_string(),
email_notifications: true,
usage_alerts: true,
announcement_notifications: true,
}
}
}
#[derive(Clone, Default)]
pub(crate) struct GatewayDataState {
config: GatewayDataConfig,
backends: Option<DataBackends>,
auth_api_key_reader: Option<Arc<dyn AuthApiKeyReadRepository>>,
auth_api_key_writer: Option<Arc<dyn AuthApiKeyWriteRepository>>,
auth_module_reader: Option<Arc<dyn AuthModuleReadRepository>>,
auth_module_writer: Option<Arc<dyn AuthModuleWriteRepository>>,
announcement_reader: Option<Arc<dyn AnnouncementReadRepository>>,
announcement_writer: Option<Arc<dyn AnnouncementWriteRepository>>,
management_token_reader: Option<Arc<dyn ManagementTokenReadRepository>>,
management_token_writer: Option<Arc<dyn ManagementTokenWriteRepository>>,
oauth_provider_reader: Option<Arc<dyn OAuthProviderReadRepository>>,
oauth_provider_writer: Option<Arc<dyn OAuthProviderWriteRepository>>,
proxy_node_reader: Option<Arc<dyn ProxyNodeReadRepository>>,
proxy_node_writer: Option<Arc<dyn ProxyNodeWriteRepository>>,
billing_reader: Option<Arc<dyn BillingReadRepository>>,
gemini_file_mapping_reader: Option<Arc<dyn GeminiFileMappingReadRepository>>,
gemini_file_mapping_writer: Option<Arc<dyn GeminiFileMappingWriteRepository>>,
global_model_reader: Option<Arc<dyn GlobalModelReadRepository>>,
global_model_writer: Option<Arc<dyn GlobalModelWriteRepository>>,
minimal_candidate_selection_reader: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
request_candidate_reader: Option<Arc<dyn RequestCandidateReadRepository>>,
request_candidate_writer: Option<Arc<dyn RequestCandidateWriteRepository>>,
provider_catalog_reader: Option<Arc<dyn ProviderCatalogReadRepository>>,
provider_catalog_writer: Option<Arc<dyn ProviderCatalogWriteRepository>>,
provider_quota_reader: Option<Arc<dyn ProviderQuotaReadRepository>>,
provider_quota_writer: Option<Arc<dyn ProviderQuotaWriteRepository>>,
usage_reader: Option<Arc<dyn UsageReadRepository>>,
usage_writer: Option<Arc<dyn UsageWriteRepository>>,
user_reader: Option<Arc<dyn UserReadRepository>>,
user_preferences: Option<Arc<RwLock<BTreeMap<String, StoredUserPreferenceRecord>>>>,
usage_worker_runner: Option<RedisStreamRunner>,
video_task_reader: Option<Arc<dyn VideoTaskReadRepository>>,
video_task_writer: Option<Arc<dyn VideoTaskWriteRepository>>,
wallet_reader: Option<Arc<dyn WalletReadRepository>>,
wallet_writer: Option<Arc<dyn WalletWriteRepository>>,
shadow_result_reader: Option<Arc<dyn ShadowResultReadRepository>>,
shadow_result_writer: Option<Arc<dyn ShadowResultWriteRepository>>,
system_config_values: Option<Arc<RwLock<BTreeMap<String, StoredSystemConfigEntry>>>>,
}
impl fmt::Debug for GatewayDataState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("GatewayDataState")
.field("config", &self.config)
.field("has_backends", &self.backends.is_some())
.field(
"has_auth_api_key_reader",
&self.auth_api_key_reader.is_some(),
)
.field(
"has_auth_api_key_writer",
&self.auth_api_key_writer.is_some(),
)
.field("has_auth_module_reader", &self.auth_module_reader.is_some())
.field("has_auth_module_writer", &self.auth_module_writer.is_some())
.field(
"has_announcement_reader",
&self.announcement_reader.is_some(),
)
.field(
"has_announcement_writer",
&self.announcement_writer.is_some(),
)
.field(
"has_management_token_reader",
&self.management_token_reader.is_some(),
)
.field(
"has_management_token_writer",
&self.management_token_writer.is_some(),
)
.field(
"has_oauth_provider_reader",
&self.oauth_provider_reader.is_some(),
)
.field(
"has_oauth_provider_writer",
&self.oauth_provider_writer.is_some(),
)
.field("has_proxy_node_reader", &self.proxy_node_reader.is_some())
.field("has_proxy_node_writer", &self.proxy_node_writer.is_some())
.field("has_billing_reader", &self.billing_reader.is_some())
.field(
"has_gemini_file_mapping_reader",
&self.gemini_file_mapping_reader.is_some(),
)
.field(
"has_gemini_file_mapping_writer",
&self.gemini_file_mapping_writer.is_some(),
)
.field(
"has_global_model_reader",
&self.global_model_reader.is_some(),
)
.field(
"has_global_model_writer",
&self.global_model_writer.is_some(),
)
.field(
"has_minimal_candidate_selection_reader",
&self.minimal_candidate_selection_reader.is_some(),
)
.field(
"has_request_candidate_reader",
&self.request_candidate_reader.is_some(),
)
.field(
"has_request_candidate_writer",
&self.request_candidate_writer.is_some(),
)
.field(
"has_provider_catalog_reader",
&self.provider_catalog_reader.is_some(),
)
.field(
"has_provider_catalog_writer",
&self.provider_catalog_writer.is_some(),
)
.field(
"has_provider_quota_reader",
&self.provider_quota_reader.is_some(),
)
.field(
"has_provider_quota_writer",
&self.provider_quota_writer.is_some(),
)
.field("has_usage_reader", &self.usage_reader.is_some())
.field("has_usage_writer", &self.usage_writer.is_some())
.field("has_user_preferences", &self.user_preferences.is_some())
.field(
"has_usage_worker_runner",
&self.usage_worker_runner.is_some(),
)
.field("has_video_task_reader", &self.video_task_reader.is_some())
.field("has_video_task_writer", &self.video_task_writer.is_some())
.field("has_wallet_reader", &self.wallet_reader.is_some())
.field("has_wallet_writer", &self.wallet_writer.is_some())
.field(
"has_shadow_result_reader",
&self.shadow_result_reader.is_some(),
)
.field(
"has_shadow_result_writer",
&self.shadow_result_writer.is_some(),
)
.field(
"has_system_config_values",
&self.system_config_values.is_some(),
)
.finish()
}
}
mod auth;
mod catalog;
mod core;
mod models;
mod runtime;
#[cfg(test)]
mod testing;

250
crates/aether-gateway/src/data/state/models.rs Normal file
View File

@@ -0,0 +1,250 @@
use super::*;
impl GatewayDataState {
pub(crate) async fn list_minimal_candidate_selection_rows(
&self,
api_format: &str,
global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
match &self.minimal_candidate_selection_reader {
Some(repository) => {
repository
.list_for_exact_api_format_and_global_model(api_format, global_model_name)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_minimal_candidate_selection_rows_for_api_format(
&self,
api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
match &self.minimal_candidate_selection_reader {
Some(repository) => repository.list_for_exact_api_format(api_format).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_public_global_models(
&self,
query: &PublicGlobalModelQuery,
) -> Result<StoredPublicGlobalModelPage, DataLayerError> {
match &self.global_model_reader {
Some(repository) => repository.list_public_models(query).await,
None => Ok(StoredPublicGlobalModelPage {
items: Vec::new(),
total: 0,
}),
}
}
pub(crate) async fn get_public_global_model_by_name(
&self,
model_name: &str,
) -> Result<Option<StoredPublicGlobalModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => repository.get_public_model_by_name(model_name).await,
None => Ok(None),
}
}
pub(crate) async fn list_public_catalog_models(
&self,
query: &PublicCatalogModelListQuery,
) -> Result<Vec<StoredPublicCatalogModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => repository.list_public_catalog_models(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn search_public_catalog_models(
&self,
query: &PublicCatalogModelSearchQuery,
) -> Result<Vec<StoredPublicCatalogModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => repository.search_public_catalog_models(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_admin_provider_models(
&self,
query: &AdminProviderModelListQuery,
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => repository.list_admin_provider_models(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_admin_global_models(
&self,
query: &AdminGlobalModelListQuery,
) -> Result<StoredAdminGlobalModelPage, DataLayerError> {
match &self.global_model_reader {
Some(repository) => repository.list_admin_global_models(query).await,
None => Ok(StoredAdminGlobalModelPage {
items: Vec::new(),
total: 0,
}),
}
}
pub(crate) async fn get_admin_provider_model(
&self,
provider_id: &str,
model_id: &str,
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => {
repository
.get_admin_provider_model(provider_id, model_id)
.await
}
None => Ok(None),
}
}
pub(crate) async fn list_admin_provider_available_source_models(
&self,
provider_id: &str,
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => {
repository
.list_admin_provider_available_source_models(provider_id)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn get_admin_global_model_by_id(
&self,
global_model_id: &str,
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => {
repository
.get_admin_global_model_by_id(global_model_id)
.await
}
None => Ok(None),
}
}
pub(crate) async fn get_admin_global_model_by_name(
&self,
model_name: &str,
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => repository.get_admin_global_model_by_name(model_name).await,
None => Ok(None),
}
}
pub(crate) async fn list_admin_provider_models_by_global_model_id(
&self,
global_model_id: &str,
) -> Result<Vec<StoredAdminProviderModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => {
repository
.list_admin_provider_models_by_global_model_id(global_model_id)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn create_admin_provider_model(
&self,
record: &UpsertAdminProviderModelRecord,
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
match &self.global_model_writer {
Some(repository) => repository.create_admin_provider_model(record).await,
None => Ok(None),
}
}
pub(crate) async fn update_admin_provider_model(
&self,
record: &UpsertAdminProviderModelRecord,
) -> Result<Option<StoredAdminProviderModel>, DataLayerError> {
match &self.global_model_writer {
Some(repository) => repository.update_admin_provider_model(record).await,
None => Ok(None),
}
}
pub(crate) async fn delete_admin_provider_model(
&self,
provider_id: &str,
model_id: &str,
) -> Result<bool, DataLayerError> {
match &self.global_model_writer {
Some(repository) => {
repository
.delete_admin_provider_model(provider_id, model_id)
.await
}
None => Ok(false),
}
}
pub(crate) async fn create_admin_global_model(
&self,
record: &CreateAdminGlobalModelRecord,
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
match &self.global_model_writer {
Some(repository) => repository.create_admin_global_model(record).await,
None => Ok(None),
}
}
pub(crate) async fn update_admin_global_model(
&self,
record: &UpdateAdminGlobalModelRecord,
) -> Result<Option<StoredAdminGlobalModel>, DataLayerError> {
match &self.global_model_writer {
Some(repository) => repository.update_admin_global_model(record).await,
None => Ok(None),
}
}
pub(crate) async fn delete_admin_global_model(
&self,
global_model_id: &str,
) -> Result<bool, DataLayerError> {
match &self.global_model_writer {
Some(repository) => repository.delete_admin_global_model(global_model_id).await,
None => Ok(false),
}
}
pub(crate) async fn list_provider_model_stats(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderModelStats>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => repository.list_provider_model_stats(provider_ids).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_active_global_model_ids_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderActiveGlobalModel>, DataLayerError> {
match &self.global_model_reader {
Some(repository) => {
repository
.list_active_global_model_ids_by_provider_ids(provider_ids)
.await
}
None => Ok(Vec::new()),
}
}
}

625
crates/aether-gateway/src/data/state/runtime.rs Normal file
View File

@@ -0,0 +1,625 @@
use super::*;
fn is_missing_shadow_results_relation_error(error: &DataLayerError) -> bool {
error
.to_string()
.contains("relation \"gateway_shadow_results\" does not exist")
}
impl GatewayDataState {
pub(crate) async fn list_announcements(
&self,
query: &AnnouncementListQuery,
) -> Result<StoredAnnouncementPage, DataLayerError> {
match &self.announcement_reader {
Some(repository) => repository.list_announcements(query).await,
None => Ok(StoredAnnouncementPage::default()),
}
}
pub(crate) async fn find_announcement_by_id(
&self,
announcement_id: &str,
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
match &self.announcement_reader {
Some(repository) => repository.find_by_id(announcement_id).await,
None => Ok(None),
}
}
pub(crate) async fn count_unread_active_announcements(
&self,
user_id: &str,
now_unix_secs: u64,
) -> Result<u64, DataLayerError> {
match &self.announcement_reader {
Some(repository) => {
repository
.count_unread_active_announcements(user_id, now_unix_secs)
.await
}
None => Ok(0),
}
}
pub(crate) async fn create_announcement(
&self,
record: CreateAnnouncementRecord,
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
match &self.announcement_writer {
Some(repository) => repository.create_announcement(record).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn update_announcement(
&self,
record: UpdateAnnouncementRecord,
) -> Result<Option<StoredAnnouncement>, DataLayerError> {
match &self.announcement_writer {
Some(repository) => repository.update_announcement(record).await,
None => Ok(None),
}
}
pub(crate) async fn delete_announcement(
&self,
announcement_id: &str,
) -> Result<bool, DataLayerError> {
match &self.announcement_writer {
Some(repository) => repository.delete_announcement(announcement_id).await,
None => Ok(false),
}
}
pub(crate) async fn mark_announcement_as_read(
&self,
user_id: &str,
announcement_id: &str,
read_at_unix_secs: u64,
) -> Result<bool, DataLayerError> {
match &self.announcement_writer {
Some(repository) => {
repository
.mark_announcement_as_read(user_id, announcement_id, read_at_unix_secs)
.await
}
None => Ok(false),
}
}
pub(crate) async fn find_video_task(
&self,
key: VideoTaskLookupKey<'_>,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
match &self.video_task_reader {
Some(repository) => repository.find(key).await,
None => Ok(None),
}
}
pub(crate) async fn list_video_task_page(
&self,
filter: &VideoTaskQueryFilter,
offset: usize,
limit: usize,
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
match &self.video_task_reader {
Some(repository) => repository.list_page(filter, offset, limit).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn count_video_tasks(
&self,
filter: &VideoTaskQueryFilter,
) -> Result<u64, DataLayerError> {
match &self.video_task_reader {
Some(repository) => repository.count(filter).await,
None => Ok(0),
}
}
pub(crate) async fn count_video_tasks_by_status(
&self,
filter: &VideoTaskQueryFilter,
) -> Result<Vec<VideoTaskStatusCount>, DataLayerError> {
match &self.video_task_reader {
Some(repository) => repository.count_by_status(filter).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn count_distinct_video_task_users(
&self,
filter: &VideoTaskQueryFilter,
) -> Result<u64, DataLayerError> {
match &self.video_task_reader {
Some(repository) => repository.count_distinct_users(filter).await,
None => Ok(0),
}
}
pub(crate) async fn top_video_task_models(
&self,
filter: &VideoTaskQueryFilter,
limit: usize,
) -> Result<Vec<VideoTaskModelCount>, DataLayerError> {
match &self.video_task_reader {
Some(repository) => repository.top_models(filter, limit).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn count_video_tasks_created_since(
&self,
filter: &VideoTaskQueryFilter,
created_since_unix_secs: u64,
) -> Result<u64, DataLayerError> {
match &self.video_task_reader {
Some(repository) => {
repository
.count_created_since(filter, created_since_unix_secs)
.await
}
None => Ok(0),
}
}
pub(crate) async fn upsert_video_task(
&self,
task: UpsertVideoTask,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
match &self.video_task_writer {
Some(repository) => repository.upsert(task).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn update_active_video_task(
&self,
task: UpsertVideoTask,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
match &self.video_task_writer {
Some(repository) => repository.update_if_active(task).await,
None => Ok(None),
}
}
pub(crate) async fn claim_due_video_tasks(
&self,
now_unix_secs: u64,
claim_until_unix_secs: u64,
limit: usize,
) -> Result<Vec<StoredVideoTask>, DataLayerError> {
match &self.video_task_writer {
Some(repository) => {
repository
.claim_due(now_unix_secs, claim_until_unix_secs, limit)
.await
}
None => Ok(Vec::new()),
}
}
pub(crate) async fn find_wallet(
&self,
key: WalletLookupKey<'_>,
) -> Result<Option<StoredWalletSnapshot>, DataLayerError> {
match &self.wallet_reader {
Some(repository) => repository.find(key).await,
None => Ok(None),
}
}
pub(crate) async fn list_wallets_by_api_key_ids(
&self,
api_key_ids: &[String],
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
match &self.wallet_reader {
Some(repository) => repository.list_wallets_by_api_key_ids(api_key_ids).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_wallets_by_user_ids(
&self,
user_ids: &[String],
) -> Result<Vec<StoredWalletSnapshot>, DataLayerError> {
match &self.wallet_reader {
Some(repository) => repository.list_wallets_by_user_ids(user_ids).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
match &self.wallet_writer {
Some(repository) => repository.settle_usage(input).await,
None => Ok(None),
}
}
pub(crate) async fn reset_due_provider_quotas(
&self,
now_unix_secs: u64,
) -> Result<usize, DataLayerError> {
match &self.provider_quota_writer {
Some(repository) => repository.reset_due(now_unix_secs).await,
None => Ok(0),
}
}
pub(crate) async fn find_provider_quota_by_provider_id(
&self,
provider_id: &str,
) -> Result<Option<StoredProviderQuotaSnapshot>, DataLayerError> {
match &self.provider_quota_reader {
Some(repository) => repository.find_by_provider_id(provider_id).await,
None => Ok(None),
}
}
#[allow(dead_code)]
pub(crate) async fn upsert_usage(
&self,
usage: UpsertUsageRecord,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
match &self.usage_writer {
Some(repository) => repository.upsert(usage).await.map(Some),
None => Ok(None),
}
}
pub(in crate::gateway) async fn find_request_usage_by_request_id(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.find_by_request_id(request_id).await,
None => Ok(None),
}
}
pub(in crate::gateway) async fn find_request_usage_by_id(
&self,
usage_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.find_by_id(usage_id).await,
None => Ok(None),
}
}
pub(crate) async fn list_usage_audits(
&self,
query: &aether_data::repository::usage::UsageAuditListQuery,
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.list_usage_audits(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_recent_usage_audits(
&self,
user_id: Option<&str>,
limit: usize,
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.list_recent_usage_audits(user_id, limit).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn summarize_usage_total_tokens_by_api_key_ids(
&self,
api_key_ids: &[String],
) -> Result<std::collections::BTreeMap<String, u64>, DataLayerError> {
match &self.usage_reader {
Some(repository) => {
repository
.summarize_total_tokens_by_api_key_ids(api_key_ids)
.await
}
None => Ok(std::collections::BTreeMap::new()),
}
}
pub(crate) async fn list_users_by_ids(
&self,
user_ids: &[String],
) -> Result<Vec<StoredUserSummary>, DataLayerError> {
match &self.user_reader {
Some(repository) => repository.list_users_by_ids(user_ids).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_export_users(
&self,
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
match &self.user_reader {
Some(repository) => repository.list_export_users().await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_export_users_page(
&self,
query: &aether_data::repository::users::UserExportListQuery,
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
match &self.user_reader {
Some(repository) => repository.list_export_users_page(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn summarize_export_users(
&self,
) -> Result<aether_data::repository::users::UserExportSummary, DataLayerError> {
match &self.user_reader {
Some(repository) => repository.summarize_export_users().await,
None => Ok(aether_data::repository::users::UserExportSummary::default()),
}
}
pub(crate) async fn find_export_user_by_id(
&self,
user_id: &str,
) -> Result<Option<StoredUserExportRow>, DataLayerError> {
match &self.user_reader {
Some(repository) => repository.find_export_user_by_id(user_id).await,
None => Ok(None),
}
}
pub(crate) async fn list_non_admin_export_users(
&self,
) -> Result<Vec<StoredUserExportRow>, DataLayerError> {
match &self.user_reader {
Some(repository) => repository.list_non_admin_export_users().await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_user_auth_by_ids(
&self,
user_ids: &[String],
) -> Result<Vec<StoredUserAuthRecord>, DataLayerError> {
match &self.user_reader {
Some(repository) => repository.list_user_auth_by_ids(user_ids).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn summarize_provider_usage_since(
&self,
provider_id: &str,
since_unix_secs: u64,
) -> Result<StoredProviderUsageSummary, DataLayerError> {
match &self.usage_reader {
Some(repository) => {
repository
.summarize_provider_usage_since(provider_id, since_unix_secs)
.await
}
None => Ok(StoredProviderUsageSummary::default()),
}
}
pub(crate) fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
self.usage_worker_runner.clone()
}
pub(in crate::gateway) async fn find_billing_model_context(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
match &self.billing_reader {
Some(repository) => {
repository
.find_model_context(provider_id, provider_api_key_id, global_model_name)
.await
}
None => Ok(None),
}
}
pub(crate) async fn read_request_candidate_trace(
&self,
request_id: &str,
attempted_only: bool,
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
read_request_candidate_trace(self, request_id, attempted_only).await
}
pub(crate) async fn read_decision_trace(
&self,
request_id: &str,
attempted_only: bool,
) -> Result<Option<DecisionTrace>, DataLayerError> {
read_decision_trace(self, request_id, attempted_only).await
}
pub(crate) async fn read_request_usage_audit(
&self,
request_id: &str,
) -> Result<Option<RequestUsageAudit>, DataLayerError> {
read_request_usage_audit(self, request_id).await
}
pub(crate) async fn read_request_audit_bundle(
&self,
request_id: &str,
attempted_only: bool,
now_unix_secs: u64,
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
read_request_audit_bundle(self, request_id, attempted_only, now_unix_secs).await
}
#[allow(dead_code)]
pub(crate) async fn read_provider_transport_snapshot(
&self,
provider_id: &str,
endpoint_id: &str,
key_id: &str,
) -> Result<Option<GatewayProviderTransportSnapshot>, DataLayerError> {
read_provider_transport_snapshot(self, provider_id, endpoint_id, key_id).await
}
pub(crate) async fn read_minimal_candidate_selection(
&self,
api_format: &str,
global_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&StoredGatewayAuthApiKeySnapshot>,
) -> Result<Vec<GatewayMinimalCandidateSelectionCandidate>, DataLayerError> {
read_minimal_candidate_selection(
self,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
)
.await
}
pub(crate) async fn read_video_task_response(
&self,
route_family: Option<&str>,
request_path: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
read_video_task_response(self, route_family, request_path).await
}
#[cfg(test)]
pub(crate) async fn write_shadow_result(
&self,
result: aether_data::repository::shadow_results::UpsertShadowResult,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
match &self.shadow_result_writer {
Some(repository) => match repository.upsert(result).await {
Ok(stored) => Ok(Some(stored)),
Err(err) if is_missing_shadow_results_relation_error(&err) => Ok(None),
Err(err) => Err(err),
},
None => Ok(None),
}
}
pub(crate) async fn record_shadow_result_sample(
&self,
sample: RecordShadowResultSample,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
let Some(writer) = &self.shadow_result_writer else {
return Ok(None);
};
let existing = match &self.shadow_result_reader {
Some(reader) => {
match reader
.find(ShadowResultLookupKey::TraceFingerprint {
trace_id: &sample.trace_id,
request_fingerprint: &sample.request_fingerprint,
})
.await
{
Ok(existing) => existing,
Err(err) if is_missing_shadow_results_relation_error(&err) => return Ok(None),
Err(err) => return Err(err),
}
}
None => None,
};
let merged = merge_shadow_result_sample(existing.as_ref(), sample);
match writer.upsert(merged).await {
Ok(stored) => Ok(Some(stored)),
Err(err) if is_missing_shadow_results_relation_error(&err) => Ok(None),
Err(err) => Err(err),
}
}
pub(crate) async fn list_recent_shadow_results(
&self,
limit: usize,
) -> Result<Vec<StoredShadowResult>, DataLayerError> {
match &self.shadow_result_reader {
Some(repository) => match repository.list_recent(limit).await {
Ok(results) => Ok(results),
Err(err) if is_missing_shadow_results_relation_error(&err) => Ok(Vec::new()),
Err(err) => Err(err),
},
None => Ok(Vec::new()),
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_data::repository::users::{InMemoryUserReadRepository, StoredUserExportRow};
use aether_data::DataLayerError;
use super::{is_missing_shadow_results_relation_error, GatewayDataState};
#[tokio::test]
async fn lists_non_admin_export_users_from_user_reader() {
let repository = Arc::new(InMemoryUserReadRepository::seed_export_users(vec![
StoredUserExportRow::new(
"user-1".to_string(),
Some("alice@example.com".to_string()),
true,
"alice".to_string(),
Some("hash".to_string()),
"user".to_string(),
"local".to_string(),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
Some(60),
Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}})),
true,
)
.expect("user export row should build"),
]));
let state = GatewayDataState::with_user_reader_for_tests(repository);
let rows = state
.list_non_admin_export_users()
.await
.expect("export users should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].username, "alice");
assert!(rows[0].email_verified);
assert_eq!(rows[0].password_hash.as_deref(), Some("hash"));
assert_eq!(rows[0].allowed_models, Some(vec!["gpt-4.1".to_string()]));
assert_eq!(
rows[0].model_capability_settings,
Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}}))
);
}
#[test]
fn detects_missing_shadow_results_relation_error_text() {
let missing_relation = DataLayerError::UnexpectedValue(
"postgres error: error returned from database: relation \"gateway_shadow_results\" does not exist"
.to_string(),
);
let other = DataLayerError::UnexpectedValue(
"postgres error: error returned from database: relation \"usage\" does not exist"
.to_string(),
);
assert!(is_missing_shadow_results_relation_error(&missing_relation));
assert!(!is_missing_shadow_results_relation_error(&other));
}
}

2777
crates/aether-gateway/src/data/state/testing.rs Normal file
View File

File diff suppressed because it is too large Load Diff

358
crates/aether-gateway/src/data/tests.rs
View File

@@ -1,8 +1,13 @@
use std::sync::Arc;
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,
};
@@ -13,13 +18,16 @@ use aether_data::repository::provider_catalog::{
use aether_data::repository::shadow_results::{
InMemoryShadowResultRepository, RecordShadowResultSample, ShadowResultLookupKey,
ShadowResultMatchStatus, ShadowResultReadRepository, ShadowResultSampleOrigin,
UpsertShadowResult,
ShadowResultWriteRepository, StoredShadowResult, UpsertShadowResult,
};
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredRequestUsageAudit};
use aether_data::repository::video_tasks::{
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus,
VideoTaskWriteRepository,
};
use aether_data::DataLayerError;
use async_trait::async_trait;
use serde_json::json;
use super::{GatewayDataConfig, GatewayDataState};
use crate::gateway::AppState;
@@ -31,8 +39,11 @@ fn disabled_gateway_data_state_has_no_backends() {
assert!(!state.has_backends());
assert!(!state.has_auth_api_key_reader());
assert!(!state.has_minimal_candidate_selection_reader());
assert!(!state.has_request_candidate_reader());
assert!(!state.has_provider_catalog_reader());
assert!(!state.has_proxy_node_reader());
assert!(!state.has_proxy_node_writer());
assert!(!state.has_usage_reader());
assert!(!state.has_video_task_reader());
assert!(!state.has_shadow_result_reader());
@@ -49,8 +60,11 @@ async fn postgres_gateway_data_state_builds_video_task_reader() {
assert!(state.has_backends());
assert!(state.has_auth_api_key_reader());
assert!(state.has_minimal_candidate_selection_reader());
assert!(state.has_request_candidate_reader());
assert!(state.has_provider_catalog_reader());
assert!(state.has_proxy_node_reader());
assert!(state.has_proxy_node_writer());
assert!(state.has_usage_reader());
assert!(state.has_video_task_reader());
assert!(state.has_shadow_result_reader());
@@ -64,19 +78,41 @@ async fn data_state_find_uses_configured_read_repository() {
.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("short-task-1".to_string()),
request_id: "request-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
username: Some("user".to_string()),
api_key_name: Some("primary".to_string()),
external_task_id: Some("ext-task-1".to_string()),
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("provider-key-1".to_string()),
client_api_format: Some("openai:video".to_string()),
provider_api_format: Some("openai:video".to_string()),
format_converted: false,
model: Some("sora-2".to_string()),
prompt: Some("hello".to_string()),
original_request_body: Some(json!({"prompt": "hello"})),
duration_seconds: Some(4),
resolution: Some("720p".to_string()),
aspect_ratio: Some("16:9".to_string()),
size: Some("1280x720".to_string()),
status: VideoTaskStatus::Queued,
progress_percent: 0,
progress_message: None,
retry_count: 0,
poll_interval_seconds: 10,
next_poll_at_unix_secs: Some(100),
poll_count: 0,
max_poll_count: 360,
created_at_unix_secs: 100,
submitted_at_unix_secs: Some(100),
completed_at_unix_secs: None,
updated_at_unix_secs: 100,
error_code: None,
error_message: None,
video_url: None,
request_metadata: None,
})
.await
.expect("upsert should succeed");
@@ -107,8 +143,11 @@ async fn app_state_wires_gateway_data_state_from_config() {
assert!(state.data.has_backends());
assert!(state.data.has_auth_api_key_reader());
assert!(state.data.has_minimal_candidate_selection_reader());
assert!(state.data.has_request_candidate_reader());
assert!(state.data.has_provider_catalog_reader());
assert!(state.data.has_proxy_node_reader());
assert!(state.data.has_proxy_node_writer());
assert!(state.data.has_usage_reader());
assert!(state.data.has_video_task_reader());
assert!(state.data.has_shadow_result_reader());
@@ -242,6 +281,85 @@ fn sample_request_usage(request_id: &str) -> StoredRequestUsageAudit {
.expect("usage should build")
}
#[derive(Default)]
struct MissingShadowResultsRelationRepository;
fn missing_shadow_results_relation_error() -> DataLayerError {
DataLayerError::UnexpectedValue(
"postgres error: error returned from database: relation \"gateway_shadow_results\" does not exist"
.to_string(),
)
}
#[async_trait]
impl ShadowResultReadRepository for MissingShadowResultsRelationRepository {
async fn find(
&self,
_key: ShadowResultLookupKey<'_>,
) -> Result<Option<StoredShadowResult>, DataLayerError> {
Err(missing_shadow_results_relation_error())
}
async fn list_recent(&self, _limit: usize) -> Result<Vec<StoredShadowResult>, DataLayerError> {
Err(missing_shadow_results_relation_error())
}
}
#[async_trait]
impl ShadowResultWriteRepository for MissingShadowResultsRelationRepository {
async fn upsert(
&self,
_result: UpsertShadowResult,
) -> Result<StoredShadowResult, DataLayerError> {
Err(missing_shadow_results_relation_error())
}
}
fn sample_minimal_candidate_selection_row(
provider_id: &str,
provider_name: &str,
provider_priority: i32,
key_id: &str,
key_name: &str,
key_internal_priority: i32,
) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: provider_id.to_string(),
provider_name: provider_name.to_string(),
provider_type: "custom".to_string(),
provider_priority,
provider_is_active: true,
endpoint_id: format!("endpoint-{provider_id}"),
endpoint_api_format: "openai:chat".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: key_id.to_string(),
key_name: key_name.to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:chat".to_string()]),
key_allowed_models: None,
key_capabilities: Some(serde_json::json!({"cache_1h": true})),
key_internal_priority,
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 3})),
model_id: format!("model-{provider_id}"),
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-4.1".to_string(),
global_model_mappings: Some(vec!["gpt-4\\.1-.*".to_string()]),
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-4.1-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-4.1-canary".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]),
model_supports_streaming: None,
model_is_active: true,
model_is_available: true,
}
}
#[tokio::test]
async fn data_state_reads_decision_trace_with_provider_catalog_metadata() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
@@ -407,6 +525,148 @@ async fn data_state_reads_request_audit_bundle_from_multiple_readers() {
);
}
#[tokio::test]
async fn data_state_reads_decrypted_provider_transport_snapshot() {
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
.expect("api key ciphertext should build");
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}",
)
.expect("auth config ciphertext should build");
let provider = sample_provider_catalog_provider().with_transport_fields(
true,
false,
true,
Some(32),
Some(3),
Some(serde_json::json!({"url":"http://provider-proxy"})),
Some(20.0),
Some(8.0),
Some(serde_json::json!({"region":"global"})),
);
let endpoint = sample_provider_catalog_endpoint()
.with_transport_fields(
"https://api.openai.com".to_string(),
Some(serde_json::json!([{"action":"set","key":"x-test","value":"1"}])),
Some(serde_json::json!([{"action":"drop","path":"stream"}])),
Some(2),
Some("/v1/chat/completions".to_string()),
Some(serde_json::json!({"api_version":"v1"})),
Some(serde_json::json!({"allow":["openai:chat"]})),
Some(serde_json::json!({"url":"http://endpoint-proxy"})),
)
.expect("endpoint transport should build");
let key = sample_provider_catalog_key()
.with_transport_fields(
Some(serde_json::json!(["openai:chat", "openai:cli"])),
encrypted_api_key,
Some(encrypted_auth_config),
Some(serde_json::json!({"openai:chat": 0.8})),
Some(serde_json::json!({"openai:chat": 1})),
Some(serde_json::json!(["gpt-4.1", "gpt-4.1-mini"])),
Some(1_800_000_000),
Some(serde_json::json!({"node_id":"proxy-node-1"})),
Some(serde_json::json!({"tls_profile":"chrome_136"})),
)
.expect("key transport should build");
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let state = GatewayDataState::with_provider_transport_reader_for_tests(
repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
);
let snapshot = state
.read_provider_transport_snapshot("provider-1", "endpoint-1", "provider-key-1")
.await
.expect("snapshot read should succeed")
.expect("snapshot should exist");
assert_eq!(snapshot.provider.name, "OpenAI");
assert_eq!(snapshot.endpoint.base_url, "https://api.openai.com");
assert_eq!(
snapshot.key.api_formats,
Some(vec!["openai:chat".to_string(), "openai:cli".to_string()])
);
assert_eq!(snapshot.key.decrypted_api_key, "sk-live-openai");
assert_eq!(
snapshot.key.decrypted_auth_config.as_deref(),
Some("{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}")
);
}
#[tokio::test]
async fn data_state_reads_minimal_candidate_selection_with_auth_filters() {
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_minimal_candidate_selection_row(
"provider-2",
"OtherProvider",
20,
"key-2",
"key-two",
20,
),
sample_minimal_candidate_selection_row(
"provider-1",
"OpenAI",
10,
"key-1",
"key-one",
10,
),
StoredMinimalCandidateSelectionRow {
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 4})),
key_allowed_models: Some(vec!["gpt-4.1-edge".to_string()]),
..sample_minimal_candidate_selection_row(
"provider-1",
"OpenAI",
10,
"key-3",
"key-three",
30,
)
},
]));
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_auth_snapshot("api-key-1", "user-1"),
)]));
let state = GatewayDataState::with_minimal_candidate_selection_and_auth_for_tests(
candidate_selection_repository,
auth_repository,
);
let auth_snapshot = state
.read_auth_api_key_snapshot("user-1", "api-key-1", 150)
.await
.expect("auth snapshot should read")
.expect("auth snapshot should exist");
let selection = state
.read_minimal_candidate_selection("openai:chat", "gpt-4.1", false, Some(&auth_snapshot))
.await
.expect("selection should read");
assert_eq!(selection.len(), 2);
assert_eq!(selection[0].provider_id, "provider-1");
assert_eq!(selection[0].selected_provider_model_name, "gpt-4.1-canary");
assert_eq!(selection[0].mapping_matched_model, None);
assert_eq!(selection[1].key_id, "key-3");
assert_eq!(
selection[1].selected_provider_model_name,
"gpt-4.1-edge".to_string()
);
assert_eq!(
selection[1].mapping_matched_model,
Some("gpt-4.1-edge".to_string())
);
}
#[tokio::test]
async fn maps_openai_video_task_repository_row_into_read_response() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
@@ -414,19 +674,41 @@ async fn maps_openai_video_task_repository_row_into_read_response() {
.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("short-task-1".to_string()),
request_id: "request-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
username: Some("user".to_string()),
api_key_name: Some("primary".to_string()),
external_task_id: Some("ext-task-1".to_string()),
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("provider-key-1".to_string()),
client_api_format: Some("openai:video".to_string()),
provider_api_format: Some("openai:video".to_string()),
format_converted: false,
model: Some("sora-2".to_string()),
prompt: Some("hello".to_string()),
original_request_body: Some(json!({"prompt": "hello"})),
duration_seconds: Some(4),
resolution: Some("720p".to_string()),
aspect_ratio: Some("16:9".to_string()),
size: Some("1280x720".to_string()),
status: VideoTaskStatus::Processing,
progress_percent: 45,
progress_message: Some("working".to_string()),
retry_count: 0,
poll_interval_seconds: 10,
next_poll_at_unix_secs: Some(120),
poll_count: 1,
max_poll_count: 360,
created_at_unix_secs: 100,
submitted_at_unix_secs: Some(100),
completed_at_unix_secs: None,
updated_at_unix_secs: 120,
error_code: None,
error_message: None,
video_url: None,
request_metadata: None,
})
.await
.expect("upsert should succeed");
@@ -451,19 +733,55 @@ async fn maps_gemini_video_task_repository_row_into_read_response() {
.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("localshort123".to_string()),
request_id: "request-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
username: Some("user".to_string()),
api_key_name: Some("primary".to_string()),
external_task_id: Some("operations/ext-task-1".to_string()),
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("provider-key-1".to_string()),
client_api_format: Some("gemini:video".to_string()),
provider_api_format: Some("gemini:video".to_string()),
format_converted: false,
model: Some("veo-3".to_string()),
prompt: Some("hello".to_string()),
original_request_body: Some(json!({"prompt": "hello"})),
duration_seconds: Some(8),
resolution: Some("720p".to_string()),
aspect_ratio: Some("16:9".to_string()),
size: Some("720p".to_string()),
status: VideoTaskStatus::Completed,
progress_percent: 100,
progress_message: None,
retry_count: 0,
poll_interval_seconds: 10,
next_poll_at_unix_secs: None,
poll_count: 4,
max_poll_count: 360,
created_at_unix_secs: 100,
submitted_at_unix_secs: Some(100),
completed_at_unix_secs: Some(120),
updated_at_unix_secs: 120,
error_code: None,
error_message: None,
video_url: None,
request_metadata: Some(json!({
"rust_local_snapshot": {
"metadata": {
"generateVideoResponse": {
"generatedSamples": [
{
"video": {
"uri": "/v1beta/files/aev_localshort123:download?alt=media"
}
}
]
}
}
}
})),
})
.await
.expect("upsert should succeed");
@@ -604,6 +922,44 @@ async fn data_state_lists_recent_shadow_results_from_reader() {
assert_eq!(recent[0].request_id.as_deref(), Some("req-shadow-1"));
}
#[tokio::test]
async fn data_state_ignores_missing_shadow_results_relation_when_recording_sample() {
let repository = Arc::new(MissingShadowResultsRelationRepository);
let state = GatewayDataState::with_shadow_result_repository_for_tests(repository);
let recorded = state
.record_shadow_result_sample(RecordShadowResultSample {
trace_id: "trace-legacy".to_string(),
request_fingerprint: "fp-legacy".to_string(),
request_id: Some("req-legacy".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Rust,
result_digest: "digest-legacy".to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs: 100,
})
.await
.expect("missing shadow result table should be ignored");
assert!(recorded.is_none());
}
#[tokio::test]
async fn data_state_returns_empty_shadow_results_when_relation_missing() {
let repository = Arc::new(MissingShadowResultsRelationRepository);
let state = GatewayDataState::with_shadow_result_repository_for_tests(repository);
let recent = state
.list_recent_shadow_results(5)
.await
.expect("missing shadow result table should be ignored");
assert!(recent.is_empty());
}
fn sample_request_candidate(
id: &str,
request_id: &str,

20
crates/aether-gateway/src/data/usage.rs
View File

@@ -1,20 +0,0 @@
use aether_data::repository::usage::StoredRequestUsageAudit;
use aether_data::DataLayerError;
use super::state::GatewayDataState;
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub(crate) struct RequestUsageAudit {
#[serde(flatten)]
pub(crate) usage: StoredRequestUsageAudit,
}
pub(crate) async fn read_request_usage_audit(
state: &GatewayDataState,
request_id: &str,
) -> Result<Option<RequestUsageAudit>, DataLayerError> {
Ok(state
.find_request_usage_by_request_id(request_id)
.await?
.map(|usage| RequestUsageAudit { usage }))
}