fix: harden concurrency limits and high-RPM runtime paths

Bound request, stream, queue, and shutdown resource lifetimes. Reduce scheduler and Redis hot-path work and isolate database maintenance. Include regression coverage, load probes, and concurrency audit results.

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

@@ -34,7 +34,7 @@ use sqlx::{
     PgPool, Postgres, QueryBuilder, Row,
 };
 use std::collections::{BTreeMap, BTreeSet};
-use std::io::Write;
+use std::io::{BufWriter, Write};
 use uuid::Uuid;
 
 use crate::{
@@ -58,6 +58,9 @@ use aether_data_contracts::repository::usage::{
 use aether_data_contracts::DataLayerError;
 
 pub mod cleanup;
+mod preparation;
+
+use preparation::prepare_usage_in_background;
 
 // Legacy inline body columns on public.usage are deprecated. Keep the threshold at zero so
 // newly captured bodies always spill to usage_body_blobs and resolve through usage_http_audits.
@@ -2121,12 +2124,9 @@ impl PreparedPendingUsage {
             ));
         }
 
-        // Keep the capture input separate from the accounting row.  The persistence sanitizer
-        // intentionally removes HTTP bodies/headers/states, but the pending batch still needs
-        // those values to populate the canonical audit/blob tables.
-        let capture_usage = usage.clone();
-        let usage = sanitize_usage_for_persistence(usage);
-        let prepared = prepare_usage_upsert_context(&capture_usage)?;
+        // Prepare captures before the accounting sanitizer removes HTTP bodies/headers/states.
+        let (usage, prepared) = prepare_usage_for_persistence(usage);
+        let prepared = prepared?;
         let input_tokens = usage
             .input_tokens
             .map(to_i32)
@@ -8450,10 +8450,10 @@ ORDER BY "usage".user_id ASC
         usage: UpsertUsageRecord,
     ) -> Result<StoredRequestUsageAudit, DataLayerError> {
         usage.validate()?;
-        // `usage` is the sanitized accounting projection; prepare the auxiliary capture and
-        // snapshots from the original event so typed `none` markers can clear prior facts.
-        let capture_usage = usage.clone();
-        let usage = sanitize_usage_for_persistence(usage);
+        // Move the event before cloning or compressing captures, and do not hold a connection
+        // while preparing them. Stale lifecycle updates still ignore preparation errors below.
+        let (usage, prepared) =
+            prepare_usage_in_background(move || Ok(prepare_usage_for_persistence(usage))).await?;
         self.tx_runner
             .run_read_write(|tx| {
                 Box::pin(async move {
@@ -8519,7 +8519,7 @@ ORDER BY "usage".user_id ASC
                         clear_provider_request_body,
                         clear_response_body,
                         clear_client_response_body,
-                    } = prepare_usage_upsert_context(&capture_usage)?;
+                    } = prepared?;
                     let capture_update_allowed = recovers_terminal_failure
                         || usage_capture_update_allowed(
                             previous_usage.as_ref().map(|stored| {
@@ -8938,33 +8938,36 @@ ORDER BY "usage".user_id ASC
             return Ok(());
         }
 
-        let mut request_id_counts = BTreeMap::<String, usize>::new();
-        for usage in &usages {
-            *request_id_counts
-                .entry(usage.request_id.clone())
-                .or_default() += 1;
-        }
-
-        // Duplicate request IDs must retain the caller's exact sequential merge order. They are
-        // uncommon in lifecycle batches, so keep them on the canonical single-row path.
-        let mut batch_rows = Vec::<(usize, PreparedPendingUsage)>::new();
-        let mut fallback_rows = Vec::<(usize, UpsertUsageRecord)>::new();
-        for (sequence, usage) in usages.into_iter().enumerate() {
-            let original_usage = usage.clone();
-            let prepared = PreparedPendingUsage::try_from_usage(usage)?;
-            if request_id_counts
-                .get(&prepared.usage.request_id)
-                .copied()
-                .unwrap_or_default()
-                == 1
-            {
-                batch_rows.push((sequence, prepared));
-            } else {
-                // Preserve capture markers for the canonical fallback; that path performs the
-                // sanitized bind only after preparing the auxiliary audit/blob state.
-                fallback_rows.push((sequence, original_usage));
+        let (batch_rows, mut fallback_rows) = prepare_usage_in_background(move || {
+            let mut request_id_counts = BTreeMap::<String, usize>::new();
+            for usage in &usages {
+                *request_id_counts
+                    .entry(usage.request_id.clone())
+                    .or_default() += 1;
             }
-        }
+
+            // Duplicate request IDs must retain the caller's exact sequential merge order.
+            let mut batch_rows = Vec::<(usize, PreparedPendingUsage)>::new();
+            let mut fallback_rows = Vec::<(usize, UpsertUsageRecord)>::new();
+            for (sequence, usage) in usages.into_iter().enumerate() {
+                let duplicate = request_id_counts
+                    .get(&usage.request_id)
+                    .copied()
+                    .unwrap_or_default()
+                    > 1;
+                let original_usage = duplicate.then(|| usage.clone());
+                let prepared = PreparedPendingUsage::try_from_usage(usage)?;
+                if let Some(original_usage) = original_usage {
+                    // Preserve capture markers for the canonical fallback, including validation
+                    // of every row before starting the batch transaction.
+                    fallback_rows.push((sequence, original_usage));
+                } else {
+                    batch_rows.push((sequence, prepared));
+                }
+            }
+            Ok((batch_rows, fallback_rows))
+        })
+        .await?;
 
         let mut inserted_request_ids = BTreeSet::<String>::new();
         if !batch_rows.is_empty() {
@@ -10294,7 +10297,7 @@ RETURNING
 
     pub async fn rebuild_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
         self.tx_runner
-            .run_read_write(|tx| {
+            .run(crate::PostgresTransactionOptions::maintenance(), |tx| {
                 Box::pin(async move {
                     sqlx::query(RESET_API_KEY_USAGE_STATS_SQL)
                         .execute(&mut **tx)
@@ -10313,7 +10316,7 @@ RETURNING
 
     pub async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
         self.tx_runner
-            .run_read_write(|tx| {
+            .run(crate::PostgresTransactionOptions::maintenance(), |tx| {
                 Box::pin(async move {
                     sqlx::query(RESET_PROVIDER_API_KEY_USAGE_STATS_SQL)
                         .execute(&mut **tx)
@@ -12359,24 +12362,35 @@ fn prepare_usage_body_storage(value: Option<&Value>) -> Result<UsageBodyStorage,
             detached_blob_bytes: None,
         });
     };
-    let bytes = serde_json::to_vec(value).map_err(|err| {
-        DataLayerError::UnexpectedValue(format!("failed to serialize usage json: {err}"))
-    })?;
-    if bytes.len() == MAX_INLINE_USAGE_BODY_BYTES {
-        return Ok(UsageBodyStorage {
-            inline_json: Some(String::from_utf8(bytes).map_err(|err| {
-                DataLayerError::UnexpectedValue(format!(
-                    "failed to encode inline usage body as utf-8: {err}"
-                ))
-            })?),
-            detached_blob_bytes: None,
-        });
-    }
-
     let mut encoder = GzEncoder::new(Vec::new(), Compression::new(6));
-    encoder.write_all(&bytes).map_err(|err| {
-        DataLayerError::UnexpectedValue(format!("failed to compress usage json: {err}"))
-    })?;
+    if MAX_INLINE_USAGE_BODY_BYTES == 0 {
+        // Coalesce serde's punctuation/escape writes without allocating a full JSON buffer.
+        let mut writer = BufWriter::with_capacity(8 * 1024, &mut encoder);
+        serde_json::to_writer(&mut writer, value).map_err(|err| {
+            let operation = if err.is_io() { "compress" } else { "serialize" };
+            DataLayerError::UnexpectedValue(format!("failed to {operation} usage json: {err}"))
+        })?;
+        writer.into_inner().map_err(|err| {
+            DataLayerError::UnexpectedValue(format!("failed to compress usage json: {err}"))
+        })?;
+    } else {
+        let bytes = serde_json::to_vec(value).map_err(|err| {
+            DataLayerError::UnexpectedValue(format!("failed to serialize usage json: {err}"))
+        })?;
+        if bytes.len() == MAX_INLINE_USAGE_BODY_BYTES {
+            return Ok(UsageBodyStorage {
+                inline_json: Some(String::from_utf8(bytes).map_err(|err| {
+                    DataLayerError::UnexpectedValue(format!(
+                        "failed to encode inline usage body as utf-8: {err}"
+                    ))
+                })?),
+                detached_blob_bytes: None,
+            });
+        }
+        encoder.write_all(&bytes).map_err(|err| {
+            DataLayerError::UnexpectedValue(format!("failed to compress usage json: {err}"))
+        })?;
+    }
     let detached_blob_bytes = encoder.finish().map_err(|err| {
         DataLayerError::UnexpectedValue(format!("failed to finish usage json compression: {err}"))
     })?;
@@ -12429,11 +12443,40 @@ fn project_usage_request_metadata(
     }
 }
 
+fn prepare_usage_for_persistence(
+    mut usage: UpsertUsageRecord,
+) -> (
+    UpsertUsageRecord,
+    Result<PreparedUsageUpsert, DataLayerError>,
+) {
+    // Capture controls and accounting metadata have different sanitizers. Move the
+    // large payloads out before copying the metadata needed by both projections.
+    let request_body = usage.request_body.take();
+    let provider_request_body = usage.provider_request_body.take();
+    let response_body = usage.response_body.take();
+    let client_response_body = usage.client_response_body.take();
+    let request_headers = usage.request_headers.take();
+    let provider_request_headers = usage.provider_request_headers.take();
+    let response_headers = usage.response_headers.take();
+    let client_response_headers = usage.client_response_headers.take();
+    let mut capture = usage.clone();
+    capture.request_body = request_body;
+    capture.provider_request_body = provider_request_body;
+    capture.response_body = response_body;
+    capture.client_response_body = client_response_body;
+    capture.request_headers = request_headers;
+    capture.provider_request_headers = provider_request_headers;
+    capture.response_headers = response_headers;
+    capture.client_response_headers = client_response_headers;
+    capture.capture_retention = std::mem::take(&mut usage.capture_retention);
+    let capture = sanitize_usage_capture_controls_for_persistence(capture);
+    let prepared = prepare_usage_upsert_context(&capture);
+    (sanitize_usage_for_persistence(usage), prepared)
+}
+
 fn prepare_usage_upsert_context(
     usage: &UpsertUsageRecord,
 ) -> Result<PreparedUsageUpsert, DataLayerError> {
-    let usage = sanitize_usage_capture_controls_for_persistence(usage.clone());
-    let usage = &usage;
     let replace_client_request_body_facts = request_body_capture_replaces_derived_facts(
         usage.request_body.as_ref(),
         usage.request_body_state,

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

@@ -0,0 +1,327 @@
+use std::sync::{Arc, OnceLock};
+use std::time::Duration;
+
+use aether_data_contracts::DataLayerError;
+use tokio::sync::Semaphore;
+
+static USAGE_PREPARATION_EXECUTOR: OnceLock<UsagePreparationExecutor> = OnceLock::new();
+
+pub(super) async fn prepare_usage_in_background<T: Send + 'static>(
+    prepare: impl FnOnce() -> Result<T, DataLayerError> + Send + 'static,
+) -> Result<T, DataLayerError> {
+    USAGE_PREPARATION_EXECUTOR
+        .get_or_init(|| {
+            UsagePreparationExecutor::new(4, 32, Duration::from_secs(1), Duration::from_secs(30))
+        })
+        .run(prepare)
+        .await
+}
+
+struct UsagePreparationExecutor {
+    workers: Arc<Semaphore>,
+    admitted: Arc<Semaphore>,
+    queue_timeout: Duration,
+    execution_timeout: Duration,
+}
+
+impl UsagePreparationExecutor {
+    fn new(
+        workers: usize,
+        admitted: usize,
+        queue_timeout: Duration,
+        execution_timeout: Duration,
+    ) -> Self {
+        Self {
+            workers: Arc::new(Semaphore::new(workers)),
+            admitted: Arc::new(Semaphore::new(admitted)),
+            queue_timeout,
+            execution_timeout,
+        }
+    }
+
+    async fn run<T: Send + 'static>(
+        &self,
+        prepare: impl FnOnce() -> Result<T, DataLayerError> + Send + 'static,
+    ) -> Result<T, DataLayerError> {
+        // Bound both running work and callers retaining input while waiting for a worker.
+        // These limits count tasks, not bytes in the caller's original usage records.
+        let admitted = self.admitted.clone().try_acquire_owned().map_err(|_| {
+            DataLayerError::TimedOut("usage preparation capacity exhausted".to_string())
+        })?;
+        let worker = tokio::time::timeout(self.queue_timeout, self.workers.clone().acquire_owned())
+            .await
+            .map_err(|_| {
+                DataLayerError::TimedOut(
+                    "timed out waiting for usage preparation worker".to_string(),
+                )
+            })?
+            .map_err(|_| {
+                DataLayerError::TimedOut("usage preparation workers unavailable".to_string())
+            })?;
+
+        // The closure owns both permits even if its caller times out or is cancelled. It
+        // prepares input only; detached completion must never begin a database transaction.
+        let mut task = tokio::task::spawn_blocking(move || {
+            let _admitted = admitted;
+            let _worker = worker;
+            prepare()
+        });
+        match tokio::time::timeout(self.execution_timeout, &mut task).await {
+            Ok(result) => result.map_err(|error| {
+                DataLayerError::TimedOut(format!("usage preparation worker failed: {error}"))
+            })?,
+            Err(_) => {
+                // This cancels work still queued in Tokio; running blocking work keeps its
+                // permits until it actually exits, since abort cannot stop a blocking thread.
+                task.abort();
+                Err(DataLayerError::TimedOut(
+                    "timed out preparing usage storage".to_string(),
+                ))
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use std::sync::mpsc;
+
+    use super::*;
+
+    fn executor(
+        admitted: usize,
+        queue_timeout: Duration,
+        execution_timeout: Duration,
+    ) -> Arc<UsagePreparationExecutor> {
+        Arc::new(UsagePreparationExecutor::new(
+            1,
+            admitted,
+            queue_timeout,
+            execution_timeout,
+        ))
+    }
+
+    async fn wait_for_worker_release(executor: &UsagePreparationExecutor) {
+        tokio::time::timeout(Duration::from_secs(2), async {
+            while executor.workers.available_permits() != 1
+                || executor.admitted.available_permits() == 0
+            {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("finished blocking work should release its permits");
+    }
+
+    #[tokio::test(flavor = "current_thread")]
+    async fn preparation_runs_off_the_runtime_thread_and_preserves_errors() {
+        let executor = executor(1, Duration::from_secs(1), Duration::from_secs(2));
+        let runtime_thread = std::thread::current().id();
+        executor
+            .run(move || {
+                assert_ne!(std::thread::current().id(), runtime_thread);
+                Ok(())
+            })
+            .await
+            .expect("preparation should succeed");
+
+        let error = executor
+            .run(|| Err::<(), _>(DataLayerError::InvalidInput("bad usage".to_string())))
+            .await
+            .expect_err("input errors must reach the caller");
+        assert!(matches!(error, DataLayerError::InvalidInput(message) if message == "bad usage"));
+    }
+
+    #[tokio::test]
+    async fn saturated_admission_rejects_work_without_running_it() {
+        let executor = executor(1, Duration::from_secs(1), Duration::from_secs(2));
+        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
+        let (release_tx, release_rx) = mpsc::channel();
+        let first_executor = executor.clone();
+        let first = tokio::spawn(async move {
+            first_executor
+                .run(move || {
+                    let _ = started_tx.send(());
+                    let _ = release_rx.recv();
+                    Ok(())
+                })
+                .await
+        });
+        started_rx.await.expect("first job should start");
+
+        let error = executor
+            .run(|| -> Result<(), DataLayerError> { panic!("rejected work must not execute") })
+            .await
+            .expect_err("admission should fail immediately");
+        assert!(matches!(error, DataLayerError::TimedOut(message) if message.contains("capacity")));
+        release_tx
+            .send(())
+            .expect("first job should still be alive");
+        first.await.unwrap().unwrap();
+    }
+
+    #[tokio::test]
+    async fn blocking_worker_failure_is_retryable_and_releases_capacity() {
+        let executor = executor(1, Duration::from_secs(1), Duration::from_secs(2));
+        let error = executor
+            .run(|| -> Result<(), DataLayerError> { panic!("simulated worker failure") })
+            .await
+            .expect_err("worker failure must reach the caller");
+        assert!(
+            matches!(error, DataLayerError::TimedOut(message) if message.contains("worker failed"))
+        );
+        executor
+            .run(|| Ok(()))
+            .await
+            .expect("failed workers should release capacity");
+    }
+
+    #[tokio::test]
+    async fn waiting_for_a_worker_has_a_deadline_and_never_starts_expired_work() {
+        let executor = executor(2, Duration::from_millis(20), Duration::from_secs(2));
+        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
+        let (release_tx, release_rx) = mpsc::channel();
+        let first_executor = executor.clone();
+        let first = tokio::spawn(async move {
+            first_executor
+                .run(move || {
+                    let _ = started_tx.send(());
+                    let _ = release_rx.recv();
+                    Ok(())
+                })
+                .await
+        });
+        started_rx.await.expect("first job should start");
+        let ran = Arc::new(AtomicBool::new(false));
+        let work_ran = ran.clone();
+        let error = executor
+            .run(move || {
+                work_ran.store(true, Ordering::SeqCst);
+                Ok(())
+            })
+            .await
+            .expect_err("the queued job should time out");
+        assert!(matches!(error, DataLayerError::TimedOut(message) if message.contains("waiting")));
+        assert_eq!(executor.admitted.available_permits(), 1);
+        release_tx
+            .send(())
+            .expect("first job should still be alive");
+        first.await.unwrap().unwrap();
+        assert!(!ran.load(Ordering::SeqCst));
+    }
+
+    #[tokio::test]
+    async fn cancellation_keeps_permits_until_running_blocking_work_exits() {
+        let executor = executor(1, Duration::from_secs(1), Duration::from_secs(2));
+        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
+        let (release_tx, release_rx) = mpsc::channel();
+        let first_executor = executor.clone();
+        let first = tokio::spawn(async move {
+            first_executor
+                .run(move || {
+                    let _ = started_tx.send(());
+                    let _ = release_rx.recv();
+                    Ok(())
+                })
+                .await
+        });
+        started_rx.await.expect("first job should start");
+        first.abort();
+        assert!(first.await.unwrap_err().is_cancelled());
+        assert_eq!(executor.workers.available_permits(), 0);
+        assert!(matches!(
+            executor.run(|| Ok(())).await,
+            Err(DataLayerError::TimedOut(_))
+        ));
+        release_tx
+            .send(())
+            .expect("blocking work should outlive cancellation");
+        wait_for_worker_release(&executor).await;
+        executor
+            .run(|| Ok(()))
+            .await
+            .expect("the executor should recover");
+    }
+
+    #[tokio::test]
+    async fn cancellation_keeps_capture_budget_until_blocking_input_is_dropped() {
+        use aether_data_contracts::repository::usage::{
+            usage_json_heap_estimate, UpsertUsageRecord, UsageCaptureMemoryBudget,
+        };
+
+        let mut usage: UpsertUsageRecord = serde_json::from_value(serde_json::json!({
+            "request_id": "req-cancelled-preparation",
+            "provider_name": "test",
+            "model": "test",
+            "status": "completed",
+            "billing_status": "pending",
+            "updated_at_unix_secs": 100,
+            "request_body": {"content": "retained".repeat(1024)}
+        }))
+        .unwrap();
+        let bytes = std::mem::size_of::<serde_json::Value>()
+            + usage_json_heap_estimate(usage.request_body.as_ref().unwrap());
+        let budget = Arc::new(UsageCaptureMemoryBudget::new(bytes));
+        assert!(usage.capture_retention.reserve(Arc::clone(&budget), bytes));
+        let executor = executor(1, Duration::from_secs(1), Duration::from_secs(2));
+        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
+        let (release_tx, release_rx) = mpsc::channel();
+        let first_executor = Arc::clone(&executor);
+        let first = tokio::spawn(async move {
+            first_executor
+                .run(move || {
+                    let _ = started_tx.send(());
+                    let _ = release_rx.recv();
+                    drop(usage);
+                    Ok(())
+                })
+                .await
+        });
+        started_rx.await.unwrap();
+        first.abort();
+        assert!(first.await.unwrap_err().is_cancelled());
+        assert_eq!(budget.retained_bytes(), bytes);
+        release_tx.send(()).unwrap();
+        wait_for_worker_release(&executor).await;
+        assert_eq!(budget.retained_bytes(), 0);
+    }
+
+    #[tokio::test]
+    async fn execution_timeout_keeps_permits_until_running_blocking_work_exits() {
+        let executor = executor(1, Duration::from_secs(1), Duration::from_millis(20));
+        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
+        let (release_tx, release_rx) = mpsc::channel();
+        let first_executor = executor.clone();
+        let first = tokio::spawn(async move {
+            first_executor
+                .run(move || {
+                    let _ = started_tx.send(());
+                    let _ = release_rx.recv();
+                    Ok(())
+                })
+                .await
+        });
+        started_rx.await.expect("first job should start");
+        let error = first
+            .await
+            .unwrap()
+            .expect_err("running work should time out");
+        assert!(
+            matches!(error, DataLayerError::TimedOut(message) if message.contains("preparing"))
+        );
+        assert_eq!(executor.workers.available_permits(), 0);
+        assert!(matches!(
+            executor.run(|| Ok(())).await,
+            Err(DataLayerError::TimedOut(_))
+        ));
+        release_tx
+            .send(())
+            .expect("blocking work should outlive timeout");
+        wait_for_worker_release(&executor).await;
+        executor
+            .run(|| Ok(()))
+            .await
+            .expect("the executor should recover");
+    }
+}

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

@@ -8,7 +8,7 @@ use super::{
     attach_usage_routing_snapshot_metadata, attach_usage_settlement_pricing_snapshot_metadata,
     clear_previous_request_body_facts, inflate_usage_json_value,
     prepare_request_metadata_for_body_storage, prepare_usage_body_storage,
-    prepare_usage_upsert_context, push_postgres_usage_websocket_filter,
+    prepare_usage_for_persistence, push_postgres_usage_websocket_filter,
     request_body_capture_replaces_derived_facts, resolved_read_usage_body_ref,
     resolved_write_usage_body_ref, split_dashboard_daily_aggregate_range,
     split_dashboard_hourly_aggregate_range, usage_body_capture_state_for_storage, usage_body_ref,
@@ -39,6 +39,7 @@ fn fast_clear_usage_record(
     terminal_service_tier: Option<&str>,
 ) -> UpsertUsageRecord {
     UpsertUsageRecord {
+        capture_retention: Default::default(),
         request_id: request_id.to_string(),
         user_id: None,
         api_key_id: None,
@@ -204,7 +205,26 @@ async fn live_full_http_capture_round_trips_for_direct_and_batch_writes() {
     let repository = SqlxUsageReadRepository::new(factory.connect_lazy().unwrap());
     crate::run_migrations(repository.pool()).await.unwrap();
 
-    for batch in [false, true] {
+    for write_mode in 0..3 {
+        use aether_data_contracts::repository::usage::{
+            usage_json_heap_estimate, UsageCaptureMemoryBudget,
+        };
+        let batch = write_mode != 0;
+        let budget = Arc::new(UsageCaptureMemoryBudget::new(4 * 1024 * 1024));
+        let retain_capture = |usage: &mut UpsertUsageRecord| {
+            let bytes = [
+                usage.request_body.as_ref(),
+                usage.provider_request_body.as_ref(),
+                usage.response_body.as_ref(),
+                usage.client_response_body.as_ref(),
+            ]
+            .into_iter()
+            .flatten()
+            .map(|body| std::mem::size_of::<serde_json::Value>() + usage_json_heap_estimate(body))
+            .sum();
+            assert!(usage.capture_retention.reserve(Arc::clone(&budget), bytes));
+            bytes
+        };
         let request_id = format!("req-full-capture-{}", uuid::Uuid::new_v4().simple());
         let now_unix_secs = Utc::now().timestamp() as u64;
         let mut pending = fast_clear_usage_record(
@@ -225,14 +245,18 @@ async fn live_full_http_capture_round_trips_for_direct_and_batch_writes() {
         pending.response_body_state = Some(UsageBodyCaptureState::Inline);
         pending.client_response_body = Some(json!("pending client response"));
         pending.client_response_body_state = Some(UsageBodyCaptureState::Inline);
+        let pending_bytes = retain_capture(&mut pending);
         if batch {
-            repository
-                .upsert_pending_many(vec![pending.clone()])
-                .await
-                .unwrap();
+            let records = if write_mode == 2 {
+                vec![pending.clone(), pending.clone()]
+            } else {
+                vec![pending.clone()]
+            };
+            repository.upsert_pending_many(records).await.unwrap();
         } else {
             repository.upsert(pending.clone()).await.unwrap();
         }
+        assert_eq!(budget.retained_bytes(), pending_bytes);
         for (field, expected) in [
             (UsageBodyField::RequestBody, pending.request_body.as_ref()),
             (
@@ -274,7 +298,10 @@ async fn live_full_http_capture_round_trips_for_direct_and_batch_writes() {
         terminal.response_body_state = Some(UsageBodyCaptureState::Inline);
         terminal.client_response_body = Some(json!({"output": "final response"}));
         terminal.client_response_body_state = Some(UsageBodyCaptureState::Inline);
+        let terminal_bytes = retain_capture(&mut terminal);
         repository.upsert(terminal.clone()).await.unwrap();
+        assert_eq!(budget.retained_bytes(), pending_bytes + terminal_bytes);
+        assert_eq!(budget.downgraded_total(), 0);
 
         let stored = repository
             .find_by_request_id_shallow(&request_id)
@@ -330,6 +357,9 @@ async fn live_full_http_capture_round_trips_for_direct_and_batch_writes() {
             .execute(repository.pool())
             .await
             .unwrap();
+        drop(pending);
+        drop(terminal);
+        assert_eq!(budget.retained_bytes(), 0);
     }
 }
 
@@ -2411,6 +2441,7 @@ async fn validates_upsert_before_hitting_database() {
     let repository = SqlxUsageReadRepository::new(pool);
     let result = repository
         .upsert(UpsertUsageRecord {
+            capture_retention: Default::default(),
             request_id: "".to_string(),
             user_id: None,
             api_key_id: None,
@@ -4324,6 +4355,96 @@ fn prepare_usage_body_storage_compresses_large_payloads() {
     );
 }
 
+#[test]
+fn prepare_usage_body_storage_streams_json_shapes_into_compatible_gzip() {
+    for payload in [
+        serde_json::Value::Null,
+        json!(false),
+        json!(42),
+        json!(["quoted\"text", "line\nbreak", "\u{4e2d}\u{6587}", null]),
+        json!({
+            "content": "escaped\n\"\\value".repeat(32 * 1024),
+            "nested": {"values": [true, null, 1.25, -7]}
+        }),
+    ] {
+        let storage = prepare_usage_body_storage(Some(&payload)).expect("body should compress");
+        assert!(storage.inline_json.is_none());
+        let compressed = storage
+            .detached_blob_bytes
+            .expect("body should be detached");
+        assert_eq!(
+            inflate_usage_json_value(&compressed).expect("body should remain readable"),
+            payload
+        );
+    }
+}
+
+#[test]
+fn managed_capture_preparation_moves_bodies_without_a_second_reservation() {
+    use aether_data_contracts::repository::usage::{
+        sanitize_usage_for_persistence, usage_json_heap_estimate, UsageCaptureMemoryBudget,
+    };
+
+    let mut usage = fast_clear_usage_record(
+        "req-managed-capture",
+        "managed-capture",
+        100,
+        true,
+        UsageBodyCaptureState::Inline,
+        Some("priority"),
+    );
+    let bodies = [
+        json!({"messages": [{"role": "user", "content": "request".repeat(4096)}]}),
+        json!({"input": "provider request".repeat(4096), "service_tier": "priority"}),
+        json!({"output": "provider response".repeat(4096)}),
+        json!({"output": "client response".repeat(4096)}),
+    ];
+    usage.request_body = Some(bodies[0].clone());
+    usage.provider_request_body = Some(bodies[1].clone());
+    usage.response_body = Some(bodies[2].clone());
+    usage.client_response_body = Some(bodies[3].clone());
+    usage.request_body_state = Some(UsageBodyCaptureState::Inline);
+    usage.response_body_state = Some(UsageBodyCaptureState::Inline);
+    usage.client_response_body_state = Some(UsageBodyCaptureState::Inline);
+    usage.request_headers = Some(json!({"content-type": "application/json"}));
+    usage.cache_read_input_tokens = Some(0);
+    usage.total_cost_usd = Some(0.25);
+    usage.actual_total_cost_usd = Some(0.125);
+    let expected_accounting = sanitize_usage_for_persistence(usage.clone());
+    let bytes = [
+        usage.request_body.as_ref(),
+        usage.provider_request_body.as_ref(),
+        usage.response_body.as_ref(),
+        usage.client_response_body.as_ref(),
+    ]
+    .into_iter()
+    .flatten()
+    .map(|body| std::mem::size_of::<serde_json::Value>() + usage_json_heap_estimate(body))
+    .sum();
+    let budget = Arc::new(UsageCaptureMemoryBudget::new(bytes));
+    assert!(usage.capture_retention.reserve(Arc::clone(&budget), bytes));
+
+    let (accounting, prepared) = prepare_usage_for_persistence(usage);
+    let prepared = prepared.expect("managed capture should prepare without cloning bodies");
+    assert_eq!(budget.retained_bytes(), 0);
+    assert_eq!(budget.downgraded_total(), 0);
+    assert_eq!(accounting, expected_accounting);
+    for (storage, expected) in [
+        prepared.request_body_storage,
+        prepared.provider_request_body_storage,
+        prepared.response_body_storage,
+        prepared.client_response_body_storage,
+    ]
+    .into_iter()
+    .zip(bodies)
+    {
+        assert_eq!(
+            inflate_usage_json_value(storage.detached_blob_bytes.as_deref().unwrap()).unwrap(),
+            expected
+        );
+    }
+}
+
 #[test]
 fn usage_body_capture_state_for_storage_marks_detached_bodies_as_reference() {
     let payload = json!({"message": "hello"});
@@ -4413,7 +4534,8 @@ fn explicit_none_capture_drops_residual_body_ref_and_incoming_fast_metadata_befo
         "provider_request_body_ref": "usage://request/req-none-residual/provider_request_body"
     }));
 
-    let prepared = prepare_usage_upsert_context(&usage).expect("usage should prepare");
+    let (_, prepared) = prepare_usage_for_persistence(usage);
+    let prepared = prepared.expect("usage should prepare");
     assert!(prepared.clear_provider_request_body);
     assert!(!prepared.provider_request_body_storage.has_detached_blob());
     assert_eq!(prepared.http_audit_refs.provider_request_body_ref, None);
@@ -4805,6 +4927,7 @@ fn attach_usage_http_audit_body_refs_adds_missing_metadata_without_overwriting_e
 fn usage_routing_snapshot_from_usage_only_activates_for_routing_metadata() {
     let snapshot = usage_routing_snapshot_from_usage(
         &UpsertUsageRecord {
+            capture_retention: Default::default(),
             request_id: "req-123".to_string(),
             user_id: None,
             api_key_id: None,
@@ -4904,6 +5027,7 @@ fn usage_routing_snapshot_from_usage_only_activates_for_routing_metadata() {
 
     let empty_snapshot = usage_routing_snapshot_from_usage(
         &UpsertUsageRecord {
+            capture_retention: Default::default(),
             request_id: "req-124".to_string(),
             user_id: None,
             api_key_id: None,
@@ -4982,6 +5106,7 @@ fn usage_routing_snapshot_from_usage_only_activates_for_routing_metadata() {
 fn usage_routing_snapshot_from_usage_prefers_typed_routing_fields_without_metadata() {
     let snapshot = usage_routing_snapshot_from_usage(
         &UpsertUsageRecord {
+            capture_retention: Default::default(),
             request_id: "req-typed-routing-1".to_string(),
             user_id: None,
             api_key_id: None,
@@ -5116,6 +5241,7 @@ fn attach_usage_routing_snapshot_metadata_adds_missing_keys_without_overwriting_
 fn usage_settlement_pricing_snapshot_from_usage_extracts_typed_billing_fields() {
     let snapshot = usage_settlement_pricing_snapshot_from_usage(
         &UpsertUsageRecord {
+            capture_retention: Default::default(),
             request_id: "req-125".to_string(),
             user_id: None,
             api_key_id: None,