Preserve usage data in system imports

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

@@ -629,6 +629,36 @@ impl AppState {
         Ok(summary)
     }
 
+    pub(crate) async fn export_admin_system_usage_aggregates(
+        &self,
+    ) -> Result<aether_data::repository::system::AdminSystemUsageAggregateSnapshot, GatewayError>
+    {
+        self.data
+            .export_admin_system_usage_aggregates()
+            .await
+            .map_err(|err| GatewayError::Internal(err.to_string()))
+    }
+
+    pub(crate) async fn import_admin_system_usage_aggregates(
+        &self,
+        snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot,
+        user_id_map: &std::collections::BTreeMap<String, String>,
+        api_key_id_map: &std::collections::BTreeMap<String, String>,
+        mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode,
+    ) -> Result<aether_data::repository::system::AdminSystemUsageAggregateImportSummary, GatewayError>
+    {
+        self.data
+            .import_admin_system_usage_aggregates(snapshot, user_id_map, api_key_id_map, mode)
+            .await
+            .map_err(|err| match err {
+                aether_data::DataLayerError::InvalidInput(detail) => GatewayError::Client {
+                    status: http::StatusCode::BAD_REQUEST,
+                    message: detail,
+                },
+                other => GatewayError::Internal(other.to_string()),
+            })
+    }
+
     pub(crate) async fn run_admin_system_cleanup_once(
         &self,
     ) -> Result<crate::maintenance::AdminSystemCleanupSummary, GatewayError> {

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

@@ -369,6 +369,25 @@ impl AppState {
         Ok(api_key)
     }
 
+    pub(crate) async fn set_api_key_usage_totals(
+        &self,
+        api_key_id: &str,
+        total_requests: u64,
+        total_tokens: u64,
+        total_cost_usd: f64,
+    ) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
+    {
+        let api_key = self
+            .data
+            .set_api_key_usage_totals(api_key_id, total_requests, total_tokens, total_cost_usd)
+            .await
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        if api_key.is_some() {
+            self.invalidate_auth_context_cache();
+        }
+        Ok(api_key)
+    }
+
     pub(crate) async fn set_standalone_api_key_feature_settings(
         &self,
         api_key_id: &str,