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: 引入 aether-runtime/cache/data/http/testkit 基础 crate，完善并发门控与审计系统

Browse Source 
新增 crate：
- aether-runtime: 服务运行时基础设施（并发门控、分布式并发、指标、队列、优雅关闭、tracing）
- aether-cache: 通用 TTL 缓存与命名空间抽象
- aether-data: 数据访问层（PostgreSQL/Redis 后端、repository 模式）
- aether-http: HTTP 客户端封装（重试、配置）
- aether-testkit: 集成测试工具集（gateway/executor/hub/proxy fixture、等待、负载测试）

gateway 扩展：
- 引入 audit 模块（shadow 执行审计、决策链路追踪、请求审计 bundle）
- 引入 cache 模块（AuthContext 缓存、direct-plan bypass 缓存）
- 引入 data 模块（auth/candidates/config/usage/video_tasks 数据访问）
- 集成 ConcurrencyGate/DistributedConcurrencyGate 请求门控
- 新增本地 auth 拒绝、过载响应构建器
- 补充 control/auth_cache/video/concurrency 集成测试

aether-proxy 扩展：
- AppState 集成 stream_gate / distributed_stream_gate 并发门控
- 新增 ProxyAdmissionError 及准入拒绝流程
- stream_handler 补充门控饱和/不可用场景测试
- 配置与注册客户端逻辑完善

aether-hub 扩展：
- main.rs 引入运行时初始化、指标端点、健康检查
- local_relay 重构为 lib.rs 暴露公共接口
This commit is contained in:
fawney19
2026-03-24 15:12:56 +08:00
parent eaf8475f9e
commit b5a0070023
157 changed files with 22097 additions and 448 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -0,0 +1,155 @@
use aether_data::repository::auth::{AuthApiKeyLookupKey, StoredAuthApiKeySnapshot};
use aether_data::DataLayerError;
use super::state::GatewayDataState;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub(crate) struct StoredGatewayAuthApiKeySnapshot {
pub(crate) user_id: String,
pub(crate) username: String,
pub(crate) email: Option<String>,
pub(crate) user_role: String,
pub(crate) user_auth_source: String,
pub(crate) user_is_active: bool,
pub(crate) user_is_deleted: bool,
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>>,
pub(crate) api_key_id: String,
pub(crate) api_key_name: Option<String>,
pub(crate) api_key_is_active: bool,
pub(crate) api_key_is_locked: bool,
pub(crate) api_key_is_standalone: bool,
pub(crate) api_key_rate_limit: Option<i32>,
pub(crate) api_key_concurrent_limit: Option<i32>,
pub(crate) api_key_expires_at_unix_secs: Option<u64>,
pub(crate) api_key_allowed_providers: Option<Vec<String>>,
pub(crate) api_key_allowed_api_formats: Option<Vec<String>>,
pub(crate) api_key_allowed_models: Option<Vec<String>>,
pub(crate) currently_usable: bool,
}
impl StoredGatewayAuthApiKeySnapshot {
fn from_stored(snapshot: StoredAuthApiKeySnapshot, now_unix_secs: u64) -> Self {
let currently_usable = snapshot.is_currently_usable(now_unix_secs);
Self {
user_id: snapshot.user_id,
username: snapshot.username,
email: snapshot.email,
user_role: snapshot.user_role,
user_auth_source: snapshot.user_auth_source,
user_is_active: snapshot.user_is_active,
user_is_deleted: snapshot.user_is_deleted,
user_allowed_providers: snapshot.user_allowed_providers,
user_allowed_api_formats: snapshot.user_allowed_api_formats,
user_allowed_models: snapshot.user_allowed_models,
api_key_id: snapshot.api_key_id,
api_key_name: snapshot.api_key_name,
api_key_is_active: snapshot.api_key_is_active,
api_key_is_locked: snapshot.api_key_is_locked,
api_key_is_standalone: snapshot.api_key_is_standalone,
api_key_rate_limit: snapshot.api_key_rate_limit,
api_key_concurrent_limit: snapshot.api_key_concurrent_limit,
api_key_expires_at_unix_secs: snapshot.api_key_expires_at_unix_secs,
api_key_allowed_providers: snapshot.api_key_allowed_providers,
api_key_allowed_api_formats: snapshot.api_key_allowed_api_formats,
api_key_allowed_models: snapshot.api_key_allowed_models,
currently_usable,
}
}
}
pub(crate) async fn read_auth_api_key_snapshot(
state: &GatewayDataState,
user_id: &str,
api_key_id: &str,
now_unix_secs: u64,
) -> Result<Option<StoredGatewayAuthApiKeySnapshot>, DataLayerError> {
let snapshot = state
.find_auth_api_key_snapshot(AuthApiKeyLookupKey::UserApiKeyIds {
user_id,
api_key_id,
})
.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 aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use std::sync::Arc;
fn sample_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(200),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
)
.expect("snapshot should build")
}
#[tokio::test]
async fn reads_trusted_auth_snapshot_and_derives_usability() {
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".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(&state, "user-1", "key-1", 150)
.await
.expect("read should succeed")
.expect("snapshot should exist");
assert_eq!(
snapshot,
StoredGatewayAuthApiKeySnapshot {
user_id: "user-1".to_string(),
username: "alice".to_string(),
email: Some("alice@example.com".to_string()),
user_role: "user".to_string(),
user_auth_source: "local".to_string(),
user_is_active: true,
user_is_deleted: false,
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()]),
api_key_id: "key-1".to_string(),
api_key_name: Some("default".to_string()),
api_key_is_active: true,
api_key_is_locked: false,
api_key_is_standalone: false,
api_key_rate_limit: Some(60),
api_key_concurrent_limit: Some(5),
api_key_expires_at_unix_secs: Some(200),
api_key_allowed_providers: Some(vec!["openai".to_string()]),
api_key_allowed_api_formats: Some(vec!["openai:chat".to_string()]),
api_key_allowed_models: Some(vec!["gpt-4.1".to_string()]),
currently_usable: true,
}
);
}
}

213
crates/aether-gateway/src/data/candidates.rs Normal file
View File

@@ -0,0 +1,213 @@
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
use aether_data::DataLayerError;
use super::state::GatewayDataState;
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum RequestCandidateFinalStatus {
Success,
Failed,
Cancelled,
Streaming,
Pending,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub(crate) struct RequestCandidateTrace {
pub(crate) request_id: String,
pub(crate) total_candidates: usize,
pub(crate) final_status: RequestCandidateFinalStatus,
pub(crate) total_latency_ms: u64,
pub(crate) candidates: Vec<StoredRequestCandidate>,
}
pub(crate) async fn read_request_candidate_trace(
state: &GatewayDataState,
request_id: &str,
attempted_only: bool,
) -> Result<Option<RequestCandidateTrace>, DataLayerError> {
let all_candidates = state
.list_request_candidates_by_request_id(request_id)
.await?;
if all_candidates.is_empty() {
return Ok(None);
}
let candidates = if attempted_only {
all_candidates
.iter()
.filter(|candidate| {
candidate
.status
.is_attempted(candidate.started_at_unix_secs)
})
.cloned()
.collect::<Vec<_>>()
} else {
all_candidates.clone()
};
let total_latency_ms = candidates
.iter()
.filter(|candidate| {
matches!(
candidate.status,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
) && candidate.latency_ms.is_some()
})
.map(|candidate| candidate.latency_ms.unwrap_or(0))
.sum();
let final_status_source = if attempted_only && candidates.is_empty() {
&all_candidates
} else {
&candidates
};
Ok(Some(RequestCandidateTrace {
request_id: request_id.to_string(),
total_candidates: candidates.len(),
final_status: derive_final_status(final_status_source),
total_latency_ms,
candidates,
}))
}
fn derive_final_status(candidates: &[StoredRequestCandidate]) -> RequestCandidateFinalStatus {
let has_success = candidates.iter().any(|candidate| {
candidate.status == RequestCandidateStatus::Success
|| matches!(candidate.status_code, Some(status_code) if (200..300).contains(&status_code))
});
if has_success {
return RequestCandidateFinalStatus::Success;
}
if candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Streaming)
{
return RequestCandidateFinalStatus::Streaming;
}
if candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Pending)
{
return RequestCandidateFinalStatus::Pending;
}
let has_cancelled = candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Cancelled);
let has_failed = candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Failed);
if has_cancelled && !has_failed {
return RequestCandidateFinalStatus::Cancelled;
}
RequestCandidateFinalStatus::Failed
}
#[cfg(test)]
mod tests {
use super::super::GatewayDataState;
use super::{derive_final_status, read_request_candidate_trace, RequestCandidateFinalStatus};
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
};
use std::sync::Arc;
fn sample_candidate(
id: &str,
request_id: &str,
candidate_index: i32,
status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>,
latency_ms: Option<i32>,
status_code: Option<i32>,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
candidate_index,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
status,
None,
false,
status_code,
None,
None,
latency_ms,
Some(1),
None,
None,
100 + i64::from(candidate_index),
started_at_unix_secs,
started_at_unix_secs.map(|value| value + 1),
)
.expect("candidate should build")
}
#[test]
fn derive_final_status_prefers_success() {
let candidates = vec![sample_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Success,
Some(100),
Some(25),
Some(200),
)];
assert_eq!(
derive_final_status(&candidates),
RequestCandidateFinalStatus::Success
);
}
#[tokio::test]
async fn read_request_candidate_trace_filters_attempted_rows() {
let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Pending,
None,
None,
None,
),
sample_candidate(
"cand-2",
"req-1",
1,
RequestCandidateStatus::Failed,
Some(101),
Some(33),
Some(502),
),
]));
let state = GatewayDataState::with_request_candidate_reader_for_tests(repository);
let trace = read_request_candidate_trace(&state, "req-1", true)
.await
.expect("trace should succeed")
.expect("trace should exist");
assert_eq!(trace.total_candidates, 1);
assert_eq!(trace.candidates[0].id, "cand-2");
assert_eq!(trace.final_status, RequestCandidateFinalStatus::Failed);
assert_eq!(trace.total_latency_ms, 33);
}
}

41
crates/aether-gateway/src/data/config.rs Normal file
View File

@@ -0,0 +1,41 @@
use aether_data::postgres::PostgresPoolConfig;
use aether_data::DataLayerConfig;
#[derive(Debug, Clone, Default)]
pub struct GatewayDataConfig {
postgres: Option<PostgresPoolConfig>,
}
impl GatewayDataConfig {
pub fn disabled() -> Self {
Self::default()
}
pub fn from_postgres_config(postgres: PostgresPoolConfig) -> Self {
Self {
postgres: Some(postgres),
}
}
pub fn from_postgres_url(database_url: impl Into<String>, require_ssl: bool) -> Self {
let mut postgres = PostgresPoolConfig::default();
postgres.database_url = database_url.into();
postgres.require_ssl = require_ssl;
Self::from_postgres_config(postgres)
}
pub fn postgres(&self) -> Option<&PostgresPoolConfig> {
self.postgres.as_ref()
}
pub fn is_enabled(&self) -> bool {
self.postgres.is_some()
}
pub fn to_data_layer_config(&self) -> DataLayerConfig {
DataLayerConfig {
postgres: self.postgres.clone(),
redis: None,
}
}
}

267
crates/aether-gateway/src/data/decision_trace.rs Normal file
View File

@@ -0,0 +1,267 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data::repository::candidates::StoredRequestCandidate;
use aether_data::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::DataLayerError;
use super::candidates::RequestCandidateFinalStatus;
use super::state::GatewayDataState;
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub(crate) struct DecisionTraceCandidate {
#[serde(flatten)]
pub(crate) candidate: StoredRequestCandidate,
pub(crate) provider_name: Option<String>,
pub(crate) provider_website: Option<String>,
pub(crate) provider_type: Option<String>,
pub(crate) endpoint_api_format: Option<String>,
pub(crate) endpoint_api_family: Option<String>,
pub(crate) endpoint_kind: Option<String>,
pub(crate) provider_key_name: Option<String>,
pub(crate) provider_key_auth_type: Option<String>,
pub(crate) provider_key_capabilities: Option<serde_json::Value>,
pub(crate) provider_key_is_active: Option<bool>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
pub(crate) struct DecisionTrace {
pub(crate) request_id: String,
pub(crate) total_candidates: usize,
pub(crate) final_status: RequestCandidateFinalStatus,
pub(crate) total_latency_ms: u64,
pub(crate) candidates: Vec<DecisionTraceCandidate>,
}
pub(crate) async fn read_decision_trace(
state: &GatewayDataState,
request_id: &str,
attempted_only: bool,
) -> Result<Option<DecisionTrace>, DataLayerError> {
let Some(trace) = state
.read_request_candidate_trace(request_id, attempted_only)
.await?
else {
return Ok(None);
};
let provider_ids = unique_ids(
trace
.candidates
.iter()
.filter_map(|item| item.provider_id.as_ref()),
);
let endpoint_ids = unique_ids(
trace
.candidates
.iter()
.filter_map(|item| item.endpoint_id.as_ref()),
);
let key_ids = unique_ids(
trace
.candidates
.iter()
.filter_map(|item| item.key_id.as_ref()),
);
let provider_map = state
.list_provider_catalog_providers_by_ids(&provider_ids)
.await?
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
let endpoint_map = state
.list_provider_catalog_endpoints_by_ids(&endpoint_ids)
.await?
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
let key_map = state
.list_provider_catalog_keys_by_ids(&key_ids)
.await?
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
Ok(Some(DecisionTrace {
request_id: trace.request_id,
total_candidates: trace.total_candidates,
final_status: trace.final_status,
total_latency_ms: trace.total_latency_ms,
candidates: trace
.candidates
.into_iter()
.map(|candidate| enrich_candidate(candidate, &provider_map, &endpoint_map, &key_map))
.collect(),
}))
}
fn enrich_candidate(
candidate: StoredRequestCandidate,
provider_map: &BTreeMap<String, StoredProviderCatalogProvider>,
endpoint_map: &BTreeMap<String, StoredProviderCatalogEndpoint>,
key_map: &BTreeMap<String, StoredProviderCatalogKey>,
) -> DecisionTraceCandidate {
let provider = candidate
.provider_id
.as_ref()
.and_then(|provider_id| provider_map.get(provider_id));
let endpoint = candidate
.endpoint_id
.as_ref()
.and_then(|endpoint_id| endpoint_map.get(endpoint_id));
let provider_key = candidate
.key_id
.as_ref()
.and_then(|key_id| key_map.get(key_id));
DecisionTraceCandidate {
provider_name: provider.map(|item| item.name.clone()),
provider_website: provider.and_then(|item| item.website.clone()),
provider_type: provider.map(|item| item.provider_type.clone()),
endpoint_api_format: endpoint.map(|item| item.api_format.clone()),
endpoint_api_family: endpoint.and_then(|item| item.api_family.clone()),
endpoint_kind: endpoint.and_then(|item| item.endpoint_kind.clone()),
provider_key_name: provider_key
.map(|item| item.name.clone())
.or_else(|| candidate.api_key_name.clone()),
provider_key_auth_type: provider_key.map(|item| item.auth_type.clone()),
provider_key_capabilities: provider_key.and_then(|item| item.capabilities.clone()),
provider_key_is_active: provider_key.map(|item| item.is_active),
candidate,
}
}
fn unique_ids<'a>(items: impl Iterator<Item = &'a String>) -> Vec<String> {
items
.cloned()
.collect::<BTreeSet<_>>()
.into_iter()
.collect()
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use super::{read_decision_trace, DecisionTrace, DecisionTraceCandidate};
use crate::gateway::data::candidates::RequestCandidateFinalStatus;
use crate::gateway::data::GatewayDataState;
fn sample_candidate(request_id: &str) -> StoredRequestCandidate {
StoredRequestCandidate::new(
"cand-1".to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Failed,
None,
false,
Some(502),
Some("bad_gateway".to_string()),
Some("upstream failed".to_string()),
Some(37),
Some(1),
None,
Some(serde_json::json!({"cache_1h": true})),
100,
Some(101),
Some(102),
)
.expect("candidate should build")
}
fn sample_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"OpenAI".to_string(),
Some("https://openai.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
}
fn sample_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
}
fn sample_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"provider-key-1".to_string(),
"provider-1".to_string(),
"prod-key".to_string(),
"api_key".to_string(),
Some(serde_json::json!({"cache_1h": true})),
true,
)
.expect("key should build")
}
#[tokio::test]
async fn enriches_request_candidate_trace_with_provider_catalog_metadata() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_candidate("req-1"),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let state = GatewayDataState::with_decision_trace_readers_for_tests(
request_candidates,
provider_catalog,
);
let trace = read_decision_trace(&state, "req-1", true)
.await
.expect("trace should read")
.expect("trace should exist");
assert_eq!(
trace,
DecisionTrace {
request_id: "req-1".to_string(),
total_candidates: 1,
final_status: RequestCandidateFinalStatus::Failed,
total_latency_ms: 37,
candidates: vec![DecisionTraceCandidate {
candidate: sample_candidate("req-1"),
provider_name: Some("OpenAI".to_string()),
provider_website: Some("https://openai.com".to_string()),
provider_type: Some("custom".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
provider_key_name: Some("prod-key".to_string()),
provider_key_auth_type: Some("api_key".to_string()),
provider_key_capabilities: Some(serde_json::json!({"cache_1h": true})),
provider_key_is_active: Some(true),
}],
}
);
}
}

80
crates/aether-gateway/src/data/gemini.rs Normal file
View File

@@ -0,0 +1,80 @@
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
use serde_json::json;
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
pub(super) fn map_gemini_video_task_to_read_response(
task: StoredVideoTask,
) -> LocalVideoTaskReadResponse {
match task.status {
VideoTaskStatus::Cancelled => LocalVideoTaskReadResponse {
status_code: 404,
body_json: json!({"detail": "Video task was cancelled"}),
},
VideoTaskStatus::Deleted => LocalVideoTaskReadResponse {
status_code: 404,
body_json: json!({"detail": "Video task not found"}),
},
VideoTaskStatus::Completed => LocalVideoTaskReadResponse {
status_code: 200,
body_json: build_gemini_completed_body(task),
},
VideoTaskStatus::Failed | VideoTaskStatus::Expired => LocalVideoTaskReadResponse {
status_code: 200,
body_json: build_gemini_failed_body(task),
},
_ => LocalVideoTaskReadResponse {
status_code: 200,
body_json: build_gemini_pending_body(task),
},
}
}
fn build_gemini_completed_body(task: StoredVideoTask) -> serde_json::Value {
let operation_name = operation_name(&task);
let short_id = task.short_id.unwrap_or_default();
json!({
"name": operation_name,
"done": true,
"response": {
"generateVideoResponse": {
"generatedSamples": [
{
"video": {
"uri": format!("/v1beta/files/aev_{short_id}:download?alt=media"),
"mimeType": "video/mp4"
}
}
]
}
}
})
}
fn build_gemini_failed_body(task: StoredVideoTask) -> serde_json::Value {
json!({
"name": operation_name(&task),
"done": true,
"error": {
"code": task.error_code.unwrap_or_else(|| "UNKNOWN".to_string()),
"message": task
.error_message
.unwrap_or_else(|| "Video generation failed".to_string()),
}
})
}
fn build_gemini_pending_body(task: StoredVideoTask) -> serde_json::Value {
json!({
"name": operation_name(&task),
"done": false,
"metadata": {}
})
}
fn operation_name(task: &StoredVideoTask) -> String {
let model = task.model.clone().unwrap_or_else(|| "unknown".to_string());
let short_id = task.short_id.clone().unwrap_or_else(|| task.id.clone());
format!("models/{model}/operations/{short_id}")
}

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

@@ -0,0 +1,21 @@
mod auth;
mod candidates;
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 auth::StoredGatewayAuthApiKeySnapshot;
pub(crate) use candidates::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;

69
crates/aether-gateway/src/data/openai.rs Normal file
View File

@@ -0,0 +1,69 @@
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
use serde_json::{json, Value};
use crate::gateway::video_tasks::LocalVideoTaskReadResponse;
pub(super) fn map_openai_video_task_to_read_response(
task: StoredVideoTask,
) -> LocalVideoTaskReadResponse {
match task.status {
VideoTaskStatus::Cancelled => LocalVideoTaskReadResponse {
status_code: 404,
body_json: json!({"detail": "Video task was cancelled"}),
},
VideoTaskStatus::Deleted => LocalVideoTaskReadResponse {
status_code: 404,
body_json: json!({"detail": "Video task not found"}),
},
status => LocalVideoTaskReadResponse {
status_code: 200,
body_json: build_openai_video_task_body(task, status),
},
}
}
fn build_openai_video_task_body(task: StoredVideoTask, status: VideoTaskStatus) -> Value {
let mut body = json!({
"id": task.id,
"object": "video",
"status": map_openai_video_status(status),
"progress": task.progress_percent,
"created_at": task.created_at_unix_secs,
});
if let Some(model) = task.model {
body["model"] = Value::String(model);
}
if let Some(prompt) = task.prompt {
body["prompt"] = Value::String(prompt);
}
if let Some(size) = task.size {
body["size"] = Value::String(size);
}
if let Some(video_url) = task.video_url {
body["video_url"] = Value::String(video_url);
}
if matches!(
status,
VideoTaskStatus::Failed | VideoTaskStatus::Expired | VideoTaskStatus::Cancelled
) {
body["error"] = json!({
"code": task.error_code.unwrap_or_else(|| "unknown".to_string()),
"message": task
.error_message
.unwrap_or_else(|| "Video generation failed".to_string()),
});
}
body
}
fn map_openai_video_status(status: VideoTaskStatus) -> &'static str {
match status {
VideoTaskStatus::Pending | VideoTaskStatus::Submitted | VideoTaskStatus::Queued => "queued",
VideoTaskStatus::Processing => "processing",
VideoTaskStatus::Completed => "completed",
VideoTaskStatus::Failed | VideoTaskStatus::Cancelled | VideoTaskStatus::Expired => "failed",
VideoTaskStatus::Deleted => "deleted",
}
}

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

@@ -0,0 +1,56 @@
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 Normal file
View File

@@ -0,0 +1,454 @@
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),
}
}
}

683
crates/aether-gateway/src/data/tests.rs Normal file
View File

@@ -0,0 +1,683 @@
use std::sync::Arc;
use aether_data::repository::auth::{
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::candidates::{
InMemoryRequestCandidateRepository, RequestCandidateStatus, StoredRequestCandidate,
};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use aether_data::repository::shadow_results::{
InMemoryShadowResultRepository, RecordShadowResultSample, ShadowResultLookupKey,
ShadowResultMatchStatus, ShadowResultReadRepository, ShadowResultSampleOrigin,
UpsertShadowResult,
};
use aether_data::repository::usage::{InMemoryUsageReadRepository, StoredRequestUsageAudit};
use aether_data::repository::video_tasks::{
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskLookupKey, VideoTaskStatus,
VideoTaskWriteRepository,
};
use super::{GatewayDataConfig, GatewayDataState};
use crate::gateway::AppState;
#[test]
fn disabled_gateway_data_state_has_no_backends() {
let state = GatewayDataState::from_config(GatewayDataConfig::disabled())
.expect("disabled config should build");
assert!(!state.has_backends());
assert!(!state.has_auth_api_key_reader());
assert!(!state.has_request_candidate_reader());
assert!(!state.has_provider_catalog_reader());
assert!(!state.has_usage_reader());
assert!(!state.has_video_task_reader());
assert!(!state.has_shadow_result_reader());
assert!(!state.has_shadow_result_writer());
}
#[tokio::test]
async fn postgres_gateway_data_state_builds_video_task_reader() {
let state = GatewayDataState::from_config(GatewayDataConfig::from_postgres_url(
"postgres://localhost/aether",
false,
))
.expect("postgres-backed state should build");
assert!(state.has_backends());
assert!(state.has_auth_api_key_reader());
assert!(state.has_request_candidate_reader());
assert!(state.has_provider_catalog_reader());
assert!(state.has_usage_reader());
assert!(state.has_video_task_reader());
assert!(state.has_shadow_result_reader());
assert!(state.has_shadow_result_writer());
}
#[tokio::test]
async fn data_state_find_uses_configured_read_repository() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("short-task-1".to_string()),
user_id: Some("user-1".to_string()),
external_task_id: Some("ext-task-1".to_string()),
provider_api_format: Some("openai:video".to_string()),
model: Some("sora-2".to_string()),
prompt: Some("hello".to_string()),
size: Some("1280x720".to_string()),
status: VideoTaskStatus::Queued,
progress_percent: 0,
created_at_unix_secs: 100,
updated_at_unix_secs: 100,
error_code: None,
error_message: None,
video_url: None,
})
.await
.expect("upsert should succeed");
let state = GatewayDataState::with_video_task_reader_for_tests(repository);
let task = state
.find_video_task(VideoTaskLookupKey::Id("task-1"))
.await
.expect("find should succeed");
assert_eq!(task.expect("task should exist").id, "task-1");
}
#[tokio::test]
async fn app_state_wires_gateway_data_state_from_config() {
let state = AppState::new_with_executor(
"http://127.0.0.1:18084",
Some("http://127.0.0.1:18085".to_string()),
Some("http://127.0.0.1:18086".to_string()),
)
.expect("app state should build")
.with_data_config(GatewayDataConfig::from_postgres_url(
"postgres://localhost/aether",
false,
))
.expect("data config should wire");
assert!(state.data.has_backends());
assert!(state.data.has_auth_api_key_reader());
assert!(state.data.has_request_candidate_reader());
assert!(state.data.has_provider_catalog_reader());
assert!(state.data.has_usage_reader());
assert!(state.data.has_video_task_reader());
assert!(state.data.has_shadow_result_reader());
assert!(state.data.has_shadow_result_writer());
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(200),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
)
.expect("auth snapshot should build")
}
#[tokio::test]
async fn data_state_reads_auth_api_key_snapshot_from_reader() {
let repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_auth_snapshot("key-1", "user-1"),
)]));
let state = GatewayDataState::with_auth_api_key_reader_for_tests(repository);
let snapshot = state
.read_auth_api_key_snapshot("user-1", "key-1", 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.username, "alice");
assert_eq!(
snapshot.api_key_allowed_models,
Some(vec!["gpt-4.1".to_string()])
);
assert!(snapshot.currently_usable);
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"OpenAI".to_string(),
Some("https://openai.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"provider-key-1".to_string(),
"provider-1".to_string(),
"prod-key".to_string(),
"api_key".to_string(),
Some(serde_json::json!({"cache_1h": true})),
true,
)
.expect("key should build")
}
fn sample_request_usage(request_id: &str) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
"usage-1".to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
"OpenAI".to_string(),
"gpt-4.1".to_string(),
Some("gpt-4.1-mini".to_string()),
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
true,
false,
120,
40,
160,
0.24,
0.36,
Some(200),
None,
None,
Some(450),
Some(120),
"completed".to_string(),
"settled".to_string(),
100,
101,
Some(102),
)
.expect("usage should build")
}
#[tokio::test]
async fn data_state_reads_decision_trace_with_provider_catalog_metadata() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Failed,
None,
false,
Some(502),
None,
None,
Some(37),
Some(1),
None,
Some(serde_json::json!({"cache_1h": true})),
100,
Some(101),
Some(102),
)
.expect("candidate should build"),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let state = GatewayDataState::with_decision_trace_readers_for_tests(
request_candidates,
provider_catalog,
);
let trace = state
.read_decision_trace("req-1", true)
.await
.expect("trace should read")
.expect("trace should exist");
assert_eq!(trace.request_id, "req-1");
assert_eq!(trace.total_candidates, 1);
assert_eq!(trace.candidates[0].provider_name.as_deref(), Some("OpenAI"));
assert_eq!(
trace.candidates[0].endpoint_api_format.as_deref(),
Some("openai:chat")
);
assert_eq!(
trace.candidates[0].provider_key_auth_type.as_deref(),
Some("api_key")
);
assert_eq!(
trace.candidates[0].provider_key_capabilities,
Some(serde_json::json!({"cache_1h": true}))
);
}
#[tokio::test]
async fn data_state_reads_request_usage_audit_from_reader() {
let repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_request_usage("req-usage-1"),
]));
let state = GatewayDataState::with_usage_reader_for_tests(repository);
let usage = state
.read_request_usage_audit("req-usage-1")
.await
.expect("read should succeed")
.expect("usage should exist");
assert_eq!(usage.usage.request_id, "req-usage-1");
assert_eq!(usage.usage.provider_name, "OpenAI");
assert_eq!(usage.usage.total_tokens, 160);
assert_eq!(usage.usage.total_cost_usd, 0.24);
assert!(usage.usage.has_format_conversion);
}
#[tokio::test]
async fn data_state_reads_request_audit_bundle_from_multiple_readers() {
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some("hash-1".to_string()),
sample_auth_snapshot("api-key-1", "user-1"),
)]));
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
StoredRequestCandidate::new(
"cand-1".to_string(),
"req-usage-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(37),
Some(1),
None,
Some(serde_json::json!({"cache_1h": true})),
100,
Some(101),
Some(102),
)
.expect("candidate should build"),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
sample_request_usage("req-usage-1"),
]));
let state = GatewayDataState::with_request_audit_readers_for_tests(
auth_repository,
request_candidates,
provider_catalog,
usage_repository,
);
let bundle = state
.read_request_audit_bundle("req-usage-1", true, 150)
.await
.expect("bundle should read")
.expect("bundle should exist");
assert_eq!(bundle.request_id, "req-usage-1");
assert_eq!(
bundle
.usage
.as_ref()
.and_then(|usage| usage.usage.target_model.as_deref()),
Some("gpt-4.1-mini")
);
assert_eq!(
bundle
.decision_trace
.as_ref()
.and_then(|trace| trace.candidates.first())
.and_then(|candidate| candidate.provider_name.as_deref()),
Some("OpenAI")
);
assert_eq!(
bundle
.auth_snapshot
.as_ref()
.map(|snapshot| snapshot.currently_usable),
Some(true)
);
}
#[tokio::test]
async fn maps_openai_video_task_repository_row_into_read_response() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("short-task-1".to_string()),
user_id: Some("user-1".to_string()),
external_task_id: Some("ext-task-1".to_string()),
provider_api_format: Some("openai:video".to_string()),
model: Some("sora-2".to_string()),
prompt: Some("hello".to_string()),
size: Some("1280x720".to_string()),
status: VideoTaskStatus::Processing,
progress_percent: 45,
created_at_unix_secs: 100,
updated_at_unix_secs: 120,
error_code: None,
error_message: None,
video_url: None,
})
.await
.expect("upsert should succeed");
let state = GatewayDataState::with_video_task_reader_for_tests(repository);
let response = state
.read_video_task_response(Some("openai"), "/v1/videos/task-1")
.await
.expect("read should succeed")
.expect("read response should exist");
assert_eq!(response.status_code, 200);
assert_eq!(response.body_json["id"], "task-1");
assert_eq!(response.body_json["status"], "processing");
assert_eq!(response.body_json["created_at"], 100);
}
#[tokio::test]
async fn maps_gemini_video_task_repository_row_into_read_response() {
let repository = Arc::new(InMemoryVideoTaskRepository::default());
repository
.upsert(UpsertVideoTask {
id: "task-1".to_string(),
short_id: Some("localshort123".to_string()),
user_id: Some("user-1".to_string()),
external_task_id: Some("operations/ext-task-1".to_string()),
provider_api_format: Some("gemini:video".to_string()),
model: Some("veo-3".to_string()),
prompt: Some("hello".to_string()),
size: Some("720p".to_string()),
status: VideoTaskStatus::Completed,
progress_percent: 100,
created_at_unix_secs: 100,
updated_at_unix_secs: 120,
error_code: None,
error_message: None,
video_url: None,
})
.await
.expect("upsert should succeed");
let state = GatewayDataState::with_video_task_reader_for_tests(repository);
let response = state
.read_video_task_response(
Some("gemini"),
"/v1beta/models/veo-3/operations/localshort123",
)
.await
.expect("read should succeed")
.expect("read response should exist");
assert_eq!(response.status_code, 200);
assert_eq!(
response.body_json["name"],
"models/veo-3/operations/localshort123"
);
assert_eq!(response.body_json["done"], true);
}
#[tokio::test]
async fn data_state_write_uses_configured_shadow_result_writer() {
let repository = Arc::new(InMemoryShadowResultRepository::default());
let state = GatewayDataState::with_shadow_result_writer_for_tests(repository.clone());
let written = state
.write_shadow_result(UpsertShadowResult {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
rust_result_digest: Some("rust-digest".to_string()),
python_result_digest: None,
match_status: ShadowResultMatchStatus::Pending,
status_code: Some(200),
error_message: None,
created_at_unix_secs: 100,
updated_at_unix_secs: 100,
})
.await
.expect("write should succeed");
assert!(written.is_some());
let stored = repository
.find(ShadowResultLookupKey::TraceFingerprint {
trace_id: "trace-1",
request_fingerprint: "fp-1",
})
.await
.expect("find should succeed");
assert_eq!(
stored.expect("stored result should exist").match_status,
ShadowResultMatchStatus::Pending
);
}
#[tokio::test]
async fn data_state_records_shadow_result_samples_and_merges_match_status() {
let repository = Arc::new(InMemoryShadowResultRepository::default());
let state = GatewayDataState::with_shadow_result_repository_for_tests(repository);
let first = state
.record_shadow_result_sample(RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Rust,
result_digest: "digest-1".to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs: 100,
})
.await
.expect("first record should succeed")
.expect("first stored result should exist");
assert_eq!(first.match_status, ShadowResultMatchStatus::Pending);
let second = state
.record_shadow_result_sample(RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Python,
result_digest: "digest-1".to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs: 200,
})
.await
.expect("second record should succeed")
.expect("second stored result should exist");
assert_eq!(second.match_status, ShadowResultMatchStatus::Match);
assert_eq!(second.created_at_unix_secs, 100);
assert_eq!(second.updated_at_unix_secs, 200);
assert_eq!(second.request_id.as_deref(), Some("req-1"));
}
#[tokio::test]
async fn data_state_lists_recent_shadow_results_from_reader() {
let repository = Arc::new(InMemoryShadowResultRepository::default());
let state = GatewayDataState::with_shadow_result_repository_for_tests(repository.clone());
state
.record_shadow_result_sample(RecordShadowResultSample {
trace_id: "trace-1".to_string(),
request_fingerprint: "fp-1".to_string(),
request_id: Some("req-shadow-1".to_string()),
route_family: Some("openai".to_string()),
route_kind: Some("chat".to_string()),
candidate_id: None,
origin: ShadowResultSampleOrigin::Rust,
result_digest: "digest-1".to_string(),
status_code: Some(200),
error_message: None,
recorded_at_unix_secs: 100,
})
.await
.expect("record should succeed");
let recent = state
.list_recent_shadow_results(5)
.await
.expect("list recent should succeed");
assert_eq!(recent.len(), 1);
assert_eq!(recent[0].trace_id, "trace-1");
assert_eq!(recent[0].request_id.as_deref(), Some("req-shadow-1"));
}
fn sample_request_candidate(
id: &str,
request_id: &str,
candidate_index: i32,
status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>,
latency_ms: Option<i32>,
status_code: Option<i32>,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
candidate_index,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
status,
None,
false,
status_code,
None,
None,
latency_ms,
Some(1),
None,
None,
100 + i64::from(candidate_index),
started_at_unix_secs,
started_at_unix_secs.map(|value| value + 1),
)
.expect("candidate should build")
}
#[tokio::test]
async fn data_state_reads_request_candidate_trace_from_reader() {
let repository = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_request_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Pending,
None,
None,
None,
),
sample_request_candidate(
"cand-2",
"req-1",
1,
RequestCandidateStatus::Success,
Some(101),
Some(42),
Some(200),
),
]));
let state = GatewayDataState::with_request_candidate_reader_for_tests(repository);
let trace = state
.read_request_candidate_trace("req-1", true)
.await
.expect("trace should succeed")
.expect("trace should exist");
assert_eq!(trace.request_id, "req-1");
assert_eq!(trace.total_candidates, 1);
assert_eq!(
trace.final_status,
super::candidates::RequestCandidateFinalStatus::Success
);
assert_eq!(trace.total_latency_ms, 42);
assert_eq!(trace.candidates[0].id, "cand-2");
}

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

@@ -0,0 +1,20 @@
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 }))
}

65
crates/aether-gateway/src/data/video_tasks.rs Normal file
View File

@@ -0,0 +1,65 @@
use aether_data::repository::video_tasks::VideoTaskLookupKey;
use aether_data::DataLayerError;
use super::gemini::map_gemini_video_task_to_read_response;
use super::openai::map_openai_video_task_to_read_response;
use super::state::GatewayDataState;
use crate::gateway::video_tasks::{
extract_gemini_short_id_from_path, extract_openai_task_id_from_path, LocalVideoTaskReadResponse,
};
pub(super) async fn read_video_task_response(
state: &GatewayDataState,
route_family: Option<&str>,
request_path: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
match route_family {
Some("openai") => read_openai_video_task_response(state, request_path).await,
Some("gemini") => read_gemini_video_task_response(state, request_path).await,
_ => Ok(None),
}
}
async fn read_openai_video_task_response(
state: &GatewayDataState,
request_path: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
let Some(task_id) = extract_openai_task_id_from_path(request_path) else {
return Ok(None);
};
let Some(task) = state
.find_video_task(VideoTaskLookupKey::Id(task_id))
.await?
else {
return Ok(None);
};
if !matches!(task.provider_api_format.as_deref(), Some("openai:video")) {
return Ok(None);
}
Ok(Some(map_openai_video_task_to_read_response(task)))
}
async fn read_gemini_video_task_response(
state: &GatewayDataState,
request_path: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
let Some(short_id) = extract_gemini_short_id_from_path(request_path) else {
return Ok(None);
};
let Some(task) = state
.find_video_task(VideoTaskLookupKey::ShortId(short_id))
.await?
else {
return Ok(None);
};
if !matches!(task.provider_api_format.as_deref(), Some("gemini:video")) {
return Ok(None);
}
Ok(Some(map_gemini_video_task_to_read_response(task)))
}