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

Rootfix usage counter outbox

Browse Source
This commit is contained in:
fawney19
2026-05-19 00:30:24 +08:00
parent 19c8688eb1
commit 5eed329916
60 changed files with 4910 additions and 362 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,15 +1,15 @@
use super::{
AuthApiKeyLookupKey, CreateManagementTokenRecord, DataLayerError, GatewayAuthApiKeySnapshot,
GatewayDataState, ManagementTokenListQuery, ProxyNodeHeartbeatMutation,
ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation, ProxyNodeRegistrationMutation,
ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation, ProxyNodeTunnelStatusMutation,
RegenerateManagementTokenSecret, StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
StoredLdapModuleConfig, StoredManagementToken, StoredManagementTokenListPage,
StoredManagementTokenWithUser, StoredOAuthProviderConfig, StoredOAuthProviderModuleConfig,
StoredProxyFleetMetricsBucket, StoredProxyNode, StoredProxyNodeEvent,
StoredProxyNodeMetricsBucket, StoredUserAuthRecord, StoredUserOAuthLinkSummary,
StoredUserPreferenceRecord, StoredUserSessionRecord, StoredWalletSnapshot,
UpdateManagementTokenRecord, UpsertOAuthProviderConfigRecord,
GatewayDataState, ManagementTokenCounterDelta, ManagementTokenListQuery, ProxyNodeCounterDelta,
ProxyNodeHeartbeatMutation, ProxyNodeManualCreateMutation, ProxyNodeManualUpdateMutation,
ProxyNodeRegistrationMutation, ProxyNodeRemoteConfigMutation, ProxyNodeTrafficMutation,
ProxyNodeTunnelStatusMutation, RegenerateManagementTokenSecret, StoredAuthApiKeyExportRecord,
StoredAuthApiKeySnapshot, StoredLdapModuleConfig, StoredManagementToken,
StoredManagementTokenListPage, StoredManagementTokenWithUser, StoredOAuthProviderConfig,
StoredOAuthProviderModuleConfig, StoredProxyFleetMetricsBucket, StoredProxyNode,
StoredProxyNodeEvent, StoredProxyNodeMetricsBucket, StoredUserAuthRecord,
StoredUserOAuthLinkSummary, StoredUserPreferenceRecord, StoredUserSessionRecord,
StoredWalletSnapshot, UpdateManagementTokenRecord, UpsertOAuthProviderConfigRecord,
};
use crate::LocalMutationOutcome;
use aether_data::repository::auth::{
@@ -1117,6 +1117,20 @@ impl GatewayDataState {
token_id: &str,
last_used_ip: Option<&str>,
) -> Result<Option<StoredManagementToken>, DataLayerError> {
if let Some(repository) = &self.usage_writer {
let enqueued = repository
.enqueue_management_token_counter_delta(ManagementTokenCounterDelta {
token_id: token_id.to_string(),
usage_count_delta: 1,
last_used_at_unix_secs: Some(chrono::Utc::now().timestamp().max(0) as u64),
last_used_ip: last_used_ip.map(ToOwned::to_owned),
})
.await?;
if enqueued {
return Ok(None);
}
}
match &self.management_token_writer {
Some(repository) => {
repository
@@ -1278,6 +1292,21 @@ impl GatewayDataState {
&self,
mutation: &ProxyNodeTrafficMutation,
) -> Result<bool, DataLayerError> {
if let Some(repository) = &self.usage_writer {
let enqueued = repository
.enqueue_proxy_node_counter_delta(ProxyNodeCounterDelta {
node_id: mutation.node_id.clone(),
total_requests_delta: mutation.total_requests_delta,
failed_requests_delta: mutation.failed_requests_delta,
dns_failures_delta: mutation.dns_failures_delta,
stream_errors_delta: mutation.stream_errors_delta,
})
.await?;
if enqueued {
return Ok(true);
}
}
match &self.proxy_node_writer {
Some(repository) => repository.record_traffic(mutation).await,
None => Ok(false),

401
apps/aether-gateway/src/data/state/candidate_cache.rs Normal file
View File

@@ -0,0 +1,401 @@
use std::collections::HashSet;
use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use aether_cache::ExpiringMap;
use aether_data::DataLayerError;
use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
};
use async_trait::async_trait;
use tokio::sync::Notify;
const CANDIDATE_SELECTION_CACHE_TTL: Duration = Duration::from_secs(5);
const CANDIDATE_SELECTION_CACHE_MAX_ENTRIES: usize = 4096;
pub(super) struct CachedMinimalCandidateSelectionReadRepository {
inner: Arc<dyn MinimalCandidateSelectionReadRepository>,
entries: ExpiringMap<CandidateSelectionCacheKey, Vec<StoredMinimalCandidateSelectionRow>>,
inflight: Mutex<HashSet<CandidateSelectionCacheKey>>,
inflight_notify: Notify,
epoch: AtomicU64,
}
impl CachedMinimalCandidateSelectionReadRepository {
pub(super) fn new(inner: Arc<dyn MinimalCandidateSelectionReadRepository>) -> Self {
Self {
inner,
entries: ExpiringMap::new(),
inflight: Mutex::new(HashSet::new()),
inflight_notify: Notify::new(),
epoch: AtomicU64::new(0),
}
}
async fn get_or_load<F, Fut>(
&self,
key: CandidateSelectionCacheKey,
load: F,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>
where
F: Fn() -> Fut,
Fut: Future<Output = Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>>,
{
if let Some(rows) = self.entries.get_fresh(&key, CANDIDATE_SELECTION_CACHE_TTL) {
return Ok(rows);
}
loop {
let notified = self.inflight_notify.notified();
match self.register_inflight(&key) {
InflightRegistration::Bypass => return load().await,
InflightRegistration::Follower => {
notified.await;
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,
);
}
}
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
}
}
Err(_) => InflightRegistration::Bypass,
}
}
fn finish_inflight(&self, key: &CandidateSelectionCacheKey) {
if let Ok(mut inflight) = self.inflight.lock() {
inflight.remove(key);
}
self.inflight_notify.notify_waiters();
}
fn clear(&self) {
self.epoch.fetch_add(1, Ordering::AcqRel);
self.entries.clear();
}
}
enum InflightRegistration {
Leader,
Follower,
Bypass,
}
#[async_trait]
impl MinimalCandidateSelectionReadRepository for CachedMinimalCandidateSelectionReadRepository {
fn clear_local_cache(&self) {
self.clear();
self.inner.clear_local_cache();
}
async fn list_for_exact_api_format(
&self,
api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let key = CandidateSelectionCacheKey::ApiFormat {
api_format: normalize_api_format_key(api_format),
};
self.get_or_load(key, || self.inner.list_for_exact_api_format(api_format))
.await
}
async fn list_for_exact_api_format_and_global_model(
&self,
api_format: &str,
global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let key = CandidateSelectionCacheKey::ApiFormatAndGlobalModel {
api_format: normalize_api_format_key(api_format),
global_model_name: global_model_name.to_string(),
};
self.get_or_load(key, || {
self.inner
.list_for_exact_api_format_and_global_model(api_format, global_model_name)
})
.await
}
async fn list_for_exact_api_format_and_requested_model(
&self,
api_format: &str,
requested_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let key = CandidateSelectionCacheKey::ApiFormatAndRequestedModel {
api_format: normalize_api_format_key(api_format),
requested_model_name: requested_model_name.to_string(),
};
self.get_or_load(key, || {
self.inner
.list_for_exact_api_format_and_requested_model(api_format, requested_model_name)
})
.await
}
async fn list_for_exact_api_format_and_requested_model_page(
&self,
query: &StoredRequestedModelCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let key = CandidateSelectionCacheKey::RequestedModelPage {
api_format: normalize_api_format_key(&query.api_format),
requested_model_name: query.requested_model_name.clone(),
offset: query.offset,
limit: query.limit,
};
self.get_or_load(key, || {
self.inner
.list_for_exact_api_format_and_requested_model_page(query)
})
.await
}
async fn list_pool_key_rows_for_group(
&self,
query: &StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let key = CandidateSelectionCacheKey::PoolKeyRowsForGroup {
api_format: normalize_api_format_key(&query.api_format),
provider_id: query.provider_id.clone(),
endpoint_id: query.endpoint_id.clone(),
model_id: query.model_id.clone(),
selected_provider_model_name: query.selected_provider_model_name.clone(),
order: CandidateSelectionPoolOrderKey::from(&query.order),
offset: query.offset,
limit: query.limit,
};
self.get_or_load(key, || self.inner.list_pool_key_rows_for_group(query))
.await
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let key = CandidateSelectionCacheKey::PoolKeyRowsForGroupKeyIds {
api_format: normalize_api_format_key(&query.api_format),
provider_id: query.provider_id.clone(),
endpoint_id: query.endpoint_id.clone(),
model_id: query.model_id.clone(),
selected_provider_model_name: query.selected_provider_model_name.clone(),
key_ids: query.key_ids.clone(),
};
self.get_or_load(key, || {
self.inner.list_pool_key_rows_for_group_key_ids(query)
})
.await
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum CandidateSelectionCacheKey {
ApiFormat {
api_format: String,
},
ApiFormatAndGlobalModel {
api_format: String,
global_model_name: String,
},
ApiFormatAndRequestedModel {
api_format: String,
requested_model_name: String,
},
RequestedModelPage {
api_format: String,
requested_model_name: String,
offset: u32,
limit: u32,
},
PoolKeyRowsForGroup {
api_format: String,
provider_id: String,
endpoint_id: String,
model_id: String,
selected_provider_model_name: String,
order: CandidateSelectionPoolOrderKey,
offset: u32,
limit: u32,
},
PoolKeyRowsForGroupKeyIds {
api_format: String,
provider_id: String,
endpoint_id: String,
model_id: String,
selected_provider_model_name: String,
key_ids: Vec<String>,
},
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum CandidateSelectionPoolOrderKey {
InternalPriority,
Lru,
CacheAffinity,
SingleAccount,
LoadBalance { seed: String },
}
impl From<&StoredPoolKeyCandidateOrder> for CandidateSelectionPoolOrderKey {
fn from(order: &StoredPoolKeyCandidateOrder) -> Self {
match order {
StoredPoolKeyCandidateOrder::InternalPriority => Self::InternalPriority,
StoredPoolKeyCandidateOrder::Lru => Self::Lru,
StoredPoolKeyCandidateOrder::CacheAffinity => Self::CacheAffinity,
StoredPoolKeyCandidateOrder::SingleAccount => Self::SingleAccount,
StoredPoolKeyCandidateOrder::LoadBalance { seed } => {
Self::LoadBalance { seed: seed.clone() }
}
}
}
}
fn normalize_api_format_key(api_format: &str) -> String {
crate::ai_serving::normalize_api_format_alias(api_format.trim())
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::AtomicUsize;
struct StubCandidateSelectionRepository {
calls: AtomicUsize,
delay: Duration,
}
impl StubCandidateSelectionRepository {
fn new(delay: Duration) -> Self {
Self {
calls: AtomicUsize::new(0),
delay,
}
}
fn calls(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
async fn load(&self) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.calls.fetch_add(1, Ordering::SeqCst);
if !self.delay.is_zero() {
tokio::time::sleep(self.delay).await;
}
Ok(Vec::new())
}
}
#[async_trait]
impl MinimalCandidateSelectionReadRepository for StubCandidateSelectionRepository {
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(
Duration::from_millis(25),
));
let cache = Arc::new(CachedMinimalCandidateSelectionReadRepository::new(
inner.clone(),
));
let mut tasks = Vec::new();
for _ in 0..16 {
let cache = cache.clone();
tasks.push(tokio::spawn(async move {
cache.list_for_exact_api_format("openai").await.unwrap();
}));
}
for task in tasks {
task.await.unwrap();
}
assert_eq!(inner.calls(), 1);
}
#[tokio::test]
async fn candidate_selection_cache_clear_invalidates_entries() {
let inner = Arc::new(StubCandidateSelectionRepository::new(Duration::ZERO));
let cache = CachedMinimalCandidateSelectionReadRepository::new(inner.clone());
cache.list_for_exact_api_format("openai").await.unwrap();
cache.list_for_exact_api_format("openai").await.unwrap();
assert_eq!(inner.calls(), 1);
cache.clear_local_cache();
cache.list_for_exact_api_format("openai").await.unwrap();
assert_eq!(inner.calls(), 2);
}
}

24
apps/aether-gateway/src/data/state/catalog.rs
View File

@@ -1,10 +1,10 @@
use super::{
DataLayerError, GatewayDataState, GeminiFileMappingListQuery, GeminiFileMappingStats,
ProviderCatalogKeyListQuery, PublicHealthStatusCount, PublicHealthTimelineBucket,
StoredGeminiFileMapping, StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats,
StoredProviderCatalogProvider, StoredRequestCandidate, UpsertGeminiFileMappingRecord,
UpsertRequestCandidateRecord,
ApiKeyLastUsedDelta, DataLayerError, GatewayDataState, GeminiFileMappingListQuery,
GeminiFileMappingStats, ProviderCatalogKeyListQuery, PublicHealthStatusCount,
PublicHealthTimelineBucket, StoredGeminiFileMapping, StoredGeminiFileMappingListPage,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, StoredRequestCandidate,
UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
};
impl GatewayDataState {
@@ -121,6 +121,18 @@ impl GatewayDataState {
&self,
api_key_id: &str,
) -> Result<bool, DataLayerError> {
if let Some(repository) = &self.usage_writer {
let enqueued = repository
.enqueue_api_key_last_used_delta(ApiKeyLastUsedDelta {
api_key_id: api_key_id.to_string(),
last_used_at_unix_secs: chrono::Utc::now().timestamp().max(0) as u64,
})
.await?;
if enqueued {
return Ok(true);
}
}
match &self.auth_api_key_writer {
Some(repository) => repository.touch_last_used_at(api_key_id).await,
None => Ok(false),

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

@@ -1,4 +1,5 @@
use aether_data::{DataBackends, DataLayerError, DatabaseDriver};
use aether_data_contracts::repository::candidate_selection::MinimalCandidateSelectionReadRepository;
use aether_runtime_state::RuntimeQueueStore;
use std::sync::Arc;
@@ -84,7 +85,17 @@ impl GatewayDataState {
let gemini_file_mapping_reader = backends.read().gemini_file_mappings();
let global_model_reader = backends.read().global_models();
let global_model_writer = backends.write().global_models();
let minimal_candidate_selection_reader = backends.read().minimal_candidate_selection();
let minimal_candidate_selection_reader =
backends
.read()
.minimal_candidate_selection()
.map(|repository| {
Arc::new(
super::candidate_cache::CachedMinimalCandidateSelectionReadRepository::new(
repository,
),
) as Arc<dyn MinimalCandidateSelectionReadRepository>
});
let request_candidate_reader = backends.read().request_candidates();
let request_candidate_writer = backends.write().request_candidates();
let gemini_file_mapping_writer = backends.write().gemini_file_mappings();
@@ -259,6 +270,12 @@ impl GatewayDataState {
self.minimal_candidate_selection_reader.is_some()
}
pub(crate) fn clear_minimal_candidate_selection_cache(&self) {
if let Some(repository) = &self.minimal_candidate_selection_reader {
repository.clear_local_cache();
}
}
pub(crate) fn has_request_candidate_reader(&self) -> bool {
self.request_candidate_reader.is_some()
}
@@ -329,6 +346,10 @@ impl GatewayDataState {
self.usage_writer.is_some()
}
pub(crate) fn has_usage_counter_flush_backend(&self) -> bool {
self.has_usage_writer() && self.database_driver() == Some(DatabaseDriver::Postgres)
}
pub(crate) fn has_usage_worker_queue(&self) -> bool {
self.usage_worker_queue.is_some()
}

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

@@ -12,7 +12,9 @@ use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::repository::usage::{
ProxyNodeCounterDelta, StoredRequestUsageAudit, UpsertUsageRecord,
};
use aether_data_contracts::repository::video_tasks::{StoredVideoTask, VideoTaskLookupKey};
use aether_runtime_state::RuntimeQueueStore;
use aether_usage_runtime::{
@@ -284,6 +286,21 @@ impl aether_usage_runtime::ManualProxyNodeCounter for GatewayDataState {
failed_delta: i64,
latency_ms: Option<i64>,
) -> Result<(), DataLayerError> {
if let Some(repository) = &self.usage_writer {
let enqueued = repository
.enqueue_proxy_node_counter_delta(ProxyNodeCounterDelta {
node_id: node_id.to_string(),
total_requests_delta: total_delta,
failed_requests_delta: failed_delta,
dns_failures_delta: 0,
stream_errors_delta: 0,
})
.await?;
if enqueued {
return Ok(());
}
}
match &self.proxy_node_writer {
Some(repository) => {
repository

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

@@ -132,8 +132,9 @@ use aether_data_contracts::repository::settlement::{
SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
UpsertUsageRecord, UsageReadRepository, UsageWriteRepository,
ApiKeyLastUsedDelta, ManagementTokenCounterDelta, PendingUsageCleanupSummary,
ProxyNodeCounterDelta, StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord,
UsageReadRepository, UsageWriteRepository,
};
use aether_data_contracts::repository::video_tasks::{
StoredVideoTask, UpsertVideoTask, VideoTaskLookupKey, VideoTaskModelCount,
@@ -310,6 +311,7 @@ impl fmt::Debug for GatewayDataState {
}
mod auth;
mod candidate_cache;
mod catalog;
mod core;
mod integrations;

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

@@ -38,7 +38,7 @@ use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
StoredProviderApiKeyWindowUsageSummary, StoredUsageDailySummary, UsageAuditListQuery,
UsageCleanupExecutionMode, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow,
UsageDailyHeatmapQuery,
UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageDailyHeatmapQuery,
};
use aether_runtime_state::RuntimeQueueStore;
use aether_video_tasks_core::read_data_backed_video_task_response;
@@ -954,6 +954,31 @@ impl GatewayDataState {
}
}
pub(crate) async fn flush_usage_counter_deltas(
&self,
batch_size: usize,
) -> Result<UsageCounterFlushSummary, DataLayerError> {
match &self.usage_writer {
Some(repository) => repository.flush_usage_counter_deltas(batch_size).await,
None => Ok(UsageCounterFlushSummary::default()),
}
}
pub(crate) async fn cleanup_processed_usage_counter_deltas(
&self,
cutoff_unix_secs: u64,
batch_size: usize,
) -> Result<usize, DataLayerError> {
match &self.usage_writer {
Some(repository) => {
repository
.cleanup_processed_usage_counter_deltas(cutoff_unix_secs, batch_size)
.await
}
None => Ok(0),
}
}
pub(crate) async fn cleanup_stale_pending_requests(
&self,
cutoff_unix_secs: u64,
@@ -1119,6 +1144,15 @@ impl GatewayDataState {
}
}
pub(crate) async fn read_usage_counter_health(
&self,
) -> Result<UsageCounterHealthSnapshot, DataLayerError> {
match &self.usage_reader {
Some(repository) => repository.read_usage_counter_health().await,
None => Ok(UsageCounterHealthSnapshot::default()),
}
}
pub(crate) async fn summarize_usage_totals_by_user_ids(
&self,
user_ids: &[String],

9
apps/aether-gateway/src/handlers/admin/observability/monitoring/activity.rs
View File

@@ -6,6 +6,7 @@ use super::route_filters::{
};
use crate::constants::INTERNAL_GATEWAY_PATH_PREFIXES;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::build_admin_usage_counter_health_payload;
use crate::GatewayError;
use aether_admin::observability::monitoring::{
admin_monitoring_bad_request_response, admin_monitoring_user_behavior_user_id_from_path,
@@ -189,6 +190,13 @@ pub(super) async fn build_admin_monitoring_system_status_response(
)
.unwrap_or(usize::MAX);
let tunnel = state.tunnel.stats();
let usage_counter_snapshot = state
.data
.read_usage_counter_health()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let usage_counter =
build_admin_usage_counter_health_payload(&usage_counter_snapshot, now_unix_secs);
Ok(build_admin_monitoring_system_status_payload_response(
now,
@@ -206,5 +214,6 @@ pub(super) async fn build_admin_monitoring_system_status_response(
tunnel.active_streams,
INTERNAL_GATEWAY_PATH_PREFIXES,
recent_errors,
usage_counter,
))
}

30
apps/aether-gateway/src/handlers/admin/observability/stats/analytics_routes.rs
View File

@@ -1,7 +1,9 @@
use super::range::{build_comparison_range, parse_bounded_u32};
use super::resolve_admin_usage_time_range;
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::admin::shared::{query_param_optional_bool, query_param_value};
use crate::handlers::admin::shared::{
build_admin_usage_counter_health_payload, query_param_optional_bool, query_param_value,
};
use crate::GatewayError;
use aether_admin::observability::stats::{
admin_stats_bad_request_response, admin_stats_comparison_empty_response,
@@ -21,6 +23,22 @@ use aether_data_contracts::repository::usage::{
};
use axum::{body::Body, http, response::Response};
async fn build_usage_counter_health_payload(
state: &AdminAppState<'_>,
) -> Result<serde_json::Value, GatewayError> {
let now_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
let snapshot = state
.as_ref()
.data
.read_usage_counter_health()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(build_admin_usage_counter_health_payload(
&snapshot,
now_unix_secs,
))
}
fn usage_summary_to_admin_stats_aggregate(
summary: &aether_data_contracts::repository::usage::StoredUsageAuditSummary,
) -> AdminStatsAggregate {
@@ -211,13 +229,18 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
Ok(value) => u64::from(value.unwrap_or(10_000)),
Err(detail) => return Ok(Some(admin_stats_bad_request_response(detail))),
};
let usage_counter = build_usage_counter_health_payload(state).await?;
if !state.has_usage_data_reader() {
return Ok(Some(admin_stats_provider_performance_empty_response()));
return Ok(Some(admin_stats_provider_performance_empty_response(
usage_counter,
)));
}
let Some((created_from_unix_secs, created_until_unix_secs)) = time_range.to_unix_bounds()
else {
return Ok(Some(admin_stats_provider_performance_empty_response()));
return Ok(Some(admin_stats_provider_performance_empty_response(
usage_counter,
)));
};
let performance = state
.summarize_usage_provider_performance(&UsageProviderPerformanceQuery {
@@ -240,6 +263,7 @@ pub(super) async fn maybe_build_local_admin_stats_analytics_response(
.await?;
return Ok(Some(build_admin_stats_provider_performance_response(
&performance,
usage_counter,
)));
}

2
apps/aether-gateway/src/handlers/admin/shared/mod.rs
View File

@@ -1,10 +1,12 @@
mod paths;
mod payloads;
mod proxy_errors;
mod usage_counter;
pub(crate) use self::paths::*;
pub(crate) use self::payloads::*;
pub(crate) use self::proxy_errors::build_proxy_error_response;
pub(crate) use self::usage_counter::build_admin_usage_counter_health_payload;
pub(crate) use crate::handlers::shared::{
attach_admin_audit_response, build_admin_provider_key_response,
decrypt_catalog_secret_with_fallbacks, default_provider_key_status_snapshot,

85
apps/aether-gateway/src/handlers/admin/shared/usage_counter.rs Normal file
View File

@@ -0,0 +1,85 @@
use aether_data_contracts::repository::usage::UsageCounterHealthSnapshot;
use serde_json::{json, Value};
pub(crate) fn build_admin_usage_counter_health_payload(
snapshot: &UsageCounterHealthSnapshot,
now_unix_secs: u64,
) -> Value {
let oldest_pending_age_secs = snapshot
.oldest_pending_created_at_unix_secs
.map(|created_at| now_unix_secs.saturating_sub(created_at));
let status = match (snapshot.pending_rows, oldest_pending_age_secs) {
(0, _) => "idle",
(_, Some(age)) if age >= 60 => "backlogged",
_ => "catching_up",
};
json!({
"status": status,
"outbox_pending_rows": snapshot.pending_rows,
"outbox_processed_rows": snapshot.processed_rows,
"oldest_pending_created_at_unix_secs": snapshot.oldest_pending_created_at_unix_secs,
"oldest_pending_age_secs": oldest_pending_age_secs,
"latest_processed_at_unix_secs": snapshot.latest_processed_at_unix_secs,
"pending_by_kind": snapshot.pending_by_kind,
})
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use aether_data_contracts::repository::usage::UsageCounterHealthSnapshot;
use serde_json::json;
use super::build_admin_usage_counter_health_payload;
#[test]
fn usage_counter_health_payload_reports_idle_without_pending_rows() {
let snapshot = UsageCounterHealthSnapshot {
processed_rows: 42,
latest_processed_at_unix_secs: Some(1_000),
..UsageCounterHealthSnapshot::default()
};
let payload = build_admin_usage_counter_health_payload(&snapshot, 1_100);
assert_eq!(payload["status"], json!("idle"));
assert_eq!(payload["outbox_pending_rows"], json!(0));
assert_eq!(payload["outbox_processed_rows"], json!(42));
assert_eq!(payload["oldest_pending_age_secs"], json!(null));
assert_eq!(payload["latest_processed_at_unix_secs"], json!(1_000));
}
#[test]
fn usage_counter_health_payload_reports_catching_up_for_fresh_backlog() {
let mut pending_by_kind = BTreeMap::new();
pending_by_kind.insert("api_key".to_string(), 3);
let snapshot = UsageCounterHealthSnapshot {
pending_rows: 3,
oldest_pending_created_at_unix_secs: Some(1_050),
pending_by_kind,
..UsageCounterHealthSnapshot::default()
};
let payload = build_admin_usage_counter_health_payload(&snapshot, 1_100);
assert_eq!(payload["status"], json!("catching_up"));
assert_eq!(payload["oldest_pending_age_secs"], json!(50));
assert_eq!(payload["pending_by_kind"]["api_key"], json!(3));
}
#[test]
fn usage_counter_health_payload_reports_backlogged_for_old_backlog() {
let snapshot = UsageCounterHealthSnapshot {
pending_rows: 1,
oldest_pending_created_at_unix_secs: Some(1_000),
..UsageCounterHealthSnapshot::default()
};
let payload = build_admin_usage_counter_health_payload(&snapshot, 1_060);
assert_eq!(payload["status"], json!("backlogged"));
assert_eq!(payload["oldest_pending_age_secs"], json!(60));
}
}

11
apps/aether-gateway/src/handlers/admin/system/shared/settings.rs
View File

@@ -1,4 +1,5 @@
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::build_admin_usage_counter_health_payload;
use crate::handlers::shared::{system_config_bool, system_config_string};
use crate::GatewayError;
use aether_admin::system::{
@@ -114,6 +115,15 @@ pub(crate) async fn build_admin_system_stats_payload(
.filter(|provider| provider.is_active)
.count() as u64;
let stats = state.read_admin_system_stats().await?;
let now_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
let usage_counter_snapshot = state
.as_ref()
.data
.read_usage_counter_health()
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let usage_counter =
build_admin_usage_counter_health_payload(&usage_counter_snapshot, now_unix_secs);
Ok(build_admin_system_stats_payload_pure(
stats.total_users,
@@ -122,6 +132,7 @@ pub(crate) async fn build_admin_system_stats_payload(
active_providers,
stats.total_api_keys,
stats.total_requests,
usage_counter,
))
}

21
apps/aether-gateway/src/maintenance/mod.rs
View File

@@ -19,14 +19,15 @@ pub(crate) use runtime::{
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_admin_request_body_cleanup_task, start_admin_system_purge_task,
start_manual_usage_cleanup_task, start_proxy_upgrade_rollout, AccountSelfCheckRunSummary,
AdminCleanupRunRecord, AdminCleanupTaskKind, AdminStatsRebuildSummary,
AdminSystemCleanupSummary, ManualUsageCleanupError, ManualUsageCleanupMode,
ManualUsageCleanupOptions, OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary,
PoolQuotaProbeWorkerConfig, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,
spawn_usage_cleanup_worker, spawn_usage_counter_flush_worker,
spawn_wallet_daily_usage_aggregation_worker, start_admin_request_body_cleanup_task,
start_admin_system_purge_task, start_manual_usage_cleanup_task, start_proxy_upgrade_rollout,
AccountSelfCheckRunSummary, AdminCleanupRunRecord, AdminCleanupTaskKind,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, ManualUsageCleanupError,
ManualUsageCleanupMode, ManualUsageCleanupOptions, OAuthTokenRefreshRunSummary,
PoolQuotaProbeRunSummary, PoolQuotaProbeWorkerConfig, ProviderCheckinRunSummary,
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary,
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutProbeConfig,
ProxyUpgradeRolloutSkippedRestoreSummary, ProxyUpgradeRolloutStatus,
ProxyUpgradeRolloutTrackedNodeState,
};

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

@@ -47,6 +47,8 @@ mod stats_hourly;
mod tests;
#[path = "runtime/usage_cleanup.rs"]
mod usage_cleanup;
#[path = "runtime/usage_counter_flush.rs"]
mod usage_counter_flush;
#[path = "runtime/wallet_daily_usage.rs"]
mod wallet_daily_usage;
#[path = "runtime/workers.rs"]
@@ -108,6 +110,7 @@ use usage_cleanup::*;
pub(crate) use usage_cleanup::{
preview_manual_usage_cleanup, ManualUsageCleanupMode, ManualUsageCleanupOptions,
};
use usage_counter_flush::*;
use wallet_daily_usage::*;
pub(crate) use workers::*;
@@ -128,6 +131,12 @@ const PROXY_NODE_STALE_MIN_GRACE_SECS: u64 = 15;
const PROXY_NODE_STALE_MISSED_HEARTBEATS: u64 = 3;
const POOL_MONITOR_INTERVAL: Duration = Duration::from_secs(5 * 60);
const OAUTH_TOKEN_REFRESH_INTERVAL: Duration = Duration::from_secs(60);
const USAGE_COUNTER_FLUSH_INTERVAL: Duration = Duration::from_secs(1);
const USAGE_COUNTER_FLUSH_BATCH_SIZE: usize = 1_000;
const USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT: usize = 20;
const USAGE_COUNTER_DELTA_CLEANUP_INTERVAL: Duration = Duration::from_secs(60);
const USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE: usize = 5_000;
const USAGE_COUNTER_DELTA_RETENTION_SECS: u64 = 7 * 24 * 60 * 60;
const PROVIDER_CHECKIN_CONCURRENCY: usize = 3;
const PROVIDER_CHECKIN_DEFAULT_TIME: &str = "01:05";
const REQUEST_CANDIDATE_CLEANUP_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);

21
apps/aether-gateway/src/maintenance/runtime/usage_counter_flush.rs Normal file
View File

@@ -0,0 +1,21 @@
use crate::data::GatewayDataState;
use aether_data::DataLayerError;
use aether_data_contracts::repository::usage::UsageCounterFlushSummary;
pub(crate) async fn run_usage_counter_flush_once(
data: &GatewayDataState,
batch_size: usize,
) -> Result<UsageCounterFlushSummary, DataLayerError> {
data.flush_usage_counter_deltas(batch_size).await
}
pub(crate) async fn cleanup_processed_usage_counter_deltas_once(
data: &GatewayDataState,
retention_secs: u64,
batch_size: usize,
) -> Result<usize, DataLayerError> {
let now = chrono::Utc::now().timestamp().max(0) as u64;
let cutoff = now.saturating_sub(retention_secs);
data.cleanup_processed_usage_counter_deltas(cutoff, batch_size)
.await
}

68
apps/aether-gateway/src/maintenance/runtime/workers.rs
View File

@@ -7,21 +7,24 @@ use crate::data::GatewayDataState;
use crate::AppState;
use super::{
duration_until_next_daily_run, duration_until_next_db_maintenance_run,
duration_until_next_stats_aggregation_run, duration_until_next_stats_hourly_aggregation_run,
maintenance_timezone, parse_hhmm_time, perform_oauth_token_refresh_once,
provider_checkin_schedule, run_audit_cleanup_once, run_db_maintenance_once,
run_gemini_file_mapping_cleanup_once, run_pending_cleanup_once, run_pool_monitor_once,
run_provider_checkin_once, run_proxy_node_metrics_cleanup_once,
cleanup_processed_usage_counter_deltas_once, duration_until_next_daily_run,
duration_until_next_db_maintenance_run, duration_until_next_stats_aggregation_run,
duration_until_next_stats_hourly_aggregation_run, maintenance_timezone, parse_hhmm_time,
perform_oauth_token_refresh_once, provider_checkin_schedule, run_audit_cleanup_once,
run_db_maintenance_once, run_gemini_file_mapping_cleanup_once, run_pending_cleanup_once,
run_pool_monitor_once, run_provider_checkin_once, run_proxy_node_metrics_cleanup_once,
run_proxy_node_stale_cleanup_once, run_proxy_upgrade_rollout_once,
run_request_candidate_cleanup_once, run_stats_aggregation_once,
run_stats_hourly_aggregation_once, run_usage_cleanup_once,
run_stats_hourly_aggregation_once, run_usage_cleanup_once, run_usage_counter_flush_once,
run_wallet_daily_usage_aggregation_once, AUDIT_LOG_CLEANUP_INTERVAL,
GEMINI_FILE_MAPPING_CLEANUP_INTERVAL, OAUTH_TOKEN_REFRESH_INTERVAL, PENDING_CLEANUP_INTERVAL,
POOL_MONITOR_INTERVAL, PROVIDER_CHECKIN_DEFAULT_TIME, PROXY_NODE_METRICS_CLEANUP_HOUR,
PROXY_NODE_METRICS_CLEANUP_MINUTE, PROXY_NODE_STALE_SWEEP_INTERVAL,
PROXY_UPGRADE_ROLLOUT_INTERVAL, REQUEST_CANDIDATE_CLEANUP_INTERVAL, USAGE_CLEANUP_HOUR,
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
USAGE_CLEANUP_MINUTE, USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE,
USAGE_COUNTER_DELTA_CLEANUP_INTERVAL, USAGE_COUNTER_DELTA_RETENTION_SECS,
USAGE_COUNTER_FLUSH_BATCH_SIZE, USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT,
USAGE_COUNTER_FLUSH_INTERVAL, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
@@ -162,6 +165,55 @@ pub(crate) fn spawn_usage_cleanup_worker(
}))
}
pub(crate) fn spawn_usage_counter_flush_worker(
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if !data.has_usage_counter_flush_backend() {
return None;
}
Some(tokio::spawn(async move {
let mut interval = tokio::time::interval(USAGE_COUNTER_FLUSH_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
interval.tick().await;
let mut last_delta_cleanup = tokio::time::Instant::now();
loop {
let mut batches = 0_usize;
while batches < USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT {
match run_usage_counter_flush_once(&data, USAGE_COUNTER_FLUSH_BATCH_SIZE).await {
Ok(summary) if summary.rows_claimed > 0 => batches += 1,
Ok(_) => break,
Err(err) => {
log_maintenance_worker_failure("usage_counter_flush", "tick", &err);
break;
}
}
}
if batches >= USAGE_COUNTER_FLUSH_CATCH_UP_BURST_LIMIT {
tokio::task::yield_now().await;
continue;
}
if last_delta_cleanup.elapsed() >= USAGE_COUNTER_DELTA_CLEANUP_INTERVAL {
if let Err(err) = cleanup_processed_usage_counter_deltas_once(
&data,
USAGE_COUNTER_DELTA_RETENTION_SECS,
USAGE_COUNTER_DELTA_CLEANUP_BATCH_SIZE,
)
.await
{
log_maintenance_worker_failure("usage_counter_delta_cleanup", "tick", &err);
}
last_delta_cleanup = tokio::time::Instant::now();
}
interval.tick().await;
}
}))
}
pub(crate) fn spawn_provider_checkin_worker(
state: AppState,
) -> Option<tokio::task::JoinHandle<()>> {

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

@@ -56,6 +56,7 @@ use crate::maintenance::spawn_request_candidate_cleanup_worker;
use crate::maintenance::spawn_stats_aggregation_worker;
use crate::maintenance::spawn_stats_hourly_aggregation_worker;
use crate::maintenance::spawn_usage_cleanup_worker;
use crate::maintenance::spawn_usage_counter_flush_worker;
use crate::maintenance::spawn_wallet_daily_usage_aggregation_worker;
const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(3);
@@ -572,6 +573,7 @@ impl AppState {
}
pub(crate) fn invalidate_provider_routing_caches(&self) {
self.data.clear_minimal_candidate_selection_cache();
self.clear_provider_transport_snapshot_cache();
self.invalidate_scheduler_affinity_cache();
}
@@ -1144,6 +1146,10 @@ impl AppState {
crate::task_runtime::TASK_KEY_USAGE_QUEUE_WORKER,
self.usage_runtime.spawn_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_USAGE_COUNTER_FLUSH,
spawn_usage_counter_flush_worker(self.data.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_PROVIDER_QUOTA_RESET,
crate::wallet_runtime::spawn_provider_quota_reset_worker(self.data.clone()),

9
apps/aether-gateway/src/task_runtime/mod.rs
View File

@@ -18,6 +18,7 @@ use crate::{AppState, GatewayError};
pub(crate) const TASK_KEY_PROVIDER_DELETE: &str = "admin.provider.delete";
pub(crate) const TASK_KEY_PROVIDER_OAUTH_BATCH_IMPORT: &str = "admin.provider.oauth.batch_import";
pub(crate) const TASK_KEY_USAGE_QUEUE_WORKER: &str = "usage.queue.worker";
pub(crate) const TASK_KEY_USAGE_COUNTER_FLUSH: &str = "usage.counter.flush.worker";
pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
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";
@@ -70,6 +71,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_USAGE_COUNTER_FLUSH,
TaskKind::Daemon,
"daemon",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_VIDEO_TASK_POLLER,
TaskKind::Scheduled,

2
apps/aether-gateway/src/tests/control/admin/monitoring.rs
View File

@@ -273,6 +273,8 @@ async fn gateway_handles_admin_monitoring_system_status_locally_with_trusted_adm
assert_eq!(payload["tunnel"]["nodes"], json!(0));
assert_eq!(payload["tunnel"]["active_streams"], json!(0));
assert_eq!(payload["recent_errors"], json!(1));
assert_eq!(payload["usage_counter"]["status"], json!("idle"));
assert_eq!(payload["usage_counter"]["outbox_pending_rows"], json!(0));
assert!(payload["timestamp"].as_str().is_some());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

4
apps/aether-gateway/src/tests/control/admin/stats.rs
View File

@@ -963,6 +963,8 @@ async fn gateway_handles_admin_stats_provider_performance_locally_with_trusted_a
assert_eq!(payload["summary"]["p99_response_time_ms"], 1000);
assert_eq!(payload["summary"]["response_time_sample_count"], 11);
assert_eq!(payload["summary"]["slow_request_count"], 0);
assert_eq!(payload["usage_counter"]["status"], json!("idle"));
assert_eq!(payload["usage_counter"]["outbox_pending_rows"], json!(0));
assert_eq!(payload["providers"].as_array().map(Vec::len), Some(2));
assert_eq!(payload["providers"][0]["provider_id"], "provider-1");
@@ -1038,6 +1040,8 @@ async fn gateway_returns_empty_admin_stats_provider_performance_without_usage_re
);
assert_eq!(payload["providers"].as_array().map(Vec::len), Some(0));
assert_eq!(payload["timeline"].as_array().map(Vec::len), Some(0));
assert_eq!(payload["usage_counter"]["status"], json!("idle"));
assert_eq!(payload["usage_counter"]["outbox_pending_rows"], json!(0));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();

2
apps/aether-gateway/src/tests/control/admin/system.rs
View File

@@ -296,6 +296,8 @@ async fn gateway_handles_admin_system_stats_locally_with_trusted_admin_principal
assert_eq!(payload["providers"]["active"], json!(1));
assert_eq!(payload["api_keys"], json!(0));
assert_eq!(payload["requests"], json!(0));
assert_eq!(payload["usage_counter"]["status"], json!("idle"));
assert_eq!(payload["usage_counter"]["outbox_pending_rows"], json!(0));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();