Rootfix usage counter outbox

apps/aether-gateway/src/data/state/auth.rs

@@ -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),

apps/aether-gateway/src/data/state/candidate_cache.rs

@@ -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);
+    }
+}

apps/aether-gateway/src/data/state/catalog.rs

@@ -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),

apps/aether-gateway/src/data/state/core.rs

@@ -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()
     }

apps/aether-gateway/src/data/state/integrations.rs

@@ -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

apps/aether-gateway/src/data/state/mod.rs

@@ -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;

apps/aether-gateway/src/data/state/runtime.rs

@@ -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],