Improve gateway scheduling and runtime admission

crates/aether-usage-runtime/src/config.rs

@@ -14,6 +14,10 @@ pub struct UsageRuntimeConfig {
     pub reclaim_idle_ms: u64,
     pub reclaim_count: usize,
     pub reclaim_interval_ms: u64,
+    pub enqueue_retry_buffer_capacity: usize,
+    pub enqueue_retry_workers: usize,
+    pub enqueue_retry_initial_backoff_ms: u64,
+    pub enqueue_retry_max_backoff_ms: u64,
 }
 
 impl Default for UsageRuntimeConfig {
@@ -31,6 +35,10 @@ impl Default for UsageRuntimeConfig {
             reclaim_idle_ms: 30_000,
             reclaim_count: 500,
             reclaim_interval_ms: 5_000,
+            enqueue_retry_buffer_capacity: 131_072,
+            enqueue_retry_workers: 4,
+            enqueue_retry_initial_backoff_ms: 10,
+            enqueue_retry_max_backoff_ms: 1_000,
         }
     }
 }
@@ -90,6 +98,31 @@ impl UsageRuntimeConfig {
                 "usage runtime reclaim_interval_ms must be positive".to_string(),
             ));
         }
+        if self.enqueue_retry_buffer_capacity == 0 {
+            return Err(DataLayerError::InvalidConfiguration(
+                "usage runtime enqueue_retry_buffer_capacity must be positive".to_string(),
+            ));
+        }
+        if self.enqueue_retry_workers == 0 {
+            return Err(DataLayerError::InvalidConfiguration(
+                "usage runtime enqueue_retry_workers must be positive".to_string(),
+            ));
+        }
+        if self.enqueue_retry_initial_backoff_ms == 0 {
+            return Err(DataLayerError::InvalidConfiguration(
+                "usage runtime enqueue_retry_initial_backoff_ms must be positive".to_string(),
+            ));
+        }
+        if self.enqueue_retry_max_backoff_ms == 0 {
+            return Err(DataLayerError::InvalidConfiguration(
+                "usage runtime enqueue_retry_max_backoff_ms must be positive".to_string(),
+            ));
+        }
+        if self.enqueue_retry_initial_backoff_ms > self.enqueue_retry_max_backoff_ms {
+            return Err(DataLayerError::InvalidConfiguration(
+                "usage runtime enqueue retry initial backoff cannot exceed max backoff".to_string(),
+            ));
+        }
 
         Ok(())
     }

crates/aether-usage-runtime/src/request_metadata.rs

@@ -146,6 +146,8 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
     copy_number(source, target, "cache_read_price_per_1m");
     copy_number(source, target, "price_per_request");
     copy_non_null_value(source, target, "proxy");
+    copy_non_null_value(source, target, "stage_timings_ms");
+    copy_non_null_value(source, target, "db_timings_ms");
     sanitize_request_path_metadata_fields(target);
 }
 
@@ -189,6 +191,8 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
     remove_number(&mut source, target, "cache_read_price_per_1m");
     remove_number(&mut source, target, "price_per_request");
     remove_non_null_value(&mut source, target, "proxy");
+    remove_non_null_value(&mut source, target, "stage_timings_ms");
+    remove_non_null_value(&mut source, target, "db_timings_ms");
     sanitize_request_path_metadata_fields(target);
 }
 
@@ -476,8 +480,38 @@ mod tests {
         }
     }
 
+    fn sample_stage_timings_metadata() -> Value {
+        json!({
+            "stream_candidate_slot": 1,
+            "stream_provider_in_flight": 2,
+            "stream_upstream_headers": 180,
+            "stream_first_data": 8210
+        })
+    }
+
+    fn sample_db_timings_metadata() -> Value {
+        json!({
+            "query_count": 2,
+            "query_total": 950,
+            "query_max": 650,
+            "operations": {
+                "request_candidate_upsert": {"count": 1, "sum": 650, "max": 650},
+                "usage_upsert": {"count": 1, "sum": 300, "max": 300}
+            },
+            "pool": {
+                "max_checked_out": 20,
+                "max_pool_size": 20,
+                "min_idle": 0,
+                "max_connections": 20,
+                "max_usage_rate": 100.0
+            }
+        })
+    }
+
     #[test]
     fn sanitizes_request_metadata_to_allowlist() {
+        let stage_timings_ms = sample_stage_timings_metadata();
+        let db_timings_ms = sample_db_timings_metadata();
         let metadata = sanitize_usage_request_metadata(Some(json!({
             "request_id": "req-1",
             "provider_id": "provider-1",
@@ -512,6 +546,8 @@ mod tests {
             "cache_creation_price_per_1m": 3.75,
             "cache_read_price_per_1m": 0.3,
             "price_per_request": 0.02,
+            "stage_timings_ms": stage_timings_ms.clone(),
+            "db_timings_ms": db_timings_ms.clone(),
             "original_headers": {"authorization": "Bearer secret"},
             "original_request_body": {"messages": []},
             "provider_request_headers": {"authorization": "Bearer secret"},
@@ -549,7 +585,9 @@ mod tests {
                 "output_price_per_1m": 15.0,
                 "cache_creation_price_per_1m": 3.75,
                 "cache_read_price_per_1m": 0.3,
-                "price_per_request": 0.02
+                "price_per_request": 0.02,
+                "stage_timings_ms": stage_timings_ms,
+                "db_timings_ms": db_timings_ms
             })
         );
     }
@@ -657,7 +695,20 @@ mod tests {
                     "global_model_name": "gpt-5",
                     "client_ip": "203.0.113.8",
                     "user_agent": "Claude-Code/1.0",
-                    "billing_snapshot": {"status": "complete"}
+                    "billing_snapshot": {"status": "complete"},
+                    "stage_timings_ms": {
+                        "stream_candidate_slot": 0,
+                        "stream_upstream_headers": 180,
+                        "stream_first_data": 8210
+                    },
+                    "db_timings_ms": {
+                        "query_count": 1,
+                        "query_total": 42,
+                        "query_max": 42,
+                        "operations": {
+                            "auth_api_key_snapshot": {"count": 1, "sum": 42, "max": 42}
+                        }
+                    }
                 })
                 .as_object()
                 .expect("object"),
@@ -676,7 +727,20 @@ mod tests {
                 "global_model_name": "gpt-5",
                 "client_ip": "203.0.113.8",
                 "user_agent": "Claude-Code/1.0",
-                "billing_snapshot": {"status": "complete"}
+                "billing_snapshot": {"status": "complete"},
+                "stage_timings_ms": {
+                    "stream_candidate_slot": 0,
+                    "stream_upstream_headers": 180,
+                    "stream_first_data": 8210
+                },
+                "db_timings_ms": {
+                    "query_count": 1,
+                    "query_total": 42,
+                    "query_max": 42,
+                    "operations": {
+                        "auth_api_key_snapshot": {"count": 1, "sum": 42, "max": 42}
+                    }
+                }
             })
         );
     }

crates/aether-usage-runtime/src/worker.rs

@@ -8,9 +8,10 @@ use async_trait::async_trait;
 use tracing::warn;
 
 use crate::executor::spawn_on_usage_background_runtime;
+use crate::runtime::UsageBillingEventEnricher;
 use crate::{
-    build_upsert_usage_record_from_event, settle_usage_if_needed, UsageEvent, UsageQueue,
-    UsageRuntimeConfig, UsageSettlementWriter,
+    build_upsert_usage_record_from_event, settle_usage_if_needed, UsageEvent, UsageEventType,
+    UsageQueue, UsageRuntimeConfig, UsageSettlementWriter,
 };
 
 #[async_trait]
@@ -50,10 +51,17 @@ impl<T> UsageDataEventRecorder<T> {
 #[async_trait]
 impl<T> UsageEventRecorder for UsageDataEventRecorder<T>
 where
-    T: UsageRecordWriter + UsageSettlementWriter + ManualProxyNodeCounter + Send + Sync,
+    T: UsageRecordWriter
+        + UsageSettlementWriter
+        + UsageBillingEventEnricher
+        + ManualProxyNodeCounter
+        + Send
+        + Sync,
 {
     async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
-        write_event_record(self.data.as_ref(), event).await
+        let mut event = event.clone();
+        enrich_terminal_event(self.data.as_ref(), &mut event).await;
+        write_event_record(self.data.as_ref(), &event).await
     }
 }
 
@@ -276,7 +284,13 @@ pub fn build_usage_queue_worker<T>(
     config: UsageRuntimeConfig,
 ) -> Result<UsageQueueWorker, DataLayerError>
 where
-    T: UsageRecordWriter + UsageSettlementWriter + ManualProxyNodeCounter + Send + Sync + 'static,
+    T: UsageRecordWriter
+        + UsageSettlementWriter
+        + UsageBillingEventEnricher
+        + ManualProxyNodeCounter
+        + Send
+        + Sync
+        + 'static,
 {
     UsageQueueWorker::new(runner, Arc::new(UsageDataEventRecorder::new(data)), config)
 }
@@ -293,6 +307,29 @@ where
     Ok(())
 }
 
+async fn enrich_terminal_event<T>(data: &T, event: &mut UsageEvent)
+where
+    T: UsageBillingEventEnricher + Send + Sync,
+{
+    if !matches!(
+        event.event_type,
+        UsageEventType::Completed | UsageEventType::Failed | UsageEventType::Cancelled
+    ) {
+        return;
+    }
+
+    if let Err(err) = data.enrich_usage_event(event).await {
+        warn!(
+            event_name = "usage_worker_billing_enrichment_failed",
+            log_type = "event",
+            request_id = %event.request_id,
+            event_type = ?event.event_type,
+            error = %err,
+            "usage worker failed to enrich terminal usage event with billing"
+        );
+    }
+}
+
 async fn increment_manual_proxy_node_from_event<T>(data: &T, event: &UsageEvent)
 where
     T: ManualProxyNodeCounter + Send + Sync,
@@ -368,6 +405,7 @@ mod tests {
         usage_event_record_error_is_permanent, write_event_record, ManualProxyNodeCounter,
         UsageEventRecorder, UsageQueueWorker, UsageRecordWriter,
     };
+    use crate::UsageBillingEventEnricher;
     use crate::{
         UsageEvent, UsageEventData, UsageEventType, UsageRuntimeConfig, UsageSettlementWriter,
     };
@@ -376,6 +414,7 @@ mod tests {
     struct TestUsageStore {
         records: Mutex<Vec<UpsertUsageRecord>>,
         settlements: Mutex<Vec<UsageSettlementInput>>,
+        enrich_calls: Mutex<Vec<String>>,
     }
 
     #[derive(Default)]
@@ -471,6 +510,18 @@ mod tests {
         }
     }
 
+    #[async_trait]
+    impl UsageBillingEventEnricher for TestUsageStore {
+        async fn enrich_usage_event(&self, event: &mut UsageEvent) -> Result<(), DataLayerError> {
+            self.enrich_calls
+                .lock()
+                .expect("enrich calls lock")
+                .push(event.request_id.clone());
+            event.data.total_cost_usd = Some(0.456);
+            Ok(())
+        }
+    }
+
     #[async_trait]
     impl UsageEventRecorder for SelectiveFailingRecorder {
         async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
@@ -532,6 +583,52 @@ mod tests {
         assert_eq!(settlements[0].request_id, "req-worker-123");
     }
 
+    #[tokio::test]
+    async fn data_event_recorder_enriches_terminal_event_before_write() {
+        let store = Arc::new(TestUsageStore::default());
+        let recorder = super::UsageDataEventRecorder::new(Arc::clone(&store));
+        let event = sample_event();
+
+        recorder
+            .record_usage_event(&event)
+            .await
+            .expect("recorder should enrich and write usage");
+
+        assert_eq!(
+            store
+                .enrich_calls
+                .lock()
+                .expect("enrich calls lock")
+                .as_slice(),
+            ["req-worker-123"]
+        );
+        let records = store.records.lock().expect("records lock");
+        assert_eq!(records.len(), 1);
+        assert_eq!(records[0].total_cost_usd, Some(0.456));
+    }
+
+    #[tokio::test]
+    async fn data_event_recorder_skips_enrichment_for_lifecycle_event() {
+        let store = Arc::new(TestUsageStore::default());
+        let recorder = super::UsageDataEventRecorder::new(Arc::clone(&store));
+        let mut event = sample_event();
+        event.event_type = UsageEventType::Pending;
+
+        recorder
+            .record_usage_event(&event)
+            .await
+            .expect("recorder should write lifecycle usage");
+
+        assert!(store
+            .enrich_calls
+            .lock()
+            .expect("enrich calls lock")
+            .is_empty());
+        let records = store.records.lock().expect("records lock");
+        assert_eq!(records.len(), 1);
+        assert_eq!(records[0].total_cost_usd, None);
+    }
+
     #[test]
     fn usage_event_record_error_classifies_permanent_failures() {
         assert!(usage_event_record_error_is_permanent(