Merge pull request #856 from Kayphoon/fix/manual-cleanup-buffered-body

fix(admin): buffer request body for manual cleanup, smtp test, and system update routes

apps/aether-gateway/src/control/tests/admin_core.rs

@@ -1,6 +1,7 @@
 use http::Uri;
 
-use crate::control::management_token_required_permission;
+use crate::control::{management_token_required_permission, GatewayPublicRequestContext};
+use crate::handlers::shared::local_proxy_route_requires_buffered_body;
 
 use super::{classify_control_route, headers};
 
@@ -206,6 +207,10 @@ fn classifies_admin_system_maintenance_write_routes_as_admin_proxy_route() {
             "/api/admin/system/important-notification/test",
             "important_notification_test",
         ),
+        (
+            "/api/admin/system/cleanup/usage/manual",
+            "cleanup_usage_manual",
+        ),
         ("/api/admin/system/cleanup", "cleanup"),
         ("/api/admin/system/purge/config", "purge_config"),
         ("/api/admin/system/purge/users", "purge_users"),
@@ -235,6 +240,28 @@ fn classifies_admin_system_maintenance_write_routes_as_admin_proxy_route() {
             Some("admin:system")
         );
         assert!(!decision.is_execution_runtime_candidate());
+
+        if matches!(
+            expected_kind,
+            "config_import"
+                | "users_import"
+                | "data_import"
+                | "smtp_test"
+                | "important_notification_test"
+                | "cleanup_usage_manual"
+        ) {
+            let context = GatewayPublicRequestContext::from_request_parts(
+                "trace-system-maintenance-write",
+                &http::Method::POST,
+                &uri,
+                &headers,
+                Some(decision),
+            );
+            assert!(
+                local_proxy_route_requires_buffered_body(&context),
+                "POST {path} should buffer request body"
+            );
+        }
     }
 }
 
@@ -303,6 +330,20 @@ fn classifies_admin_system_update_routes_as_admin_proxy_routes() {
             Some("admin:system")
         );
         assert!(!decision.is_execution_runtime_candidate());
+
+        if matches!(expected_kind, "prepare_update" | "apply_update") {
+            let context = GatewayPublicRequestContext::from_request_parts(
+                "trace-system-update-write",
+                &method,
+                &uri,
+                &headers,
+                Some(decision),
+            );
+            assert!(
+                local_proxy_route_requires_buffered_body(&context),
+                "{method} {path} should buffer request body"
+            );
+        }
     }
 }
 

apps/aether-gateway/src/handlers/shared/request_utils.rs

@@ -275,6 +275,10 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
                 | (Some("system_manage"), http::Method::POST, Some("config_import"))
                 | (Some("system_manage"), http::Method::POST, Some("users_import"))
                 | (Some("system_manage"), http::Method::POST, Some("data_import"))
+                | (Some("system_manage"), http::Method::POST, Some("cleanup_usage_manual"))
+                | (Some("system_manage"), http::Method::POST, Some("smtp_test"))
+                | (Some("system_manage"), http::Method::POST, Some("prepare_update"))
+                | (Some("system_manage"), http::Method::POST, Some("apply_update"))
                 | (Some("system_manage"), http::Method::PUT, Some("settings_set"))
                 | (Some("system_manage"), http::Method::PUT, Some("config_set"))
                 | (Some("system_manage"), http::Method::PUT, Some("email_template_set"))
@@ -609,4 +613,27 @@ mod tests {
             "/v1/chat/completions?key=passthrough"
         );
     }
+
+    #[test]
+    fn manual_cleanup_route_requires_buffered_body() {
+        use crate::control::GatewayPublicRequestContext;
+        let mut decision = GatewayControlDecision::synthetic(
+            "/api/admin/system/cleanup/usage/manual",
+            Some("admin_proxy".to_string()),
+            Some("system_manage".to_string()),
+            Some("cleanup_usage_manual".to_string()),
+            Some("system_manage:cleanup_usage_manual".to_string()),
+        );
+        decision.route_class = Some("admin_proxy".to_string());
+        let uri: http::Uri = "/api/admin/system/cleanup/usage/manual".parse().unwrap();
+        let headers = http::HeaderMap::new();
+        let context = GatewayPublicRequestContext::from_request_parts(
+            "trace-manual-cleanup",
+            &http::Method::POST,
+            &uri,
+            &headers,
+            Some(decision),
+        );
+        assert!(super::admin_proxy_local_requires_buffered_body(&context));
+    }
 }

crates/aether-data/adapters/postgres/src/usage/cleanup.rs

@@ -687,11 +687,31 @@ async fn cleanup_usage_raw_body_fields(
     Ok(total_cleaned)
 }
 
+async fn truncate_usage_body_blobs_table(pool: &PostgresPool) -> Result<(), DataLayerError> {
+    let mut tx = pool.begin().await.map_err(postgres_error)?;
+    sqlx::query("SET LOCAL lock_timeout = '2s'")
+        .execute(&mut *tx)
+        .await
+        .map_err(postgres_error)?;
+    sqlx::query("TRUNCATE TABLE usage_body_blobs")
+        .execute(&mut *tx)
+        .await
+        .map_err(postgres_error)?;
+    tx.commit().await.map_err(postgres_error)?;
+    Ok(())
+}
+
 async fn cleanup_usage_compressed_body_fields(
     pool: &PostgresPool,
     cutoff_time: DateTime<Utc>,
     batch_size: usize,
 ) -> Result<usize, DataLayerError> {
+    if let Err(err) = truncate_usage_body_blobs_table(pool).await {
+        warn!(
+            error = %err,
+            "usage cleanup truncate usage_body_blobs table failed or timed out, falling back to batch deletion"
+        );
+    }
     let mut total_cleaned = 0usize;
     loop {
         let rows = fetch_usage_body_cleanup_rows(