refactor(workspace): enforce layered crate boundaries

This commit is contained in:
elky
2026-07-15 23:47:19 +08:00
parent a728c090a9
commit 8616fe6ee2
969 changed files with 40187 additions and 27240 deletions
@@ -423,7 +423,7 @@ mod tests {
.json_body
.as_ref()
.is_some_and(|body| body.get("stream").is_none()));
assert!(built.plan.headers.get("accept").is_none());
assert!(!built.plan.headers.contains_key("accept"));
}
#[test]
+30 -26
View File
@@ -144,34 +144,38 @@ pub(crate) fn spawn_video_task_poller(state: AppState) -> Option<JoinHandle<()>>
return None;
}
Some(tokio::spawn(async move {
let mut interval = tokio::time::interval(config.interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut deferred_since = None;
loop {
Some(crate::task_runtime::spawn_singleton_worker(
state,
crate::task_runtime::TASK_KEY_VIDEO_TASK_POLLER,
move |state| async move {
let mut interval = tokio::time::interval(config.interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
if state
.data
.should_defer_maintenance_for_database_pool_pressure(&mut deferred_since)
{
debug!(
event_name = "video_task_poller_deferred",
log_type = "event",
"gateway video task poller deferred because database pool has no idle reserve"
);
continue;
let mut deferred_since = None;
loop {
interval.tick().await;
if state
.data
.should_defer_maintenance_for_database_pool_pressure(&mut deferred_since)
{
debug!(
event_name = "video_task_poller_deferred",
log_type = "event",
"gateway video task poller deferred because database pool has no idle reserve"
);
continue;
}
if let Err(err) = poll_video_tasks_once(&state, config.batch_size).await {
warn!(
event_name = "video_task_poller_tick_failed",
log_type = "event",
error = ?err,
"gateway video task poller tick failed"
);
}
}
if let Err(err) = poll_video_tasks_once(&state, config.batch_size).await {
warn!(
event_name = "video_task_poller_tick_failed",
log_type = "event",
error = ?err,
"gateway video task poller tick failed"
);
}
}
}))
},
))
}
async fn fetch_video_task_refresh_attempt(
+13 -9
View File
@@ -20,17 +20,21 @@ pub(crate) fn spawn_s3_backup_worker(app: AppState) -> Option<JoinHandle<()>> {
return None;
}
Some(tokio::spawn(async move {
let mut interval = tokio::time::interval(S3_BACKUP_WORKER_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
loop {
Some(crate::task_runtime::spawn_singleton_worker(
app,
S3_BACKUP_WORKER_TASK_KEY,
|app| async move {
let mut interval = tokio::time::interval(S3_BACKUP_WORKER_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
if let Err(error) = run_s3_backup_schedule_tick(&app, Utc::now()).await {
warn!(error = ?error, "S3 backup schedule tick failed");
loop {
interval.tick().await;
if let Err(error) = run_s3_backup_schedule_tick(&app, Utc::now()).await {
warn!(error = ?error, "S3 backup schedule tick failed");
}
}
}
}))
},
))
}
async fn run_s3_backup_schedule_tick(
+52 -3
View File
@@ -213,6 +213,7 @@ async fn estimate_execution_plan_cost_upper_bound_usd(
);
let ttl = state.frontdoor_runtime_guards.auth_capacity_cache_ttl;
if ttl.is_zero() {
let _permit = state.acquire_auth_snapshot_load_gate().await?;
return calculate_execution_plan_cost_upper_bound(
state,
plan,
@@ -230,6 +231,7 @@ async fn estimate_execution_plan_cost_upper_bound_usd(
state
.auth_request_cost_upper_bound_cache
.get_or_load(cache_key, ttl, || async {
let _permit = state.acquire_auth_snapshot_load_gate().await?;
calculate_execution_plan_cost_upper_bound(
state,
plan,
@@ -499,9 +501,16 @@ async fn request_model_resolves_to_allowed_model(
.model_directive_policy
.resolve_reasoning(&api_format, Some(requested_model));
let routing_model = resolution.base_model().unwrap_or(requested_model);
let rows = state
.list_minimal_candidate_selection_rows_for_api_format(&api_format)
.await?;
let rows = {
// Model alias authorization runs before candidate planning, so its database read must
// participate in the same foreground DB admission budget as the rest of auth. Keep
// the permit scoped to this one read; callers do not hold this gate, and releasing it
// here avoids carrying a DB permit through pure filtering or subsequent formats.
let _permit = state.acquire_auth_snapshot_load_gate().await?;
state
.list_minimal_candidate_selection_rows_for_api_format(&api_format)
.await?
};
let matching_rows = rows
.into_iter()
.filter(|row| {
@@ -553,6 +562,7 @@ mod tests {
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
@@ -564,6 +574,7 @@ mod tests {
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::DataLayerError;
use aether_runtime::ConcurrencyGate;
use async_trait::async_trait;
use axum::body::Bytes;
use axum::http::{HeaderMap, Uri};
@@ -896,6 +907,44 @@ mod tests {
assert_eq!(rejection, None);
}
#[tokio::test]
async fn model_alias_resolution_waits_for_auth_database_gate() {
let mut state = state_with_model_mapping();
state.auth_snapshot_load_gate = Some(Arc::new(ConcurrencyGate::new(
"test_auth_model_resolution",
1,
)));
let held = state
.acquire_auth_snapshot_load_gate()
.await
.expect("auth gate acquisition should succeed")
.expect("auth gate should be configured");
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let headers = json_headers();
let body = Bytes::from_static(br#"{"model":"gpt-5.2","messages":[]}"#);
let blocked = tokio::time::timeout(
Duration::from_millis(25),
request_model_local_rejection(&state, Some(&decision), &uri, &headers, &body),
)
.await;
assert!(
blocked.is_err(),
"model alias candidate reads must wait for the auth DB gate"
);
drop(held);
let rejection = tokio::time::timeout(
Duration::from_secs(1),
request_model_local_rejection(&state, Some(&decision), &uri, &headers, &body),
)
.await
.expect("model alias resolution should resume after releasing the auth gate")
.expect("model rejection should resolve");
assert_eq!(rejection, None);
}
#[tokio::test]
async fn model_rejection_allows_cross_format_provider_mapping_to_allowed_global_model() {
let mut row = sample_row_for_api_format("gemini:generate_content");
+2 -50
View File
@@ -1,59 +1,11 @@
use axum::http::Uri;
use crate::headers::header_value_str;
use crate::{AppState, GatewayError};
use super::{resolve_control_route, GatewayControlDecision};
#[derive(Debug, Clone)]
pub(crate) struct GatewayPublicRequestContext {
pub(crate) trace_id: String,
pub(crate) request_method: http::Method,
pub(crate) request_path: String,
pub(crate) request_query_string: Option<String>,
pub(crate) request_content_type: Option<String>,
pub(crate) host_header: Option<String>,
pub(crate) control_decision: Option<GatewayControlDecision>,
}
impl GatewayPublicRequestContext {
pub(crate) fn from_request_parts(
trace_id: impl Into<String>,
method: &http::Method,
uri: &Uri,
headers: &http::HeaderMap,
control_decision: Option<GatewayControlDecision>,
) -> Self {
let request_path = if uri.path().starts_with('/') {
uri.path().to_string()
} else {
format!("/{}", uri.path())
};
let request_query_string = uri.query().map(ToOwned::to_owned);
Self {
trace_id: trace_id.into(),
request_method: method.clone(),
request_path,
request_query_string,
request_content_type: header_value_str(headers, http::header::CONTENT_TYPE.as_str()),
host_header: header_value_str(headers, http::header::HOST.as_str()),
control_decision,
}
}
pub(crate) fn request_path_and_query(&self) -> String {
if let Some(query) = self
.request_query_string
.as_deref()
.filter(|value| !value.is_empty())
{
format!("{}?{query}", self.request_path)
} else {
self.request_path.clone()
}
}
}
pub(crate) type GatewayPublicRequestContext =
aether_gateway_control::PublicRequestContext<GatewayControlDecision>;
pub(crate) async fn resolve_public_request_context(
state: &AppState,
+154
View File
@@ -89,4 +89,158 @@ impl GatewayDataConfig {
postgres: self.postgres.clone(),
}
}
pub(crate) fn split_runtime_pools(&self) -> (Self, Option<Self>) {
let configured_background_max =
std::env::var("AETHER_GATEWAY_BACKGROUND_DB_MAX_CONNECTIONS")
.ok()
.and_then(|value| value.trim().parse::<u32>().ok());
self.split_runtime_pools_with_background_max(configured_background_max)
}
/// Returns the database capacity reserved for isolated background work, when enabled.
pub fn background_database_config(&self) -> Option<SqlDatabaseConfig> {
self.split_runtime_pools()
.1
.and_then(|config| config.database)
}
fn split_runtime_pools_with_background_max(
&self,
configured_background_max: Option<u32>,
) -> (Self, Option<Self>) {
let Some(database) = self.database.as_ref() else {
return (self.clone(), None);
};
let total_max = database.pool.max_connections;
if total_max < 2
|| configured_background_max == Some(0)
|| is_private_sqlite_memory_database(database)
{
return (self.clone(), None);
}
let default_background_max = (total_max / 5).clamp(1, 8);
let background_max = configured_background_max
.unwrap_or(default_background_max)
.clamp(1, total_max.saturating_sub(1));
let foreground_max = total_max.saturating_sub(background_max).max(1);
let total_min = database.pool.min_connections.min(total_max);
let background_min = u32::from(total_min > 1).min(background_max);
// The configured minimum protects foreground readiness. Isolating background work must
// not take one of those warm foreground connections away; the background pool receives
// its own single warm connection while the combined hard maximum remains unchanged.
let foreground_min = total_min.min(foreground_max);
(
self.with_pool_limits(foreground_min, foreground_max),
Some(self.with_pool_limits(background_min, background_max)),
)
}
fn with_pool_limits(&self, min_connections: u32, max_connections: u32) -> Self {
let mut config = self.clone();
if let Some(database) = config.database.as_mut() {
database.pool.min_connections = min_connections.min(max_connections);
database.pool.max_connections = max_connections;
}
if let Some(postgres) = config.postgres.as_mut() {
postgres.min_connections = min_connections.min(max_connections);
postgres.max_connections = max_connections;
}
config
}
}
fn is_private_sqlite_memory_database(database: &aether_data::SqlDatabaseConfig) -> bool {
database.driver == aether_data::DatabaseDriver::Sqlite
&& matches!(database.url.trim(), "sqlite::memory:" | "sqlite://:memory:")
}
#[cfg(test)]
mod tests {
use super::GatewayDataConfig;
use aether_data::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
#[test]
fn runtime_pool_split_preserves_total_connection_budget() {
let config = GatewayDataConfig::from_database_config(
SqlDatabaseConfig::new(
DatabaseDriver::Postgres,
"postgres://localhost/aether",
SqlPoolConfig {
min_connections: 4,
max_connections: 20,
..SqlPoolConfig::default()
},
)
.expect("database config should be valid"),
);
let (foreground, background) = config.split_runtime_pools_with_background_max(Some(4));
let foreground = foreground.database().expect("foreground database");
let background = background
.expect("background database config")
.database()
.expect("background database")
.clone();
assert_eq!(foreground.pool.max_connections, 16);
assert_eq!(background.pool.max_connections, 4);
assert_eq!(
foreground.pool.max_connections + background.pool.max_connections,
20
);
assert_eq!(foreground.pool.min_connections, 4);
assert_eq!(background.pool.min_connections, 1);
}
#[test]
fn runtime_pool_split_can_be_disabled_or_degrade_for_single_connection() {
let mut database = SqlDatabaseConfig::sqlite_default();
database.pool.max_connections = 1;
let config = GatewayDataConfig::from_database_config(database);
assert!(config
.split_runtime_pools_with_background_max(Some(1))
.1
.is_none());
let mut database = SqlDatabaseConfig::sqlite_default();
database.pool.max_connections = 8;
let config = GatewayDataConfig::from_database_config(database);
assert!(config
.split_runtime_pools_with_background_max(Some(0))
.1
.is_none());
}
#[test]
fn runtime_pool_split_keeps_private_sqlite_memory_database_in_one_pool() {
for url in ["sqlite::memory:", "sqlite://:memory:"] {
let config = GatewayDataConfig::from_database_config(
SqlDatabaseConfig::new(
DatabaseDriver::Sqlite,
url,
SqlPoolConfig {
min_connections: 1,
max_connections: 8,
..SqlPoolConfig::default()
},
)
.expect("sqlite memory database config should be valid"),
);
let (foreground, background) = config.split_runtime_pools_with_background_max(Some(2));
assert!(background.is_none(), "private SQLite URL {url} was split");
assert_eq!(
foreground
.database()
.expect("foreground database")
.pool
.max_connections,
8
);
}
}
}
+54 -49
View File
@@ -1818,24 +1818,11 @@ impl GatewayDataState {
.effective_user_groups_for_user(&snapshot.user_id)
.await?;
let allowed_providers =
resolve_effective_list_policy(None, "unrestricted", &groups, |group| {
(
&group.allowed_providers_mode,
group.allowed_providers.clone(),
)
});
let allowed_api_formats =
resolve_effective_api_format_policy(None, "unrestricted", &groups, |group| {
(
&group.allowed_api_formats_mode,
group.allowed_api_formats.clone(),
)
});
let allowed_models =
resolve_effective_list_policy(None, "unrestricted", &groups, |group| {
(&group.allowed_models_mode, group.allowed_models.clone())
});
let GatewayUserEffectiveListPolicies {
allowed_providers,
allowed_api_formats,
allowed_models,
} = resolve_group_effective_list_policies(&groups);
let user_rate_limit = resolve_effective_rate_limit_policy(None, "system", &groups);
snapshot.user_allowed_providers = allowed_providers;
snapshot.user_allowed_api_formats = allowed_api_formats;
@@ -1860,36 +1847,7 @@ impl GatewayDataState {
} else {
Vec::new()
};
Ok(GatewayUserEffectiveListPolicies {
allowed_providers: resolve_effective_list_policy(
user.allowed_providers.clone(),
&user.allowed_providers_mode,
&groups,
|group| {
(
&group.allowed_providers_mode,
group.allowed_providers.clone(),
)
},
),
allowed_api_formats: resolve_effective_api_format_policy(
user.allowed_api_formats.clone(),
&user.allowed_api_formats_mode,
&groups,
|group| {
(
&group.allowed_api_formats_mode,
group.allowed_api_formats.clone(),
)
},
),
allowed_models: resolve_effective_list_policy(
user.allowed_models.clone(),
&user.allowed_models_mode,
&groups,
|group| (&group.allowed_models_mode, group.allowed_models.clone()),
),
})
Ok(resolve_group_effective_list_policies(&groups))
}
async fn effective_user_groups_for_user(
@@ -1979,6 +1937,35 @@ fn apply_admin_unrestricted_auth_snapshot(snapshot: &mut StoredAuthApiKeySnapsho
snapshot.api_key_concurrent_limit = None;
}
// Per-user list policy columns are retained only for legacy import/export compatibility.
// Runtime authorization and user-facing catalogs must both treat group policies as authoritative.
fn resolve_group_effective_list_policies(
groups: &[aether_data::repository::users::StoredUserGroup],
) -> GatewayUserEffectiveListPolicies {
GatewayUserEffectiveListPolicies {
allowed_providers: resolve_effective_list_policy(None, "unrestricted", groups, |group| {
(
&group.allowed_providers_mode,
group.allowed_providers.clone(),
)
}),
allowed_api_formats: resolve_effective_api_format_policy(
None,
"unrestricted",
groups,
|group| {
(
&group.allowed_api_formats_mode,
group.allowed_api_formats.clone(),
)
},
),
allowed_models: resolve_effective_list_policy(None, "unrestricted", groups, |group| {
(&group.allowed_models_mode, group.allowed_models.clone())
}),
}
}
fn resolve_effective_list_policy(
user_values: Option<Vec<String>>,
user_mode: &str,
@@ -2592,8 +2579,9 @@ mod tests {
Some("hash-user".to_string()),
snapshot,
)]));
let user = sample_auth_user("user-1", "user");
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![
sample_auth_user("user-1", "user"),
user.clone()
]));
let group = user_repository
.create_user_group(UpsertUserGroupRecord {
@@ -2638,6 +2626,23 @@ mod tests {
Some(&["claude-sonnet-4-5".to_string()][..])
);
assert_eq!(resolved.user_rate_limit, Some(30));
let catalog_policies = state
.resolve_user_effective_list_policies(&user)
.await
.expect("catalog policies should resolve");
assert_eq!(
catalog_policies.allowed_providers.as_deref(),
resolved.effective_allowed_providers()
);
assert_eq!(
catalog_policies.allowed_api_formats.as_deref(),
resolved.effective_allowed_api_formats()
);
assert_eq!(
catalog_policies.allowed_models.as_deref(),
resolved.effective_allowed_models()
);
}
#[tokio::test]
File diff suppressed because it is too large Load Diff
+44 -4
View File
@@ -38,7 +38,8 @@ use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
StoredProviderApiKeyWindowUsageSummary, StoredUsageDailySummary, UsageAuditListQuery,
UsageCleanupExecutionMode, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow,
UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageDailyHeatmapQuery,
UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot,
UsageDailyHeatmapQuery,
};
use aether_runtime_state::RuntimeQueueStore;
use aether_video_tasks_core::read_data_backed_video_task_response;
@@ -152,6 +153,13 @@ impl GatewayDataState {
}
}
pub(crate) async fn warm_database_pool(&self) -> Result<(), DataLayerError> {
match &self.backends {
Some(backends) => backends.warm_database_pool().await,
None => Ok(()),
}
}
pub(crate) async fn pending_database_backfills(
&self,
) -> Result<
@@ -198,15 +206,22 @@ impl GatewayDataState {
pub(crate) fn database_pool_summary_under_maintenance_pressure(
summary: &aether_data::DatabasePoolSummary,
) -> bool {
summary.checked_out > 0 && summary.idle <= Self::maintenance_pool_idle_reserve(summary)
summary.checked_out > 0
&& Self::database_pool_available_capacity(summary)
<= Self::maintenance_pool_idle_reserve(summary)
}
pub(crate) fn database_pool_summary_under_usage_worker_pressure(
summary: &aether_data::DatabasePoolSummary,
) -> bool {
summary.checked_out > 0
&& (summary.checked_out >= summary.max_connections as usize
|| summary.idle <= Self::usage_worker_pool_idle_reserve(summary))
&& Self::database_pool_available_capacity(summary)
<= Self::usage_worker_pool_idle_reserve(summary)
}
fn database_pool_available_capacity(summary: &aether_data::DatabasePoolSummary) -> usize {
let unopened = (summary.max_connections as usize).saturating_sub(summary.pool_size);
summary.idle.saturating_add(unopened)
}
pub(crate) fn maintenance_pool_idle_reserve(
@@ -1353,6 +1368,15 @@ impl GatewayDataState {
}
}
pub(crate) async fn read_usage_counter_pending_health(
&self,
) -> Result<UsageCounterPendingHealthSnapshot, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.read_usage_counter_pending_health().await,
None => Ok(UsageCounterPendingHealthSnapshot::default()),
}
}
pub(crate) async fn summarize_usage_totals_by_user_ids(
&self,
user_ids: &[String],
@@ -1439,6 +1463,22 @@ impl GatewayDataState {
}
}
pub(crate) async fn summarize_dashboard_stats(
&self,
query: &aether_data_contracts::repository::usage::UsageDashboardSummaryQuery,
) -> Result<
aether_data_contracts::repository::usage::StoredUsageDashboardStatsSummary,
DataLayerError,
> {
match &self.usage_reader {
Some(repository) => repository.summarize_dashboard_stats(query).await,
None => Ok(
aether_data_contracts::repository::usage::StoredUsageDashboardStatsSummary::default(
),
),
}
}
pub(crate) async fn list_dashboard_daily_breakdown(
&self,
query: &aether_data_contracts::repository::usage::UsageDashboardDailyBreakdownQuery,
+4 -2
View File
@@ -54,7 +54,7 @@ fn disabled_gateway_data_state_has_no_backends() {
#[test]
fn maintenance_pool_pressure_keeps_idle_reserve_for_foreground_work() {
let pressured = aether_data::DatabasePoolSummary {
let pool_can_still_grow = aether_data::DatabasePoolSummary {
driver: DatabaseDriver::Postgres,
checked_out: 6,
pool_size: 6,
@@ -62,7 +62,9 @@ fn maintenance_pool_pressure_keeps_idle_reserve_for_foreground_work() {
max_connections: 20,
usage_rate: 30.0,
};
assert!(GatewayDataState::database_pool_summary_under_maintenance_pressure(&pressured));
assert!(
!GatewayDataState::database_pool_summary_under_maintenance_pressure(&pool_can_still_grow)
);
let reserve_idle_left = aether_data::DatabasePoolSummary {
driver: DatabaseDriver::Postgres,
@@ -1,3 +1,5 @@
pub(crate) const MAX_ERROR_BODY_BYTES: usize = 16_384;
pub(crate) const MAX_STREAM_PREFETCH_FRAMES: usize = 5;
pub(crate) const MAX_STREAM_PREFETCH_BYTES: usize = 16_384;
//! Compatibility facade for execution resource limits.
pub(crate) use aether_gateway_execution::{
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
};
@@ -1,3 +1,5 @@
//! Compatibility facade for execution stream framing.
use std::io::Error as IoError;
use aether_contracts::StreamFrame;
@@ -6,36 +8,10 @@ use axum::body::Bytes;
use crate::GatewayError;
pub(crate) fn encode_stream_frame_ndjson(frame: &StreamFrame) -> Result<Bytes, IoError> {
let mut raw = serde_json::to_vec(frame).map_err(|err| IoError::other(err.to_string()))?;
raw.push(b'\n');
Ok(Bytes::from(raw))
aether_gateway_execution::stream::encode_stream_frame_ndjson(frame)
}
pub(crate) fn decode_stream_frame_ndjson(line: &[u8]) -> Result<StreamFrame, GatewayError> {
serde_json::from_slice(line).map_err(|err| GatewayError::Internal(err.to_string()))
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use aether_contracts::{StreamFramePayload, StreamFrameType};
use super::{decode_stream_frame_ndjson, encode_stream_frame_ndjson};
#[test]
fn ndjson_round_trip_preserves_frame() {
let frame = aether_contracts::StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code: 200,
headers: BTreeMap::from([("content-type".into(), "text/event-stream".into())]),
},
};
let raw = encode_stream_frame_ndjson(&frame).expect("frame should encode");
let decoded =
decode_stream_frame_ndjson(raw.trim_ascii_end()).expect("frame should decode");
assert_eq!(decoded, frame);
}
aether_gateway_execution::stream::decode_stream_frame_ndjson(line)
.map_err(|err| GatewayError::Internal(err.to_string()))
}
@@ -259,6 +259,11 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
"has_format_conversion",
),
slow_threshold_ms,
include_timeline: query_param_optional_bool(
request_context.query_string(),
"include_timeline",
)
.unwrap_or(true),
})
.await?;
return Ok(Some(build_admin_stats_provider_performance_response(
@@ -600,7 +600,7 @@ mod tests {
};
use crate::AppState;
use aether_runtime_state::{RedisClientConfig, RuntimeState, RuntimeStateConfig};
use aether_testkit::ManagedRedisServer;
use aether_test_support::ManagedRedisServer;
use std::collections::BTreeMap;
async fn start_managed_redis_or_skip() -> Option<ManagedRedisServer> {
@@ -2,11 +2,11 @@ use crate::api::response::build_local_http_error_response;
use crate::control::GatewayPublicRequestContext;
use crate::headers::RequestBodyNormalizationError;
use crate::{AppState, GatewayError};
use axum::body::{to_bytes, Body, Bytes};
use aether_gateway_frontdoor::{BodyBufferError, BodyBufferPolicy as FrontdoorBodyBufferPolicy};
use axum::body::{Body, Bytes};
use axum::http::{self, Response};
use std::error::Error as StdError;
use std::sync::Arc;
use std::time::{Duration, Instant};
use std::time::Duration;
use tokio::sync::Semaphore;
use tracing::{info, warn};
@@ -16,23 +16,22 @@ const REQUEST_BODY_READ_FAILED_DETAIL: &str = "Failed to read request body";
#[derive(Debug, Clone)]
pub(super) struct RequestBodyBufferPolicy {
max_bytes: u64,
read_timeout: Duration,
queue_timeout: Duration,
budget_bytes: usize,
budget: Arc<Semaphore>,
inner: FrontdoorBodyBufferPolicy,
}
impl RequestBodyBufferPolicy {
pub(super) fn from_state(state: &AppState) -> Self {
Self {
max_bytes: crate::headers::max_request_body_bytes(),
read_timeout: state.frontdoor_runtime_guards.request_body_read_timeout,
queue_timeout: state.frontdoor_runtime_guards.internal_gate_queue_budget,
budget_bytes: state
.frontdoor_runtime_guards
.request_body_buffer_budget_bytes,
budget: Arc::clone(&state.request_body_buffer_budget),
inner: FrontdoorBodyBufferPolicy::with_permit_bytes(
crate::headers::max_request_body_bytes(),
state.frontdoor_runtime_guards.request_body_read_timeout,
state.frontdoor_runtime_guards.internal_gate_queue_budget,
state
.frontdoor_runtime_guards
.request_body_buffer_budget_bytes,
crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
Arc::clone(&state.request_body_buffer_budget),
),
}
}
@@ -42,14 +41,17 @@ impl RequestBodyBufferPolicy {
.unwrap_or(usize::MAX)
.max(crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES);
Self {
max_bytes,
read_timeout,
queue_timeout: read_timeout,
budget_bytes,
budget: Arc::new(Semaphore::new(
budget_bytes.saturating_add(crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES - 1)
/ crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
)),
inner: FrontdoorBodyBufferPolicy::with_permit_bytes(
max_bytes,
read_timeout,
read_timeout,
budget_bytes,
crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
Arc::new(Semaphore::new(
budget_bytes.saturating_add(crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES - 1)
/ crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
)),
),
}
}
@@ -62,13 +64,35 @@ impl RequestBodyBufferPolicy {
budget: Arc<Semaphore>,
) -> Self {
Self {
max_bytes,
read_timeout,
queue_timeout,
budget_bytes,
budget,
inner: FrontdoorBodyBufferPolicy::with_permit_bytes(
max_bytes,
read_timeout,
queue_timeout,
budget_bytes,
crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES,
budget,
),
}
}
fn max_bytes(&self) -> u64 {
self.inner.max_bytes()
}
fn budget_bytes(&self) -> usize {
self.inner.budget_bytes()
}
fn read_timeout(&self) -> Duration {
self.inner.read_timeout()
}
async fn reserve(
&self,
headers: &http::HeaderMap,
) -> Result<aether_gateway_frontdoor::BodyBufferReservation, BodyBufferError> {
self.inner.reserve(headers).await
}
}
#[derive(Debug)]
@@ -135,30 +159,23 @@ impl RequestBodyBufferError {
}
}
fn request_body_buffer_reservation_bytes(headers: &http::HeaderMap, max_bytes: u64) -> usize {
let max_bytes = usize::try_from(max_bytes).unwrap_or(usize::MAX);
let encoded = headers
.get(http::header::CONTENT_ENCODING)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.is_some_and(|value| !value.is_empty() && !value.eq_ignore_ascii_case("identity"));
if encoded {
return max_bytes;
impl From<BodyBufferError> for RequestBodyBufferError {
fn from(error: BodyBufferError) -> Self {
match error {
BodyBufferError::TooLarge { limit_bytes } => Self::TooLarge { limit_bytes },
BodyBufferError::Overloaded {
requested_bytes,
budget_bytes,
timeout_ms,
} => Self::Overloaded {
requested_bytes,
budget_bytes,
timeout_ms,
},
BodyBufferError::Timeout { timeout_ms } => Self::Timeout { timeout_ms },
BodyBufferError::ReadFailed { message } => Self::ReadFailed { message },
}
}
headers
.get(http::header::CONTENT_LENGTH)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.trim().parse::<usize>().ok())
.map(|value| value.min(max_bytes))
.unwrap_or(max_bytes)
}
fn request_body_buffer_reservation_permits(reservation_bytes: usize) -> u32 {
let permits = reservation_bytes
.max(1)
.saturating_add(crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES - 1)
/ crate::state::REQUEST_BODY_BUFFER_PERMIT_BYTES;
u32::try_from(permits).unwrap_or(u32::MAX).max(1)
}
pub(super) async fn buffer_and_normalize_request_body(
@@ -171,33 +188,12 @@ pub(super) async fn buffer_and_normalize_request_body(
phase: &'static str,
policy: RequestBodyBufferPolicy,
) -> Result<Bytes, RequestBodyBufferError> {
if let Err(err) =
crate::headers::check_request_content_length_with_limit(headers, policy.max_bytes)
{
return Err(RequestBodyBufferError::Normalization(err));
}
let reservation = policy
.reserve(headers)
.await
.map_err(RequestBodyBufferError::from)?;
let reservation_bytes = reservation.requested_bytes();
let reservation_bytes = request_body_buffer_reservation_bytes(headers, policy.max_bytes);
let reservation_permits = request_body_buffer_reservation_permits(reservation_bytes);
let queue_timeout_ms = policy.queue_timeout.as_millis() as u64;
let _budget_permit = match tokio::time::timeout(
policy.queue_timeout,
Arc::clone(&policy.budget).acquire_many_owned(reservation_permits),
)
.await
{
Ok(Ok(permit)) => permit,
Ok(Err(_)) | Err(_) => {
return Err(RequestBodyBufferError::Overloaded {
requested_bytes: reservation_bytes,
budget_bytes: policy.budget_bytes,
timeout_ms: queue_timeout_ms,
});
}
};
let read_started_at = Instant::now();
let timeout_ms = policy.read_timeout.as_millis() as u64;
info!(
event_name = "frontdoor_request_body_buffer_started",
log_type = "event",
@@ -205,45 +201,27 @@ pub(super) async fn buffer_and_normalize_request_body(
method = %method,
path = %path_and_query,
phase,
max_body_bytes = policy.max_bytes,
max_body_bytes = policy.max_bytes(),
reserved_body_bytes = reservation_bytes,
body_buffer_budget_bytes = policy.budget_bytes,
timeout_ms,
body_buffer_budget_bytes = policy.budget_bytes(),
timeout_ms = policy.read_timeout().as_millis() as u64,
"gateway started buffering request body"
);
let body_limit = usize::try_from(policy.max_bytes).unwrap_or(usize::MAX);
let body = match tokio::time::timeout(
policy.read_timeout,
to_bytes(
request_body.take().expect(body_owner_expectation),
body_limit,
),
)
.await
{
Ok(Ok(body)) => body,
Ok(Err(err)) if request_body_collection_exceeded_limit(&err) => {
return Err(RequestBodyBufferError::TooLarge {
limit_bytes: policy.max_bytes,
});
}
Ok(Err(err)) => {
return Err(RequestBodyBufferError::ReadFailed {
message: err.to_string(),
});
}
Err(_) => {
return Err(RequestBodyBufferError::Timeout { timeout_ms });
}
};
let normalized = crate::headers::normalize_request_body_headers_and_bytes_with_limit(
headers,
body,
policy.max_bytes,
)
.map_err(RequestBodyBufferError::Normalization)?;
let buffered = reservation
.collect(request_body.take().expect(body_owner_expectation))
.await
.map_err(RequestBodyBufferError::from)?;
let elapsed_ms = buffered.elapsed().as_millis() as u64;
let normalized = buffered
.try_map(|body| {
crate::headers::normalize_request_body_headers_and_bytes_with_limit(
headers,
body,
policy.max_bytes(),
)
})
.map_err(RequestBodyBufferError::Normalization)?;
info!(
event_name = "frontdoor_request_body_buffer_completed",
log_type = "event",
@@ -252,23 +230,12 @@ pub(super) async fn buffer_and_normalize_request_body(
path = %path_and_query,
phase,
body_bytes = normalized.len(),
elapsed_ms = read_started_at.elapsed().as_millis() as u64,
elapsed_ms,
"gateway completed buffering request body"
);
Ok(normalized)
}
fn request_body_collection_exceeded_limit(error: &(dyn StdError + 'static)) -> bool {
let mut current = Some(error);
while let Some(error) = current {
if error.to_string().contains("length limit exceeded") {
return true;
}
current = error.source();
}
false
}
pub(super) fn build_request_body_buffer_error_response(
trace_id: &str,
request_context: &GatewayPublicRequestContext,
@@ -5,8 +5,8 @@ use super::{
};
use aether_data_contracts::repository::usage::{
StoredUsageCostSavingsSummary, StoredUsageDashboardDailyBreakdownRow,
StoredUsageDashboardSummary, UsageAuditAggregationGroupBy, UsageAuditAggregationQuery,
UsageCostSavingsSummaryQuery, UsageDashboardDailyBreakdownQuery,
StoredUsageDashboardStatsSummary, StoredUsageDashboardSummary, UsageAuditAggregationGroupBy,
UsageAuditAggregationQuery, UsageDashboardDailyBreakdownQuery,
UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery,
};
use axum::{
@@ -550,6 +550,39 @@ async fn dashboard_summary_for_range(
dashboard_summary_for_range_raw(state, range, user_id, error_context).await
}
async fn dashboard_stats_for_range(
state: &AppState,
range: DashboardDateRange,
user_id: Option<&str>,
error_context: &str,
) -> Result<StoredUsageDashboardStatsSummary, Response<Body>> {
let Some((created_from_unix_secs, created_until_unix_secs)) =
dashboard_range_bounds_unix(range)
else {
return Err(build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("{error_context}: invalid time range"),
false,
));
};
match state
.summarize_dashboard_stats(&UsageDashboardSummaryQuery {
created_from_unix_secs,
created_until_unix_secs,
user_id: user_id.map(ToOwned::to_owned),
})
.await
{
Ok(value) => Ok(value),
Err(err) => Err(build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("{error_context}: {err:?}"),
false,
)),
}
}
async fn dashboard_daily_breakdown_for_range(
state: &AppState,
range: DashboardDateRange,
@@ -855,28 +888,6 @@ fn dashboard_cache_savings_usd(summary: &StoredUsageCostSavingsSummary) -> f64 {
)
}
async fn dashboard_load_cache_savings(
state: &AppState,
range: DashboardDateRange,
user_id: Option<&str>,
) -> Result<f64, GatewayError> {
let Some((created_from_unix_secs, created_until_unix_secs)) =
dashboard_range_bounds_unix(range)
else {
return Ok(0.0);
};
let summary = state
.summarize_usage_cost_savings(&UsageCostSavingsSummaryQuery {
created_from_unix_secs,
created_until_unix_secs,
user_id: user_id.map(ToOwned::to_owned),
provider_name: None,
model: None,
})
.await?;
Ok(dashboard_cache_savings_usd(&summary))
}
pub(super) async fn handle_dashboard_stats_get(
state: &AppState,
request_context: &GatewayPublicRequestContext,
@@ -902,7 +913,7 @@ pub(super) async fn handle_dashboard_stats_get(
.as_deref()
.unwrap_or("");
let cache_key = format!("stats:{cache_identity}:{query_string}");
let cache_ttl = std::time::Duration::from_secs(15);
let cache_ttl = std::time::Duration::from_secs(30);
if let Some(cached) = state.dashboard_response_cache.get(&cache_key, cache_ttl) {
return Response::builder()
@@ -924,31 +935,63 @@ pub(super) async fn handle_dashboard_stats_get(
tz_offset_minutes: summary_range.tz_offset_minutes,
};
let user_filter = (!is_admin).then_some(auth.user.id.as_str());
let period_summary = match dashboard_summary_for_range(
state,
summary_range,
user_filter,
"dashboard stats lookup failed",
)
.await
{
Ok(value) => value,
Err(response) => return response,
let (period_totals, today_totals, admin_cost_savings) = if is_admin {
let (period_result, today_result) = tokio::join!(
dashboard_stats_for_range(
state,
summary_range,
user_filter,
"dashboard stats lookup failed",
),
dashboard_stats_for_range(
state,
today_range,
user_filter,
"dashboard today stats lookup failed",
),
);
let period = match period_result {
Ok(value) => value,
Err(response) => return response,
};
let today = match today_result {
Ok(value) => value,
Err(response) => return response,
};
(
dashboard_usage_totals_from_summary(&period.usage),
dashboard_usage_totals_from_summary(&today.usage),
Some((period.cost_savings, today.cost_savings)),
)
} else {
let period_summary = match dashboard_summary_for_range(
state,
summary_range,
user_filter,
"dashboard stats lookup failed",
)
.await
{
Ok(value) => value,
Err(response) => return response,
};
let today_summary = match dashboard_summary_for_range(
state,
today_range,
user_filter,
"dashboard today stats lookup failed",
)
.await
{
Ok(value) => value,
Err(response) => return response,
};
(
dashboard_usage_totals_from_summary(&period_summary),
dashboard_usage_totals_from_summary(&today_summary),
None,
)
};
let today_summary = match dashboard_summary_for_range(
state,
today_range,
user_filter,
"dashboard today stats lookup failed",
)
.await
{
Ok(value) => value,
Err(response) => return response,
};
let period_totals = dashboard_usage_totals_from_summary(&period_summary);
let today_totals = dashboard_usage_totals_from_summary(&today_summary);
let api_key_counts = match dashboard_load_api_key_counts(state, is_admin, &auth.user.id).await {
Ok(value) => value,
@@ -1030,28 +1073,10 @@ pub(super) async fn handle_dashboard_stats_get(
/ today_totals.requests as f64
* 100.0
};
let today_cost_savings =
match dashboard_load_cache_savings(state, today_range, user_filter).await {
Ok(value) => value,
Err(err) => {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("dashboard today cache savings lookup failed: {err:?}"),
false,
);
}
};
let period_cost_savings =
match dashboard_load_cache_savings(state, summary_range, user_filter).await {
Ok(value) => value,
Err(err) => {
return build_auth_error_response(
http::StatusCode::INTERNAL_SERVER_ERROR,
format!("dashboard cache savings lookup failed: {err:?}"),
false,
);
}
};
let (period_cost_savings_summary, today_cost_savings_summary) =
admin_cost_savings.unwrap_or_default();
let today_cost_savings = dashboard_cache_savings_usd(&today_cost_savings_summary);
let period_cost_savings = dashboard_cache_savings_usd(&period_cost_savings_summary);
let stats = json!([
{
"name": "今日请求 / 费用",
@@ -1248,7 +1273,7 @@ pub(super) async fn handle_dashboard_daily_stats_get(
.as_deref()
.unwrap_or("");
let cache_key = format!("daily:{cache_identity}:{query_string}");
let cache_ttl = std::time::Duration::from_secs(30);
let cache_ttl = std::time::Duration::from_secs(60);
if let Some(cached) = state.dashboard_response_cache.get(&cache_key, cache_ttl) {
return Response::builder()
+2 -73
View File
@@ -1,74 +1,3 @@
pub(crate) fn short_request_id(value: &str) -> String {
let trimmed = value.trim();
if trimmed.is_empty() {
return "-".to_string();
}
//! Compatibility facade for request identifier formatting.
if trimmed.chars().count() <= 12 {
return trimmed.to_string();
}
if looks_like_uuid(trimmed) {
return trimmed.chars().take(8).collect();
}
let prefix: String = trimmed.chars().take(6).collect();
let suffix: String = trimmed
.chars()
.rev()
.take(4)
.collect::<String>()
.chars()
.rev()
.collect();
format!("{prefix}...{suffix}")
}
fn looks_like_uuid(value: &str) -> bool {
let bytes = value.as_bytes();
if bytes.len() != 36 {
return false;
}
for (index, byte) in bytes.iter().enumerate() {
let is_hyphen = matches!(index, 8 | 13 | 18 | 23);
if is_hyphen {
if *byte != b'-' {
return false;
}
continue;
}
if !byte.is_ascii_hexdigit() {
return false;
}
}
true
}
#[cfg(test)]
mod tests {
use super::short_request_id;
#[test]
fn shortens_uuid_like_request_ids_to_prefix() {
assert_eq!(
short_request_id("d07e1e94-41b8-409f-a18a-27993ae7ecb1"),
"d07e1e94"
);
}
#[test]
fn shortens_long_named_request_ids_to_prefix_and_suffix() {
assert_eq!(
short_request_id("trace-openai-cli-stream-sync-direct-123"),
"trace-...-123"
);
}
#[test]
fn preserves_short_request_ids() {
assert_eq!(short_request_id("req-123"), "req-123");
}
}
pub(crate) use aether_gateway_frontdoor::short_request_id;
+81 -16
View File
@@ -138,6 +138,10 @@ impl NodeRoleArg {
const fn spawns_background_tasks(self) -> bool {
matches!(self, Self::All | Self::Background)
}
const fn isolates_background_database(self) -> bool {
matches!(self, Self::All)
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)]
@@ -358,6 +362,18 @@ fn usage_queue_workers_for_request_concurrency(request_concurrency: usize) -> us
workers.clamp(AUTO_USAGE_QUEUE_WORKERS_MIN, MAX_USAGE_QUEUE_WORKERS)
}
fn usage_database_config_for_role<'a>(
node_role: NodeRoleArg,
database: Option<&'a SqlDatabaseConfig>,
isolated_background_database: Option<&'a SqlDatabaseConfig>,
) -> Option<&'a SqlDatabaseConfig> {
if node_role.isolates_background_database() {
isolated_background_database.or(database)
} else {
database
}
}
fn usage_queue_worker_database_cap(
node_role: NodeRoleArg,
database: Option<&SqlDatabaseConfig>,
@@ -1279,7 +1295,7 @@ struct Args {
#[arg(
long,
env = "AETHER_RUNTIME_COMMAND_TIMEOUT_MS",
default_value_t = 1_000
default_value_t = 2_000
)]
runtime_command_timeout_ms: u64,
@@ -1710,6 +1726,18 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
runtime_redis_url.as_deref(),
runtime_backend,
)?;
let data_config = args.data.to_config();
let isolate_background_database = args.node_role.isolates_background_database();
let background_database_config = if isolate_background_database {
data_config.background_database_config()
} else {
None
};
let usage_database_config = usage_database_config_for_role(
args.node_role,
data_config.database(),
background_database_config.as_ref(),
);
let request_concurrency_limit = args
.max_in_flight_requests
.filter(|limit| *limit > 0)
@@ -1728,16 +1756,16 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
args.node_role,
Some(request_concurrency_limit),
args.distributed_request_limit,
sql_database_config.as_ref(),
usage_database_config,
);
let usage_queue_worker_max_count = args.usage.effective_queue_worker_max_count(
args.node_role,
sql_database_config.as_ref(),
usage_database_config,
usage_queue_workers,
);
let usage_worker_record_concurrency_limit = args
.usage
.effective_worker_record_concurrency_limit(args.node_role, sql_database_config.as_ref());
.effective_worker_record_concurrency_limit(args.node_role, usage_database_config);
let usage_config = args.usage.to_config(
usage_queue_workers,
usage_queue_worker_max_count,
@@ -1745,7 +1773,7 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
);
let usage_blocking_stream_lanes = args.usage.runtime_state_blocking_stream_lanes(
args.node_role,
sql_database_config.as_ref(),
usage_database_config,
usage_config.worker_max_count,
);
let runtime_state = Arc::new(
@@ -1757,7 +1785,6 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
.await
.map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err.to_string()))?,
);
let data_config = args.data.to_config();
let rate_limit_config = if matches!(args.deployment_topology, DeploymentTopologyArg::MultiNode)
{
args.rate_limit.config().with_local_fallback(false)
@@ -1859,7 +1886,7 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
let mut state = AppState::new()?
.with_runtime_state(runtime_state)
.with_data_config(data_config)?
.with_data_config_and_background_isolation(data_config, isolate_background_database)?
.with_usage_runtime_config(usage_config)?
.with_video_task_truth_source_mode(args.video_task_truth_source_mode.into());
if let Some(cors_config) = args.frontdoor.cors_config() {
@@ -1927,6 +1954,7 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
"aether-gateway data layer configured"
);
prepare_database_startup_requirements(&state, args.auto_prepare_database).await?;
state.warm_database_pools().await?;
let reset_stale_proxy_nodes = state.reset_stale_proxy_node_tunnel_statuses().await?;
if reset_stale_proxy_nodes > 0 {
info!(
@@ -1986,6 +2014,13 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
);
None
};
if state.prewarm_metric_snapshot().await {
info!("gateway metric snapshot prewarmed");
} else {
warn!(
"gateway metric snapshot prewarm did not complete; continuing with fail-open metrics"
);
}
let listen_backlog = gateway_listen_backlog(args.listen_backlog);
let listener_shards = gateway_listener_shards(args.listener_shards);
let listeners = gateway_listeners(bind_addr, listen_backlog, listener_shards)?;
@@ -2378,15 +2413,15 @@ mod tests {
automatic_gateway_request_concurrency_for_parallelism, automatic_sql_pool_config,
automatic_sql_pool_config_for_parallelism, automatic_usage_queue_workers_for_parallelism,
ensure_database_backfills_are_current, ensure_database_schema_is_current,
pending_backfills_error, pending_schema_error, resolve_healthcheck_url, Args,
DatabaseDriverArg, DeploymentTopologyArg, GatewayDataArgs, GatewayFrontdoorArgs,
GatewayLogDestinationArg, GatewayLogFormatArg, GatewayLogRotationArg, GatewayLoggingArgs,
GatewayRateLimitArgs, GatewayUsageArgs, NodeRoleArg, RuntimeBackendArg,
VideoTaskTruthSourceArg, DEFAULT_GATEWAY_HTTP2_MAX_CONCURRENT_STREAMS,
DEFAULT_GATEWAY_LISTENER_SHARDS, DEFAULT_GATEWAY_LISTEN_BACKLOG,
MAX_GATEWAY_HTTP2_MAX_CONCURRENT_STREAMS, MAX_GATEWAY_LISTENER_SHARDS,
MAX_GATEWAY_LISTEN_BACKLOG, MIN_GATEWAY_HTTP2_MAX_CONCURRENT_STREAMS,
MIN_GATEWAY_LISTEN_BACKLOG,
pending_backfills_error, pending_schema_error, resolve_healthcheck_url,
usage_database_config_for_role, Args, DatabaseDriverArg, DeploymentTopologyArg,
GatewayDataArgs, GatewayFrontdoorArgs, GatewayLogDestinationArg, GatewayLogFormatArg,
GatewayLogRotationArg, GatewayLoggingArgs, GatewayRateLimitArgs, GatewayUsageArgs,
NodeRoleArg, RuntimeBackendArg, VideoTaskTruthSourceArg,
DEFAULT_GATEWAY_HTTP2_MAX_CONCURRENT_STREAMS, DEFAULT_GATEWAY_LISTENER_SHARDS,
DEFAULT_GATEWAY_LISTEN_BACKLOG, MAX_GATEWAY_HTTP2_MAX_CONCURRENT_STREAMS,
MAX_GATEWAY_LISTENER_SHARDS, MAX_GATEWAY_LISTEN_BACKLOG,
MIN_GATEWAY_HTTP2_MAX_CONCURRENT_STREAMS, MIN_GATEWAY_LISTEN_BACKLOG,
};
use aether_data::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig};
use aether_gateway::AppState;
@@ -2695,6 +2730,36 @@ mod tests {
assert_eq!(many_cpu.max_connections, 100);
}
#[test]
fn gateway_database_pool_isolation_and_usage_capacity_follow_role() {
assert!(NodeRoleArg::All.isolates_background_database());
assert!(!NodeRoleArg::Frontdoor.isolates_background_database());
assert!(!NodeRoleArg::Background.isolates_background_database());
let database = test_database(DatabaseDriver::Postgres, 20);
let isolated_background = test_database(DatabaseDriver::Postgres, 4);
assert_eq!(
usage_database_config_for_role(
NodeRoleArg::All,
Some(&database),
Some(&isolated_background),
)
.expect("all-role usage database")
.pool
.max_connections,
4
);
for role in [NodeRoleArg::Frontdoor, NodeRoleArg::Background] {
assert_eq!(
usage_database_config_for_role(role, Some(&database), Some(&isolated_background),)
.expect("single-pool role usage database")
.pool
.max_connections,
20
);
}
}
#[test]
fn gateway_usage_queue_workers_manual_override_wins_and_is_capped() {
let mut args = test_args();
@@ -761,33 +761,38 @@ pub(crate) fn spawn_account_self_check_worker(
}
let config = AccountSelfCheckWorkerConfig::from_env();
Some(tokio::spawn(async move {
let mut interval = tokio::time::interval(config.scan_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut deferred_since = None;
loop {
Some(crate::task_runtime::spawn_singleton_worker(
state,
crate::task_runtime::TASK_KEY_ACCOUNT_SELF_CHECK,
move |state| async move {
let mut interval = tokio::time::interval(config.scan_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
if state
.data
.should_defer_maintenance_for_database_pool_pressure(&mut deferred_since)
{
debug!(
event_name = "maintenance_worker_deferred",
log_type = "ops",
worker = "account_self_check",
"gateway account self-check deferred because database pool has no idle reserve"
);
continue;
let mut deferred_since = None;
loop {
interval.tick().await;
if state
.data
.should_defer_maintenance_for_database_pool_pressure(&mut deferred_since)
{
debug!(
event_name = "maintenance_worker_deferred",
log_type = "ops",
worker = "account_self_check",
"gateway account self-check deferred because database pool has no idle reserve"
);
continue;
}
if let Err(err) = perform_account_self_check_once_with_config(&state, config).await
{
warn!(
error = ?err,
"gateway account self-check worker tick failed"
);
}
}
if let Err(err) = perform_account_self_check_once_with_config(&state, config).await {
warn!(
error = ?err,
"gateway account self-check worker tick failed"
);
}
}
}))
},
))
}
#[cfg(test)]
@@ -1701,20 +1701,24 @@ pub(crate) fn spawn_pool_quota_probe_worker(
}
let config = PoolQuotaProbeWorkerConfig::from_env();
Some(tokio::spawn(async move {
let mut interval = tokio::time::interval(config.scan_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
loop {
Some(crate::task_runtime::spawn_singleton_worker(
state,
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
move |state| async move {
let mut interval = tokio::time::interval(config.scan_interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
if let Err(err) = perform_pool_quota_probe_once_with_config(&state, config).await {
warn!(
error = ?err,
"gateway pool quota probe worker tick failed"
);
loop {
interval.tick().await;
if let Err(err) = perform_pool_quota_probe_once_with_config(&state, config).await {
warn!(
error = ?err,
"gateway pool quota probe worker tick failed"
);
}
}
}
}))
},
))
}
#[cfg(test)]
@@ -401,50 +401,54 @@ pub(crate) fn spawn_pool_score_rebuild_worker(
}
let config = PoolScoreRebuildWorkerConfig::from_env();
Some(tokio::spawn(async move {
if let Err(err) = perform_pool_score_rebuild_once_with_config(&state, config).await {
warn!(
error = ?err,
"gateway pool score rebuild initial tick failed"
);
}
let mut interval = tokio::time::interval(config.interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let mut deferred_since = None;
loop {
interval.tick().await;
if state
.data
.should_defer_maintenance_for_database_pool_pressure(&mut deferred_since)
{
debug!(
event_name = "maintenance_worker_deferred",
log_type = "ops",
worker = "pool_score_rebuild",
"gateway pool score rebuild deferred because database pool has no idle reserve"
Some(crate::task_runtime::spawn_singleton_worker(
state,
crate::task_runtime::TASK_KEY_POOL_SCORE_REBUILD,
move |state| async move {
if let Err(err) = perform_pool_score_rebuild_once_with_config(&state, config).await {
warn!(
error = ?err,
"gateway pool score rebuild initial tick failed"
);
continue;
}
match perform_pool_score_rebuild_once_with_config(&state, config).await {
Ok(summary) if summary.scores_upserted > 0 => {
info!(
providers_checked = summary.providers_checked,
providers_scored = summary.providers_scored,
keys_seen = summary.keys_seen,
scores_upserted = summary.scores_upserted,
"gateway pool score rebuild completed"
let mut interval = tokio::time::interval(config.interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
let mut deferred_since = None;
loop {
interval.tick().await;
if state
.data
.should_defer_maintenance_for_database_pool_pressure(&mut deferred_since)
{
debug!(
event_name = "maintenance_worker_deferred",
log_type = "ops",
worker = "pool_score_rebuild",
"gateway pool score rebuild deferred because database pool has no idle reserve"
);
continue;
}
Ok(_) => {}
Err(err) => {
warn!(
error = ?err,
"gateway pool score rebuild worker tick failed"
);
match perform_pool_score_rebuild_once_with_config(&state, config).await {
Ok(summary) if summary.scores_upserted > 0 => {
info!(
providers_checked = summary.providers_checked,
providers_scored = summary.providers_scored,
keys_seen = summary.keys_seen,
scores_upserted = summary.scores_upserted,
"gateway pool score rebuild completed"
);
}
Ok(_) => {}
Err(err) => {
warn!(
error = ?err,
"gateway pool score rebuild worker tick failed"
);
}
}
}
}
}))
},
))
}
#[cfg(test)]
@@ -116,7 +116,7 @@ pub(super) async fn run_proxy_upgrade_rollout_once(state: &AppState) -> Result<(
for probe in probes {
match state
.tunnel
.probe_node_url(&probe.node_id, &probe.url, probe.timeout_secs)
.probe_node_url_routed(state, &probe.node_id, &probe.url, probe.timeout_secs)
.await
{
Ok(status) if (200..300).contains(&status) => {
@@ -39,24 +39,34 @@ use super::{
#[tokio::test]
async fn spawn_audit_cleanup_worker_skips_when_postgres_unavailable() {
assert!(spawn_audit_cleanup_worker(Arc::new(GatewayDataState::disabled())).is_none());
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_audit_cleanup_worker(state).is_none());
}
#[tokio::test]
async fn spawn_db_maintenance_worker_skips_when_database_maintenance_unavailable() {
assert!(spawn_db_maintenance_worker(Arc::new(GatewayDataState::disabled())).is_none());
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_db_maintenance_worker(state).is_none());
}
#[tokio::test]
async fn spawn_pending_cleanup_worker_skips_when_usage_writer_unavailable() {
assert!(spawn_pending_cleanup_worker(Arc::new(GatewayDataState::disabled())).is_none());
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_pending_cleanup_worker(state).is_none());
}
#[tokio::test]
async fn spawn_proxy_node_stale_cleanup_worker_skips_when_proxy_nodes_unavailable() {
assert!(
spawn_proxy_node_stale_cleanup_worker(Arc::new(GatewayDataState::disabled())).is_none()
);
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_proxy_node_stale_cleanup_worker(state).is_none());
}
#[tokio::test]
@@ -95,7 +105,10 @@ async fn spawn_proxy_upgrade_rollout_worker_skips_when_system_config_unavailable
#[tokio::test]
async fn spawn_pool_monitor_worker_skips_when_postgres_unavailable() {
assert!(spawn_pool_monitor_worker(Arc::new(GatewayDataState::disabled())).is_none());
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_pool_monitor_worker(state).is_none());
}
#[tokio::test]
@@ -559,28 +572,35 @@ async fn proxy_upgrade_rollout_active_probe_advances_next_wave_after_version_con
#[tokio::test]
async fn spawn_stats_aggregation_worker_skips_when_stats_daily_backend_unavailable() {
assert!(spawn_stats_aggregation_worker(Arc::new(GatewayDataState::disabled())).is_none());
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_stats_aggregation_worker(state).is_none());
}
#[tokio::test]
async fn spawn_stats_hourly_aggregation_worker_skips_when_stats_hourly_backend_unavailable() {
assert!(
spawn_stats_hourly_aggregation_worker(Arc::new(GatewayDataState::disabled())).is_none()
);
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_stats_hourly_aggregation_worker(state).is_none());
}
#[tokio::test]
async fn spawn_usage_cleanup_worker_skips_when_usage_writer_unavailable() {
assert!(spawn_usage_cleanup_worker(Arc::new(GatewayDataState::disabled())).is_none());
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_usage_cleanup_worker(state).is_none());
}
#[tokio::test]
async fn spawn_wallet_daily_usage_aggregation_worker_skips_when_wallet_daily_usage_backend_unavailable(
) {
assert!(
spawn_wallet_daily_usage_aggregation_worker(Arc::new(GatewayDataState::disabled()))
.is_none()
);
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_wallet_daily_usage_aggregation_worker(state).is_none());
}
#[tokio::test]
File diff suppressed because it is too large Load Diff
@@ -1,670 +1,6 @@
use std::time::Instant;
//! Compatibility facade for frontdoor access logging.
use axum::body::Body;
use axum::extract::Request;
use axum::http::header::{HeaderName, HeaderValue};
use axum::http::Method;
use axum::middleware::Next;
use axum::response::Response;
use tracing::{info, trace, warn};
use crate::ai_serving::api::sanitize_request_path_and_query;
use crate::constants::{
CONTROL_REQUEST_ID_HEADER, CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
pub(crate) use aether_gateway_frontdoor::{
access_log_middleware, sanitize_access_log_path, should_downgrade_access_log,
GatewayRequestAcceptedAt, RequestLogEmitted,
};
use crate::headers::extract_or_generate_trace_id;
use crate::log_ids::short_request_id;
#[derive(Debug, Clone, Copy)]
pub(crate) struct RequestLogEmitted;
#[derive(Debug, Clone, Copy)]
pub(crate) struct GatewayRequestAcceptedAt(pub(crate) Instant);
fn is_usage_detail_path(path: &str) -> bool {
let Some(detail_id) = path.strip_prefix("/api/admin/usage/") else {
return false;
};
!detail_id.is_empty()
&& !detail_id.contains('/')
&& !matches!(detail_id, "active" | "records" | "stats" | "heatmap")
}
pub(crate) fn should_downgrade_access_log(method: &Method, path: &str) -> bool {
if method != Method::GET {
return false;
}
let normalized_path = path.split('?').next().unwrap_or(path);
matches!(
normalized_path,
"/api/admin/usage/active"
| "/api/users/me/usage/active"
| "/api/admin/usage/records"
| "/api/admin/usage/stats"
| "/api/admin/usage/aggregation/stats"
| "/api/admin/usage/heatmap"
| "/api/admin/usage/cache-affinity/interval-timeline"
| "/api/admin/usage/cache-affinity/ttl-analysis"
| "/api/admin/usage/cache-affinity/hit-analysis"
| "/api/admin/users"
| "/api/admin/monitoring/cache/stats"
| "/api/admin/monitoring/cache/model-mapping/stats"
| "/api/admin/monitoring/cache/config"
| "/api/admin/monitoring/cache/redis-keys"
| "/api/admin/monitoring/cache/affinities"
) || is_usage_detail_path(normalized_path)
|| normalized_path.starts_with("/api/admin/monitoring/trace/")
}
pub(crate) fn sanitize_access_log_path(path: &str) -> String {
sanitize_request_path_and_query(path, None).unwrap_or_else(|| "/".to_string())
}
pub(crate) async fn access_log_middleware(mut request: Request<Body>, next: Next) -> Response {
let started_at = Instant::now();
request
.extensions_mut()
.insert(GatewayRequestAcceptedAt(started_at));
let method = request.method().clone();
let raw_path = request
.uri()
.path_and_query()
.map(|value| value.as_str().to_string())
.unwrap_or_else(|| "/".to_string());
let path = sanitize_access_log_path(&raw_path);
let trace_id = extract_or_generate_trace_id(request.headers());
if !request.headers().contains_key(TRACE_ID_HEADER) {
request.headers_mut().insert(
HeaderName::from_static(TRACE_ID_HEADER),
HeaderValue::from_str(&trace_id).expect("trace id should be a valid header value"),
);
}
trace!(
event_name = "http_request_started",
log_type = "access",
status = "started",
trace_id = %trace_id,
request_id = "-",
method = %method,
path = %path,
route_class = "pending",
execution_path = "pending",
"gateway request started"
);
let mut response = next.run(request).await;
if !response.headers().contains_key(TRACE_ID_HEADER) {
response.headers_mut().insert(
HeaderName::from_static(TRACE_ID_HEADER),
HeaderValue::from_str(&trace_id).expect("trace id should be a valid header value"),
);
}
if response.extensions().get::<RequestLogEmitted>().is_none() {
let route_class = response
.headers()
.get(CONTROL_ROUTE_CLASS_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or("local");
let execution_path = response
.headers()
.get(EXECUTION_PATH_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or("local_route");
let request_id = response
.headers()
.get(CONTROL_REQUEST_ID_HEADER)
.and_then(|value| value.to_str().ok())
.filter(|value| !value.trim().is_empty())
.unwrap_or("-");
let request_id = short_request_id(request_id);
let status_code = response.status().as_u16();
let elapsed_ms = started_at.elapsed().as_millis() as u64;
if response.status().is_server_error() {
warn!(
event_name = "http_request_failed",
log_type = "access",
status = "failed",
status_code,
trace_id = %trace_id,
request_id,
method = %method,
path = %path,
route_class,
execution_path,
elapsed_ms,
"gateway request failed"
);
} else if should_downgrade_access_log(&method, &path) {
trace!(
event_name = "http_request_completed",
log_type = "access",
status = "completed",
status_code,
trace_id = %trace_id,
request_id,
method = %method,
path = %path,
route_class,
execution_path,
elapsed_ms,
"gateway completed request"
);
} else {
info!(
event_name = "http_request_completed",
log_type = "access",
status = "completed",
status_code,
trace_id = %trace_id,
request_id,
method = %method,
path = %path,
route_class,
execution_path,
elapsed_ms,
"gateway completed request"
);
}
}
response
}
#[cfg(test)]
mod tests {
use super::{access_log_middleware, sanitize_access_log_path, should_downgrade_access_log};
use crate::constants::{
CONTROL_REQUEST_ID_HEADER, CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER,
TRACE_ID_HEADER,
};
use axum::body::Body;
use axum::http::{Method, Request, Response, StatusCode};
use axum::routing::get;
use axum::Router;
use bytes::Bytes;
use futures_util::stream;
use std::sync::{Arc, Mutex};
use tower::ServiceExt;
use tracing_subscriber::filter::LevelFilter;
use tracing_subscriber::prelude::*;
#[derive(Clone, Default)]
struct SharedBuffer(Arc<Mutex<Vec<u8>>>);
struct SharedBufferWriter(Arc<Mutex<Vec<u8>>>);
impl SharedBuffer {
fn lines(&self) -> Vec<serde_json::Value> {
String::from_utf8(self.0.lock().expect("buffer should lock").clone())
.expect("buffer should contain valid utf-8")
.lines()
.filter(|line| !line.trim().is_empty())
.map(|line| serde_json::from_str(line).expect("json log line should parse"))
.collect()
}
}
impl std::io::Write for SharedBufferWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0
.lock()
.expect("buffer should lock")
.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
impl<'a> tracing_subscriber::fmt::writer::MakeWriter<'a> for SharedBuffer {
type Writer = SharedBufferWriter;
fn make_writer(&'a self) -> Self::Writer {
SharedBufferWriter(Arc::clone(&self.0))
}
}
#[test]
fn access_log_path_redacts_credential_query_values() {
assert_eq!(
sanitize_access_log_path(
"/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse&pageSize=10&token=hidden"
),
"/v1beta/models/gemini-3-flash-preview:generateContent?alt=sse&pageSize=10"
);
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_emits_sanitized_path() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::INFO),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let app = Router::new()
.route(
"/v1beta/models/gemini-3-flash-preview:generateContent",
get(|| async { Response::new(Body::empty()) }),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let _response = app
.oneshot(
Request::builder()
.uri("/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
let logs = writer.lines();
assert_eq!(logs.len(), 1);
assert_eq!(
logs[0]["path"],
"/v1beta/models/gemini-3-flash-preview:generateContent?alt=sse"
);
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_emits_completed_events_by_default() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::INFO),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let app = Router::new()
.route(
"/ok",
get(|| async {
let mut response = Response::new(Body::empty());
response.headers_mut().insert(
CONTROL_ROUTE_CLASS_HEADER,
"local".parse().expect("header should parse"),
);
response.headers_mut().insert(
EXECUTION_PATH_HEADER,
"local_route".parse().expect("header should parse"),
);
response.headers_mut().insert(
CONTROL_REQUEST_ID_HEADER,
"req-123".parse().expect("header should parse"),
);
response
}),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let response = app
.oneshot(
Request::builder()
.uri("/ok")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
assert!(response.headers().contains_key(TRACE_ID_HEADER));
let logs = writer.lines();
assert_eq!(logs.len(), 1);
assert_eq!(logs[0]["event_name"], "http_request_completed");
assert_eq!(logs[0]["status"], "completed");
assert_eq!(logs[0]["status_code"], 200);
assert_eq!(logs[0]["request_id"], "req-123");
assert_eq!(logs[0]["route_class"], "local");
assert_eq!(logs[0]["execution_path"], "local_route");
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_propagates_generated_trace_id_to_downstream_handler() {
let app = Router::new()
.route(
"/trace",
get(|headers: http::HeaderMap| async move {
Response::builder()
.status(StatusCode::OK)
.header(
"x-seen-trace-id",
headers
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.unwrap_or("-"),
)
.body(Body::empty())
.expect("response should build")
}),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let response = app
.oneshot(
Request::builder()
.uri("/trace")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
let response_trace_id = response
.headers()
.get(TRACE_ID_HEADER)
.and_then(|value| value.to_str().ok())
.expect("response trace id should exist")
.to_string();
let seen_trace_id = response
.headers()
.get("x-seen-trace-id")
.and_then(|value| value.to_str().ok())
.expect("downstream seen trace id should exist")
.to_string();
assert_eq!(seen_trace_id, response_trace_id);
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_shortens_long_request_ids() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::INFO),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let app = Router::new()
.route(
"/ok",
get(|| async {
let mut response = Response::new(Body::empty());
response.headers_mut().insert(
CONTROL_ROUTE_CLASS_HEADER,
"local".parse().expect("header should parse"),
);
response.headers_mut().insert(
EXECUTION_PATH_HEADER,
"local_route".parse().expect("header should parse"),
);
response.headers_mut().insert(
CONTROL_REQUEST_ID_HEADER,
"d07e1e94-41b8-409f-a18a-27993ae7ecb1"
.parse()
.expect("header should parse"),
);
response
}),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let _response = app
.oneshot(
Request::builder()
.uri("/ok")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
let logs = writer.lines();
assert_eq!(logs[0]["request_id"], "d07e1e94");
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_emits_failed_events_by_default_for_server_errors() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::INFO),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let app = Router::new()
.route(
"/fail",
get(|| async {
Response::builder()
.status(StatusCode::BAD_GATEWAY)
.header(CONTROL_ROUTE_CLASS_HEADER, "passthrough")
.header(EXECUTION_PATH_HEADER, "execution_runtime_sync")
.body(Body::empty())
.expect("response should build")
}),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let _response = app
.oneshot(
Request::builder()
.uri("/fail")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
let logs = writer.lines();
assert_eq!(logs.len(), 1);
assert_eq!(logs[0]["event_name"], "http_request_failed");
assert_eq!(logs[0]["status"], "failed");
assert_eq!(logs[0]["status_code"], 502);
assert_eq!(logs[0]["route_class"], "passthrough");
assert_eq!(logs[0]["execution_path"], "execution_runtime_sync");
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_treats_client_errors_as_completed_events() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::INFO),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let app = Router::new()
.route(
"/missing",
get(|| async {
Response::builder()
.status(StatusCode::UNAUTHORIZED)
.header(CONTROL_ROUTE_CLASS_HEADER, "auth")
.header(EXECUTION_PATH_HEADER, "local_auth_denied")
.body(Body::empty())
.expect("response should build")
}),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let _response = app
.oneshot(
Request::builder()
.uri("/missing")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
let logs = writer.lines();
assert_eq!(logs.len(), 1);
assert_eq!(logs[0]["event_name"], "http_request_completed");
assert_eq!(logs[0]["status"], "completed");
assert_eq!(logs[0]["status_code"], 401);
assert_eq!(logs[0]["route_class"], "auth");
assert_eq!(logs[0]["execution_path"], "local_auth_denied");
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_emits_completed_events_for_streaming_responses() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::INFO),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let app = Router::new()
.route(
"/stream",
get(|| async {
let body = Body::from_stream(stream::iter(vec![
Ok::<Bytes, std::convert::Infallible>(Bytes::from("chunk-1")),
Ok::<Bytes, std::convert::Infallible>(Bytes::from("chunk-2")),
]));
Response::builder()
.status(StatusCode::OK)
.header(CONTROL_ROUTE_CLASS_HEADER, "ai_public")
.header(EXECUTION_PATH_HEADER, "execution_runtime_stream")
.header(CONTROL_REQUEST_ID_HEADER, "req-stream")
.body(body)
.expect("response should build")
}),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let response = app
.oneshot(
Request::builder()
.uri("/stream")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let logs = writer.lines();
assert_eq!(logs.len(), 1);
assert_eq!(logs[0]["event_name"], "http_request_completed");
assert_eq!(logs[0]["status_code"], 200);
assert_eq!(logs[0]["request_id"], "req-stream");
assert_eq!(logs[0]["route_class"], "ai_public");
assert_eq!(logs[0]["execution_path"], "execution_runtime_stream");
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_downgrades_usage_active_polling_to_trace() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::TRACE),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let app = Router::new()
.route(
"/api/admin/usage/active",
get(|| async {
Response::builder()
.status(StatusCode::OK)
.header(CONTROL_ROUTE_CLASS_HEADER, "admin_proxy")
.header(EXECUTION_PATH_HEADER, "public_proxy_passthrough")
.body(Body::empty())
.expect("response should build")
}),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let _response = app
.oneshot(
Request::builder()
.uri("/api/admin/usage/active?ids=req-1")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
let logs = writer.lines();
assert_eq!(logs.len(), 2);
assert_eq!(logs[0]["level"], "TRACE");
assert_eq!(logs[0]["event_name"], "http_request_started");
assert_eq!(logs[1]["level"], "TRACE");
assert_eq!(logs[1]["event_name"], "http_request_completed");
}
#[test]
fn access_log_marks_usage_active_paths_as_high_frequency() {
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/usage/active"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/usage/active?ids=req-1"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/users/me/usage/active"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/usage/records?limit=20"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/usage/123e4567-e89b-12d3-a456-426614174000?include_bodies=false"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/monitoring/trace/req-123?attempted_only=false"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/monitoring/cache/stats"
));
assert!(!should_downgrade_access_log(
&Method::DELETE,
"/api/admin/monitoring/cache/affinity/provider/key/model/openai:responses"
));
assert!(!should_downgrade_access_log(&Method::GET, "/v1/responses"));
}
}
+2 -3
View File
@@ -1,11 +1,10 @@
mod access_log;
mod frontdoor_cors;
mod strip_cf_headers;
pub(crate) use access_log::{
access_log_middleware, sanitize_access_log_path, should_downgrade_access_log,
GatewayRequestAcceptedAt, RequestLogEmitted,
};
pub use aether_gateway_frontdoor::strip_cf_headers_middleware;
pub(crate) use aether_gateway_frontdoor::{apply_cf_header_stripping, CfConnectingIp};
pub(crate) use frontdoor_cors::frontdoor_cors_middleware;
pub use strip_cf_headers::strip_cf_headers_middleware;
pub(crate) use strip_cf_headers::{apply_cf_header_stripping, CfConnectingIp};
@@ -1,164 +0,0 @@
use axum::{extract::Request, middleware::Next, response::Response, Router};
use http::{header::HeaderName, HeaderMap};
/// Cloudflare-specific headers that are not part of the `cf-*` prefix family.
const CF_EXACT_HEADERS: &[&str] = &["cdn-loop", "true-client-ip"];
#[derive(Clone, Debug)]
pub(crate) struct CfConnectingIp(pub(crate) String);
fn should_strip_cf_header(name: &HeaderName) -> bool {
let normalized = name.as_str();
normalized.starts_with("cf-") || CF_EXACT_HEADERS.contains(&normalized)
}
fn cf_connecting_ip(headers: &HeaderMap) -> Option<String> {
headers
.get("cf-connecting-ip")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|value| value.chars().take(45).collect())
}
fn strip_cf_headers(headers: &mut HeaderMap) {
let to_remove: Vec<_> = headers
.keys()
.filter(|name| should_strip_cf_header(name))
.cloned()
.collect();
for name in to_remove {
headers.remove(name);
}
}
pub(crate) fn apply_cf_header_stripping(router: Router) -> Router {
router.layer(axum::middleware::from_fn(strip_cf_headers_middleware))
}
pub async fn strip_cf_headers_middleware(mut request: Request, next: Next) -> Response {
if let Some(client_ip) = cf_connecting_ip(request.headers()) {
request.extensions_mut().insert(CfConnectingIp(client_ip));
}
strip_cf_headers(request.headers_mut());
let mut response = next.run(request).await;
strip_cf_headers(response.headers_mut());
response
}
#[cfg(test)]
mod tests {
use axum::body::{to_bytes, Body};
use axum::routing::any;
use axum::Router;
use http::{HeaderValue, Request, Response};
use tower::ServiceExt;
use super::apply_cf_header_stripping;
#[tokio::test]
async fn strips_cf_prefixed_and_exact_headers_from_request_and_response() {
let app = apply_cf_header_stripping(Router::new().route(
"/",
any(|headers: http::HeaderMap| async move {
let leaked = headers.contains_key("cf-ipcity")
|| headers.contains_key("cf-ray")
|| headers.contains_key("cf-connecting-ip")
|| headers.contains_key("true-client-ip")
|| headers.contains_key("cdn-loop");
let mut response =
Response::new(Body::from(if leaked { "leaked" } else { "clean" }));
response.headers_mut().insert(
http::header::HeaderName::from_static("cf-ipcity"),
HeaderValue::from_static("Shanghai"),
);
response.headers_mut().insert(
http::header::HeaderName::from_static("cf-cache-status"),
HeaderValue::from_static("HIT"),
);
response.headers_mut().insert(
http::header::HeaderName::from_static("true-client-ip"),
HeaderValue::from_static("1.1.1.1"),
);
response.headers_mut().insert(
http::header::HeaderName::from_static("cdn-loop"),
HeaderValue::from_static("cloudflare"),
);
response
}),
));
let response = app
.oneshot(
Request::builder()
.uri("/")
.header("cf-ipcity", "Shanghai")
.header("cf-ray", "abc123")
.header("cf-connecting-ip", "203.0.113.10")
.header("true-client-ip", "1.1.1.1")
.header("cdn-loop", "cloudflare")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
assert!(response.headers().get("cf-ipcity").is_none());
assert!(response.headers().get("cf-cache-status").is_none());
assert!(response.headers().get("cf-connecting-ip").is_none());
assert!(response.headers().get("true-client-ip").is_none());
assert!(response.headers().get("cdn-loop").is_none());
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should be readable");
assert_eq!(body.as_ref(), b"clean");
}
#[tokio::test]
async fn preserves_non_cf_headers() {
let app = apply_cf_header_stripping(Router::new().route(
"/",
any(|headers: http::HeaderMap| async move {
let mut response = Response::new(Body::from(
headers
.get("x-custom-header")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
));
response.headers_mut().insert(
http::header::HeaderName::from_static("x-custom-response"),
HeaderValue::from_static("kept"),
);
response
}),
));
let response = app
.oneshot(
Request::builder()
.uri("/")
.header("x-custom-header", "kept")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
assert_eq!(
response
.headers()
.get("x-custom-response")
.and_then(|value| value.to_str().ok()),
Some("kept")
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should be readable");
assert_eq!(body.as_ref(), b"kept");
}
}
+25 -21
View File
@@ -34,31 +34,35 @@ pub(crate) fn spawn_model_fetch_worker(state: AppState) -> Option<tokio::task::J
return None;
}
Some(tokio::spawn(async move {
if model_fetch_startup_enabled() {
let startup_delay = model_fetch_startup_delay_seconds();
if startup_delay > 0 {
tokio::time::sleep(Duration::from_secs(startup_delay)).await;
Some(crate::task_runtime::spawn_singleton_worker(
state,
crate::task_runtime::TASK_KEY_MODEL_FETCH_WORKER,
|state| async move {
if model_fetch_startup_enabled() {
let startup_delay = model_fetch_startup_delay_seconds();
if startup_delay > 0 {
tokio::time::sleep(Duration::from_secs(startup_delay)).await;
}
if let Err(err) = run_model_fetch_cycle(&state, "startup").await {
warn!(error = ?err, "gateway model fetch startup failed");
}
} else {
info!("gateway model fetch startup disabled");
}
if let Err(err) = run_model_fetch_cycle(&state, "startup").await {
warn!(error = ?err, "gateway model fetch startup failed");
}
} else {
info!("gateway model fetch startup disabled");
}
let mut interval = tokio::time::interval(Duration::from_secs(
model_fetch_interval_minutes().saturating_mul(60),
));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
loop {
let mut interval = tokio::time::interval(Duration::from_secs(
model_fetch_interval_minutes().saturating_mul(60),
));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
if let Err(err) = run_model_fetch_cycle(&state, "tick").await {
warn!(error = ?err, "gateway model fetch tick failed");
loop {
interval.tick().await;
if let Err(err) = run_model_fetch_cycle(&state, "tick").await {
warn!(error = ?err, "gateway model fetch tick failed");
}
}
}
}))
},
))
}
pub(crate) async fn perform_model_fetch_once(
@@ -38,6 +38,7 @@ pub(crate) const SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD: &str = "scheduler_affini
pub(crate) const POOL_KEY_LEASE_KEY_REPORT_FIELD: &str = "pool_key_lease_key";
pub(crate) const POOL_KEY_LEASE_OWNER_REPORT_FIELD: &str = "pool_key_lease_owner";
pub(crate) const POOL_KEY_LEASE_TOKEN_REPORT_FIELD: &str = "pool_key_lease_token";
pub(crate) const POOL_KEY_LEASE_FENCING_REPORT_FIELD: &str = "pool_key_lease_fencing_token";
pub(crate) const POOL_KEY_LEASE_TTL_MS_REPORT_FIELD: &str = "pool_key_lease_ttl_ms";
pub(crate) fn attempt_identity_from_report_context(
@@ -92,6 +93,10 @@ pub(crate) fn insert_pool_key_lease_report_context_fields(
POOL_KEY_LEASE_TOKEN_REPORT_FIELD.to_string(),
Value::String(lease.token.clone()),
);
extra_fields.insert(
POOL_KEY_LEASE_FENCING_REPORT_FIELD.to_string(),
Value::Number(lease.fencing_token.into()),
);
extra_fields.insert(
POOL_KEY_LEASE_TTL_MS_REPORT_FIELD.to_string(),
Value::Number(lease.ttl_ms.into()),
@@ -119,11 +124,17 @@ fn pool_key_lease_from_report_context(report_context: Option<&Value>) -> Option<
.get(POOL_KEY_LEASE_TTL_MS_REPORT_FIELD)
.and_then(Value::as_u64)
.filter(|value| *value > 0)?;
let fencing_token = report_context
.get(POOL_KEY_LEASE_FENCING_REPORT_FIELD)
.and_then(Value::as_u64)
.filter(|value| *value > 0)
.unwrap_or(1);
Some(RuntimeLockLease {
key: key.to_string(),
owner: owner.to_string(),
token: token.to_string(),
fencing_token,
ttl_ms,
})
}
@@ -485,6 +496,7 @@ mod tests {
"pool_key_lease_key": "ap:provider-1:lease:key-1",
"pool_key_lease_owner": "gateway-1",
"pool_key_lease_token": "gateway-1:token-1",
"pool_key_lease_fencing_token": 7,
"pool_key_lease_ttl_ms": 900000,
})));
@@ -497,6 +509,7 @@ mod tests {
key: "ap:provider-1:lease:key-1".to_string(),
owner: "gateway-1".to_string(),
token: "gateway-1:token-1".to_string(),
fencing_token: 7,
ttl_ms: 900000,
}),
scheduler_affinity_epoch: None,
@@ -1283,7 +1283,7 @@ mod tests {
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_testkit::ManagedRedisServer;
use aether_test_support::ManagedRedisServer;
use serde_json::{json, Value};
use super::{
+1 -1
View File
@@ -4169,7 +4169,7 @@ mod tests {
use std::time::Duration;
use aether_runtime_state::{RedisClientConfig, RuntimeState};
use aether_testkit::ManagedRedisServer;
use aether_test_support::ManagedRedisServer;
use serde_json::{json, Value};
fn assert_debug_surface_hides_values(debug: &str, originals: &[&str], sentinels: &[String]) {
+14 -2
View File
@@ -2,15 +2,16 @@ use std::collections::HashMap;
use std::sync::atomic::AtomicU64;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::time::Duration;
use std::time::{Duration, Instant};
use aether_data::repository::users::StoredUserGroup;
use aether_data_contracts::repository::billing::UserDailyQuotaAvailabilityRecord;
use aether_data_contracts::repository::quota::StoredProviderQuotaSnapshot;
use aether_data_contracts::repository::usage::UsageCounterHealthSnapshot;
use aether_runtime::ConcurrencyGate;
use aether_runtime_state::{RuntimeSemaphore, RuntimeState};
use dashmap::DashMap;
use tokio::sync::{Mutex as TokioMutex, Semaphore};
use tokio::sync::{Mutex as TokioMutex, RwLock as TokioRwLock, Semaphore};
use super::super::async_task::{VideoTaskPollerConfig, VideoTaskService};
use super::super::cache::{
@@ -106,6 +107,8 @@ pub(crate) struct FrontdoorRuntimeGuardConfig {
pub(crate) upstream_target_gate_limit: Option<usize>,
}
pub(crate) const METRIC_SNAPSHOT_TTL: Duration = Duration::from_secs(2);
impl FrontdoorRuntimeGuardConfig {
pub(crate) fn from_env() -> Self {
Self {
@@ -364,6 +367,8 @@ pub struct AppState {
#[cfg(test)]
pub(crate) execution_runtime_sync_override: Option<TestExecutionRuntimeSyncOverride>,
pub(crate) data: Arc<GatewayDataState>,
pub(crate) background_data: Arc<GatewayDataState>,
pub(crate) background_data_isolated: bool,
pub(crate) runtime_state: Arc<RuntimeState>,
pub(crate) usage_runtime: Arc<usage::UsageRuntime>,
pub(crate) video_tasks: Arc<VideoTaskService>,
@@ -412,6 +417,13 @@ pub struct AppState {
pub(crate) usage_counter_flush_metrics: Arc<UsageCounterFlushRuntimeMetrics>,
pub(crate) task_supervisor_metrics: TaskSupervisorMetrics,
pub(crate) process_resource_monitor: Arc<crate::process_metrics::GatewayProcessResourceMonitor>,
pub(crate) metric_snapshot:
Arc<TokioRwLock<Option<(Instant, Vec<aether_runtime::MetricSample>)>>>,
pub(crate) metric_snapshot_refresh: Arc<TokioMutex<()>>,
pub(crate) usage_counter_exact_health_metric_snapshot:
Arc<TokioRwLock<Option<(Instant, UsageCounterHealthSnapshot)>>>,
pub(crate) usage_counter_exact_health_metric_last_attempt: Arc<StdMutex<Option<Instant>>>,
pub(crate) usage_counter_exact_health_metric_refresh: Arc<TokioMutex<()>>,
pub(crate) request_candidate_queue: Option<Arc<RequestCandidateQueueRuntime>>,
pub(crate) frontdoor_cors: Option<Arc<FrontdoorCorsConfig>>,
pub(crate) frontdoor_user_rpm: Arc<FrontdoorUserRpmLimiter>,
File diff suppressed because it is too large Load Diff
@@ -13,7 +13,7 @@ const AUTH_API_KEY_SNAPSHOT_RUNTIME_CACHE_TTL: Duration = Duration::from_secs(30
use super::super::super::{AUTH_API_KEY_LAST_USED_MAX_ENTRIES, AUTH_API_KEY_LAST_USED_TTL};
impl AppState {
async fn acquire_auth_snapshot_load_gate(
pub(crate) async fn acquire_auth_snapshot_load_gate(
&self,
) -> Result<Option<aether_runtime::ConcurrencyPermit>, GatewayError> {
let Some(gate) = self.auth_snapshot_load_gate.as_ref() else {
@@ -525,12 +525,27 @@ impl AppState {
}
let ttl = self.frontdoor_runtime_guards.auth_capacity_cache_ttl;
if ttl.is_zero() {
return self.find_user_daily_quota_availability(user_id).await;
return self
.find_user_daily_quota_availability_for_auth_uncached(user_id)
.await;
}
self.auth_daily_quota_availability_cache
.get_or_load(user_id.to_string(), ttl, || async move {
let _permit = self.acquire_auth_snapshot_load_gate().await?;
self.find_user_daily_quota_availability(user_id).await
})
.await
}
pub(crate) async fn find_user_daily_quota_availability_for_auth_uncached(
&self,
user_id: &str,
) -> Result<Option<UserDailyQuotaAvailabilityRecord>, GatewayError> {
let user_id = user_id.trim();
if user_id.is_empty() {
return Ok(None);
}
let _permit = self.acquire_auth_snapshot_load_gate().await?;
self.find_user_daily_quota_availability(user_id).await
}
}
@@ -181,6 +181,16 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn summarize_dashboard_stats(
&self,
query: &usage::UsageDashboardSummaryQuery,
) -> Result<usage::StoredUsageDashboardStatsSummary, GatewayError> {
self.data
.summarize_dashboard_stats(query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_dashboard_daily_breakdown(
&self,
query: &usage::UsageDashboardDailyBreakdownQuery,
@@ -73,15 +73,16 @@ impl AppState {
let ttl = self.frontdoor_runtime_guards.auth_capacity_cache_ttl;
if ttl.is_zero() {
return self.find_wallet(lookup).await;
return self
.read_wallet_snapshot_for_auth_uncached(user_id, api_key_id, api_key_is_standalone)
.await;
}
self.auth_wallet_snapshot_cache
.get_or_load(
cache_key,
ttl,
|| async move { self.find_wallet(lookup).await },
)
.get_or_load(cache_key, ttl, || async move {
let _permit = self.acquire_auth_snapshot_load_gate().await?;
self.find_wallet(lookup).await
})
.await
}
@@ -117,6 +118,7 @@ impl AppState {
return Ok(None);
};
let _permit = self.acquire_auth_snapshot_load_gate().await?;
self.find_wallet(lookup).await
}
+1 -1
View File
@@ -609,7 +609,7 @@ mod tests {
settings
.api_format_suffixes("openai:chat")
.expect("chat suffixes")
.into_iter()
.iter()
.collect::<Vec<_>>(),
vec!["low", "max"]
);
+116 -1
View File
@@ -9,6 +9,7 @@ use aether_runtime::task::spawn_named;
use aether_task_runtime::{RetryPolicy, TaskDefinition, TaskKind};
pub(crate) use aether_task_runtime::{TaskSupervisor, TaskSupervisorMetrics};
use serde_json::Value;
use sha2::{Digest, Sha256};
use tokio::task::JoinHandle;
use tracing::warn;
use uuid::Uuid;
@@ -26,6 +27,7 @@ pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
pub(crate) const TASK_KEY_ACCOUNT_SELF_CHECK: &str = "account.self_check.worker";
pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker";
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database";
@@ -52,6 +54,58 @@ const PROVIDER_DELETE_LOCK_TTL_SECS: u64 = 60 * 60 * 6;
const RETRY_ONCE: RetryPolicy = RetryPolicy { max_attempts: 1 };
const RETRY_THREE: RetryPolicy = RetryPolicy { max_attempts: 3 };
const BACKGROUND_TASK_RUN_ID_MAX_BYTES: usize = 64;
const WORKER_BOOT_RUN_ID_HASH_HEX_BYTES: usize = 20;
fn build_worker_boot_run_id(task_key: &str, instance_id: &str) -> String {
let full_run_id = format!("boot:{task_key}:{instance_id}");
if full_run_id.len() <= BACKGROUND_TASK_RUN_ID_MAX_BYTES {
return full_run_id;
}
let digest_hex = format!("{:x}", Sha256::digest(full_run_id.as_bytes()));
let suffix = &digest_hex[..WORKER_BOOT_RUN_ID_HASH_HEX_BYTES];
let mut prefix_bytes = BACKGROUND_TASK_RUN_ID_MAX_BYTES - 1 - WORKER_BOOT_RUN_ID_HASH_HEX_BYTES;
while !full_run_id.is_char_boundary(prefix_bytes) {
prefix_bytes -= 1;
}
format!("{}~{suffix}", &full_run_id[..prefix_bytes])
}
pub(crate) fn spawn_singleton_worker<F, Fut>(
app: AppState,
task_key: &'static str,
worker: F,
) -> JoinHandle<()>
where
F: Fn(AppState) -> Fut + Send + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
spawn_singleton_worker_with_context(app, task_key, move |app, _context| worker(app))
}
pub(crate) fn spawn_singleton_worker_with_context<F, Fut>(
app: AppState,
task_key: &'static str,
worker: F,
) -> JoinHandle<()>
where
F: Fn(AppState, aether_gateway_workers::SingletonWorkerContext) -> Fut + Send + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let runtime_state = app.runtime_state.clone();
let metrics = app.task_supervisor_metrics.clone();
let owner = app.tunnel.local_instance_id().to_string();
aether_gateway_workers::spawn_singleton_worker_with_context(
runtime_state,
metrics,
owner,
task_key,
aether_gateway_workers::SingletonWorkerConfig::default(),
move |context| worker(app.clone(), context),
)
}
const TASK_DEFINITIONS: &[TaskDefinition] = &[
TaskDefinition::new(
@@ -142,6 +196,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_QUOTA_PROBE,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_MONITOR,
TaskKind::Scheduled,
@@ -448,7 +510,7 @@ pub(crate) fn spawn_record_worker_boot(
) -> JoinHandle<()> {
spawn_named("task-runtime-record-worker-boot", async move {
let now = now_unix_secs();
let run_id = format!("boot:{}:{}", task_key, app.tunnel.local_instance_id());
let run_id = build_worker_boot_run_id(task_key, app.tunnel.local_instance_id());
let run = UpsertBackgroundTaskRun {
id: run_id.clone(),
task_key: task_key.to_string(),
@@ -719,3 +781,56 @@ pub(crate) async fn submit_provider_delete_task(
Ok(Some(task_id))
}
#[cfg(test)]
mod worker_boot_run_id_tests {
use super::{build_worker_boot_run_id, BACKGROUND_TASK_RUN_ID_MAX_BYTES};
#[test]
fn worker_boot_run_id_preserves_short_legacy_format() {
assert_eq!(
build_worker_boot_run_id("usage.queue.worker", "gateway-1"),
"boot:usage.queue.worker:gateway-1"
);
}
#[test]
fn worker_boot_run_id_preserves_exact_database_limit() {
let task_key = "task";
let instance_id = "i".repeat(BACKGROUND_TASK_RUN_ID_MAX_BYTES - "boot:task:".len());
let run_id = build_worker_boot_run_id(task_key, &instance_id);
assert_eq!(run_id.len(), BACKGROUND_TASK_RUN_ID_MAX_BYTES);
assert_eq!(run_id, format!("boot:{task_key}:{instance_id}"));
}
#[test]
fn worker_boot_run_id_compacts_oversized_values_deterministically() {
let instance_id = "gateway-instance-".repeat(8);
let first = build_worker_boot_run_id("usage.queue.worker", &instance_id);
let second = build_worker_boot_run_id("usage.queue.worker", &instance_id);
assert!(first.len() <= BACKGROUND_TASK_RUN_ID_MAX_BYTES);
assert_eq!(first, second);
assert!(first.starts_with("boot:usage.queue.worker:"));
assert!(first.contains('~'));
}
#[test]
fn worker_boot_run_id_hash_distinguishes_shared_long_prefixes() {
let shared_prefix = "gateway-instance-with-a-very-long-shared-prefix-".repeat(2);
let first = build_worker_boot_run_id("usage.queue.worker", &format!("{shared_prefix}a"));
let second = build_worker_boot_run_id("usage.queue.worker", &format!("{shared_prefix}b"));
assert_ne!(first, second);
}
#[test]
fn worker_boot_run_id_handles_unicode_at_truncation_boundary() {
let run_id = build_worker_boot_run_id("usage.queue.worker", &"网关实例".repeat(16));
assert!(run_id.len() <= BACKGROUND_TASK_RUN_ID_MAX_BYTES);
assert!(run_id.is_char_boundary(run_id.len()));
assert!(run_id.starts_with("boot:usage.queue.worker:"));
}
}
@@ -661,7 +661,7 @@ fn admin_monitoring_snapshots_stay_app_local() {
"monitoring/resilience/snapshot.rs should define AdminMonitoringResilienceSnapshot locally"
);
let data_system = read_workspace_file("crates/aether-data/src/repository/system.rs");
let data_system = read_workspace_file("crates/aether-data/runtime/src/repository/system.rs");
assert!(
!data_system.contains("AdminMonitoringCacheSnapshot")
&& !data_system.contains("AdminMonitoringResilienceSnapshot"),
@@ -58,7 +58,8 @@ fn admin_provider_root_stays_thin() {
let ops_mod = read_workspace_file("apps/aether-gateway/src/handlers/admin/provider/ops/mod.rs");
assert!(
!ops_mod.contains("pub(crate) use self::providers::admin_provider_ops_local_action_response;"),
!ops_mod
.contains("pub(crate) use self::providers::admin_provider_ops_local_action_response;"),
"handlers/admin/provider/ops/mod.rs should not re-export admin_provider_ops_local_action_response"
);
@@ -120,7 +121,7 @@ fn admin_provider_oauth_complete_dispatch_remains_thin() {
#[test]
fn postgres_provider_cleanup_preserves_usage_history() {
let postgres_provider_catalog =
read_workspace_file("crates/aether-data/src/repository/provider_catalog/postgres.rs");
read_workspace_file("crates/aether-data/adapters/postgres/src/provider_catalog.rs");
for forbidden in [
"UPDATE usage SET provider_id = NULL",
@@ -137,9 +138,9 @@ fn postgres_provider_cleanup_preserves_usage_history() {
#[test]
fn provider_cleanup_keeps_common_backends_in_sync() {
for path in [
"crates/aether-data/src/repository/provider_catalog/postgres.rs",
"crates/aether-data/src/repository/provider_catalog/mysql.rs",
"crates/aether-data/src/repository/provider_catalog/sqlite.rs",
"crates/aether-data/adapters/postgres/src/provider_catalog.rs",
"crates/aether-data/adapters/mysql/src/provider_catalog.rs",
"crates/aether-data/adapters/sqlite/src/provider_catalog.rs",
] {
let source = read_workspace_file(path);
for required in [
@@ -370,7 +371,9 @@ fn admin_provider_ops_routes_directoryized() {
}
assert!(
routes_mod.contains("pub(crate) async fn maybe_build_local_admin_provider_ops_providers_response("),
routes_mod.contains(
"pub(crate) async fn maybe_build_local_admin_provider_ops_providers_response("
),
"handlers/admin/provider/ops/providers/routes/mod.rs should keep the provider ops entry seam"
);
for forbidden in [
@@ -386,7 +389,9 @@ fn admin_provider_ops_routes_directoryized() {
);
}
assert!(
!workspace_file_exists("apps/aether-gateway/src/handlers/admin/provider/ops/providers/routes.rs"),
!workspace_file_exists(
"apps/aether-gateway/src/handlers/admin/provider/ops/providers/routes.rs"
),
"handlers/admin/provider/ops/providers/routes.rs should be removed once routes are directoryized"
);
@@ -667,8 +672,10 @@ fn admin_provider_query_and_strategy_use_specific_local_owners() {
let strategy_routes =
read_workspace_file("apps/aether-gateway/src/handlers/admin/provider/strategy/routes.rs");
assert!(
strategy_routes.contains("state\n .maybe_build_admin_provider_strategy_route_response(")
|| strategy_routes.contains("state.maybe_build_admin_provider_strategy_route_response("),
strategy_routes
.contains("state\n .maybe_build_admin_provider_strategy_route_response(")
|| strategy_routes
.contains("state.maybe_build_admin_provider_strategy_route_response("),
"handlers/admin/provider/strategy/routes.rs should delegate to request/provider route owner"
);
assert!(
@@ -686,11 +693,15 @@ fn admin_provider_query_and_strategy_use_specific_local_owners() {
);
assert!(
workspace_file_exists("apps/aether-gateway/src/handlers/admin/provider/strategy/responses.rs"),
workspace_file_exists(
"apps/aether-gateway/src/handlers/admin/provider/strategy/responses.rs"
),
"handlers/admin/provider/strategy/responses.rs should own strategy route-level shared responses"
);
assert!(
!workspace_file_exists("apps/aether-gateway/src/handlers/admin/provider/strategy/shared.rs"),
!workspace_file_exists(
"apps/aether-gateway/src/handlers/admin/provider/strategy/shared.rs"
),
"handlers/admin/provider/strategy/shared.rs should be removed once the local shared hub is narrowed"
);
}
@@ -1180,7 +1191,7 @@ fn admin_provider_write_uses_specific_local_owners() {
] {
assert!(
endpoint_keys_mutations.contains(pattern),
"handlers/admin/provider/endpoint_keys/mutations/mod.rs should expose explicit mutation owner {pattern}"
"handlers/admin/provider/endpoint_keys/mutations/mod.rs should expose explicit mutation owner {pattern}"
);
}
@@ -1299,7 +1310,9 @@ fn admin_provider_ops_actions_mod_stays_thin() {
);
}
assert!(
!workspace_file_exists("apps/aether-gateway/src/handlers/admin/provider/ops/providers/actions.rs"),
!workspace_file_exists(
"apps/aether-gateway/src/handlers/admin/provider/ops/providers/actions.rs"
),
"handlers/admin/provider/ops/providers/actions.rs should be removed once actions logic is directoryized"
);
@@ -1811,7 +1824,8 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
"handlers/admin/provider/oauth/quota/shared.rs should delegate quota metadata provider detection to aether-provider-pool"
);
assert!(
!quota_shared.contains("[\"codex\", \"kiro\", \"antigravity\", \"gemini_cli\", \"chatgpt_web\"]"),
!quota_shared
.contains("[\"codex\", \"kiro\", \"antigravity\", \"gemini_cli\", \"chatgpt_web\"]"),
"handlers/admin/provider/oauth/quota/shared.rs should not hardcode quota metadata provider list"
);
@@ -1873,9 +1887,8 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
"apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/invalid.rs",
);
assert!(
quota_codex_invalid.contains(
"use crate::handlers::admin::provider::shared::payloads::{"
) || quota_codex_invalid.contains("use aether_admin::provider::quota"),
quota_codex_invalid.contains("use crate::handlers::admin::provider::shared::payloads::{")
|| quota_codex_invalid.contains("use aether_admin::provider::quota"),
"handlers/admin/provider/oauth/quota/codex/invalid.rs should either own codex invalid-state helpers locally or delegate to aether-admin"
);
let quota_codex_plan = read_workspace_file(
@@ -1949,7 +1962,8 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
"handlers/admin/provider/oauth/quota/antigravity.rs should import common quota helpers from shared.rs"
);
assert!(
quota_antigravity.contains("use aether_provider_pool::build_antigravity_pool_quota_request;"),
quota_antigravity
.contains("use aether_provider_pool::build_antigravity_pool_quota_request;"),
"handlers/admin/provider/oauth/quota/antigravity.rs should delegate antigravity quota request construction to aether-provider-pool"
);
let quota_chatgpt_web = read_workspace_file(
@@ -2130,12 +2144,14 @@ fn admin_provider_oauth_dispatch_batch_mod_stays_thin() {
"apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/kiro_import.rs",
);
assert!(
batch_kiro_import.contains("pub(super) async fn execute_admin_provider_oauth_kiro_batch_import("),
batch_kiro_import
.contains("pub(super) async fn execute_admin_provider_oauth_kiro_batch_import("),
"handlers/admin/provider/oauth/dispatch/batch/kiro_import.rs should own the kiro batch execution owner"
);
assert!(
batch_kiro_import.contains("build_kiro_batch_import_key_name(")
|| batch_kiro_import.contains("aether_admin::provider::oauth::build_kiro_batch_import_key_name"),
|| batch_kiro_import
.contains("aether_admin::provider::oauth::build_kiro_batch_import_key_name"),
"handlers/admin/provider/oauth/dispatch/batch/kiro_import.rs should either own or delegate the kiro key-name builder"
);
assert!(
@@ -2162,7 +2178,8 @@ fn admin_provider_oauth_dispatch_batch_mod_stays_thin() {
"apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/orchestration.rs",
);
assert!(
batch_orchestration.contains("pub(in super::super) async fn handle_admin_provider_oauth_batch_import("),
batch_orchestration
.contains("pub(in super::super) async fn handle_admin_provider_oauth_batch_import("),
"handlers/admin/provider/oauth/dispatch/batch/orchestration.rs should own the direct batch import route"
);
@@ -2333,7 +2350,9 @@ fn admin_provider_oauth_device_mod_stays_thin() {
}
assert!(
!workspace_file_exists("apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device.rs"),
!workspace_file_exists(
"apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device.rs"
),
"handlers/admin/provider/oauth/dispatch/device.rs should be removed once device dispatch is directoryized"
);
}
@@ -103,12 +103,15 @@ fn admin_wrapped_state_owns_api_key_and_proxy_capabilities() {
"{path} should use AdminAppState encryption capability instead of raw state.app() encryption"
);
assert!(
!contents.contains("decrypt_catalog_secret_with_fallbacks(state.app().encryption_key(),"),
!contents
.contains("decrypt_catalog_secret_with_fallbacks(state.app().encryption_key(),"),
"{path} should use AdminAppState decryption capability instead of raw state.app().encryption_key()"
);
assert!(
!contents.contains("resolve_transport_proxy_snapshot_with_tunnel_affinity(\n state.app(),")
&& !contents.contains("resolve_transport_proxy_snapshot_with_tunnel_affinity(state.app(),"),
!contents.contains(
"resolve_transport_proxy_snapshot_with_tunnel_affinity(\n state.app(),"
) && !contents
.contains("resolve_transport_proxy_snapshot_with_tunnel_affinity(state.app(),"),
"{path} should use AdminAppState proxy capability instead of raw state.app() transport proxy resolution"
);
}
@@ -682,7 +685,8 @@ fn crate_root_exposes_real_admin_and_ai_serving_facades() {
let ai_serving_api_mod = read_workspace_file("apps/aether-gateway/src/ai_serving/api.rs");
assert!(
ai_serving_api_mod.contains("use crate::ai_serving::{is_json_request, GatewayControlDecision};"),
ai_serving_api_mod
.contains("use crate::ai_serving::{is_json_request, GatewayControlDecision};"),
"ai_serving/api.rs should depend on the crate-facing ai_serving seam instead of deep internal modules"
);
}
@@ -695,7 +699,7 @@ fn gateway_ai_serving_api_module_delegates_pure_ownership_to_format_crate() {
"ai_serving/api.rs should re-export pure ownership through aether_ai_formats::api"
);
let format_crate_api = read_workspace_file("crates/aether-ai-formats/src/api.rs");
let format_crate_api = read_workspace_file("crates/aether-ai/formats/src/api.rs");
for pattern in [
"pub use crate::contracts::{",
"pub use crate::provider_compat::kiro_stream::{",
@@ -29,8 +29,8 @@ fn ai_serving_target_structure_removes_legacy_pipeline_boundary() {
let mut violations = Vec::new();
for root in [
"apps/aether-gateway/src",
"crates/aether-ai-serving/src",
"crates/aether-ai-formats/src",
"crates/aether-ai/serving/src",
"crates/aether-ai/formats/src",
] {
for file in collect_workspace_rust_files(root) {
let relative = file
@@ -75,7 +75,7 @@ fn ai_serving_target_structure_removes_legacy_pipeline_boundary() {
#[test]
fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
let serving_manifest = read_workspace_file("crates/aether-ai-serving/Cargo.toml");
let serving_manifest = read_workspace_file("crates/aether-ai/serving/Cargo.toml");
for forbidden in ["axum", "sqlx", "redis", "aether-gateway"] {
assert!(
!serving_manifest.contains(forbidden),
@@ -84,7 +84,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
}
let mut violations = Vec::new();
for file in collect_workspace_rust_files("crates/aether-ai-serving/src") {
for file in collect_workspace_rust_files("crates/aether-ai/serving/src") {
let source = std::fs::read_to_string(&file).expect("source file should be readable");
let hits = ["AppState", "axum::", "sqlx::", "redis::"]
.iter()
@@ -101,7 +101,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
violations.join("\n")
);
let serving_attempt_loop = read_workspace_file("crates/aether-ai-serving/src/attempt_loop.rs");
let serving_attempt_loop = read_workspace_file("crates/aether-ai/serving/src/attempt_loop.rs");
for pattern in [
"pub trait AiExecutionAttempt",
"pub trait AiAttemptLoopPort",
@@ -129,7 +129,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
}
let serving_decision_path =
read_workspace_file("crates/aether-ai-serving/src/decision_path.rs");
read_workspace_file("crates/aether-ai/serving/src/decision_path.rs");
for pattern in [
"pub enum AiSyncDecisionStep",
"pub enum AiStreamDecisionStep",
@@ -168,7 +168,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
);
}
let serving_plan_payload = read_workspace_file("crates/aether-ai-serving/src/plan_payload.rs");
let serving_plan_payload = read_workspace_file("crates/aether-ai/serving/src/plan_payload.rs");
for pattern in [
"pub fn build_ai_sync_execution_plan_payload",
"pub fn build_ai_stream_execution_plan_payload",
@@ -204,7 +204,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
);
}
let serving_attempt_plan = read_workspace_file("crates/aether-ai-serving/src/attempt_plan.rs");
let serving_attempt_plan = read_workspace_file("crates/aether-ai/serving/src/attempt_plan.rs");
for pattern in [
"pub fn build_ai_execution_decision_from_plan",
"pub fn infer_ai_upstream_base_url",
@@ -238,7 +238,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
}
let serving_candidate_materialization =
read_workspace_file("crates/aether-ai-serving/src/candidate_materialization.rs");
read_workspace_file("crates/aether-ai/serving/src/candidate_materialization.rs");
for pattern in [
"pub trait AiCandidateMaterializationPort",
"pub async fn run_ai_candidate_materialization",
@@ -270,7 +270,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
}
let serving_candidate_preselection =
read_workspace_file("crates/aether-ai-serving/src/candidate_preselection.rs");
read_workspace_file("crates/aether-ai/serving/src/candidate_preselection.rs");
for pattern in [
"pub trait AiCandidatePreselectionPort",
"pub async fn run_ai_candidate_preselection",
@@ -290,7 +290,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
}
let serving_candidate_ranking =
read_workspace_file("crates/aether-ai-serving/src/candidate_ranking.rs");
read_workspace_file("crates/aether-ai/serving/src/candidate_ranking.rs");
for pattern in [
"pub trait AiCandidateRankingPort",
"pub async fn run_ai_candidate_ranking",
@@ -312,7 +312,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
}
let serving_candidate_resolution =
read_workspace_file("crates/aether-ai-serving/src/candidate_resolution.rs");
read_workspace_file("crates/aether-ai/serving/src/candidate_resolution.rs");
for pattern in [
"pub trait AiCandidateResolutionPort",
"pub async fn run_ai_candidate_resolution",
@@ -370,7 +370,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
}
let provider_transport_conversion =
read_workspace_file("crates/aether-provider-transport/src/conversion.rs");
read_workspace_file("crates/aether-provider/transport/src/conversion.rs");
for pattern in [
"pub struct CandidateTransportPolicyFacts",
"pub fn candidate_common_transport_skip_reason(",
@@ -445,7 +445,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
}
let serving_execution_path =
read_workspace_file("crates/aether-ai-serving/src/execution_path.rs");
read_workspace_file("crates/aether-ai/serving/src/execution_path.rs");
for pattern in [
"pub enum AiSyncExecutionStep",
"pub enum AiStreamExecutionStep",
@@ -488,7 +488,7 @@ fn ai_serving_crate_owns_attempt_loop_without_gateway_runtime_deps() {
#[test]
fn ai_serving_internal_dtos_use_ai_execution_names() {
let serving_dto = read_workspace_file("crates/aether-ai-serving/src/dto.rs");
let serving_dto = read_workspace_file("crates/aether-ai/serving/src/dto.rs");
for expected in [
"pub struct AiExecutionDecision",
"pub struct AiExecutionPlanPayload",
@@ -525,8 +525,8 @@ fn ai_serving_internal_dtos_use_ai_execution_names() {
"apps/aether-gateway/src/ai_serving",
"apps/aether-gateway/src/executor",
"apps/aether-gateway/src/execution_runtime",
"crates/aether-ai-serving/src",
"crates/aether-ai-formats/src",
"crates/aether-ai/serving/src",
"crates/aether-ai/formats/src",
] {
for file in collect_workspace_rust_files(root) {
let relative = file
@@ -560,7 +560,7 @@ fn ai_serving_internal_dtos_use_ai_execution_names() {
#[test]
fn ai_format_crate_stays_free_of_gateway_runtime_deps() {
for manifest_path in ["crates/aether-ai-formats/Cargo.toml"] {
for manifest_path in ["crates/aether-ai/formats/Cargo.toml"] {
let manifest = read_workspace_file(manifest_path);
for forbidden in [
"axum",
@@ -578,7 +578,7 @@ fn ai_format_crate_stays_free_of_gateway_runtime_deps() {
}
let mut violations = Vec::new();
for root in ["crates/aether-ai-formats/src"] {
for root in ["crates/aether-ai/formats/src"] {
for file in collect_workspace_rust_files(root) {
let source = std::fs::read_to_string(&file).expect("source file should be readable");
let hits = [
@@ -608,7 +608,7 @@ fn ai_format_crate_stays_free_of_gateway_runtime_deps() {
#[test]
fn aether_runtime_stays_free_of_ai_serving_policy() {
let runtime_manifest = read_workspace_file("crates/aether-runtime/Cargo.toml");
let runtime_manifest = read_workspace_file("crates/aether-runtime/base/Cargo.toml");
for forbidden in [
"aether-ai-serving",
"aether-ai-formats",
@@ -622,7 +622,7 @@ fn aether_runtime_stays_free_of_ai_serving_policy() {
}
let mut violations = Vec::new();
for file in collect_workspace_rust_files("crates/aether-runtime/src") {
for file in collect_workspace_rust_files("crates/aether-runtime/base/src") {
let source = std::fs::read_to_string(&file).expect("source file should be readable");
let hits = [
"aether_ai_serving",
@@ -824,7 +824,7 @@ fn ai_serving_routes_control_and_execution_deps_through_facades() {
);
}
let serving_attempt_plan = read_workspace_file("crates/aether-ai-serving/src/attempt_plan.rs");
let serving_attempt_plan = read_workspace_file("crates/aether-ai/serving/src/attempt_plan.rs");
for pattern in [
"pub fn take_ai_decision_plan_core(",
"pub fn take_ai_upstream_auth_pair(",
@@ -871,13 +871,15 @@ fn ai_serving_routes_control_and_execution_deps_through_facades() {
let gateway_finalize_common =
read_workspace_file("apps/aether-gateway/src/ai_serving/finalize/common.rs");
assert!(
gateway_finalize_common
.contains("prepare_local_success_response_parts as prepare_local_success_response_parts_impl"),
gateway_finalize_common.contains(
"prepare_local_success_response_parts as prepare_local_success_response_parts_impl"
),
"finalize/common.rs should delegate success response-part normalization to the format crate"
);
assert!(
gateway_finalize_common
.contains("build_local_success_background_report as build_local_success_background_report_impl"),
gateway_finalize_common.contains(
"build_local_success_background_report as build_local_success_background_report_impl"
),
"finalize/common.rs should delegate pure success background-report construction to the format crate"
);
assert!(
@@ -1009,7 +1011,7 @@ fn ai_serving_routes_provider_transport_deps_through_facade() {
"ai_serving/runtime should stay removed after facade cleanup"
);
assert!(
!workspace_file_exists("crates/aether-ai-formats/src/transport.rs"),
!workspace_file_exists("crates/aether-ai/formats/src/transport.rs"),
"aether-ai-formats should not expose a provider transport bridge"
);
@@ -1384,7 +1386,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
);
}
let serving_lib = read_workspace_file("crates/aether-ai-serving/src/lib.rs");
let serving_lib = read_workspace_file("crates/aether-ai/serving/src/lib.rs");
for forbidden in ["pub mod pool_scheduler;", "pub mod pool_scores;"] {
assert!(
!serving_lib.contains(forbidden),
@@ -1392,7 +1394,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
);
}
let provider_pool_lib = read_workspace_file("crates/aether-provider-pool/src/lib.rs");
let provider_pool_lib = read_workspace_file("crates/aether-provider/pool/src/lib.rs");
for pattern in [
"mod capability;",
"mod plan;",
@@ -1410,7 +1412,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
);
}
let provider_pool_provider = read_workspace_file("crates/aether-provider-pool/src/provider.rs");
let provider_pool_provider = read_workspace_file("crates/aether-provider/pool/src/provider.rs");
for pattern in [
"pub trait ProviderPoolAdapter",
"ProviderPoolMemberInput",
@@ -1423,7 +1425,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
);
}
let provider_pool_service = read_workspace_file("crates/aether-provider-pool/src/service.rs");
let provider_pool_service = read_workspace_file("crates/aether-provider/pool/src/service.rs");
for pattern in [
"pub struct ProviderPoolService",
"with_builtin_adapters",
@@ -1449,7 +1451,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
);
let provider_pool_providers =
read_workspace_file("crates/aether-provider-pool/src/providers/mod.rs");
read_workspace_file("crates/aether-provider/pool/src/providers/mod.rs");
for pattern in [
"pub mod default;",
"pub mod unsupported;",
@@ -1466,15 +1468,15 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
}
for (path, patterns) in [
(
"crates/aether-provider-pool/src/providers/default.rs",
"crates/aether-provider/pool/src/providers/default.rs",
vec!["DefaultProviderPoolAdapter"],
),
(
"crates/aether-provider-pool/src/providers/antigravity.rs",
"crates/aether-provider/pool/src/providers/antigravity.rs",
vec!["AntigravityProviderPoolAdapter"],
),
(
"crates/aether-provider-pool/src/providers/codex.rs",
"crates/aether-provider/pool/src/providers/codex.rs",
vec![
"CodexProviderPoolAdapter",
"recent_refresh",
@@ -1482,7 +1484,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
],
),
(
"crates/aether-provider-pool/src/providers/gemini_cli.rs",
"crates/aether-provider/pool/src/providers/gemini_cli.rs",
vec![
"GeminiCliProviderPoolAdapter",
"build_gemini_cli_pool_quota_request",
@@ -1490,11 +1492,11 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
],
),
(
"crates/aether-provider-pool/src/providers/kiro.rs",
"crates/aether-provider/pool/src/providers/kiro.rs",
vec!["KiroProviderPoolAdapter", "quota_exhausted_from_bucket"],
),
(
"crates/aether-provider-pool/src/providers/chatgpt_web.rs",
"crates/aether-provider/pool/src/providers/chatgpt_web.rs",
vec![
"ChatGptWebProviderPoolAdapter",
"build_chatgpt_web_pool_quota_request",
@@ -1504,7 +1506,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
],
),
(
"crates/aether-provider-pool/src/providers/unsupported.rs",
"crates/aether-provider/pool/src/providers/unsupported.rs",
vec![
"UnsupportedQuotaProviderPoolAdapter",
"CLAUDE_CODE_PROVIDER_POOL_ADAPTER",
@@ -1521,7 +1523,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
}
}
let provider_pool_plan = read_workspace_file("crates/aether-provider-pool/src/plan.rs");
let provider_pool_plan = read_workspace_file("crates/aether-provider/pool/src/plan.rs");
for pattern in ["normalize_provider_plan_tier", "derive_plan_tier"] {
assert!(
provider_pool_plan.contains(pattern),
@@ -1529,7 +1531,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
);
}
let provider_pool_quota = read_workspace_file("crates/aether-provider-pool/src/quota.rs");
let provider_pool_quota = read_workspace_file("crates/aether-provider/pool/src/quota.rs");
for pattern in [
"provider_pool_key_account_quota_exhausted",
"provider_pool_member_quota_snapshot",
@@ -1544,7 +1546,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
);
}
let provider_pool_presets = read_workspace_file("crates/aether-provider-pool/src/presets.rs");
let provider_pool_presets = read_workspace_file("crates/aether-provider/pool/src/presets.rs");
for pattern in [
"normalize_provider_scheduling_presets",
"build_admin_pool_scheduling_presets_payload",
@@ -1561,7 +1563,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
"plan_priority_score(",
] {
let mut violations = Vec::new();
for file in collect_workspace_rust_files("crates/aether-provider-pool/src") {
for file in collect_workspace_rust_files("crates/aether-provider/pool/src") {
let source = std::fs::read_to_string(&file).expect("source file should be readable");
if source.contains(forbidden) {
violations.push(file.display().to_string());
@@ -1578,7 +1580,7 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
#[test]
fn ai_serving_candidate_preparation_owns_shared_auth_and_mapped_model_resolution() {
let serving_candidate_preparation =
read_workspace_file("crates/aether-ai-serving/src/candidate_preparation.rs");
read_workspace_file("crates/aether-ai/serving/src/candidate_preparation.rs");
for pattern in [
"pub struct AiPreparedHeaderAuthenticatedCandidate",
"pub fn prepare_ai_header_authenticated_candidate(",
@@ -1592,7 +1594,7 @@ fn ai_serving_candidate_preparation_owns_shared_auth_and_mapped_model_resolution
);
}
let serving_lib = read_workspace_file("crates/aether-ai-serving/src/lib.rs");
let serving_lib = read_workspace_file("crates/aether-ai/serving/src/lib.rs");
for pattern in [
"pub mod candidate_preparation;",
"prepare_ai_header_authenticated_candidate",
@@ -1690,7 +1692,7 @@ fn ai_serving_candidate_materialization_owns_affinity_and_candidate_runtime_pers
);
let serving_candidate_persistence =
read_workspace_file("crates/aether-ai-serving/src/candidate_persistence.rs");
read_workspace_file("crates/aether-ai/serving/src/candidate_persistence.rs");
for pattern in [
"pub trait AiAvailableCandidatePersistencePort",
"pub async fn run_ai_available_candidate_persistence",
@@ -1837,7 +1839,7 @@ fn ai_serving_materialization_policy_owns_local_candidate_persistence_modes() {
);
let serving_candidate_persistence_policy =
read_workspace_file("crates/aether-ai-serving/src/candidate_persistence_policy.rs");
read_workspace_file("crates/aether-ai/serving/src/candidate_persistence_policy.rs");
for pattern in [
"pub enum AiCandidatePersistencePolicyKind {",
"pub struct AiCandidatePersistencePolicySpec {",
@@ -1913,9 +1915,9 @@ fn ai_serving_candidate_metadata_owns_local_execution_candidate_extra_data_shape
let candidate_metadata =
read_workspace_file("apps/aether-gateway/src/ai_serving/planner/candidate_metadata.rs");
let serving_ranking_metadata =
read_workspace_file("crates/aether-ai-serving/src/ranking_metadata.rs");
read_workspace_file("crates/aether-ai/serving/src/ranking_metadata.rs");
let serving_candidate_metadata =
read_workspace_file("crates/aether-ai-serving/src/candidate_metadata.rs");
read_workspace_file("crates/aether-ai/serving/src/candidate_metadata.rs");
for pattern in [
"pub struct AiCandidateMetadataParts<'a> {",
"pub fn build_ai_candidate_metadata(",
@@ -2015,7 +2017,7 @@ fn ai_serving_runtime_miss_owns_local_execution_miss_state_machine() {
"planner/mod.rs should wire runtime_miss helper module"
);
let serving_runtime_miss = read_workspace_file("crates/aether-ai-serving/src/runtime_miss.rs");
let serving_runtime_miss = read_workspace_file("crates/aether-ai/serving/src/runtime_miss.rs");
for pattern in [
"pub trait AiRuntimeMissDiagnosticPort",
"pub trait AiRuntimeMissDiagnosticFields",
@@ -2264,7 +2266,7 @@ fn ai_serving_video_routes_request_preparation_through_request_payload_seams() {
}
let provider_transport_video =
read_workspace_file("crates/aether-provider-transport/src/video/mod.rs");
read_workspace_file("crates/aether-provider/transport/src/video/mod.rs");
for pattern in [
"pub enum ProviderVideoCreateFamily",
"pub fn video_create_transport_unsupported_reason(",
@@ -2339,7 +2341,7 @@ fn ai_serving_files_routes_request_preparation_through_request_payload_seams() {
}
let provider_transport_files =
read_workspace_file("crates/aether-provider-transport/src/gemini_files/mod.rs");
read_workspace_file("crates/aether-provider/transport/src/gemini_files/mod.rs");
for pattern in [
"pub fn gemini_files_transport_unsupported_reason(",
"pub fn resolve_gemini_files_auth(",
@@ -2409,7 +2411,7 @@ fn ai_serving_image_routes_split_surface_normalization_and_transport_policy() {
}
let surface_image =
read_workspace_file("crates/aether-ai-formats/src/formats/openai/image/request.rs");
read_workspace_file("crates/aether-ai/formats/src/formats/openai/image/request.rs");
for pattern in [
"pub enum OpenAiImageOperation",
"pub fn is_openai_image_stream_request(",
@@ -2425,7 +2427,7 @@ fn ai_serving_image_routes_split_surface_normalization_and_transport_policy() {
}
let provider_transport_image =
read_workspace_file("crates/aether-provider-transport/src/openai_image/mod.rs");
read_workspace_file("crates/aether-provider/transport/src/openai_image/mod.rs");
for pattern in [
"pub fn openai_image_transport_unsupported_reason(",
"pub fn resolve_openai_image_auth(",
@@ -2605,7 +2607,7 @@ fn ai_serving_payload_metadata_owns_local_execution_decision_response_shape() {
"gateway payload_metadata.rs should be removed after serving extraction"
);
let decision_payload = read_workspace_file("crates/aether-ai-serving/src/decision_payload.rs");
let decision_payload = read_workspace_file("crates/aether-ai/serving/src/decision_payload.rs");
for pattern in [
"pub struct AiExecutionDecisionResponseParts {",
"pub fn build_ai_execution_decision_response(",
@@ -2651,7 +2653,7 @@ fn ai_serving_owns_pure_planner_diagnostics_and_execution_labels() {
}
let serving_failure_diagnostic =
read_workspace_file("crates/aether-ai-serving/src/failure_diagnostic.rs");
read_workspace_file("crates/aether-ai/serving/src/failure_diagnostic.rs");
for pattern in [
"pub enum CandidateFailureDiagnosticKind {",
"pub struct CandidateFailureDiagnostic {",
@@ -2667,7 +2669,7 @@ fn ai_serving_owns_pure_planner_diagnostics_and_execution_labels() {
}
let serving_request_body_diagnostics =
read_workspace_file("crates/aether-ai-serving/src/request_body_diagnostics.rs");
read_workspace_file("crates/aether-ai/serving/src/request_body_diagnostics.rs");
for pattern in [
"pub fn request_body_build_failure_extra_data(",
"pub fn same_format_provider_request_body_failure_extra_data(",
@@ -2692,7 +2694,7 @@ fn ai_serving_owns_pure_planner_diagnostics_and_execution_labels() {
"gateway standard planner should consume request-body diagnostics from aether-ai-serving"
);
let serving_dto = read_workspace_file("crates/aether-ai-serving/src/dto.rs");
let serving_dto = read_workspace_file("crates/aether-ai/serving/src/dto.rs");
for pattern in ["pub enum ExecutionStrategy", "pub enum ConversionMode"] {
assert!(
serving_dto.contains(pattern),
@@ -2726,7 +2728,7 @@ fn ai_serving_report_context_owns_local_execution_context_shape() {
);
let serving_report_context =
read_workspace_file("crates/aether-ai-serving/src/report_context.rs");
read_workspace_file("crates/aether-ai/serving/src/report_context.rs");
for pattern in [
"pub struct AiExecutionReportContextParts<'a> {",
"pub fn build_ai_execution_report_context(",
@@ -2878,7 +2880,7 @@ fn ai_serving_standard_attempts_consume_eligible_local_candidates_without_transp
}
let provider_transport_standard =
read_workspace_file("crates/aether-provider-transport/src/standard/mod.rs");
read_workspace_file("crates/aether-provider/transport/src/standard/mod.rs");
for pattern in [
"pub struct StandardProviderRequestHeadersInput",
"pub struct StandardProviderRequestHeaders",
@@ -2898,7 +2900,7 @@ fn ai_serving_standard_attempts_consume_eligible_local_candidates_without_transp
}
let provider_transport_request_url =
read_workspace_file("crates/aether-provider-transport/src/request_url/mod.rs");
read_workspace_file("crates/aether-provider/transport/src/request_url/mod.rs");
assert!(
provider_transport_request_url.contains("pub fn build_kiro_cross_format_upstream_url("),
"provider-transport request_url.rs should own Kiro cross-format URL hook"
@@ -2960,7 +2962,7 @@ fn ai_serving_standard_plan_builders_delegate_fallback_transport_policy() {
}
let provider_transport_standard =
read_workspace_file("crates/aether-provider-transport/src/standard/mod.rs");
read_workspace_file("crates/aether-provider/transport/src/standard/mod.rs");
for pattern in [
"pub enum StandardPlanFallbackAcceptPolicy",
"pub struct StandardPlanFallbackHeadersInput",
@@ -3063,7 +3065,7 @@ fn ai_serving_spec_metadata_owns_family_requested_model_and_plan_builder_routing
let spec_metadata =
read_workspace_file("apps/aether-gateway/src/ai_serving/planner/spec_metadata.rs");
let serving_surface_spec = read_workspace_file("crates/aether-ai-serving/src/surface_spec.rs");
let serving_surface_spec = read_workspace_file("crates/aether-ai/serving/src/surface_spec.rs");
for pattern in [
"pub enum AiRequestedModelFamily {",
"pub struct AiExecutionSurfaceSpecMetadata {",
@@ -3228,7 +3230,7 @@ fn ai_serving_same_format_provider_request_policy_owns_provider_type_behavior()
"apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/request/policy.rs",
);
let provider_transport_policy =
read_workspace_file("crates/aether-provider-transport/src/same_format_provider/mod.rs");
read_workspace_file("crates/aether-provider/transport/src/same_format_provider/mod.rs");
for pattern in [
"pub struct SameFormatProviderRequestBehavior {",
"pub struct SameFormatProviderRequestBodyInput",
@@ -3314,7 +3316,7 @@ fn ai_serving_decision_inputs_share_authenticated_input_helper() {
);
let serving_decision_input =
read_workspace_file("crates/aether-ai-serving/src/decision_input.rs");
read_workspace_file("crates/aether-ai/serving/src/decision_input.rs");
for pattern in [
"pub trait AiAuthenticatedDecisionInputPort",
"pub async fn run_ai_authenticated_decision_input",
@@ -3488,7 +3490,8 @@ fn ai_serving_leaf_planner_owners_route_contract_specs_through_gateway_seams() {
"apps/aether-gateway/src/ai_serving/planner/specialized/video/support.rs",
);
assert!(
specialized_video_support.contains("use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};"),
specialized_video_support
.contains("use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};"),
"planner/specialized/video/support.rs should use local video seams for LocalVideoCreate* types"
);
}
@@ -3496,10 +3499,10 @@ fn ai_serving_leaf_planner_owners_route_contract_specs_through_gateway_seams() {
#[test]
fn ai_serving_m5_moves_contracts_and_route_logic_into_format_crate() {
for path in [
"crates/aether-ai-formats/src/contracts/actions.rs",
"crates/aether-ai-formats/src/contracts/plan_kinds.rs",
"crates/aether-ai-formats/src/contracts/report_kinds.rs",
"crates/aether-ai-formats/src/formats/shared/routing.rs",
"crates/aether-ai/formats/src/contracts/actions.rs",
"crates/aether-ai/formats/src/contracts/plan_kinds.rs",
"crates/aether-ai/formats/src/contracts/report_kinds.rs",
"crates/aether-ai/formats/src/formats/shared/routing.rs",
] {
assert!(
workspace_file_exists(path),
@@ -3561,7 +3564,7 @@ fn ai_serving_m5_moves_contracts_and_route_logic_into_format_crate() {
}
let surface_route =
read_workspace_file("crates/aether-ai-formats/src/formats/shared/routing.rs");
read_workspace_file("crates/aether-ai/formats/src/formats/shared/routing.rs");
for pattern in [
"pub fn is_matching_stream_http_request(",
"is_openai_image_stream_request(parts, body_json, body_base64)",
@@ -3614,12 +3617,12 @@ fn ai_serving_m5_moves_contracts_and_route_logic_into_format_crate() {
#[test]
fn ai_serving_m5_moves_kiro_stream_helpers_into_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/provider_compat/kiro_stream.rs"),
"crates/aether-ai-formats/src/provider_compat/kiro_stream.rs should exist after kiro helper extraction"
workspace_file_exists("crates/aether-ai/formats/src/provider_compat/kiro_stream.rs"),
"crates/aether-ai/formats/src/provider_compat/kiro_stream.rs should exist after kiro helper extraction"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/provider_compat/kiro_stream/state.rs"),
"crates/aether-ai-formats/src/provider_compat/kiro_stream/state.rs should own the Kiro stream state machine"
workspace_file_exists("crates/aether-ai/formats/src/provider_compat/kiro_stream/state.rs"),
"crates/aether-ai/formats/src/provider_compat/kiro_stream/state.rs should own the Kiro stream state machine"
);
for path in [
@@ -3633,7 +3636,7 @@ fn ai_serving_m5_moves_kiro_stream_helpers_into_format_crate() {
);
}
let surface_api = read_workspace_file("crates/aether-ai-formats/src/api.rs");
let surface_api = read_workspace_file("crates/aether-ai/formats/src/api.rs");
assert!(
surface_api.contains("KiroToClaudeCliStreamState"),
"aether-ai-formats api should export KiroToClaudeCliStreamState"
@@ -3670,7 +3673,7 @@ fn ai_serving_m5_moves_kiro_stream_helpers_into_format_crate() {
#[test]
fn ai_serving_private_envelope_stream_normalizer_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/provider_compat/private_envelope.rs"),
workspace_file_exists("crates/aether-ai/formats/src/provider_compat/private_envelope.rs"),
"surface private envelope adapter should exist"
);
assert!(
@@ -3687,7 +3690,7 @@ fn ai_serving_private_envelope_stream_normalizer_is_owned_by_format_crate() {
);
let surface_private_envelope =
read_workspace_file("crates/aether-ai-formats/src/provider_compat/private_envelope.rs");
read_workspace_file("crates/aether-ai/formats/src/provider_compat/private_envelope.rs");
for expected in [
"pub struct ProviderPrivateStreamNormalizer",
"pub fn maybe_build_provider_private_stream_normalizer",
@@ -3810,11 +3813,11 @@ fn ai_serving_error_body_is_owned_by_format_finalize_module() {
"ai_serving/conversion/error.rs should stay removed"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/error_body.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/error_body.rs"),
"format error response-body helpers should live under finalize/error_body.rs"
);
assert!(
!workspace_file_exists("crates/aether-ai-formats/src/formats/conversion/error.rs"),
!workspace_file_exists("crates/aether-ai/formats/src/formats/conversion/error.rs"),
"aether-ai-formats should not keep error response-body helpers under conversion"
);
@@ -3839,7 +3842,7 @@ fn ai_serving_error_body_is_owned_by_format_finalize_module() {
#[test]
fn ai_serving_conversion_request_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/conversion/request.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/conversion/request.rs"),
"request conversion should live in aether-ai-formats"
);
assert!(
@@ -3864,7 +3867,7 @@ fn ai_serving_conversion_request_is_owned_by_format_crate() {
"gateway ai_serving/mod.rs should not keep request re-export shell after root-seam consolidation"
);
let surface_api = read_workspace_file("crates/aether-ai-formats/src/api.rs");
let surface_api = read_workspace_file("crates/aether-ai/formats/src/api.rs");
assert!(
surface_api.contains("pub use aether_ai_formats::formats::conversion::request::{"),
"format API facade should re-export request conversion directly from aether-ai-formats"
@@ -3874,7 +3877,7 @@ fn ai_serving_conversion_request_is_owned_by_format_crate() {
#[test]
fn ai_serving_conversion_response_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/conversion/response.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/conversion/response.rs"),
"response conversion should live in aether-ai-formats"
);
assert!(
@@ -3899,7 +3902,7 @@ fn ai_serving_conversion_response_is_owned_by_format_crate() {
"gateway ai_serving/mod.rs should not keep response re-export shell after root-seam consolidation"
);
let surface_api = read_workspace_file("crates/aether-ai-formats/src/api.rs");
let surface_api = read_workspace_file("crates/aether-ai/formats/src/api.rs");
assert!(
surface_api.contains("pub use aether_ai_formats::formats::conversion::response::{"),
"format API facade should re-export response conversion directly from aether-ai-formats"
@@ -3909,17 +3912,17 @@ fn ai_serving_conversion_response_is_owned_by_format_crate() {
#[test]
fn ai_format_crate_owns_conversion_and_surface_facade() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/conversion"),
workspace_file_exists("crates/aether-ai/formats/src/formats/conversion"),
"aether-ai-formats should own the conversion directory"
);
let surface_lib = read_workspace_file("crates/aether-ai-formats/src/lib.rs");
let surface_lib = read_workspace_file("crates/aether-ai/formats/src/lib.rs");
assert!(
surface_lib.contains("pub mod protocol;"),
"aether-ai-formats lib.rs should expose the protocol module"
);
let surface_api = read_workspace_file("crates/aether-ai-formats/src/api.rs");
let surface_api = read_workspace_file("crates/aether-ai/formats/src/api.rs");
for pattern in [
"pub use aether_ai_formats::{",
"pub use aether_ai_formats::formats::conversion::request::{",
@@ -4003,12 +4006,12 @@ fn ai_serving_finalize_standard_sync_response_converters_are_owned_by_format_cra
#[test]
fn ai_serving_finalize_stream_engine_is_owned_by_format_crate() {
for path in [
"crates/aether-ai-formats/src/formats/shared/sse.rs",
"crates/aether-ai-formats/src/formats/shared/stream_core/common.rs",
"crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs",
"crates/aether-ai-formats/src/formats/openai/chat/stream.rs",
"crates/aether-ai-formats/src/formats/claude/messages/stream.rs",
"crates/aether-ai-formats/src/formats/gemini/generate_content/stream.rs",
"crates/aether-ai/formats/src/formats/shared/sse.rs",
"crates/aether-ai/formats/src/formats/shared/stream_core/common.rs",
"crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs",
"crates/aether-ai/formats/src/formats/openai/chat/stream.rs",
"crates/aether-ai/formats/src/formats/claude/messages/stream.rs",
"crates/aether-ai/formats/src/formats/gemini/generate_content/stream.rs",
] {
assert!(
workspace_file_exists(path),
@@ -4044,7 +4047,7 @@ fn ai_serving_finalize_stream_engine_is_owned_by_format_crate() {
}
let surface_format_matrix = read_workspace_file(
"crates/aether-ai-formats/src/formats/shared/stream_core/format_matrix.rs",
"crates/aether-ai/formats/src/formats/shared/stream_core/format_matrix.rs",
);
for pattern in [
"pub struct StreamingStandardFormatMatrix",
@@ -4068,16 +4071,16 @@ fn ai_serving_finalize_stream_engine_is_owned_by_format_crate() {
#[test]
fn ai_serving_finalize_standard_sync_products_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/sync_products.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/sync_products.rs"),
"finalize sync_products should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/sync_to_stream.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/sync_to_stream.rs"),
"finalize sync-to-stream bridge should live in aether-ai-formats"
);
let surface_sync_products =
read_workspace_file("crates/aether-ai-formats/src/formats/shared/sync_products.rs");
read_workspace_file("crates/aether-ai/formats/src/formats/shared/sync_products.rs");
for expected in [
"pub fn maybe_build_standard_cross_format_sync_product_from_normalized_payload(",
"pub fn maybe_build_standard_same_format_sync_body_from_normalized_payload(",
@@ -4167,9 +4170,8 @@ fn ai_serving_finalize_standard_sync_products_are_owned_by_format_crate() {
"apps/aether-gateway/src/ai_serving/finalize/internal/sync_finalize.rs",
);
assert!(
gateway_internal_sync.contains(
"maybe_build_standard_sync_finalize_product_from_normalized_payload"
),
gateway_internal_sync
.contains("maybe_build_standard_sync_finalize_product_from_normalized_payload"),
"gateway internal/sync_finalize.rs should delegate normalized standard sync finalize dispatch to aether-ai-formats"
);
assert!(
@@ -4204,7 +4206,7 @@ fn ai_serving_finalize_standard_sync_products_are_owned_by_format_crate() {
}
let surface_openai_image_stream =
read_workspace_file("crates/aether-ai-formats/src/formats/openai/image/stream.rs");
read_workspace_file("crates/aether-ai/formats/src/formats/openai/image/stream.rs");
for expected in [
"pub fn maybe_build_openai_image_sync_finalize_product(",
"pub struct OpenAiImageSyncFinalizeProduct",
@@ -4219,7 +4221,7 @@ fn ai_serving_finalize_standard_sync_products_are_owned_by_format_crate() {
}
let surface_sync_to_stream =
read_workspace_file("crates/aether-ai-formats/src/formats/shared/sync_to_stream.rs");
read_workspace_file("crates/aether-ai/formats/src/formats/shared/sync_to_stream.rs");
for expected in [
"pub fn maybe_bridge_standard_sync_json_to_stream(",
"pub struct SyncToStreamBridgeOutcome",
@@ -4256,11 +4258,11 @@ fn ai_serving_finalize_standard_sync_products_are_owned_by_format_crate() {
#[test]
fn ai_serving_finalize_stream_rewrite_matrix_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/stream_rewrite.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/stream_rewrite.rs"),
"finalize stream rewrite matrix should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/openai/image/stream.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/openai/image/stream.rs"),
"OpenAI image stream rewrite state should live in aether-ai-formats"
);
@@ -4309,7 +4311,7 @@ fn ai_serving_finalize_stream_rewrite_matrix_is_owned_by_format_crate() {
#[test]
fn ai_serving_planner_common_parser_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/request.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/request.rs"),
"planner/common pure parser should exist in aether-ai-formats"
);
@@ -4325,8 +4327,9 @@ fn ai_serving_planner_common_parser_is_owned_by_format_crate() {
"gateway planner/common.rs should delegate body parsing through the ai_serving root seam"
);
assert!(
gateway_common_runtime
.contains("force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl"),
gateway_common_runtime.contains(
"force_upstream_streaming_for_provider as force_upstream_streaming_for_provider_impl"
),
"gateway planner/common.rs should delegate upstream streaming policy through the ai_serving root seam"
);
for forbidden in [
@@ -4404,7 +4407,8 @@ fn ai_serving_planner_common_parser_is_owned_by_format_crate() {
);
assert!(
openai_chat_diagnostic.contains("set_local_runtime_miss_diagnostic_reason(")
|| openai_chat_diagnostic.contains("set_local_runtime_candidate_evaluation_diagnostic("),
|| openai_chat_diagnostic
.contains("set_local_runtime_candidate_evaluation_diagnostic("),
"openai chat diagnostic.rs should delegate miss diagnostic handling through planner/runtime_miss.rs"
);
assert!(
@@ -4415,7 +4419,7 @@ fn ai_serving_planner_common_parser_is_owned_by_format_crate() {
#[test]
fn ai_serving_root_owns_shared_gemini_request_path_parser() {
let serving_surface_spec = read_workspace_file("crates/aether-ai-serving/src/surface_spec.rs");
let serving_surface_spec = read_workspace_file("crates/aether-ai/serving/src/surface_spec.rs");
assert!(
serving_surface_spec.contains("pub fn extract_ai_gemini_model_from_path("),
"aether-ai-serving should own shared gemini request-path parsing"
@@ -4423,9 +4427,8 @@ fn ai_serving_root_owns_shared_gemini_request_path_parser() {
let ai_serving_mod = read_workspace_file("apps/aether-gateway/src/ai_serving/mod.rs");
assert!(
ai_serving_mod.contains(
"extract_ai_gemini_model_from_path as extract_gemini_model_from_path"
),
ai_serving_mod
.contains("extract_ai_gemini_model_from_path as extract_gemini_model_from_path"),
"ai_serving/mod.rs should expose shared gemini request-path parsing through the serving seam"
);
@@ -4452,7 +4455,7 @@ fn ai_serving_root_owns_shared_gemini_request_path_parser() {
#[test]
fn ai_serving_planner_standard_normalize_is_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/standard_normalize.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/standard_normalize.rs"),
"planner/standard/normalize should live in aether-ai-formats"
);
@@ -4515,7 +4518,7 @@ fn ai_serving_planner_standard_normalize_is_owned_by_format_crate() {
#[test]
fn ai_serving_openai_helpers_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/openai/shared.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/openai/shared.rs"),
"planner/openai helper owner should exist in aether-ai-formats"
);
@@ -4548,17 +4551,17 @@ fn ai_serving_openai_helpers_are_owned_by_format_crate() {
#[test]
fn ai_serving_standard_matrix_delegates_format_conversion_to_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/request_matrix.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/request_matrix.rs"),
"planner/matrix facade should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/standard_matrix.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/standard_matrix.rs"),
"format standard request-body planner should live in aether-ai-formats"
);
for path in [
"crates/aether-ai-formats/src/protocol/canonical.rs",
"crates/aether-ai-formats/src/formats/matrix.rs",
"crates/aether-ai-formats/src/formats/registry.rs",
"crates/aether-ai/formats/src/protocol/canonical.rs",
"crates/aether-ai/formats/src/formats/matrix.rs",
"crates/aether-ai/formats/src/formats/registry.rs",
] {
assert!(
workspace_file_exists(path),
@@ -4566,7 +4569,7 @@ fn ai_serving_standard_matrix_delegates_format_conversion_to_format_crate() {
);
}
let surface_matrix =
read_workspace_file("crates/aether-ai-formats/src/formats/shared/standard_matrix.rs");
read_workspace_file("crates/aether-ai/formats/src/formats/shared/standard_matrix.rs");
assert!(
surface_matrix.contains("use aether_ai_formats::formats::registry::{")
&& surface_matrix.contains("convert_request")
@@ -4618,26 +4621,26 @@ fn ai_serving_standard_matrix_delegates_format_conversion_to_format_crate() {
#[test]
fn ai_serving_standard_family_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/family.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/family.rs"),
"planner/standard/family pure spec owner should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/claude/messages/chat_spec.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/claude/messages/chat_spec.rs"),
"planner/standard/claude/chat pure spec resolver should live in aether-ai-formats"
);
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/claude/messages/cli_spec.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/claude/messages/cli_spec.rs"),
"planner/standard/claude/cli pure spec resolver should live in aether-ai-formats"
);
assert!(
workspace_file_exists(
"crates/aether-ai-formats/src/formats/gemini/generate_content/chat_spec.rs"
"crates/aether-ai/formats/src/formats/gemini/generate_content/chat_spec.rs"
),
"planner/standard/gemini/chat pure spec resolver should live in aether-ai-formats"
);
assert!(
workspace_file_exists(
"crates/aether-ai-formats/src/formats/gemini/generate_content/cli_spec.rs"
"crates/aether-ai/formats/src/formats/gemini/generate_content/cli_spec.rs"
),
"planner/standard/gemini/cli pure spec resolver should live in aether-ai-formats"
);
@@ -4705,7 +4708,7 @@ fn ai_serving_standard_family_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_same_format_provider_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/passthrough.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/passthrough.rs"),
"planner/passthrough/provider pure spec owner should live in aether-ai-formats"
);
@@ -4756,7 +4759,7 @@ fn ai_serving_same_format_provider_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_passthrough_provider_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/passthrough.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/passthrough.rs"),
"planner/passthrough/provider pure spec owner should live in aether-ai-formats"
);
@@ -4801,7 +4804,7 @@ fn ai_serving_passthrough_provider_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_specialized_files_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/gemini/files/spec.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/gemini/files/spec.rs"),
"planner/specialized/files pure spec owner should live in aether-ai-formats"
);
@@ -4831,12 +4834,12 @@ fn ai_serving_specialized_files_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_specialized_video_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/shared/video.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/shared/video.rs"),
"planner/specialized/video shared spec seam should live in aether-ai-formats"
);
for path in [
"crates/aether-ai-formats/src/formats/openai/video/spec.rs",
"crates/aether-ai-formats/src/formats/gemini/video/spec.rs",
"crates/aether-ai/formats/src/formats/openai/video/spec.rs",
"crates/aether-ai/formats/src/formats/gemini/video/spec.rs",
] {
assert!(
workspace_file_exists(path),
@@ -4868,7 +4871,7 @@ fn ai_serving_specialized_video_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_openai_responses_specs_are_owned_by_format_crate() {
assert!(
workspace_file_exists("crates/aether-ai-formats/src/formats/openai/responses/spec.rs"),
workspace_file_exists("crates/aether-ai/formats/src/formats/openai/responses/spec.rs"),
"planner/standard/openai_responses pure spec owner should live in aether-ai-formats"
);
@@ -4908,9 +4911,9 @@ fn ai_serving_openai_responses_specs_are_owned_by_format_crate() {
#[test]
fn ai_serving_legacy_api_format_names_stay_out_of_primary_paths() {
for path in [
"crates/aether-ai-formats/src/contracts/plan_kinds.rs",
"crates/aether-ai-formats/src/formats/shared/routing.rs",
"crates/aether-ai-formats/src/formats/openai/responses/spec.rs",
"crates/aether-ai/formats/src/contracts/plan_kinds.rs",
"crates/aether-ai/formats/src/formats/shared/routing.rs",
"crates/aether-ai/formats/src/formats/openai/responses/spec.rs",
"apps/aether-gateway/src/ai_serving/planner/decision/control_plan.rs",
"apps/aether-gateway/src/execution_runtime/fallback.rs",
] {
@@ -4934,7 +4937,7 @@ fn ai_serving_legacy_api_format_names_stay_out_of_primary_paths() {
}
}
let registry = read_workspace_file("crates/aether-ai-formats/src/formats/registry.rs");
let registry = read_workspace_file("crates/aether-ai/formats/src/formats/registry.rs");
let implementation = registry
.split("#[cfg(test)]")
.next()
@@ -4969,12 +4972,12 @@ fn retired_api_format_occurrences_are_whitelisted() {
"apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs",
"apps/aether-gateway/src/handlers/admin/request/system/import.rs",
"apps/aether-gateway/src/tests/control/admin/system_import.rs",
"crates/aether-ai-formats/src/formats/id.rs",
"crates/aether-ai-formats/src/formats/matrix.rs",
"crates/aether-ai-formats/src/formats/registry.rs",
"crates/aether-data/src/migrate.rs",
"crates/aether-data/src/lifecycle/migrate/tests.rs",
"crates/aether-usage-runtime/src/report.rs",
"crates/aether-ai/formats/src/formats/id.rs",
"crates/aether-ai/formats/src/formats/matrix.rs",
"crates/aether-ai/formats/src/formats/registry.rs",
"crates/aether-data/runtime/src/migrate.rs",
"crates/aether-data/runtime/src/lifecycle/migrate/tests.rs",
"crates/aether-usage/runtime/src/report.rs",
"frontend/src/api/endpoints/types/__tests__/api-format.spec.ts",
];
let allowed = allowed_paths
@@ -119,6 +119,23 @@ pub(super) fn workspace_file_exists(root_relative_path: &str) -> bool {
.exists()
}
pub(super) fn workspace_files_with_extension(
root_relative_path: &str,
extension: &str,
) -> Vec<PathBuf> {
let root = Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../..")
.join(root_relative_path);
let mut files = fs::read_dir(root)
.expect("workspace directory should be readable")
.filter_map(Result::ok)
.map(|entry| entry.path())
.filter(|path| path.extension().and_then(|value| value.to_str()) == Some(extension))
.collect::<Vec<_>>();
files.sort();
files
}
pub(super) fn collect_workspace_rust_files(root_relative_path: &str) -> Vec<PathBuf> {
let root = Path::new(env!("CARGO_MANIFEST_DIR"))
.join("../..")
@@ -179,3 +196,4 @@ mod ai_serving;
mod runtime_and_security;
mod sql_and_data;
mod usage;
mod workspace_tiers;
@@ -161,11 +161,11 @@ fn runtime_state_owns_redis_runtime_boundaries() {
"crates/aether-admin/src",
"crates/aether-billing/src",
"crates/aether-model-fetch/src",
"crates/aether-provider-pool/src",
"crates/aether-runtime/src",
"crates/aether-task-runtime/src",
"crates/aether-usage-runtime/src",
"crates/aether-provider-transport/src",
"crates/aether-provider/pool/src",
"crates/aether-runtime/base/src",
"crates/aether-task/runtime/src",
"crates/aether-usage/runtime/src",
"crates/aether-provider/transport/src",
"crates/aether-wallet/src",
] {
for path in collect_workspace_rust_files(root) {
@@ -199,11 +199,11 @@ fn runtime_state_owns_redis_runtime_boundaries() {
"crates/aether-admin/Cargo.toml",
"crates/aether-billing/Cargo.toml",
"crates/aether-model-fetch/Cargo.toml",
"crates/aether-provider-pool/Cargo.toml",
"crates/aether-provider-transport/Cargo.toml",
"crates/aether-runtime/Cargo.toml",
"crates/aether-task-runtime/Cargo.toml",
"crates/aether-usage-runtime/Cargo.toml",
"crates/aether-provider/pool/Cargo.toml",
"crates/aether-provider/transport/Cargo.toml",
"crates/aether-runtime/base/Cargo.toml",
"crates/aether-task/runtime/Cargo.toml",
"crates/aether-usage/runtime/Cargo.toml",
"crates/aether-wallet/Cargo.toml",
] {
let cargo = read_workspace_file(manifest);
@@ -220,7 +220,7 @@ fn runtime_state_owns_redis_runtime_boundaries() {
);
let mut runtime_state_violations = Vec::new();
for path in collect_workspace_rust_files("crates/aether-runtime-state/src") {
for path in collect_workspace_rust_files("crates/aether-runtime/state/src") {
if path
.components()
.any(|component| component.as_os_str() == "redis")
@@ -244,13 +244,13 @@ fn runtime_state_owns_redis_runtime_boundaries() {
}
assert!(
runtime_state_violations.is_empty(),
"only crates/aether-runtime-state/src/redis may depend on the redis crate directly:\n{}",
"only crates/aether-runtime/state/src/redis may depend on the redis crate directly:\n{}",
runtime_state_violations.join("\n")
);
let mut runtime_connection_violations = Vec::new();
for path in collect_workspace_rust_files("crates/aether-runtime-state/src") {
if path.ends_with("crates/aether-runtime-state/src/redis/client.rs") {
for path in collect_workspace_rust_files("crates/aether-runtime/state/src") {
if path.ends_with("crates/aether-runtime/state/src/redis/client.rs") {
continue;
}
let source = production_workspace_source(&path);
@@ -267,17 +267,17 @@ fn runtime_state_owns_redis_runtime_boundaries() {
#[test]
fn aether_data_stays_free_of_redis_runtime_backends() {
let cargo = read_workspace_file("crates/aether-data/Cargo.toml");
let cargo = read_workspace_file("crates/aether-data/runtime/Cargo.toml");
assert!(
!cargo.contains("redis.workspace"),
"aether-data should not depend on redis; runtime Redis belongs to aether-runtime-state"
);
for removed_path in [
"crates/aether-data/src/backend/redis.rs",
"crates/aether-data/src/backend/locks.rs",
"crates/aether-data/src/backend/workers.rs",
"crates/aether-data/src/driver/redis/mod.rs",
"crates/aether-data/runtime/src/backend/redis.rs",
"crates/aether-data/runtime/src/backend/locks.rs",
"crates/aether-data/runtime/src/backend/workers.rs",
"crates/aether-data/runtime/src/driver/redis/mod.rs",
] {
assert!(
!workspace_file_exists(removed_path),
@@ -285,7 +285,7 @@ fn aether_data_stays_free_of_redis_runtime_backends() {
);
}
for path in collect_workspace_rust_files("crates/aether-data/src") {
for path in collect_workspace_rust_files("crates/aether-data/runtime/src") {
let source = production_workspace_source(&path);
for forbidden in [
"pub mod redis",
@@ -328,7 +328,7 @@ fn gateway_request_candidate_trace_type_is_owned_by_aether_data_contracts() {
}
let candidate_types =
read_workspace_file("crates/aether-data-contracts/src/repository/candidates/types.rs");
read_workspace_file("crates/aether-data/contracts/src/repository/candidates/types.rs");
for pattern in [
"pub enum RequestCandidateFinalStatus",
"pub struct RequestCandidateTrace",
@@ -366,7 +366,7 @@ fn gateway_decision_trace_type_is_owned_by_aether_data_contracts() {
}
let candidate_types =
read_workspace_file("crates/aether-data-contracts/src/repository/candidates/types.rs");
read_workspace_file("crates/aether-data/contracts/src/repository/candidates/types.rs");
for pattern in [
"pub struct DecisionTraceCandidate",
"pub struct DecisionTrace",
@@ -911,7 +911,8 @@ fn gateway_request_audit_bundle_type_is_owned_by_aether_data() {
);
}
let audit_types = read_workspace_file("crates/aether-data/src/repository/audit.rs");
let audit_types =
read_workspace_file("crates/aether-data/runtime/src/repository/audit/types.rs");
for pattern in [
"pub struct RequestAuditBundle",
"pub trait RequestAuditReader",
@@ -1434,7 +1435,7 @@ fn provider_transport_cache_helpers_live_in_shared_crate() {
"state/mod.rs should import refresh detection from shared provider transport"
);
let transport_cache = read_workspace_file("crates/aether-provider-transport/src/cache.rs");
let transport_cache = read_workspace_file("crates/aether-provider/transport/src/cache.rs");
for pattern in [
"pub struct ProviderTransportSnapshotCacheKey",
"pub fn provider_transport_snapshot_looks_refreshed(",
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,349 @@
use super::{collect_workspace_rust_files, read_workspace_file};
fn assert_manifest_excludes(manifest_path: &str, forbidden: &[&str]) {
let manifest = read_workspace_file(manifest_path);
let violations = forbidden
.iter()
.filter(|dependency| manifest.contains(**dependency))
.copied()
.collect::<Vec<_>>();
assert!(
violations.is_empty(),
"{manifest_path} crosses its dependency tier through: {}",
violations.join(", ")
);
}
#[test]
fn pure_policy_crates_do_not_depend_on_runtime_adapters() {
let pure_manifests = [
"crates/aether-admission-core/Cargo.toml",
"crates/aether-provider/core/Cargo.toml",
"crates/aether-task/core/Cargo.toml",
"crates/aether-usage/core/Cargo.toml",
];
let forbidden = [
"axum",
"sqlx",
"redis",
"reqwest",
"wreq",
"tokio",
"aether-data =",
"aether-gateway",
];
for manifest in pure_manifests {
assert_manifest_excludes(manifest, &forbidden);
}
}
#[test]
fn database_adapters_are_independent_driver_boundaries() {
let adapters = [
(
"crates/aether-data/adapters/postgres/Cargo.toml",
"features = [\"postgres\"",
["features = [\"mysql\"", "features = [\"sqlite\""],
),
(
"crates/aether-data/adapters/mysql/Cargo.toml",
"features = [\"mysql\"",
["features = [\"postgres\"", "features = [\"sqlite\""],
),
(
"crates/aether-data/adapters/sqlite/Cargo.toml",
"features = [\"sqlite\"",
["features = [\"postgres\"", "features = [\"mysql\""],
),
];
for (manifest_path, expected_driver, other_drivers) in adapters {
let manifest = read_workspace_file(manifest_path);
assert!(manifest.contains("aether-data-contracts.workspace = true"));
assert!(manifest.contains(expected_driver));
assert_manifest_excludes(
manifest_path,
&[other_drivers[0], other_drivers[1], "aether-gateway", "axum"],
);
}
}
#[test]
fn data_facade_preserves_legacy_driver_paths_without_owning_driver_code() {
for (path, adapter) in [
(
"crates/aether-data/runtime/src/driver/postgres.rs",
"aether_data_postgres",
),
(
"crates/aether-data/runtime/src/driver/mysql.rs",
"aether_data_mysql",
),
(
"crates/aether-data/runtime/src/driver/sqlite.rs",
"aether_data_sqlite",
),
] {
let source = read_workspace_file(path);
assert!(
source.contains(&format!("pub use {adapter}::*;")),
"{path} should remain a thin compatibility facade"
);
assert!(!source.contains("sqlx::"));
}
}
#[test]
fn gateway_runtime_components_keep_focused_dependency_surfaces() {
assert_manifest_excludes(
"crates/aether-gateway/frontdoor/Cargo.toml",
&[
"aether-data",
"aether-provider-transport",
"aether-gateway-workers",
"sqlx",
"redis",
],
);
assert_manifest_excludes(
"crates/aether-gateway/workers/Cargo.toml",
&[
"axum",
"aether-gateway-frontdoor",
"aether-provider-transport",
],
);
assert_manifest_excludes(
"crates/aether-gateway/execution/Cargo.toml",
&[
"axum",
"sqlx",
"redis",
"aether-data",
"aether-gateway-frontdoor",
"aether-gateway-workers",
],
);
assert_manifest_excludes(
"crates/aether-gateway/control/Cargo.toml",
&[
"sqlx",
"redis",
"reqwest",
"aether-data",
"aether-provider-transport",
"aether-gateway-workers",
],
);
assert_manifest_excludes(
"crates/aether-gateway/tunnel/Cargo.toml",
&[
"axum",
"sqlx",
"redis",
"reqwest",
"wreq",
"aether-data",
"aether-provider-transport",
"aether-gateway-workers",
],
);
assert_manifest_excludes(
"crates/aether-testing/loadtools/Cargo.toml",
&[
"aether-gateway",
"aether-testkit",
"aether-data",
"axum",
"redis",
],
);
}
#[test]
fn tunnel_binary_uses_shared_tunnel_boundary_without_gateway_runtime_dependency() {
let manifest = read_workspace_file("apps/aether-tunnel/Cargo.toml");
let dependencies = manifest
.split_once("[dependencies]")
.expect("tunnel manifest should declare dependencies")
.1
.split("[dev-dependencies]")
.next()
.expect("normal dependency section should exist");
assert!(dependencies.contains("aether-gateway-tunnel.workspace = true"));
assert!(!dependencies.contains("aether-gateway.workspace = true"));
let protocol_facade = read_workspace_file("apps/aether-tunnel/src/tunnel/protocol.rs");
assert!(protocol_facade.contains("aether_gateway_tunnel::protocol::*"));
}
#[test]
fn data_facade_defaults_to_postgres_and_gateway_selects_all_drivers_explicitly() {
let data_manifest = read_workspace_file("crates/aether-data/runtime/Cargo.toml");
assert!(data_manifest.contains("default = [\"postgres\"]"));
for dependency in [
"aether-data-postgres = { workspace = true, optional = true }",
"aether-data-mysql = { workspace = true, optional = true }",
"aether-data-sqlite = { workspace = true, optional = true }",
] {
assert!(
data_manifest.contains(dependency),
"aether-data should keep {dependency} optional"
);
}
assert!(
!data_manifest.contains("features = [\"postgres\", \"mysql\", \"sqlite\"\"]"),
"aether-data must not unconditionally enable every sqlx driver"
);
let gateway_manifest = read_workspace_file("apps/aether-gateway/Cargo.toml");
assert!(gateway_manifest
.contains("aether-data = { workspace = true, features = [\"all-drivers\"] }"));
let data_lib = read_workspace_file("crates/aether-data/runtime/src/lib.rs");
for backend in ["PostgresBackend", "MysqlBackend", "SqliteBackend"] {
assert!(
data_lib.contains(&format!("pub use backend::{backend};")),
"aether-data should expose enabled backends symmetrically at its facade root"
);
}
}
#[test]
fn data_query_helpers_belong_to_adapters_not_the_runtime_facade() {
let data_manifest = read_workspace_file("crates/aether-data/runtime/Cargo.toml");
assert!(
!data_manifest.contains("aether-data-query.workspace = true"),
"aether-data should not keep a direct query-helper dependency after SQL repositories move to adapters"
);
for adapter_manifest in [
"crates/aether-data/adapters/postgres/Cargo.toml",
"crates/aether-data/adapters/mysql/Cargo.toml",
"crates/aether-data/adapters/sqlite/Cargo.toml",
] {
let manifest = read_workspace_file(adapter_manifest);
assert!(
manifest.contains("aether-data-query.workspace = true"),
"{adapter_manifest} should own its query-helper dependency"
);
}
let query_helpers = read_workspace_file("crates/aether-data/query/src/lib.rs");
for dialect in ["Postgres", "MySql", "Sqlite"] {
assert!(
query_helpers.contains(dialect),
"aether-data-query should render the {dialect} dialect"
);
}
}
#[test]
fn sql_adapters_centralize_error_mapping_boilerplate() {
for (adapter, driver) in [
("aether-data-mysql", "mysql"),
("aether-data-sqlite", "sqlite"),
] {
let root = format!("crates/aether-data/adapters/{driver}/src");
let files = collect_workspace_rust_files(&root);
let trait_owners = files
.iter()
.filter(|path| {
std::fs::read_to_string(path)
.expect("adapter source should be readable")
.contains("trait SqlResultExt<T>")
})
.collect::<Vec<_>>();
assert_eq!(
trait_owners.len(),
1,
"{adapter} should have exactly one SqlResultExt owner, found: {trait_owners:?}"
);
assert_eq!(
trait_owners[0].file_name().and_then(|name| name.to_str()),
Some("error.rs"),
"{adapter} should keep SQL error mapping in src/error.rs"
);
let lib = read_workspace_file(&format!("crates/aether-data/adapters/{driver}/src/lib.rs"));
assert!(lib.contains("mod error;"));
}
}
#[test]
fn gateway_tunnel_protocol_path_is_a_thin_compatibility_facade() {
let source = read_workspace_file("apps/aether-gateway/src/tunnel/embedded/protocol.rs");
assert_eq!(
source.trim(),
"pub use aether_gateway_tunnel::embedded::protocol::*;"
);
}
#[test]
fn frontdoor_owns_bounded_request_body_buffering() {
let frontdoor = read_workspace_file("crates/aether-gateway/frontdoor/src/body.rs");
assert!(frontdoor.contains("acquire_many_owned"));
assert!(frontdoor.contains("to_bytes(body, body_limit)"));
assert!(frontdoor.contains("BodyBufferReservation"));
let gateway = read_workspace_file("apps/aether-gateway/src/handlers/proxy/body_buffer.rs");
assert!(gateway.contains("FrontdoorBodyBufferPolicy"));
assert!(!gateway.contains("acquire_many_owned"));
assert!(!gateway.contains("request_body_collection_exceeded_limit"));
}
#[test]
fn benchmark_binaries_are_outside_the_reusable_testkit() {
assert!(
collect_workspace_rust_files("crates/aether-testing/testkit/src/bin").is_empty(),
"aether-testkit must not own benchmark binaries"
);
assert!(
!collect_workspace_rust_files("crates/aether-testing/loadtools/src/bin").is_empty(),
"standalone load tools should live in aether-loadtools"
);
assert!(
!collect_workspace_rust_files("crates/aether-testing/integration/src/bin").is_empty(),
"gateway-backed scenarios should live in aether-integration-tests"
);
}
#[test]
fn testkit_gateway_harness_is_opt_in() {
let testkit_manifest = read_workspace_file("crates/aether-testing/testkit/Cargo.toml");
assert!(
testkit_manifest.contains("default = []"),
"aether-testkit should keep the default feature set dependency-light"
);
assert!(
testkit_manifest
.contains("gateway = [\"dep:aether-gateway\", \"dep:aether-runtime-state\"]"),
"gateway harnesses should be behind the explicit gateway feature"
);
assert!(
testkit_manifest.contains("postgres = [\"dep:aether-data\", \"dep:sqlx\"]"),
"Postgres schema helpers should be behind the explicit postgres feature"
);
assert!(testkit_manifest.contains(
"aether-gateway = { workspace = true, features = [\"testkit\"], optional = true }"
));
let testkit_lib = read_workspace_file("crates/aether-testing/testkit/src/lib.rs");
for module in ["execution_runtime", "gateway", "tunnel"] {
assert!(
testkit_lib.contains(&format!("#[cfg(feature = \"gateway\")]\nmod {module};")),
"aether-testkit::{module} should be feature-gated"
);
}
assert!(
testkit_lib.contains("#[cfg(feature = \"postgres\")]\nmod postgres;"),
"the Postgres helper should be feature-gated"
);
let integration_manifest = read_workspace_file("crates/aether-testing/integration/Cargo.toml");
assert!(integration_manifest
.contains("aether-testkit = { workspace = true, features = [\"gateway\", \"postgres\"] }"));
}
+10 -9
View File
@@ -310,15 +310,15 @@ fn gateway_exposes_request_concurrency_metrics() {
}
async fn gateway_exposes_request_concurrency_metrics_impl() {
let gateway = build_router_with_state(
AppState::new()
.expect("gateway state should build")
.with_request_concurrency_limit(3)
.with_distributed_request_concurrency_gate(memory_runtime_semaphore(
"gateway_requests_distributed",
5,
)),
);
let state = AppState::new()
.expect("gateway state should build")
.with_request_concurrency_limit(3)
.with_distributed_request_concurrency_gate(memory_runtime_semaphore(
"gateway_requests_distributed",
5,
));
assert!(state.prewarm_metric_snapshot().await);
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
@@ -437,6 +437,7 @@ async fn gateway_exposes_fallback_metrics_impl() {
Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS),
GatewayFallbackReason::LocalExecutionPathRequired,
);
assert!(state.prewarm_metric_snapshot().await);
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
@@ -10,7 +10,7 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
use aether_runtime_state::{RedisClientConfig, RuntimeState};
use aether_testkit::ManagedRedisServer;
use aether_test_support::ManagedRedisServer;
use axum::body::to_bytes;
use axum::body::Body;
use axum::routing::{any, get, post};
@@ -1009,6 +1009,26 @@ async fn gateway_handles_admin_stats_provider_performance_locally_with_trusted_a
payload["timeline"][1]["avg_first_byte_time_ms"],
serde_json::Value::Null
);
let without_timeline_response = admin_request(reqwest::Client::new().get(format!(
"{gateway_url}/api/admin/stats/performance/providers?start_date=2024-03-21&end_date=2024-03-21&granularity=hour&limit=2&tz_offset_minutes=0&include_timeline=false"
)))
.send()
.await
.expect("request without timeline should succeed");
assert_eq!(without_timeline_response.status(), StatusCode::OK);
let without_timeline_payload: serde_json::Value = without_timeline_response
.json()
.await
.expect("json body without timeline should parse");
assert_eq!(without_timeline_payload["summary"], payload["summary"]);
assert_eq!(without_timeline_payload["providers"], payload["providers"]);
assert_eq!(
without_timeline_payload["timeline"]
.as_array()
.map(Vec::len),
Some(0)
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
@@ -8240,7 +8240,9 @@ async fn gateway_returns_service_unavailable_for_users_me_management_token_write
#[tokio::test]
async fn gateway_handles_users_me_providers_locally_without_proxying_upstream() {
let now = Utc::now();
let user = sample_auth_user(now);
let mut user = sample_auth_user(now);
user.allowed_providers = Some(vec!["claude".to_string()]);
user.allowed_providers_mode = "specific".to_string();
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
@@ -8300,6 +8302,27 @@ async fn gateway_handles_users_me_providers_locally_without_proxying_upstream()
]),
);
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![user]));
let group = user_repository
.create_user_group(UpsertUserGroupRecord {
name: "OpenAI only".to_string(),
description: None,
priority: 0,
allowed_providers: Some(vec!["openai".to_string()]),
allowed_providers_mode: "specific".to_string(),
allowed_api_formats: None,
allowed_api_formats_mode: "unrestricted".to_string(),
allowed_models: None,
allowed_models_mode: "unrestricted".to_string(),
rate_limit: None,
rate_limit_mode: "system".to_string(),
})
.await
.expect("group should create")
.expect("group should exist");
user_repository
.add_user_to_group(&group.id, "user-auth-1")
.await
.expect("group membership should create");
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_builder(|| {
@@ -9697,13 +9720,13 @@ async fn gateway_handles_users_me_available_models_locally_without_proxying_upst
}
#[tokio::test]
async fn gateway_filters_users_me_available_models_by_group_policy_and_hides_model_mappings() {
async fn gateway_refreshes_users_me_available_models_after_group_assignment() {
let now = Utc::now();
let mut user = sample_auth_user(now);
user.allowed_providers = None;
user.allowed_providers_mode = "unrestricted".to_string();
user.allowed_models = None;
user.allowed_models_mode = "unrestricted".to_string();
// Keep the legacy gpt-5 personal policy from sample_auth_user. Group policies are the
// authority now, so this stale field must not narrow the user catalog.
let access_token = build_test_auth_token(
"access",
serde_json::Map::from_iter([
@@ -9754,31 +9777,25 @@ async fn gateway_filters_users_me_available_models_by_group_policy_and_hides_mod
.await
.expect("group should create")
.expect("group should exist");
user_repository
.add_user_to_group(&group.id, "user-auth-1")
.await
.expect("group membership should create");
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_builder(|| {
let data_state = crate::data::GatewayDataState::with_global_model_reader_for_tests(
global_model_repository,
)
let data_state =
crate::data::GatewayDataState::with_global_model_reader_for_tests(global_model_repository)
.with_user_reader(Arc::clone(&user_repository));
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_auth_sessions_for_tests([sample_auth_session(
"user-auth-1",
"session-users-me-group-models",
"device-users-me-group-models",
"refresh-token-placeholder",
now,
)])
})
.await;
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
.with_auth_sessions_for_tests([sample_auth_session(
"user-auth-1",
"session-users-me-group-models",
"device-users-me-group-models",
"refresh-token-placeholder",
now,
)]);
let mutation_state = state.clone();
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_builder(|| state).await;
let response = reqwest::Client::new()
let client = reqwest::Client::new();
let response = client
.get(format!("{gateway_url}/api/users/me/available-models"))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-users-me-group-models")
@@ -9787,6 +9804,24 @@ async fn gateway_filters_users_me_available_models_by_group_policy_and_hides_mod
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["total"], 2);
mutation_state
.replace_user_groups_for_user("user-auth-1", std::slice::from_ref(&group.id))
.await
.expect("group membership should create");
let response = client
.get(format!("{gateway_url}/api/users/me/available-models"))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-users-me-group-models")
.header("user-agent", "AetherTest/1.0")
.send()
.await
.expect("request should succeed after group assignment");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let models = payload["models"]
@@ -9926,6 +9961,27 @@ async fn gateway_returns_service_unavailable_for_users_me_available_models_witho
.expect("active global model ref should build")]),
);
let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![user]));
let group = user_repository
.create_user_group(UpsertUserGroupRecord {
name: "OpenAI provider only".to_string(),
description: None,
priority: 0,
allowed_providers: Some(vec!["openai".to_string()]),
allowed_providers_mode: "specific".to_string(),
allowed_api_formats: None,
allowed_api_formats_mode: "unrestricted".to_string(),
allowed_models: None,
allowed_models_mode: "unrestricted".to_string(),
rate_limit: None,
rate_limit_mode: "system".to_string(),
})
.await
.expect("group should create")
.expect("group should exist");
user_repository
.add_user_to_group(&group.id, "user-auth-1")
.await
.expect("group membership should create");
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
start_auth_gateway_with_builder(|| {
let data_state = crate::data::GatewayDataState::with_global_model_reader_for_tests(
@@ -159,6 +159,15 @@ fn relay_header_timeout(meta: &protocol::RequestMeta) -> Duration {
Duration::from_millis(resolve_tunnel_request_timeouts(meta).first_byte_ms)
}
fn is_rollout_probe_request(headers: &HeaderMap, forwarded_by_gateway: bool) -> bool {
forwarded_by_gateway
&& headers
.get(crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_HEADER)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.is_some_and(|value| value == crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_VALUE)
}
pub async fn relay_request(
Path(node_id): Path<String>,
State(state): State<AppState>,
@@ -171,6 +180,7 @@ pub async fn relay_request(
.and_then(|value| value.to_str().ok())
.map(str::trim)
.is_some_and(|value| !value.is_empty());
let rollout_probe = is_rollout_probe_request(request.headers(), forwarded_by_gateway);
if !addr.ip().is_loopback() && !forwarded_by_gateway {
return tunnel_error_response(
StatusCode::FORBIDDEN,
@@ -348,12 +358,16 @@ pub async fn relay_request(
);
}
};
if let Err(error) = record_proxy_upgrade_traffic_success(state.data.as_ref(), &node_id).await {
warn!(
node_id = %node_id,
error = %error,
"failed to record proxy upgrade traffic confirmation"
);
if !rollout_probe {
if let Err(error) =
record_proxy_upgrade_traffic_success(state.data.as_ref(), &node_id).await
{
warn!(
node_id = %node_id,
error = %error,
"failed to record proxy upgrade traffic confirmation"
);
}
}
let Some(mut body_rx) = stream.take_body_receiver() else {
@@ -462,8 +476,8 @@ mod tests {
use super::super::hub::ProxyConn;
use super::super::{protocol, AppState, ConnConfig, ControlPlaneClient};
use super::{
relay_header_timeout, relay_request, Body, Request, SocketAddr, StatusCode,
TUNNEL_ERROR_HEADER,
is_rollout_probe_request, relay_header_timeout, relay_request, Body, HeaderMap, Request,
SocketAddr, StatusCode, TUNNEL_ERROR_HEADER,
};
use crate::data::GatewayDataState;
use crate::maintenance::start_proxy_upgrade_rollout;
@@ -482,6 +496,20 @@ mod tests {
use std::time::Duration;
use tokio::sync::watch;
#[test]
fn rollout_probe_marker_is_only_trusted_from_a_forwarding_gateway() {
let mut headers = HeaderMap::new();
headers.insert(
crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_HEADER,
crate::tunnel::TUNNEL_RELAY_ROLLOUT_PROBE_VALUE
.parse()
.expect("probe marker should be a valid header"),
);
assert!(!is_rollout_probe_request(&headers, false));
assert!(is_rollout_probe_request(&headers, true));
}
fn test_app_state() -> AppState {
AppState::new(
ControlPlaneClient::disabled(),
+3 -119
View File
@@ -6,6 +6,9 @@ mod proxy_conn;
use std::sync::Arc;
use aether_gateway_tunnel::{
resolve_proxy_max_streams, resolve_proxy_node_name, resolve_proxy_protocol_version,
};
use aether_runtime::{
hold_admission_permit_until, prometheus_response, service_up_sample, AdmissionPermit,
ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot, MetricKind, MetricLabel, MetricSample,
@@ -17,7 +20,6 @@ use axum::http::HeaderMap;
use axum::response::{IntoResponse, Json};
use axum::routing::{get, post};
use axum::Router;
use base64::Engine as _;
use dashmap::DashMap;
use tracing::warn;
@@ -335,121 +337,3 @@ pub async fn ws_proxy(
})
.into_response()
}
fn resolve_proxy_max_streams(headers: &HeaderMap, fallback: usize) -> usize {
headers
.get("x-tunnel-max-streams")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<usize>().ok())
.unwrap_or(fallback)
.clamp(1, 2048)
}
fn resolve_proxy_node_name(headers: &HeaderMap, node_id: &str) -> String {
if let Some(decoded) = headers
.get(aether_contracts::tunnel::TUNNEL_NODE_NAME_B64_HEADER)
.and_then(|value| value.to_str().ok())
.and_then(|value| {
base64::engine::general_purpose::URL_SAFE_NO_PAD
.decode(value.trim())
.ok()
})
.and_then(|bytes| String::from_utf8(bytes).ok())
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty() && value.chars().count() <= 100)
{
return decoded;
}
headers
.get("x-node-name")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(node_id)
.to_string()
}
fn resolve_proxy_protocol_version(headers: &HeaderMap) -> u8 {
headers
.get(aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER)
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u8>().ok())
.filter(|value| *value >= 1)
.map(|value| value.min(aether_contracts::tunnel::CURRENT_TUNNEL_PROTOCOL_VERSION))
.unwrap_or(1)
}
#[cfg(test)]
mod tests {
use axum::http::{HeaderMap, HeaderValue};
use base64::Engine as _;
use super::{
resolve_proxy_max_streams, resolve_proxy_node_name, resolve_proxy_protocol_version,
};
#[test]
fn proxy_max_streams_honors_small_advertised_capacity() {
let mut headers = HeaderMap::new();
headers.insert("x-tunnel-max-streams", HeaderValue::from_static("8"));
assert_eq!(resolve_proxy_max_streams(&headers, 128), 8);
}
#[test]
fn proxy_max_streams_caps_unreasonably_large_capacity() {
let mut headers = HeaderMap::new();
headers.insert("x-tunnel-max-streams", HeaderValue::from_static("9999"));
assert_eq!(resolve_proxy_max_streams(&headers, 128), 2048);
}
#[test]
fn proxy_protocol_version_defaults_to_v1_when_header_missing() {
let headers = HeaderMap::new();
assert_eq!(resolve_proxy_protocol_version(&headers), 1);
}
#[test]
fn proxy_protocol_version_reads_advertised_version() {
let mut headers = HeaderMap::new();
headers.insert(
aether_contracts::tunnel::TUNNEL_PROTOCOL_VERSION_HEADER,
HeaderValue::from_static("2"),
);
assert_eq!(resolve_proxy_protocol_version(&headers), 2);
}
#[test]
fn proxy_node_name_reads_legacy_ascii_header() {
let mut headers = HeaderMap::new();
headers.insert("x-node-name", HeaderValue::from_static("edge-1"));
assert_eq!(resolve_proxy_node_name(&headers, "node-1"), "edge-1");
}
#[test]
fn proxy_node_name_decodes_base64_header() {
let mut headers = HeaderMap::new();
let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode("日本节点");
headers.insert(
aether_contracts::tunnel::TUNNEL_NODE_NAME_B64_HEADER,
HeaderValue::from_str(&encoded).expect("encoded header value should parse"),
);
assert_eq!(resolve_proxy_node_name(&headers, "node-1"), "日本节点");
}
#[test]
fn proxy_node_name_falls_back_to_node_id_for_invalid_base64() {
let mut headers = HeaderMap::new();
headers.insert(
aether_contracts::tunnel::TUNNEL_NODE_NAME_B64_HEADER,
HeaderValue::from_static("not valid"),
);
assert_eq!(resolve_proxy_node_name(&headers, "node-1"), "node-1");
}
}
@@ -1,23 +1 @@
use bytes::Bytes;
pub use aether_contracts::tunnel::{
decode_payload, encode_connection_close, encode_frame, encode_goaway, encode_goaway_v3,
encode_hello, encode_load_report, encode_ping, encode_pong, encode_reset_stream,
encode_settings, encode_stream_error, encode_window_update, frame_payload_by_header,
ConnectionClosePayload, FrameHeader, GoAwayPayload, HelloPayload, LoadReportPayload,
RequestMeta, ResetStreamPayload, ResponseMeta, SettingsPayload, WindowUpdatePayload,
CONNECTION_CLOSE, FLAG_END_STREAM, FLAG_GZIP_COMPRESSED, GOAWAY, HEADER_SIZE, HEARTBEAT_ACK,
HEARTBEAT_DATA, HELLO, LOAD_REPORT, PING, PONG, REQUEST_BODY, REQUEST_HEADERS, RESET_STREAM,
RESPONSE_BODY, RESPONSE_HEADERS, SETTINGS, STREAM_END, STREAM_ERROR, WINDOW_UPDATE,
};
pub fn compress_payload(payload: &[u8]) -> Result<(Vec<u8>, u8), std::io::Error> {
let (compressed, flags) =
aether_contracts::tunnel::compress_payload(Bytes::copy_from_slice(payload));
Ok((compressed.to_vec(), flags))
}
pub fn raw_payload(payload: &[u8]) -> (Vec<u8>, u8) {
let (payload, flags) = aether_contracts::tunnel::raw_payload(Bytes::copy_from_slice(payload));
(payload.to_vec(), flags)
}
pub use aether_gateway_tunnel::embedded::protocol::*;
+205 -61
View File
@@ -14,6 +14,7 @@ use aether_contracts::tunnel::{
use aether_data::repository::proxy_nodes::{
ProxyNodeHeartbeatMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode,
};
use aether_gateway_tunnel::EmbeddedTunnelDefaults;
use aether_runtime::MetricSample;
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
use async_stream::stream;
@@ -36,6 +37,11 @@ use super::error::GatewayError;
use super::headers::{extract_or_generate_trace_id, should_skip_request_header};
use super::AppState;
pub(crate) use aether_gateway_tunnel::{
is_tunnel_heartbeat_path, is_tunnel_node_status_path, TunnelAttachmentRecord,
DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES, DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES, PROXY_TUNNEL_PATH,
TUNNEL_HEARTBEAT_PATH, TUNNEL_NODE_STATUS_PATH, TUNNEL_RELAY_PATH_PATTERN, TUNNEL_ROUTE_FAMILY,
};
pub(crate) use embedded::DirectRelayResponse;
pub(crate) use embedded::ProxyConn as TunnelProxyConn;
pub use embedded::{
@@ -44,24 +50,14 @@ pub use embedded::{
ControlPlaneClient as TunnelControlPlaneClient,
};
pub(crate) const PROXY_TUNNEL_PATH: &str = "/api/internal/proxy-tunnel";
pub(crate) const TUNNEL_HEARTBEAT_PATH: &str = "/api/internal/tunnel/heartbeat";
pub(crate) const TUNNEL_NODE_STATUS_PATH: &str = "/api/internal/tunnel/node-status";
pub(crate) const TUNNEL_RELAY_PATH_PATTERN: &str = "/api/internal/tunnel/relay/{node_id}";
pub(crate) const TUNNEL_ROUTE_FAMILY: &str = "tunnel_manage";
const DEFAULT_PROXY_IDLE_TIMEOUT_MS: u64 = 0;
const DEFAULT_PING_INTERVAL_MS: u64 = 15_000;
const DEFAULT_MAX_STREAMS: usize = 2048;
const DEFAULT_OUTBOUND_QUEUE_CAPACITY: usize = 512;
const DEFAULT_ATTACHMENT_TTL_SECS: u64 = 90;
const DEFAULT_OWNER_RELAY_BODY_LIMIT_BYTES: usize = 5_242_880;
const DEFAULT_TUNNEL_PROBE_BODY_LIMIT_BYTES: usize = 64 * 1024;
const TUNNEL_ATTACHMENT_KEY_PREFIX: &str = "tunnel.attachments.";
const TUNNEL_ATTACHMENT_REDIS_KEY_PREFIX: &str = "tunnel:attachments:";
const TUNNEL_INSTANCE_ID_ENV: &str = "AETHER_GATEWAY_INSTANCE_ID";
const TUNNEL_RELAY_BASE_URL_ENV: &str = "AETHER_TUNNEL_RELAY_BASE_URL";
const TUNNEL_ATTACHMENT_TTL_ENV: &str = "AETHER_TUNNEL_ATTACHMENT_TTL_SECS";
pub(crate) const TUNNEL_RELAY_ROLLOUT_PROBE_HEADER: &str = "x-aether-tunnel-rollout-probe";
pub(crate) const TUNNEL_RELAY_ROLLOUT_PROBE_VALUE: &str = "1";
pub(crate) async fn send_owner_forward_request(
request: reqwest::RequestBuilder,
@@ -118,14 +114,6 @@ pub(crate) struct TunnelInstanceIdentity {
attachment_ttl_secs: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct TunnelAttachmentRecord {
pub(crate) gateway_instance_id: String,
pub(crate) relay_base_url: String,
pub(crate) conn_count: usize,
pub(crate) observed_at_unix_secs: u64,
}
#[derive(Debug, Clone)]
pub(crate) struct TunnelAttachmentDirectory {
identity: Arc<TunnelInstanceIdentity>,
@@ -263,11 +251,7 @@ impl TunnelAttachmentDirectory {
let Some(record) = self.read_attachment_record(data, node_id).await? else {
return Ok(None);
};
let is_expired = record
.observed_at_unix_secs
.saturating_add(self.identity.attachment_ttl_secs)
< current_unix_secs();
if is_expired || record.relay_base_url.trim().is_empty() {
if !record.is_routable(current_unix_secs(), self.identity.attachment_ttl_secs) {
return Ok(None);
}
Ok(Some(record))
@@ -281,7 +265,7 @@ impl TunnelAttachmentDirectory {
let Some(record) = self.read_attachment_record(data, node_id).await? else {
return Ok(());
};
if record.gateway_instance_id == self.identity.instance_id {
if record.is_owned_by(&self.identity.instance_id) {
self.delete_attachment_record(data, node_id).await?;
}
Ok(())
@@ -462,15 +446,16 @@ impl EmbeddedTunnelState {
data: Arc<GatewayDataState>,
attachment_directory: TunnelAttachmentDirectory,
) -> Self {
let defaults = EmbeddedTunnelDefaults::default();
Self {
inner: TunnelAppState::new(
build_embedded_control_plane(Arc::clone(&data), attachment_directory.clone()),
ConnConfig {
ping_interval: Duration::from_millis(DEFAULT_PING_INTERVAL_MS),
idle_timeout: Duration::from_millis(DEFAULT_PROXY_IDLE_TIMEOUT_MS),
outbound_queue_capacity: DEFAULT_OUTBOUND_QUEUE_CAPACITY,
ping_interval: defaults.ping_interval,
idle_timeout: defaults.proxy_idle_timeout,
outbound_queue_capacity: defaults.outbound_queue_capacity,
},
DEFAULT_MAX_STREAMS,
defaults.max_streams,
)
.with_data(data),
attachment_directory,
@@ -531,6 +516,53 @@ impl EmbeddedTunnelState {
.status)
}
pub(crate) async fn probe_node_url_routed(
&self,
state: &AppState,
node_id: &str,
url: &str,
timeout_secs: u64,
) -> Result<u16, String> {
if self.has_local_proxy(node_id) {
return self.probe_node_url(node_id, url, timeout_secs).await;
}
let Some(owner) = self
.lookup_attachment_owner(state.data.as_ref(), node_id)
.await?
else {
return self.probe_node_url(node_id, url, timeout_secs).await;
};
if owner.gateway_instance_id == self.local_instance_id() {
self.clear_local_attachment_if_stale(state.data.as_ref(), node_id)
.await?;
return self.probe_node_url(node_id, url, timeout_secs).await;
}
let timeout_secs = timeout_secs.clamp(5, 60);
let owner_url = build_owner_relay_url(&owner.relay_base_url, node_id)
.map_err(|error| format!("invalid owner tunnel probe URL: {error:?}"))?;
let payload = encode_tunnel_relay_envelope(&build_tunnel_probe_meta(url, timeout_secs))?;
let response = state
.owner_forward_client
.post(owner_url)
.header(TUNNEL_RELAY_FORWARDED_BY_HEADER, self.local_instance_id())
.header(
TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
owner.gateway_instance_id.as_str(),
)
.header(
TUNNEL_RELAY_ROLLOUT_PROBE_HEADER,
TUNNEL_RELAY_ROLLOUT_PROBE_VALUE,
)
.timeout(Duration::from_secs(timeout_secs))
.body(payload)
.send()
.await
.map_err(|error| format!("owner tunnel probe failed: {error}"))?;
Ok(response.status().as_u16())
}
pub(crate) async fn probe_node_url_with_response(
&self,
node_id: &str,
@@ -538,21 +570,7 @@ impl EmbeddedTunnelState {
timeout_secs: u64,
) -> Result<TunnelProbeResponse, String> {
let timeout_secs = timeout_secs.clamp(5, 60);
let meta = tunnel_protocol::RequestMeta {
provider_id: None,
endpoint_id: None,
key_id: None,
method: "GET".to_string(),
url: url.trim().to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: timeout_secs,
follow_redirects: Some(false),
http1_only: false,
transport_profile: None,
};
let meta = build_tunnel_probe_meta(url, timeout_secs);
let stream = self.inner.hub.open_local_stream(node_id, &meta).await?;
let stream_id = stream.id;
let result = async {
@@ -623,6 +641,35 @@ impl EmbeddedTunnelState {
}
}
fn build_tunnel_probe_meta(url: &str, timeout_secs: u64) -> tunnel_protocol::RequestMeta {
tunnel_protocol::RequestMeta {
provider_id: None,
endpoint_id: None,
key_id: None,
method: "GET".to_string(),
url: url.trim().to_string(),
headers: HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: timeout_secs,
follow_redirects: Some(false),
http1_only: false,
transport_profile: None,
}
}
fn encode_tunnel_relay_envelope(meta: &tunnel_protocol::RequestMeta) -> Result<Vec<u8>, String> {
let meta = serde_json::to_vec(meta)
.map_err(|error| format!("failed to encode tunnel probe metadata: {error}"))?;
let meta_len = u32::try_from(meta.len())
.map_err(|_| "tunnel probe metadata exceeds relay envelope limit".to_string())?;
let mut payload = Vec::with_capacity(4usize.saturating_add(meta.len()));
payload.extend_from_slice(&meta_len.to_be_bytes());
payload.extend_from_slice(&meta);
Ok(payload)
}
impl Default for EmbeddedTunnelState {
fn default() -> Self {
Self::new()
@@ -631,11 +678,15 @@ impl Default for EmbeddedTunnelState {
impl fmt::Debug for EmbeddedTunnelState {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let defaults = EmbeddedTunnelDefaults::default();
f.debug_struct("EmbeddedTunnelState")
.field("proxy_idle_timeout_ms", &DEFAULT_PROXY_IDLE_TIMEOUT_MS)
.field("ping_interval_ms", &DEFAULT_PING_INTERVAL_MS)
.field("max_streams", &DEFAULT_MAX_STREAMS)
.field("outbound_queue_capacity", &DEFAULT_OUTBOUND_QUEUE_CAPACITY)
.field(
"proxy_idle_timeout_ms",
&defaults.proxy_idle_timeout.as_millis(),
)
.field("ping_interval_ms", &defaults.ping_interval.as_millis())
.field("max_streams", &defaults.max_streams)
.field("outbound_queue_capacity", &defaults.outbound_queue_capacity)
.field(
"instance_id",
&self.attachment_directory.local_instance_id(),
@@ -714,14 +765,6 @@ pub(crate) async fn relay_request(
.into_response())
}
pub(crate) fn is_tunnel_heartbeat_path(path: &str) -> bool {
path == TUNNEL_HEARTBEAT_PATH
}
pub(crate) fn is_tunnel_node_status_path(path: &str) -> bool {
path == TUNNEL_NODE_STATUS_PATH
}
fn build_embedded_control_plane(
data: Arc<GatewayDataState>,
attachment_directory: TunnelAttachmentDirectory,
@@ -1134,17 +1177,25 @@ fn parse_embedded_tunnel_heartbeat_request(
#[cfg(test)]
mod tests {
use super::{
apply_embedded_tunnel_heartbeat, apply_embedded_tunnel_node_status, current_unix_secs,
prepare_owner_relay_request_body, tunnel_attachment_key, GatewayDataState,
apply_embedded_tunnel_heartbeat, apply_embedded_tunnel_node_status,
build_tunnel_probe_meta, current_unix_secs, encode_tunnel_relay_envelope,
prepare_owner_relay_request_body, tunnel_attachment_key, AppState, GatewayDataState,
TunnelAttachmentDirectory, TunnelAttachmentRecord,
};
use aether_contracts::tunnel::{
try_decode_tunnel_relay_request_meta, TUNNEL_RELAY_FORWARDED_BY_HEADER,
TUNNEL_RELAY_OWNER_INSTANCE_HEADER,
};
use aether_data::repository::proxy_nodes::{
InMemoryProxyNodeRepository, ProxyNodeReadRepository, StoredProxyNode,
};
use axum::body::Body;
use axum::body::{Body, Bytes};
use axum::http::{HeaderMap, StatusCode};
use axum::routing::post;
use axum::Router;
use serde_json::json;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::sync::{Arc, Mutex};
fn sample_proxy_node(node_id: &str) -> StoredProxyNode {
StoredProxyNode::new(
@@ -1183,6 +1234,99 @@ mod tests {
)
}
#[test]
fn routed_tunnel_probe_builds_a_valid_owner_relay_envelope() {
let meta = build_tunnel_probe_meta("https://probe.example/health", 7);
let envelope = encode_tunnel_relay_envelope(&meta).expect("probe should encode");
let (decoded, body_offset) = try_decode_tunnel_relay_request_meta(&envelope)
.expect("probe envelope should decode")
.expect("probe envelope should contain complete metadata");
assert_eq!(decoded.method, "GET");
assert_eq!(decoded.url, "https://probe.example/health");
assert_eq!(decoded.timeout, 7);
assert_eq!(body_offset, envelope.len());
}
#[tokio::test]
async fn routed_tunnel_probe_forwards_to_the_attachment_owner() {
let captured = Arc::new(Mutex::new(None::<(HeaderMap, Bytes)>));
let captured_for_route = Arc::clone(&captured);
let app = Router::new().route(
"/api/internal/tunnel/relay/{node_id}",
post(move |headers: HeaderMap, body: Bytes| {
let captured = Arc::clone(&captured_for_route);
async move {
*captured.lock().expect("capture lock") = Some((headers, body));
StatusCode::NO_CONTENT
}
}),
);
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("owner listener should bind");
let owner_base_url = format!("http://{}", listener.local_addr().expect("owner address"));
let server = tokio::spawn(async move {
axum::serve(listener, app)
.await
.expect("owner server should run");
});
let owner = TunnelAttachmentRecord {
gateway_instance_id: "gateway-b".to_string(),
relay_base_url: owner_base_url,
conn_count: 1,
observed_at_unix_secs: current_unix_secs(),
};
let data = GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
tunnel_attachment_key("node-remote"),
serde_json::to_value(owner).expect("owner record should serialize"),
)]);
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(data)
.with_tunnel_identity("gateway-a", Some("http://gateway-a.internal"));
let status = state
.tunnel
.probe_node_url_routed(&state, "node-remote", "https://probe.example/health", 5)
.await
.expect("remote owner probe should succeed");
assert_eq!(status, StatusCode::NO_CONTENT.as_u16());
let (headers, body) = captured
.lock()
.expect("capture lock")
.take()
.expect("owner should receive the probe");
assert_eq!(
headers
.get(TUNNEL_RELAY_FORWARDED_BY_HEADER)
.and_then(|value| value.to_str().ok()),
Some("gateway-a")
);
assert_eq!(
headers
.get(TUNNEL_RELAY_OWNER_INSTANCE_HEADER)
.and_then(|value| value.to_str().ok()),
Some("gateway-b")
);
assert_eq!(
headers
.get(super::TUNNEL_RELAY_ROLLOUT_PROBE_HEADER)
.and_then(|value| value.to_str().ok()),
Some(super::TUNNEL_RELAY_ROLLOUT_PROBE_VALUE)
);
let (meta, body_offset) = try_decode_tunnel_relay_request_meta(&body)
.expect("owner probe envelope should decode")
.expect("owner probe metadata should be complete");
assert_eq!(meta.url, "https://probe.example/health");
assert_eq!(body_offset, body.len());
server.abort();
let _ = server.await;
}
#[tokio::test]
async fn owner_relay_body_preparation_rejects_invalid_metadata() {
let mut envelope = Vec::new();
@@ -55,11 +55,16 @@ async fn resolve_wallet_auth_gate_with_cache(
None => WalletAccessDecision::wallet_unavailable(None),
};
if !auth_snapshot.api_key_is_standalone {
if let Some(quota) = state
.find_user_daily_quota_availability(&auth_snapshot.user_id)
.await?
.filter(|quota| quota.has_active_daily_quota)
{
let quota = if use_cache {
state
.find_user_daily_quota_availability_for_auth(&auth_snapshot.user_id)
.await?
} else {
state
.find_user_daily_quota_availability_for_auth_uncached(&auth_snapshot.user_id)
.await?
};
if let Some(quota) = quota.filter(|quota| quota.has_active_daily_quota) {
let has_remaining_quota = quota.remaining_usd > DAILY_QUOTA_EPSILON_USD;
if decision.failure == Some(WalletAccessFailure::BalanceDenied) && has_remaining_quota {
return Ok(Some(WalletAccessDecision::allowed(Some(
@@ -106,6 +111,7 @@ fn map_wallet_snapshot(snapshot: &StoredWalletSnapshot) -> WalletSnapshot {
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
@@ -113,6 +119,7 @@ mod tests {
BillingReadRepository, StoredBillingModelContext, UserDailyQuotaAvailabilityRecord,
};
use aether_data_contracts::DataLayerError;
use aether_runtime::ConcurrencyGate;
use aether_wallet::{WalletAccessFailure, WalletLimitMode, WalletSnapshot, WalletStatus};
use async_trait::async_trait;
@@ -257,6 +264,41 @@ mod tests {
assert_eq!(decision.remaining, Some(4.0));
}
#[tokio::test]
async fn disabled_auth_capacity_cache_still_gates_wallet_reads() {
let mut state = state_with_wallet_and_quota(empty_user_wallet(), None);
let mut guard_config = (*state.frontdoor_runtime_guards).clone();
guard_config.auth_capacity_cache_ttl = Duration::ZERO;
state = state.with_frontdoor_runtime_guard_config_for_tests(guard_config);
state.auth_snapshot_load_gate =
Some(Arc::new(ConcurrencyGate::new("test_auth_wallet_load", 1)));
let held = state
.acquire_auth_snapshot_load_gate()
.await
.expect("auth gate acquisition should succeed")
.expect("auth gate should be configured");
let blocked = tokio::time::timeout(
Duration::from_millis(25),
state.read_wallet_snapshot_for_auth("user-1", "api-key-1", false),
)
.await;
assert!(
blocked.is_err(),
"zero-TTL wallet reads must wait for the auth DB gate"
);
drop(held);
let wallet = tokio::time::timeout(
Duration::from_secs(1),
state.read_wallet_snapshot_for_auth("user-1", "api-key-1", false),
)
.await
.expect("wallet read should resume after releasing the auth gate")
.expect("wallet read should succeed");
assert!(wallet.is_some());
}
#[tokio::test]
async fn admin_wallet_recharge_invalidates_cached_auth_capacity_state() {
let wallet = empty_user_wallet();
+27 -19
View File
@@ -1,9 +1,9 @@
use std::sync::Arc;
use std::time::Duration;
use tracing::warn;
use crate::data::GatewayDataState;
use crate::AppState;
const QUOTA_RESET_INTERVAL: Duration = Duration::from_secs(60 * 60);
@@ -18,26 +18,31 @@ pub(crate) async fn reset_due_provider_quotas_once(
}
pub(crate) fn spawn_provider_quota_reset_worker(
data: Arc<GatewayDataState>,
app: AppState,
) -> Option<tokio::task::JoinHandle<()>> {
if !data.has_provider_quota_writer() {
if !app.data.has_provider_quota_writer() {
return None;
}
Some(tokio::spawn(async move {
if let Err(err) = reset_due_provider_quotas_once(&data).await {
warn!(error = %err, "gateway provider quota reset startup failed");
}
let mut interval = tokio::time::interval(QUOTA_RESET_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
loop {
interval.tick().await;
Some(crate::task_runtime::spawn_singleton_worker(
app,
crate::task_runtime::TASK_KEY_PROVIDER_QUOTA_RESET,
|app| async move {
let data = app.data;
if let Err(err) = reset_due_provider_quotas_once(&data).await {
warn!(error = %err, "gateway provider quota reset tick failed");
warn!(error = %err, "gateway provider quota reset startup failed");
}
}
}))
let mut interval = tokio::time::interval(QUOTA_RESET_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
loop {
interval.tick().await;
if let Err(err) = reset_due_provider_quotas_once(&data).await {
warn!(error = %err, "gateway provider quota reset tick failed");
}
}
},
))
}
#[cfg(test)]
@@ -52,6 +57,7 @@ mod tests {
use super::{reset_due_provider_quotas_once, spawn_provider_quota_reset_worker};
use crate::data::GatewayDataState;
use crate::AppState;
#[tokio::test]
async fn resets_due_provider_quotas_from_runtime() {
@@ -98,10 +104,12 @@ mod tests {
)
.expect("quota should build"),
]));
let data = Arc::new(GatewayDataState::with_provider_quota_repository_for_tests(
repository.clone(),
));
let handle = spawn_provider_quota_reset_worker(data).expect("worker should spawn");
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::with_provider_quota_repository_for_tests(
repository.clone(),
));
let handle = spawn_provider_quota_reset_worker(state).expect("worker should spawn");
let stored = tokio::time::timeout(Duration::from_secs(1), async {
loop {