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

fix: harden frontdoor and usage ingestion

Browse Source
This commit is contained in:
fawney19
2026-05-22 23:57:38 +08:00
parent 6ef6cbade2
commit d18b13a91a
25 changed files with 1363 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
apps/aether-gateway/src/constants.rs
View File

@@ -26,6 +26,8 @@ pub(crate) const EXECUTION_PATH_EXECUTION_RUNTIME_STREAM: &str = "execution_runt
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_SYNC: &str = "control_execute_sync";
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_STREAM: &str = "control_execute_stream";
pub(crate) const EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS: &str = "local_execution_runtime_miss";
pub(crate) const EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT: &str =
"local_execution_planning_timeout";
pub(crate) const EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED: &str =
"local_api_key_concurrency_limited";
pub(crate) const API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS: u64 = 150;

355
apps/aether-gateway/src/data/state/candidate_cache.rs
View File

@@ -1,4 +1,4 @@
use std::collections::HashSet;
use std::collections::HashMap;
use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
@@ -13,15 +13,26 @@ use aether_data_contracts::repository::candidate_selection::{
};
use async_trait::async_trait;
use tokio::sync::Notify;
use tokio::time::timeout;
use tracing::warn;
const CANDIDATE_SELECTION_CACHE_TTL: Duration = Duration::from_secs(5);
const CANDIDATE_SELECTION_CACHE_MAX_ENTRIES: usize = 4096;
#[cfg(not(test))]
const CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT: Duration = Duration::from_secs(10);
#[cfg(test)]
const CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT: Duration = Duration::from_millis(50);
#[cfg(not(test))]
const CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_secs(10);
#[cfg(test)]
const CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_millis(50);
pub(super) struct CachedMinimalCandidateSelectionReadRepository {
inner: Arc<dyn MinimalCandidateSelectionReadRepository>,
entries: ExpiringMap<CandidateSelectionCacheKey, Vec<StoredMinimalCandidateSelectionRow>>,
inflight: Mutex<HashSet<CandidateSelectionCacheKey>>,
inflight: Mutex<HashMap<CandidateSelectionCacheKey, u64>>,
inflight_notify: Notify,
next_inflight_token: AtomicU64,
epoch: AtomicU64,
}
@@ -30,8 +41,9 @@ impl CachedMinimalCandidateSelectionReadRepository {
Self {
inner,
entries: ExpiringMap::new(),
inflight: Mutex::new(HashSet::new()),
inflight: Mutex::new(HashMap::new()),
inflight_notify: Notify::new(),
next_inflight_token: AtomicU64::new(1),
epoch: AtomicU64::new(0),
}
}
@@ -52,67 +64,169 @@ impl CachedMinimalCandidateSelectionReadRepository {
loop {
let notified = self.inflight_notify.notified();
match self.register_inflight(&key) {
InflightRegistration::Bypass => return load().await,
InflightRegistration::Bypass => {
return load_candidate_selection_rows_with_timeout(&key, load()).await;
}
InflightRegistration::Follower => {
notified.await;
if timeout(CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT, notified)
.await
.is_err()
{
self.expire_inflight(&key);
}
if let Some(rows) = self.entries.get_fresh(&key, CANDIDATE_SELECTION_CACHE_TTL)
{
return Ok(rows);
}
continue;
}
InflightRegistration::Leader => {}
}
let load_epoch = self.epoch.load(Ordering::Acquire);
let result = load().await;
if let Ok(rows) = &result {
if load_epoch == self.epoch.load(Ordering::Acquire) {
self.entries.insert(
key.clone(),
rows.clone(),
CANDIDATE_SELECTION_CACHE_TTL,
CANDIDATE_SELECTION_CACHE_MAX_ENTRIES,
);
InflightRegistration::Leader(token) => {
let mut guard = InflightGuard::new(self, key.clone(), token);
let load_epoch = self.epoch.load(Ordering::Acquire);
let result = load_candidate_selection_rows_with_timeout(&key, load()).await;
if let Ok(rows) = &result {
if load_epoch == self.epoch.load(Ordering::Acquire) {
self.entries.insert(
key.clone(),
rows.clone(),
CANDIDATE_SELECTION_CACHE_TTL,
CANDIDATE_SELECTION_CACHE_MAX_ENTRIES,
);
}
}
guard.finish();
return result;
}
}
self.finish_inflight(&key);
return result;
}
}
fn register_inflight(&self, key: &CandidateSelectionCacheKey) -> InflightRegistration {
match self.inflight.lock() {
Ok(mut inflight) => {
if inflight.insert(key.clone()) {
InflightRegistration::Leader
} else {
InflightRegistration::Follower
if inflight.contains_key(key) {
return InflightRegistration::Follower;
}
let token = self.next_inflight_token.fetch_add(1, Ordering::AcqRel);
inflight.insert(key.clone(), token);
InflightRegistration::Leader(token)
}
Err(_) => InflightRegistration::Bypass,
}
}
fn finish_inflight(&self, key: &CandidateSelectionCacheKey) {
fn finish_inflight(&self, key: &CandidateSelectionCacheKey, token: u64) {
let mut removed = false;
if let Ok(mut inflight) = self.inflight.lock() {
inflight.remove(key);
if inflight.get(key).copied() == Some(token) {
inflight.remove(key);
removed = true;
}
}
if removed {
self.inflight_notify.notify_waiters();
}
}
fn expire_inflight(&self, key: &CandidateSelectionCacheKey) {
let mut removed = false;
if let Ok(mut inflight) = self.inflight.lock() {
removed = inflight.remove(key).is_some();
}
if removed {
warn!(
event_name = "candidate_selection_cache_inflight_expired",
log_type = "ops",
cache_key = ?key,
wait_timeout_ms = CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT.as_millis() as u64,
"gateway candidate selection cache expired stale inflight load"
);
self.inflight_notify.notify_waiters();
}
self.inflight_notify.notify_waiters();
}
fn clear(&self) {
self.epoch.fetch_add(1, Ordering::AcqRel);
self.entries.clear();
let mut cleared_inflight = false;
if let Ok(mut inflight) = self.inflight.lock() {
cleared_inflight = !inflight.is_empty();
inflight.clear();
}
if cleared_inflight {
warn!(
event_name = "candidate_selection_cache_inflight_cleared",
log_type = "ops",
"gateway candidate selection cache cleared in-flight loads"
);
self.inflight_notify.notify_waiters();
}
}
}
enum InflightRegistration {
Leader,
Leader(u64),
Follower,
Bypass,
}
struct InflightGuard<'a> {
cache: &'a CachedMinimalCandidateSelectionReadRepository,
key: Option<CandidateSelectionCacheKey>,
token: u64,
}
impl<'a> InflightGuard<'a> {
fn new(
cache: &'a CachedMinimalCandidateSelectionReadRepository,
key: CandidateSelectionCacheKey,
token: u64,
) -> Self {
Self {
cache,
key: Some(key),
token,
}
}
fn finish(&mut self) {
if let Some(key) = self.key.take() {
self.cache.finish_inflight(&key, self.token);
}
}
}
impl Drop for InflightGuard<'_> {
fn drop(&mut self) {
self.finish();
}
}
async fn load_candidate_selection_rows_with_timeout<Fut>(
key: &CandidateSelectionCacheKey,
load: Fut,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>
where
Fut: Future<Output = Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>>,
{
match timeout(CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT, load).await {
Ok(result) => result,
Err(_) => {
warn!(
event_name = "candidate_selection_cache_load_timeout",
log_type = "ops",
cache_key = ?key,
timeout_ms = CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT.as_millis() as u64,
"gateway candidate selection cache load timed out"
);
Err(DataLayerError::TimedOut(format!(
"candidate selection cache load exceeded {}ms for {key:?}",
CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT.as_millis()
)))
}
}
}
#[async_trait]
impl MinimalCandidateSelectionReadRepository for CachedMinimalCandidateSelectionReadRepository {
fn clear_local_cache(&self) {
@@ -286,6 +400,7 @@ fn normalize_api_format_key(api_format: &str) -> String {
#[cfg(test)]
mod tests {
use super::*;
use std::future::pending;
use std::sync::atomic::AtomicUsize;
struct StubCandidateSelectionRepository {
@@ -361,6 +476,77 @@ mod tests {
}
}
struct FirstLoadPendingThenFastRepository {
calls: AtomicUsize,
}
impl FirstLoadPendingThenFastRepository {
fn new() -> Self {
Self {
calls: AtomicUsize::new(0),
}
}
fn calls(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
async fn load(&self) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let call = self.calls.fetch_add(1, Ordering::SeqCst);
if call == 0 {
pending::<()>().await;
}
Ok(Vec::new())
}
}
#[async_trait]
impl MinimalCandidateSelectionReadRepository for FirstLoadPendingThenFastRepository {
async fn list_for_exact_api_format(
&self,
_api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_for_exact_api_format_and_global_model(
&self,
_api_format: &str,
_global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_for_exact_api_format_and_requested_model(
&self,
_api_format: &str,
_requested_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_for_exact_api_format_and_requested_model_page(
&self,
_query: &StoredRequestedModelCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_pool_key_rows_for_group(
&self,
_query: &StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
_query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
}
#[tokio::test]
async fn candidate_selection_cache_coalesces_concurrent_loads() {
let inner = Arc::new(StubCandidateSelectionRepository::new(
@@ -398,4 +584,117 @@ mod tests {
cache.list_for_exact_api_format("openai").await.unwrap();
assert_eq!(inner.calls(), 2);
}
#[tokio::test]
async fn candidate_selection_cache_releases_inflight_when_leader_is_cancelled() {
let inner = Arc::new(FirstLoadPendingThenFastRepository::new());
let cache = Arc::new(CachedMinimalCandidateSelectionReadRepository::new(
inner.clone(),
));
let leader_cache = cache.clone();
let leader = tokio::spawn(async move {
leader_cache
.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
)
.await
});
tokio::time::sleep(Duration::from_millis(10)).await;
leader.abort();
let _ = leader.await;
tokio::time::timeout(
Duration::from_millis(200),
cache.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
),
)
.await
.expect("cancelled leader must not leave a permanent inflight wait")
.unwrap();
assert_eq!(inner.calls(), 2);
}
#[tokio::test]
async fn candidate_selection_cache_times_out_and_clears_stuck_load() {
let inner = Arc::new(FirstLoadPendingThenFastRepository::new());
let cache = CachedMinimalCandidateSelectionReadRepository::new(inner.clone());
let err = cache
.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
)
.await
.unwrap_err();
assert!(matches!(err, DataLayerError::TimedOut(_)));
cache
.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
)
.await
.unwrap();
assert_eq!(inner.calls(), 2);
}
#[tokio::test]
async fn candidate_selection_cache_clear_releases_inflight_waiters() {
let inner = Arc::new(FirstLoadPendingThenFastRepository::new());
let cache = Arc::new(CachedMinimalCandidateSelectionReadRepository::new(
inner.clone(),
));
let leader_cache = cache.clone();
let leader = tokio::spawn(async move {
leader_cache
.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
)
.await
});
tokio::time::sleep(Duration::from_millis(10)).await;
cache.clear_local_cache();
tokio::time::timeout(
Duration::from_millis(200),
cache.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
),
)
.await
.expect("cache clear must release stale inflight waiters")
.unwrap();
assert_eq!(inner.calls(), 2);
leader.abort();
let _ = leader.await;
}
}

63
apps/aether-gateway/src/error.rs
View File

@@ -11,12 +11,42 @@ use crate::insert_header_if_missing;
#[derive(Debug)]
pub(crate) enum GatewayError {
UpstreamUnavailable { trace_id: String, message: String },
ControlUnavailable { trace_id: String, message: String },
Client { status: StatusCode, message: String },
UpstreamUnavailable {
trace_id: String,
message: String,
},
ControlUnavailable {
trace_id: String,
message: String,
},
LocalExecutionPlanningTimeout {
trace_id: String,
phase: &'static str,
timeout_ms: u64,
},
Client {
status: StatusCode,
message: String,
},
Internal(String),
}
impl GatewayError {
pub(crate) fn into_message(self) -> String {
match self {
Self::UpstreamUnavailable { message, .. }
| Self::ControlUnavailable { message, .. }
| Self::Client { message, .. }
| Self::Internal(message) => message,
Self::LocalExecutionPlanningTimeout {
phase, timeout_ms, ..
} => {
format!("local execution planning timed out in {phase} after {timeout_ms}ms")
}
}
}
}
impl IntoResponse for GatewayError {
fn into_response(self) -> Response<Body> {
match self {
@@ -56,6 +86,33 @@ impl IntoResponse for GatewayError {
);
response
}
Self::LocalExecutionPlanningTimeout {
trace_id,
phase,
timeout_ms,
} => {
warn!(
trace_id = %trace_id,
phase,
timeout_ms,
"gateway local execution planning timed out"
);
let body = Json(json!({
"error": {
"message": "gateway local execution planning timed out",
"trace_id": trace_id,
}
}));
let mut response = (StatusCode::GATEWAY_TIMEOUT, body).into_response();
let _ =
insert_header_if_missing(response.headers_mut(), TRACE_ID_HEADER, &trace_id);
let _ = insert_header_if_missing(
response.headers_mut(),
GATEWAY_HEADER,
"rust-phase3b",
);
response
}
Self::Client { status, message } => (
status,
Json(json!({

101
apps/aether-gateway/src/executor/candidate_loop.rs
View File

@@ -125,7 +125,16 @@ where
decision,
plan_kind,
};
run_dynamic_attempt_loop(&port, &mut source).await
run_dynamic_attempt_loop(
&port,
&mut source,
trace_id,
plan_kind,
state
.frontdoor_runtime_guards
.local_execution_planning_timeout,
)
.await
}
.instrument(span)
.await
@@ -281,7 +290,16 @@ where
decision,
plan_kind,
};
run_dynamic_attempt_loop(&port, &mut source).await
run_dynamic_attempt_loop(
&port,
&mut source,
trace_id,
plan_kind,
state
.frontdoor_runtime_guards
.local_execution_planning_timeout,
)
.await
}
.instrument(span)
.await
@@ -290,6 +308,9 @@ where
async fn run_dynamic_attempt_loop<Port, Source, Attempt>(
port: &Port,
source: &mut Source,
trace_id: &str,
plan_kind: &str,
planning_timeout: Duration,
) -> Result<LocalExecutionRequestOutcome, GatewayError>
where
Port: AiAttemptLoopPort<
@@ -303,7 +324,9 @@ where
{
let mut last_attempted = None;
while let Some(attempt) = source.next_execution_attempt().await? {
while let Some(attempt) =
next_execution_attempt_with_timeout(source, trace_id, plan_kind, planning_timeout).await?
{
last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context()));
if let Some(response) = port.execute_attempt(&attempt).await? {
let remaining = source.drain_execution_attempts().await?;
@@ -322,6 +345,37 @@ where
))
}
async fn next_execution_attempt_with_timeout<Source, Attempt>(
source: &mut Source,
trace_id: &str,
plan_kind: &str,
planning_timeout: Duration,
) -> Result<Option<Attempt>, GatewayError>
where
Source: LocalExecutionAttemptSource<Attempt>,
{
match timeout(planning_timeout, source.next_execution_attempt()).await {
Ok(result) => result,
Err(_) => {
let timeout_ms = planning_timeout.as_millis() as u64;
warn!(
event_name = "local_execution_candidate_planning_timeout",
log_type = "ops",
trace_id,
plan_kind,
timeout_ms,
phase = "next_execution_attempt",
"gateway timed out while planning the next local execution candidate"
);
Err(GatewayError::LocalExecutionPlanningTimeout {
trace_id: trace_id.to_string(),
phase: "next_execution_attempt",
timeout_ms,
})
}
}
}
struct StreamAttemptLoopPort<'a> {
state: &'a AppState,
trace_id: &'a str,
@@ -632,6 +686,20 @@ mod tests {
}
}
struct PendingAttemptSource;
#[async_trait]
impl LocalExecutionAttemptSource<()> for PendingAttemptSource {
async fn next_execution_attempt(&mut self) -> Result<Option<()>, GatewayError> {
std::future::pending::<()>().await;
Ok(None)
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<()>, GatewayError> {
Ok(Vec::new())
}
}
fn test_plan(timeouts: Option<ExecutionTimeouts>) -> ExecutionPlan {
ExecutionPlan {
request_id: "req_watchdog".to_string(),
@@ -656,6 +724,33 @@ mod tests {
}
}
#[tokio::test]
async fn next_execution_attempt_times_out_instead_of_waiting_forever() {
let mut source = PendingAttemptSource;
let err = next_execution_attempt_with_timeout(
&mut source,
"trace-planning-timeout",
"openai_responses_sync",
Duration::from_millis(5),
)
.await
.expect_err("pending candidate planning should time out");
match err {
GatewayError::LocalExecutionPlanningTimeout {
trace_id,
phase,
timeout_ms,
} => {
assert_eq!(trace_id, "trace-planning-timeout");
assert_eq!(phase, "next_execution_attempt");
assert_eq!(timeout_ms, 5);
}
other => panic!("unexpected error: {other:?}"),
}
}
fn test_report_context() -> serde_json::Value {
json!({
"request_id": "req_watchdog",

7
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
View File

@@ -317,12 +317,7 @@ pub(super) async fn execute_provider_quota_plan(
match state.execute_execution_runtime_sync_plan(None, &plan).await {
Ok(result) => Ok(ProviderQuotaExecutionOutcome::Response(result)),
Err(err) => {
let error = match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
};
let error = err.into_message();
let proxy_node_id = plan
.proxy
.as_ref()

7
apps/aether-gateway/src/handlers/admin/provider/ops/providers/verify/request.rs
View File

@@ -301,12 +301,7 @@ fn admin_provider_ops_decode_response_bytes(
}
fn admin_provider_ops_gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
error.into_message()
}
pub(super) fn admin_provider_ops_verify_execution_error_message(error: &str) -> String {

69
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/reads.rs
View File

@@ -16,7 +16,12 @@ use aether_runtime_state::{DataLayerError, RuntimeState};
use futures_util::future::join_all;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
use tracing::{info, warn};
const DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 512;
const MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 10_000;
const POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV: &str =
"AETHER_GATEWAY_ADMIN_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT";
fn current_unix_secs() -> u64 {
SystemTime::now()
@@ -29,6 +34,20 @@ fn should_load_active_probe_members(pool_config: &AdminProviderPoolConfig) -> bo
pool_config.probing_enabled
}
fn pool_runtime_window_metric_key_limit() -> usize {
std::env::var(POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|value| *value > 0)
.unwrap_or(DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
.clamp(1, MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
}
fn bounded_runtime_window_metric_key_ids(key_ids: &[String], limit: usize) -> &[String] {
let end = key_ids.len().min(limit.max(1));
&key_ids[..end]
}
pub(crate) async fn read_admin_provider_pool_cooldown_counts(
runtime: &RuntimeState,
provider_ids: &[String],
@@ -54,8 +73,21 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
) -> AdminProviderPoolRuntimeState {
let mut state = AdminProviderPoolRuntimeState::default();
let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
let cost_keys = pool_cost_keys(provider_id, key_ids);
let latency_keys = pool_latency_keys(provider_id, key_ids);
let metric_key_limit = pool_runtime_window_metric_key_limit();
let metric_key_ids = bounded_runtime_window_metric_key_ids(key_ids, metric_key_limit);
if metric_key_ids.len() < key_ids.len() {
info!(
event_name = "admin_pool_runtime_window_metrics_truncated",
log_type = "event",
provider_id,
total_key_count = key_ids.len(),
scanned_key_count = metric_key_ids.len(),
metric_key_limit,
"gateway limited admin pool runtime cost/latency window reads"
);
}
let cost_keys = pool_cost_keys(provider_id, metric_key_ids);
let latency_keys = pool_latency_keys(provider_id, metric_key_ids);
let sticky_sessions_enabled = pool_config.sticky_session_ttl_seconds > 0
&& admin_provider_pool_cache_affinity_enabled(pool_config);
@@ -179,7 +211,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
.map(|cost_key| runtime.score_range_by_min(cost_key, cost_window_start)),
)
.await;
for (key_id, members) in key_ids.iter().zip(cost_results) {
for (key_id, members) in metric_key_ids.iter().zip(cost_results) {
let total = members
.unwrap_or_default()
.iter()
@@ -197,7 +229,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
.map(|latency_key| runtime.score_range_by_min(latency_key, latency_window_start)),
)
.await;
for (key_id, members) in key_ids.iter().zip(latency_results) {
for (key_id, members) in metric_key_ids.iter().zip(latency_results) {
let samples = members
.unwrap_or_default()
.iter()
@@ -265,3 +297,30 @@ pub(crate) async fn read_admin_provider_pool_key_cooldown_reason(
.kv_get(&pool_cooldown_key(provider_id, key_id))
.await
}
#[cfg(test)]
mod tests {
use super::bounded_runtime_window_metric_key_ids;
#[test]
fn runtime_window_metric_key_ids_are_bounded() {
let key_ids = vec![
"key-1".to_string(),
"key-2".to_string(),
"key-3".to_string(),
];
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 2);
assert_eq!(bounded, &key_ids[..2]);
}
#[test]
fn runtime_window_metric_key_ids_keep_at_least_one_key() {
let key_ids = vec!["key-1".to_string(), "key-2".to_string()];
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 0);
assert_eq!(bounded, &key_ids[..1]);
}
}

7
apps/aether-gateway/src/handlers/admin/request/provider/oauth.rs
View File

@@ -662,10 +662,5 @@ fn admin_provider_oauth_decode_response_bytes(
}
fn admin_provider_oauth_gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
error.into_message()
}

7
apps/aether-gateway/src/handlers/internal/gateway_helpers.rs
View File

@@ -354,12 +354,7 @@ pub(crate) async fn maybe_build_internal_finalize_video_response(
}
pub(crate) fn gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
error.into_message()
}
pub(crate) fn build_internal_tunnel_heartbeat_ack(

442
apps/aether-gateway/src/handlers/proxy/mod.rs
View File

@@ -21,13 +21,13 @@ use crate::constants::{
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED,
EXECUTION_PATH_LOCAL_AUTH_DENIED, EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED,
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, EXECUTION_PATH_LOCAL_INVALID_REQUEST,
EXECUTION_PATH_LOCAL_OVERLOADED, EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED,
EXECUTION_PATH_LOCAL_RATE_LIMITED, EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND,
EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH, EXECUTION_RUNTIME_LOOP_GUARD_HEADER,
FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER, FORWARDED_PROTO_HEADER, GATEWAY_HEADER,
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
EXECUTION_PATH_LOCAL_INVALID_REQUEST, EXECUTION_PATH_LOCAL_OVERLOADED,
EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED, EXECUTION_PATH_LOCAL_RATE_LIMITED,
EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND, EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH,
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER,
FORWARDED_PROTO_HEADER, GATEWAY_HEADER, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
TRACE_ID_HEADER, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
TRUSTED_AUTH_BALANCE_HEADER, TRUSTED_AUTH_USER_ID_HEADER, TUNNEL_AFFINITY_FORWARDED_BY_HEADER,
TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER,
};
@@ -65,7 +65,11 @@ use axum::extract::{ConnectInfo, Request, State};
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
use futures_util::StreamExt;
use sha2::{Digest, Sha256};
use std::{collections::BTreeMap, time::Instant};
use std::{
collections::BTreeMap,
error::Error as StdError,
time::{Duration, Instant},
};
use tracing::{debug, info, warn};
const OPENAI_CHAT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
@@ -89,23 +93,192 @@ const LOCAL_EXECUTION_LOOP_DETECTED_DETAIL: &str =
"Gateway detected an execution runtime request loop back into the local frontdoor";
const AUTH_API_KEY_CONCURRENCY_LIMIT_REACHED_DETAIL: &str =
"当前 API Key 并发请求数已达上限，请稍后重试";
const REQUEST_BODY_READ_TIMEOUT_DETAIL: &str =
"Request body read timed out before the gateway could route the request";
const REQUEST_BODY_READ_FAILED_DETAIL: &str = "Failed to read request body";
const LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL: &str =
"当前 AI 请求在本地执行规划阶段超时，请稍后重试";
const EXECUTION_PATH_TUNNEL_AFFINITY_FORWARD: &str = "tunnel_affinity_forward";
const MANAGEMENT_TOKEN_PREFIX: &str = "ae-";
const LEGACY_MANAGEMENT_TOKEN_PREFIX: &str = "ae_";
fn build_request_body_normalization_error_response(
#[derive(Debug, Clone, Copy)]
struct RequestBodyBufferPolicy {
max_bytes: u64,
read_timeout: Duration,
}
impl RequestBodyBufferPolicy {
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,
}
}
#[cfg(test)]
fn for_tests(max_bytes: u64, read_timeout: Duration) -> Self {
Self {
max_bytes,
read_timeout,
}
}
}
#[derive(Debug)]
enum RequestBodyBufferError {
Normalization(RequestBodyNormalizationError),
TooLarge { limit_bytes: u64 },
Timeout { timeout_ms: u64 },
ReadFailed { message: String },
}
impl RequestBodyBufferError {
fn http_status(&self) -> http::StatusCode {
match self {
Self::Normalization(error) => error.http_status(),
Self::TooLarge { .. } => http::StatusCode::PAYLOAD_TOO_LARGE,
Self::Timeout { .. } => http::StatusCode::REQUEST_TIMEOUT,
Self::ReadFailed { .. } => http::StatusCode::BAD_REQUEST,
}
}
fn client_message(&self) -> String {
match self {
Self::Normalization(error) => error.client_message(),
Self::TooLarge { limit_bytes } => format!("Request body exceeds {limit_bytes} bytes"),
Self::Timeout { .. } => REQUEST_BODY_READ_TIMEOUT_DETAIL.to_string(),
Self::ReadFailed { .. } => REQUEST_BODY_READ_FAILED_DETAIL.to_string(),
}
}
fn reason(&self) -> &'static str {
match self {
Self::Normalization(error) => match error {
RequestBodyNormalizationError::UnsupportedContentEncoding(_) => {
"unsupported_content_encoding"
}
RequestBodyNormalizationError::DecodeFailed { .. } => "decode_failed",
RequestBodyNormalizationError::DecompressedBodyTooLarge { .. } => {
"decompressed_body_too_large"
}
RequestBodyNormalizationError::RequestBodyTooLarge { .. } => {
"request_body_too_large"
}
},
Self::TooLarge { .. } => "request_body_too_large",
Self::Timeout { .. } => "request_body_read_timeout",
Self::ReadFailed { .. } => "request_body_read_failed",
}
}
}
async fn buffer_and_normalize_request_body(
request_body: &mut Option<Body>,
headers: &mut http::HeaderMap,
body_owner_expectation: &'static str,
trace_id: &str,
method: &http::Method,
path_and_query: &str,
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 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",
trace_id,
method = %method,
path = %path_and_query,
phase,
max_body_bytes = policy.max_bytes,
timeout_ms,
"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)?;
info!(
event_name = "frontdoor_request_body_buffer_completed",
log_type = "event",
trace_id,
method = %method,
path = %path_and_query,
phase,
body_bytes = normalized.len(),
elapsed_ms = read_started_at.elapsed().as_millis() as u64,
"gateway completed request body buffering"
);
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
}
fn build_request_body_buffer_error_response(
trace_id: &str,
request_context: &GatewayPublicRequestContext,
error: &RequestBodyNormalizationError,
error: &RequestBodyBufferError,
) -> Result<Response<Body>, GatewayError> {
warn!(
event_name = "frontdoor_request_body_normalization_failed",
event_name = "frontdoor_request_body_buffer_failed",
log_type = "ops",
trace_id,
method = %request_context.request_method,
path = %request_context.request_path_and_query(),
error = %error,
"gateway rejected request with invalid encoded body"
status_code = error.http_status().as_u16(),
reason = error.reason(),
detail = %error.client_message(),
read_error = match error {
RequestBodyBufferError::ReadFailed { message } => message.as_str(),
_ => "",
},
"gateway rejected request body before local execution planning"
);
build_local_http_error_response(
trace_id,
@@ -115,36 +288,16 @@ fn build_request_body_normalization_error_response(
)
}
async fn buffer_and_normalize_request_body(
request_body: &mut Option<Body>,
headers: &mut http::HeaderMap,
body_owner_expectation: &'static str,
) -> Result<Result<Bytes, RequestBodyNormalizationError>, GatewayError> {
if let Err(err) = crate::headers::check_request_content_length(headers) {
return Ok(Err(err));
}
let body = to_bytes(
request_body.take().expect(body_owner_expectation),
usize::MAX,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(crate::headers::normalize_request_body_headers_and_bytes(
headers, body,
))
}
fn finalize_request_body_normalization_rejection(
fn finalize_request_body_buffer_rejection(
state: &AppState,
request_context: &GatewayPublicRequestContext,
remote_addr: &std::net::SocketAddr,
started_at: &std::time::Instant,
trace_id: &str,
request_permit: Option<aether_runtime::AdmissionPermit>,
error: &RequestBodyNormalizationError,
error: &RequestBodyBufferError,
) -> Result<Response<Body>, GatewayError> {
let response =
build_request_body_normalization_error_response(trace_id, request_context, error)?;
let response = build_request_body_buffer_error_response(trace_id, request_context, error)?;
Ok(finalize_gateway_response_with_context(
state,
response,
@@ -156,6 +309,59 @@ fn finalize_request_body_normalization_rejection(
))
}
fn local_execution_planning_timeout_parts(error: &GatewayError) -> Option<(&'static str, u64)> {
match error {
GatewayError::LocalExecutionPlanningTimeout {
phase, timeout_ms, ..
} => Some((*phase, *timeout_ms)),
_ => None,
}
}
fn finalize_local_execution_planning_timeout(
state: &AppState,
request_context: &GatewayPublicRequestContext,
remote_addr: &std::net::SocketAddr,
started_at: &std::time::Instant,
trace_id: &str,
request_permit: Option<aether_runtime::AdmissionPermit>,
control_decision: Option<&GatewayControlDecision>,
phase: &'static str,
timeout_ms: u64,
) -> Result<Response<Body>, GatewayError> {
warn!(
event_name = "frontdoor_local_execution_planning_timeout",
log_type = "ops",
trace_id,
method = %request_context.request_method,
path = %request_context.request_path_and_query(),
route_family = control_decision
.and_then(|decision| decision.route_family.as_deref())
.unwrap_or("-"),
route_kind = control_decision
.and_then(|decision| decision.route_kind.as_deref())
.unwrap_or("-"),
phase,
timeout_ms,
"gateway failed local execution before a candidate could be selected"
);
let response = build_local_http_error_response(
trace_id,
control_decision,
http::StatusCode::GATEWAY_TIMEOUT,
LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL,
)?;
Ok(finalize_gateway_response_with_context(
state,
response,
remote_addr,
request_context,
EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT,
started_at,
request_permit,
))
}
fn local_execution_outcome_label(outcome: &LocalExecutionRequestOutcome) -> &'static str {
match outcome {
LocalExecutionRequestOutcome::Responded(_) => "responded",
@@ -991,16 +1197,22 @@ pub(crate) async fn proxy_request(
}
let mut request_body = Some(body);
let local_proxy_body = if local_proxy_route_requires_buffered_body(&request_context) {
let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state);
let body = buffer_and_normalize_request_body(
&mut request_body,
&mut parts.headers,
"local proxy body buffering should own request body",
&trace_id,
&parts.method,
&request_context.request_path_and_query(),
"local_proxy",
body_buffer_policy,
)
.await?;
.await;
match body {
Ok(body) => Some(body),
Err(err) => {
return finalize_request_body_normalization_rejection(
return finalize_request_body_buffer_rejection(
&state,
&request_context,
&remote_addr,
@@ -1171,16 +1383,22 @@ pub(crate) async fn proxy_request(
&& request_enables_control_execute(&parts.headers);
let buffered_body = if should_buffer_body {
let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state);
let body = buffer_and_normalize_request_body(
&mut request_body,
&mut parts.headers,
"buffered auth/execution runtime path should own request body",
&trace_id,
&parts.method,
&request_context.request_path_and_query(),
"auth_execution",
body_buffer_policy,
)
.await?;
.await;
match body {
Ok(body) => Some(body),
Err(err) => {
return finalize_request_body_normalization_rejection(
return finalize_request_body_buffer_rejection(
&state,
&request_context,
&remote_addr,
@@ -1335,14 +1553,34 @@ pub(crate) async fn proxy_request(
let stream_request = request_wants_stream(&request_context, &parts.headers, buffered_body);
let mut local_execution_exhaustion = None;
if stream_request {
let stream_outcome = maybe_execute_stream_request(
let stream_outcome = match maybe_execute_stream_request(
&state,
&parts,
buffered_body,
&trace_id,
control_decision,
)
.await?;
.await
{
Ok(outcome) => outcome,
Err(err) => {
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err)
{
return finalize_local_execution_planning_timeout(
&state,
&request_context,
&remote_addr,
&started_at,
&trace_id,
request_permit.take(),
control_decision,
phase,
timeout_ms,
);
}
return Err(err);
}
};
debug!(
event_name = "proxy_stream_local_execute_outcome",
log_type = "debug",
@@ -1380,9 +1618,34 @@ pub(crate) async fn proxy_request(
LocalExecutionRequestOutcome::NoPath => {}
}
}
match maybe_execute_sync_request(&state, &parts, buffered_body, &trace_id, control_decision)
.await?
let sync_outcome = match maybe_execute_sync_request(
&state,
&parts,
buffered_body,
&trace_id,
control_decision,
)
.await
{
Ok(outcome) => outcome,
Err(err) => {
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err) {
return finalize_local_execution_planning_timeout(
&state,
&request_context,
&remote_addr,
&started_at,
&trace_id,
request_permit.take(),
control_decision,
phase,
timeout_ms,
);
}
return Err(err);
}
};
match sync_outcome {
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
let execution_runtime_response = restore_redacted_sync_execution_response(
execution_runtime_response,
@@ -1406,15 +1669,35 @@ pub(crate) async fn proxy_request(
LocalExecutionRequestOutcome::NoPath => {}
}
if parts.method != http::Method::POST {
match maybe_execute_stream_request(
let stream_outcome = match maybe_execute_stream_request(
&state,
&parts,
buffered_body,
&trace_id,
control_decision,
)
.await?
.await
{
Ok(outcome) => outcome,
Err(err) => {
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err)
{
return finalize_local_execution_planning_timeout(
&state,
&request_context,
&remote_addr,
&started_at,
&trace_id,
request_permit.take(),
control_decision,
phase,
timeout_ms,
);
}
return Err(err);
}
};
match stream_outcome {
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
let execution_runtime_response = restore_redacted_stream_execution_response(
execution_runtime_response,
@@ -1977,14 +2260,17 @@ fn local_execution_runtime_miss_route_detail(
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::{
api_key_remote_ip_allowed, diagnostic_is_auth_api_key_concurrency_limited,
local_execution_runtime_miss_detail, restore_redacted_stream_execution_response,
restore_redacted_sync_execution_response, GatewayControlDecision,
LocalExecutionRuntimeMissDiagnostic,
api_key_remote_ip_allowed, buffer_and_normalize_request_body,
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response,
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError,
RequestBodyBufferPolicy,
};
use axum::body::{to_bytes, Body};
use axum::http::{header, Response};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::{header, HeaderMap, Method, Response};
use serde_json::json;
#[test]
@@ -2112,6 +2398,58 @@ mod tests {
assert!(!message.contains(&sentinel));
}
#[tokio::test]
async fn request_body_buffer_rejects_chunked_body_when_limit_is_exceeded() {
let mut body = Some(Body::from(Bytes::from_static(b"abcdef")));
let mut headers = HeaderMap::new();
let err = buffer_and_normalize_request_body(
&mut body,
&mut headers,
"test owns body",
"trace-body-large",
&Method::POST,
"/v1/responses",
"test",
RequestBodyBufferPolicy::for_tests(5, Duration::from_secs(1)),
)
.await
.expect_err("body exceeding the ingress limit should fail");
assert!(matches!(
err,
RequestBodyBufferError::TooLarge { limit_bytes: 5 }
));
}
#[tokio::test]
async fn request_body_buffer_times_out_instead_of_waiting_forever() {
let stream = async_stream::stream! {
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(b"{"));
std::future::pending::<()>().await;
};
let mut body = Some(Body::from_stream(stream));
let mut headers = HeaderMap::new();
let err = buffer_and_normalize_request_body(
&mut body,
&mut headers,
"test owns body",
"trace-body-timeout",
&Method::POST,
"/v1/responses",
"test",
RequestBodyBufferPolicy::for_tests(1024, Duration::from_millis(5)),
)
.await
.expect_err("body buffering should time out");
assert!(matches!(
err,
RequestBodyBufferError::Timeout { timeout_ms: 5 }
));
}
#[test]
fn runtime_miss_detail_returns_model_specific_stream_message_when_candidates_are_unavailable() {
let decision = GatewayControlDecision::synthetic(

7
apps/aether-gateway/src/handlers/public/support/announcements/shared.rs
View File

@@ -136,12 +136,7 @@ pub(super) fn announcements_internal_error_response(detail: impl Into<String>) -
}
pub(super) fn announcements_internal_detail(err: GatewayError) -> String {
match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
err.into_message()
}
pub(super) fn parse_optional_rfc3339_unix_secs(

98
apps/aether-gateway/src/headers.rs
View File

@@ -21,6 +21,10 @@ static MAX_REQUEST_BODY_BYTES: LazyLock<u64> = LazyLock::new(|| {
.saturating_mul(1024 * 1024)
});
pub(crate) fn max_request_body_bytes() -> u64 {
*MAX_REQUEST_BODY_BYTES
}
pub(crate) fn extract_or_generate_trace_id(headers: &http::HeaderMap) -> String {
header_value_str(headers, TRACE_ID_HEADER).unwrap_or_else(|| Uuid::new_v4().to_string())
}
@@ -246,9 +250,21 @@ impl std::error::Error for RequestBodyNormalizationError {}
pub(crate) fn normalize_request_body_headers_and_bytes(
headers: &mut http::HeaderMap,
body_bytes: Bytes,
) -> Result<Bytes, RequestBodyNormalizationError> {
normalize_request_body_headers_and_bytes_with_limit(
headers,
body_bytes,
max_request_body_bytes(),
)
}
pub(crate) fn normalize_request_body_headers_and_bytes_with_limit(
headers: &mut http::HeaderMap,
body_bytes: Bytes,
limit_bytes: u64,
) -> Result<Bytes, RequestBodyNormalizationError> {
let body_was_encoded = !request_content_encodings(headers).is_empty();
let decoded = decoded_request_body_bytes(headers, body_bytes.as_ref())?;
let decoded = decoded_request_body_bytes_with_limit(headers, body_bytes.as_ref(), limit_bytes)?;
if !body_was_encoded {
return Ok(body_bytes);
}
@@ -264,7 +280,13 @@ pub(crate) fn normalize_request_body_headers_and_bytes(
pub(crate) fn check_request_content_length(
headers: &http::HeaderMap,
) -> Result<(), RequestBodyNormalizationError> {
let limit = *MAX_REQUEST_BODY_BYTES;
check_request_content_length_with_limit(headers, max_request_body_bytes())
}
pub(crate) fn check_request_content_length_with_limit(
headers: &http::HeaderMap,
limit: u64,
) -> Result<(), RequestBodyNormalizationError> {
let declared = header_value_str(headers, http::header::CONTENT_LENGTH.as_str())
.and_then(|value| value.trim().parse::<u64>().ok());
if declared.is_some_and(|value| value > limit) {
@@ -276,10 +298,17 @@ pub(crate) fn check_request_content_length(
pub(crate) fn decoded_request_body_bytes<'a>(
headers: &http::HeaderMap,
body_bytes: &'a [u8],
) -> Result<Cow<'a, [u8]>, RequestBodyNormalizationError> {
decoded_request_body_bytes_with_limit(headers, body_bytes, max_request_body_bytes())
}
pub(crate) fn decoded_request_body_bytes_with_limit<'a>(
headers: &http::HeaderMap,
body_bytes: &'a [u8],
limit: u64,
) -> Result<Cow<'a, [u8]>, RequestBodyNormalizationError> {
let encodings = request_content_encodings(headers);
if encodings.is_empty() {
let limit = *MAX_REQUEST_BODY_BYTES;
if body_bytes.len() as u64 > limit {
return Err(RequestBodyNormalizationError::RequestBodyTooLarge { limit_bytes: limit });
}
@@ -288,7 +317,7 @@ pub(crate) fn decoded_request_body_bytes<'a>(
let mut decoded = body_bytes.to_vec();
for encoding in encodings.iter().rev() {
decoded = decode_single_request_body(encoding, decoded.as_slice())?;
decoded = decode_single_request_body_with_limit(encoding, decoded.as_slice(), limit)?;
}
Ok(Cow::Owned(decoded))
}
@@ -310,11 +339,19 @@ fn request_content_encodings(headers: &http::HeaderMap) -> Vec<String> {
fn decode_single_request_body(
encoding: &str,
body_bytes: &[u8],
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
decode_single_request_body_with_limit(encoding, body_bytes, max_request_body_bytes())
}
fn decode_single_request_body_with_limit(
encoding: &str,
body_bytes: &[u8],
limit: u64,
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
match encoding {
"gzip" | "x-gzip" => decode_gzip_body(encoding, body_bytes),
"deflate" => decode_deflate_body(encoding, body_bytes),
"zstd" => decode_zstd_body(encoding, body_bytes),
"gzip" | "x-gzip" => decode_gzip_body_with_limit(encoding, body_bytes, limit),
"deflate" => decode_deflate_body_with_limit(encoding, body_bytes, limit),
"zstd" => decode_zstd_body_with_limit(encoding, body_bytes, limit),
_ => Err(RequestBodyNormalizationError::UnsupportedContentEncoding(
encoding.to_string(),
)),
@@ -324,27 +361,43 @@ fn decode_single_request_body(
fn decode_gzip_body(
encoding: &str,
body_bytes: &[u8],
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
decode_gzip_body_with_limit(encoding, body_bytes, max_request_body_bytes())
}
fn decode_gzip_body_with_limit(
encoding: &str,
body_bytes: &[u8],
limit: u64,
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
let mut decoder = GzDecoder::new(body_bytes);
read_request_decoder_to_end(encoding, &mut decoder)
read_request_decoder_to_end_with_limit(encoding, &mut decoder, limit)
}
fn decode_deflate_body(
encoding: &str,
body_bytes: &[u8],
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
decode_deflate_body_with_limit(encoding, body_bytes, max_request_body_bytes())
}
fn decode_deflate_body_with_limit(
encoding: &str,
body_bytes: &[u8],
limit: u64,
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
let mut zlib_decoder = ZlibDecoder::new(body_bytes);
match read_request_decoder_to_end(encoding, &mut zlib_decoder) {
match read_request_decoder_to_end_with_limit(encoding, &mut zlib_decoder, limit) {
Ok(decoded) => Ok(decoded),
Err(err @ RequestBodyNormalizationError::DecompressedBodyTooLarge { .. }) => Err(err),
Err(zlib_error) => {
let mut raw_decoder = DeflateDecoder::new(body_bytes);
read_request_decoder_to_end(encoding, &mut raw_decoder).map_err(|raw_error| {
RequestBodyNormalizationError::DecodeFailed {
read_request_decoder_to_end_with_limit(encoding, &mut raw_decoder, limit).map_err(
|raw_error| RequestBodyNormalizationError::DecodeFailed {
encoding: encoding.to_string(),
reason: format!("{zlib_error}; raw deflate fallback failed: {raw_error}"),
}
})
},
)
}
}
}
@@ -352,6 +405,14 @@ fn decode_deflate_body(
fn decode_zstd_body(
encoding: &str,
body_bytes: &[u8],
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
decode_zstd_body_with_limit(encoding, body_bytes, max_request_body_bytes())
}
fn decode_zstd_body_with_limit(
encoding: &str,
body_bytes: &[u8],
limit: u64,
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
let mut decoder = zstd::stream::read::Decoder::new(body_bytes).map_err(|err| {
RequestBodyNormalizationError::DecodeFailed {
@@ -359,14 +420,21 @@ fn decode_zstd_body(
reason: err.to_string(),
}
})?;
read_request_decoder_to_end(encoding, &mut decoder)
read_request_decoder_to_end_with_limit(encoding, &mut decoder, limit)
}
fn read_request_decoder_to_end(
encoding: &str,
decoder: &mut impl Read,
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
let limit = *MAX_REQUEST_BODY_BYTES;
read_request_decoder_to_end_with_limit(encoding, decoder, max_request_body_bytes())
}
fn read_request_decoder_to_end_with_limit(
encoding: &str,
decoder: &mut impl Read,
limit: u64,
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
let mut limited = decoder.take(limit.saturating_add(1));
let mut out = Vec::new();
limited

7
apps/aether-gateway/src/maintenance/runtime/account_self_check.rs
View File

@@ -551,12 +551,7 @@ fn endpoint_for_self_check(
}
fn gateway_error_message(err: GatewayError) -> String {
match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
err.into_message()
}
fn update_summary_from_outcome(

7
apps/aether-gateway/src/oauth/http_executor.rs
View File

@@ -157,10 +157,5 @@ fn decode_response_bytes(bytes: &[u8], content_encoding: Option<&str>) -> Option
}
fn gateway_error_to_oauth_error(error: GatewayError) -> OAuthError {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => OAuthError::Transport(message),
}
OAuthError::Transport(error.into_message())
}

58
apps/aether-gateway/src/state/app.rs
View File

@@ -23,6 +23,17 @@ use super::{
LocalProviderDeleteTaskState, ProviderTransportSnapshotCacheKey,
};
const DEFAULT_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 120_000;
const MIN_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 1_000;
const MAX_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 600_000;
const REQUEST_BODY_READ_TIMEOUT_MS_ENV: &str = "AETHER_GATEWAY_REQUEST_BODY_READ_TIMEOUT_MS";
const DEFAULT_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 30_000;
const MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 500;
const MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 120_000;
const LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV: &str =
"AETHER_GATEWAY_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS";
#[cfg(test)]
type TestExecutionRuntimeSyncOverrideFn = dyn Fn(
&aether_contracts::ExecutionPlan,
@@ -43,6 +54,52 @@ impl std::fmt::Debug for TestExecutionRuntimeSyncOverride {
}
}
#[derive(Debug, Clone)]
pub(crate) struct FrontdoorRuntimeGuardConfig {
pub(crate) request_body_read_timeout: Duration,
pub(crate) local_execution_planning_timeout: Duration,
}
impl FrontdoorRuntimeGuardConfig {
pub(crate) fn from_env() -> Self {
Self {
request_body_read_timeout: env_duration_ms(
REQUEST_BODY_READ_TIMEOUT_MS_ENV,
DEFAULT_REQUEST_BODY_READ_TIMEOUT_MS,
MIN_REQUEST_BODY_READ_TIMEOUT_MS,
MAX_REQUEST_BODY_READ_TIMEOUT_MS,
),
local_execution_planning_timeout: env_duration_ms(
LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV,
DEFAULT_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
),
}
}
#[cfg(test)]
pub(crate) fn for_tests(
request_body_read_timeout: Duration,
local_execution_planning_timeout: Duration,
) -> Self {
Self {
request_body_read_timeout,
local_execution_planning_timeout,
}
}
}
fn env_duration_ms(key: &str, default_ms: u64, min_ms: u64, max_ms: u64) -> Duration {
let ms = std::env::var(key)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.filter(|value| *value > 0)
.unwrap_or(default_ms)
.clamp(min_ms, max_ms);
Duration::from_millis(ms)
}
#[derive(Debug, Clone)]
pub struct AppState {
#[cfg(test)]
@@ -54,6 +111,7 @@ pub struct AppState {
pub(crate) usage_runtime: Arc<usage::UsageRuntime>,
pub(crate) video_tasks: Arc<VideoTaskService>,
pub(crate) video_task_poller: Option<VideoTaskPollerConfig>,
pub(crate) frontdoor_runtime_guards: Arc<FrontdoorRuntimeGuardConfig>,
pub(crate) request_gate: Option<Arc<ConcurrencyGate>>,
pub(crate) distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
pub(crate) client: reqwest::Client,

5
apps/aether-gateway/src/state/core.rs
View File

@@ -21,7 +21,9 @@ use aether_runtime_state::{
};
use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS;
use super::{AppState, FrontdoorCorsConfig, LocalExecutionRuntimeMissDiagnostic};
use super::{
AppState, FrontdoorCorsConfig, FrontdoorRuntimeGuardConfig, LocalExecutionRuntimeMissDiagnostic,
};
use super::super::async_task::{
spawn_video_task_poller, VideoTaskPollerConfig, VideoTaskService, VideoTaskTruthSourceMode,
@@ -227,6 +229,7 @@ impl AppState {
VideoTaskTruthSourceMode::PythonSyncReport,
)),
video_task_poller: None,
frontdoor_runtime_guards: Arc::new(FrontdoorRuntimeGuardConfig::from_env()),
request_gate: None,
distributed_request_gate: None,
client,

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

@@ -60,12 +60,7 @@ impl provider_transport::VideoTaskTransportSnapshotLookup for AppState {
) -> Result<Option<GatewayProviderTransportSnapshot>, String> {
self.read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
.await
.map_err(|err| match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
})
.map_err(GatewayError::into_message)
}
}
@@ -77,12 +72,7 @@ impl ModelFetchTransportRuntime for AppState {
) -> Result<Option<LocalResolvedOAuthRequestAuth>, String> {
AppState::resolve_local_oauth_request_auth(self, transport)
.await
.map_err(|err| match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
})
.map_err(GatewayError::into_message)
}
async fn resolve_model_fetch_proxy(
@@ -99,12 +89,7 @@ impl ModelFetchTransportRuntime for AppState {
) -> Result<ExecutionResult, String> {
execution_runtime::execute_execution_runtime_sync_plan(self, None, plan)
.await
.map_err(|err| match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
})
.map_err(GatewayError::into_message)
}
}

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

@@ -27,6 +27,7 @@ pub(crate) use self::admin_types::{
UserDailyQuotaAvailabilityRecord, UserPlanEntitlementRecord,
};
pub use self::app::AppState;
pub(crate) use self::app::FrontdoorRuntimeGuardConfig;
pub(crate) use self::cache::{
CachedProviderTransportSnapshot, AUTH_API_KEY_LAST_USED_MAX_ENTRIES,
AUTH_API_KEY_LAST_USED_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES,

7
apps/aether-gateway/src/state/oauth.rs
View File

@@ -1448,12 +1448,7 @@ impl AppState {
.map_err(
|err| provider_transport::LocalOAuthRefreshError::InvalidResponse {
provider_type,
message: match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
},
message: err.into_message(),
},
)?;
let response_body_text = local_oauth_execution_body_text(&result);

10
apps/aether-gateway/src/state/testing.rs
View File

@@ -11,7 +11,7 @@ use aether_data_contracts::repository::video_tasks::{
};
use serde_json::json;
use super::{AppState, GatewayDataState};
use super::{AppState, FrontdoorRuntimeGuardConfig, GatewayDataState};
use crate::{provider_transport, usage};
#[cfg(test)]
@@ -31,6 +31,14 @@ impl AppState {
self
}
pub(crate) fn with_frontdoor_runtime_guard_config_for_tests(
mut self,
config: FrontdoorRuntimeGuardConfig,
) -> Self {
self.frontdoor_runtime_guards = Arc::new(config);
self
}
pub(crate) fn with_tunnel_identity_for_tests(
mut self,
instance_id: &str,

21
crates/aether-data/migrations/postgres/20260522000000_decouple_usage_identity_foreign_keys.sql Normal file
View File

@@ -0,0 +1,21 @@
-- Usage is a historical fact table. Terminal usage events can arrive after
-- mutable catalog/auth rows have been disabled or deleted, so these snapshot
-- identity columns must not make ingestion depend on current dimension rows.
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_provider_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_provider_endpoint_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_provider_api_key_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_api_key_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_user_id_fkey;
ALTER TABLE ONLY public.usage
DROP CONSTRAINT IF EXISTS usage_wallet_id_fkey;

76
crates/aether-data/schema/bootstrap/postgres/005_foreign_keys.sql
View File

@@ -583,80 +583,12 @@ END $mig$;
--
-- Name: usage usage_api_key_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
-- Usage is a historical fact table. Snapshot identity columns such as
-- user_id/api_key_id/provider_endpoint_id/provider_api_key_id/wallet_id
-- intentionally do not carry foreign keys because terminal usage events may
-- arrive after those mutable dimension rows have been disabled or deleted.
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES public.api_keys(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_provider_api_key_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_provider_api_key_id_fkey FOREIGN KEY (provider_api_key_id) REFERENCES public.provider_api_keys(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_provider_endpoint_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_provider_endpoint_id_fkey FOREIGN KEY (provider_endpoint_id) REFERENCES public.provider_endpoints(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_user_id_fkey FOREIGN KEY (user_id) REFERENCES public.users(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: usage usage_wallet_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.usage
ADD CONSTRAINT usage_wallet_id_fkey FOREIGN KEY (wallet_id) REFERENCES public.wallets(id) ON DELETE SET NULL;
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: user_model_usage_counts user_model_usage_counts_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
--

2
crates/aether-data/src/lifecycle/bootstrap/postgres.rs
View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260520010000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260522000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

30
crates/aether-data/src/lifecycle/migrate/tests.rs
View File

@@ -312,6 +312,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260519130000,
20260520000000,
20260520010000,
20260522000000,
]
);
}
@@ -393,6 +394,34 @@ fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role()
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("usage_count bigint DEFAULT 0 NOT NULL"));
}
#[test]
fn usage_identity_foreign_keys_are_decoupled_for_historical_ingestion() {
let migration = POSTGRES_MIGRATOR
.iter()
.find(|migration| migration.version == 20260522000000)
.expect("usage identity foreign key decoupling migration should be embedded");
for constraint in [
"usage_provider_id_fkey",
"usage_provider_endpoint_id_fkey",
"usage_provider_api_key_id_fkey",
"usage_api_key_id_fkey",
"usage_user_id_fkey",
"usage_wallet_id_fkey",
] {
assert!(
migration
.sql
.contains(format!("DROP CONSTRAINT IF EXISTS {constraint}").as_str()),
"migration should drop {constraint}"
);
assert!(
!EMPTY_DATABASE_SNAPSHOT_SQL.contains(format!("ADD CONSTRAINT {constraint}").as_str()),
"fresh bootstrap snapshot should not recreate {constraint}"
);
}
}
#[test]
fn empty_database_snapshot_sql_includes_payment_gateway_and_plans() {
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("payment_provider character varying(64)"));
@@ -1179,6 +1208,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260519130000,
20260520000000,
20260520010000,
20260522000000,
]
);
}

205
crates/aether-usage-runtime/src/worker.rs
View File

@@ -193,16 +193,83 @@ impl UsageQueueWorker {
let event = match UsageEvent::from_stream_fields(&entry.fields) {
Ok(event) => event,
Err(err) => {
warn!(
event_name = "usage_worker_entry_decode_dead_lettered",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
entry_id = %entry.id,
error = %err,
"usage worker moved malformed queue entry to dead letter"
);
self.queue.push_dead_letter(entry, &err.to_string()).await?;
return Ok(true);
}
};
self.recorder.record_usage_event(&event).await?;
Ok(true)
match self.recorder.record_usage_event(&event).await {
Ok(()) => Ok(true),
Err(err) if usage_event_record_error_is_permanent(&err) => {
warn!(
event_name = "usage_worker_entry_record_dead_lettered",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
entry_id = %entry.id,
request_id = %event.request_id,
event_type = ?event.event_type,
provider_name = %event.data.provider_name,
model = %event.data.model,
api_format = event.data.api_format.as_deref().unwrap_or(""),
provider_id = event.data.provider_id.as_deref().unwrap_or(""),
provider_endpoint_id = event.data.provider_endpoint_id.as_deref().unwrap_or(""),
provider_api_key_id = event.data.provider_api_key_id.as_deref().unwrap_or(""),
error = %err,
"usage worker moved non-retryable usage event to dead letter"
);
self.queue.push_dead_letter(entry, &err.to_string()).await?;
Ok(true)
}
Err(err) => {
warn!(
event_name = "usage_worker_entry_record_retryable_failed",
log_type = "ops",
worker_consumer = %self.consumer,
worker_group = %self.config.consumer_group,
entry_id = %entry.id,
request_id = %event.request_id,
event_type = ?event.event_type,
provider_name = %event.data.provider_name,
model = %event.data.model,
api_format = event.data.api_format.as_deref().unwrap_or(""),
provider_id = event.data.provider_id.as_deref().unwrap_or(""),
provider_endpoint_id = event.data.provider_endpoint_id.as_deref().unwrap_or(""),
provider_api_key_id = event.data.provider_api_key_id.as_deref().unwrap_or(""),
error = %err,
"usage worker will retry usage event after record failure"
);
Err(err)
}
}
}
}
fn usage_event_record_error_is_permanent(err: &DataLayerError) -> bool {
match err {
DataLayerError::InvalidConfiguration(_)
| DataLayerError::InvalidInput(_)
| DataLayerError::UnexpectedValue(_) => true,
DataLayerError::Postgres(message) | DataLayerError::Sql(message) => {
database_error_is_known_permanent(message)
}
DataLayerError::Redis(_) | DataLayerError::TimedOut(_) => false,
}
}
fn database_error_is_known_permanent(message: &str) -> bool {
message.contains("SQLSTATE 23503") || message.contains("violates foreign key constraint")
}
pub fn build_usage_queue_worker<T>(
runner: Arc<dyn RuntimeQueueStore>,
data: Arc<T>,
@@ -287,16 +354,23 @@ fn consumer_name() -> String {
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use std::sync::{Arc, Mutex};
use aether_data_contracts::repository::settlement::{
StoredUsageSettlement, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeQueueStore, RuntimeState};
use async_trait::async_trait;
use super::{write_event_record, ManualProxyNodeCounter, UsageRecordWriter};
use crate::{UsageEvent, UsageEventData, UsageEventType, UsageSettlementWriter};
use super::{
usage_event_record_error_is_permanent, write_event_record, ManualProxyNodeCounter,
UsageEventRecorder, UsageQueueWorker, UsageRecordWriter,
};
use crate::{
UsageEvent, UsageEventData, UsageEventType, UsageRuntimeConfig, UsageSettlementWriter,
};
#[derive(Default)]
struct TestUsageStore {
@@ -304,6 +378,11 @@ mod tests {
settlements: Mutex<Vec<UsageSettlementInput>>,
}
#[derive(Default)]
struct SelectiveFailingRecorder {
calls: Mutex<Vec<String>>,
}
#[async_trait]
impl UsageRecordWriter for TestUsageStore {
async fn upsert_usage_record(
@@ -392,6 +471,22 @@ mod tests {
}
}
#[async_trait]
impl UsageEventRecorder for SelectiveFailingRecorder {
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
self.calls
.lock()
.expect("calls lock")
.push(event.request_id.clone());
if event.request_id == "req-worker-poison" {
return Err(DataLayerError::UnexpectedValue(
"permanent test error".to_string(),
));
}
Ok(())
}
}
fn sample_event() -> UsageEvent {
UsageEvent::new(
UsageEventType::Completed,
@@ -436,4 +531,104 @@ mod tests {
assert_eq!(settlements.len(), 1);
assert_eq!(settlements[0].request_id, "req-worker-123");
}
#[test]
fn usage_event_record_error_classifies_permanent_failures() {
assert!(usage_event_record_error_is_permanent(
&DataLayerError::UnexpectedValue("bad payload".to_string())
));
assert!(usage_event_record_error_is_permanent(
&DataLayerError::Postgres(
"error returned from database: violates foreign key constraint (SQLSTATE 23503)"
.to_string()
)
));
assert!(!usage_event_record_error_is_permanent(
&DataLayerError::Redis("connection refused".to_string())
));
assert!(!usage_event_record_error_is_permanent(
&DataLayerError::TimedOut("postgres acquire".to_string())
));
}
#[tokio::test]
async fn process_entries_dead_letters_permanent_record_error_and_continues() {
let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let queue_runner: Arc<dyn RuntimeQueueStore> = runner.clone();
let recorder = Arc::new(SelectiveFailingRecorder::default());
let config = UsageRuntimeConfig {
enabled: true,
stream_key: "usage:test:worker:events".to_string(),
consumer_group: "usage:test:worker:group".to_string(),
dlq_stream_key: "usage:test:worker:dlq".to_string(),
consumer_batch_size: 10,
consumer_block_ms: 1,
..UsageRuntimeConfig::default()
};
let worker = UsageQueueWorker::new(queue_runner, recorder.clone(), config)
.expect("worker should build");
worker
.queue
.ensure_consumer_group()
.await
.expect("group should initialize");
let mut poison = sample_event();
poison.request_id = "req-worker-poison".to_string();
let mut ok = sample_event();
ok.request_id = "req-worker-ok".to_string();
worker
.queue
.enqueue(&poison)
.await
.expect("poison event should enqueue");
worker
.queue
.enqueue(&ok)
.await
.expect("ok event should enqueue");
let entries = worker
.queue
.read_group(&worker.consumer)
.await
.expect("events should read");
assert_eq!(entries.len(), 2);
worker
.process_entries(entries)
.await
.expect("permanent failure should not block batch");
assert_eq!(
recorder.calls.lock().expect("calls lock").as_slice(),
["req-worker-poison", "req-worker-ok"]
);
runner
.ensure_consumer_group(
"usage:test:worker:dlq",
"usage:test:worker:dlq-group",
"0-0",
)
.await
.expect("dlq group should initialize");
let dlq_entries = runner
.read_group(
"usage:test:worker:dlq",
"usage:test:worker:dlq-group",
"usage-test-dlq-consumer",
10,
Some(1),
)
.await
.expect("dlq should read");
assert_eq!(dlq_entries.len(), 1);
let payload = dlq_entries[0]
.fields
.get("payload")
.expect("dlq payload should exist");
assert!(payload.contains("req-worker-poison"));
assert!(payload.contains("permanent test error"));
}
}