From 66f127c4048b29291fa2088d1fd6e9816dfbfcf7 Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 16 Jul 2026 13:21:12 +0800 Subject: [PATCH 1/9] feat(usage): durable minute rollup so per-user cost survives ledger pruning --- crates/server/src/main.rs | 2 + crates/state/src/store.rs | 379 ++++++++++++++++++++++++++++++++++---- crates/task/src/lib.rs | 24 +++ docs/governance.md | 8 + docs/observability.md | 3 + 5 files changed, 377 insertions(+), 39 deletions(-) diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 77f5106..2e7e149 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -84,6 +84,7 @@ async fn main() -> anyhow::Result<()> { // daily quota reset; governance is a preserved seam, so this survives reloads let quota_task = gw_task::spawn_quota_reset(state.clone(), gw_task::DAILY); let purge_task = gw_task::spawn_content_purge(state.clone(), gw_task::PURGE_PERIOD); + let rollup_task = gw_task::spawn_usage_rollup(state.clone(), gw_task::ROLLUP_PERIOD); let distributed_batches = state.store.distributed_batches(); let transport = select_transport()?; @@ -196,6 +197,7 @@ async fn main() -> anyhow::Result<()> { quota_task.abort(); purge_task.abort(); + rollup_task.abort(); tracing::info!("gw drained and exiting"); Ok(()) } diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 08abe30..aa2536d 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -27,6 +27,14 @@ pub const MAX_METERED_TOKENS: i64 = 1_000_000_000; /// approximate by at most this many rows, saving a round-trip per billing). const LEDGER_PRUNE_EVERY: usize = 64; +/// Usage-rollup bucket width. +const ROLLUP_BUCKET_SECS: i64 = 60; + +/// Each rollup advance recomputes this trailing window whole, so late rows and +/// missed ticks self-heal. Size `ledger_max_rows` to hold at least this much +/// traffic, or a recomputed bucket could undercount already-pruned rows. +const ROLLUP_BACKFILL_SECS: i64 = 20 * 60; + /// One billing entry (recorded locally only; no reporting upstream). #[derive(Debug, Clone, serde::Serialize)] pub struct BillingRecord { @@ -230,6 +238,61 @@ pub struct UserUsageRow { pub vendor_cost_micros: i64, } +impl UserUsageRow { + fn zero(user_id: String, model: String) -> Self { + Self { + user_id, + model, + requests: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + cost_micros: 0, + vendor_cost_micros: 0, + } + } + + /// Fold `o`'s counters into self (saturating). + fn absorb(&mut self, o: &UserUsageRow) { + self.requests = self.requests.saturating_add(o.requests); + self.prompt_tokens = self.prompt_tokens.saturating_add(o.prompt_tokens); + self.completion_tokens = self.completion_tokens.saturating_add(o.completion_tokens); + self.total_tokens = self.total_tokens.saturating_add(o.total_tokens); + self.cost_micros = self.cost_micros.saturating_add(o.cost_micros); + self.vendor_cost_micros = self.vendor_cost_micros.saturating_add(o.vendor_cost_micros); + } +} + +/// One raw ledger row as a single-request usage line. +fn usage_of(r: &BillingRecord) -> UserUsageRow { + UserUsageRow { + user_id: r.user_id.clone(), + model: r.model.clone(), + requests: 1, + prompt_tokens: r.prompt_tokens, + completion_tokens: r.completion_tokens, + total_tokens: r.total_tokens, + cost_micros: r.cost_micros, + vendor_cost_micros: r.vendor_cost_micros, + } +} + +/// Fold grouped usage rows into `map` by (user, model). +fn fold_user_usage( + map: &mut std::collections::BTreeMap<(String, String), UserUsageRow>, + rows: impl IntoIterator, +) { + for r in rows { + map.entry((r.user_id.clone(), r.model.clone())) + .or_insert_with(|| UserUsageRow::zero(r.user_id.clone(), r.model.clone())) + .absorb(&r); + } +} + +fn bucket_floor(ts: i64) -> i64 { + ts - ts.rem_euclid(ROLLUP_BUCKET_SECS) +} + /// A content-safety outcome, recorded WITHOUT the offending text — only which /// key/user/rule fired and what the gateway did, so hits are queryable per /// ak/tenant without retaining prompt content. @@ -267,7 +330,8 @@ pub struct AdminAudit { pub actor: String, /// The presented scope: "global" | "tenant". pub scope: String, - /// The mutation: "key_create" | "key_patch" | "key_delete" | "config_publish" | "reload". + /// The mutation: "key_create" | "key_patch" | "key_delete" | + /// "config_publish" | "reload" | "content_erase". pub action: String, /// The object acted on (an ak, a config version, …). pub target: String, @@ -283,9 +347,11 @@ pub trait Store: Send + Sync + std::fmt::Debug { async fn ledger_snapshot(&self, limit: usize) -> GResult<(usize, Vec)>; /// Usage rolled up by (tenant, requested model), sorted. async fn ledger_usage(&self, tenant: Option<&str>) -> GResult>; - /// Precise per-user cost over `[since, until]` (unix secs), grouped by - /// (user, requested model); optional tenant/user filter. The billing-period - /// query behind per-user invoicing. + /// Per-user cost over `[since, until]` (unix secs), grouped by (user, + /// requested model); optional tenant/user filter. The billing-period query + /// behind per-user invoicing. Served from the minute rollup for rolled + /// minutes (where the bounds are minute-aligned) plus the raw ledger tail, + /// so the result survives `ledger_max_rows` pruning. async fn usage_by_user( &self, tenant: Option<&str>, @@ -293,6 +359,12 @@ pub trait Store: Send + Sync + std::fmt::Debug { since: i64, until: i64, ) -> GResult>; + /// Roll completed minutes of the ledger into durable `usage_rollup` + /// buckets: every bucket in the trailing backfill window is recomputed + /// from the raw rows and upserted — never deleted — so the periodic task + /// is idempotent, self-heals missed ticks, and a bucket outlives the raw + /// rows it summarizes. Returns the buckets written. + async fn usage_rollup_advance(&self, now_epoch_secs: i64) -> GResult; /// Append a content-safety event (no prompt text retained). async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()>; @@ -398,6 +470,9 @@ pub trait Store: Send + Sync + std::fmt::Debug { #[derive(Debug, Default)] pub struct MemoryStore { records: Mutex>, + /// Minute buckets keyed by (minute, tenant, user, model); see + /// [`Store::usage_rollup_advance`]. + rollup: Mutex>, sec_events: Mutex>, audit: Mutex>, content: Mutex>, @@ -485,40 +560,70 @@ impl Store for MemoryStore { since: i64, until: i64, ) -> GResult> { + let mut map = std::collections::BTreeMap::new(); + let watermark = { + let rollup = self + .rollup + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for ((minute, t, u, _), row) in rollup.iter() { + if *minute >= bucket_floor(since) + && *minute <= bucket_floor(until) + && tenant.is_none_or(|f| f == t) + && user.is_none_or(|f| f == u) + { + fold_user_usage(&mut map, [row.clone()]); + } + } + rollup + .keys() + .next_back() + .map_or(0, |k| k.0 + ROLLUP_BUCKET_SECS) + }; let records = self .records .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - let mut rollup: std::collections::BTreeMap<(String, String), UserUsageRow> = - std::collections::BTreeMap::new(); - for r in records.iter() { - if tenant.is_some_and(|t| t != r.tenant) - || user.is_some_and(|u| u != r.user_id) - || r.created_at_epoch_secs < since - || r.created_at_epoch_secs > until + let tail = records.iter().filter(|r| { + r.created_at_epoch_secs >= since.max(watermark) + && r.created_at_epoch_secs <= until + && tenant.is_none_or(|t| t == r.tenant) + && user.is_none_or(|u| u == r.user_id) + }); + fold_user_usage(&mut map, tail.map(usage_of)); + Ok(map.into_values().collect()) + } + + async fn usage_rollup_advance(&self, now: i64) -> GResult { + let hi = bucket_floor(now); + let lo = hi - ROLLUP_BACKFILL_SECS; + let mut fresh = std::collections::BTreeMap::new(); + { + let records = self + .records + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + for r in records + .iter() + .filter(|r| r.created_at_epoch_secs >= lo && r.created_at_epoch_secs < hi) { - continue; + fresh + .entry(( + bucket_floor(r.created_at_epoch_secs), + r.tenant.clone(), + r.user_id.clone(), + r.model.clone(), + )) + .or_insert_with(|| UserUsageRow::zero(r.user_id.clone(), r.model.clone())) + .absorb(&usage_of(r)); } - let e = rollup - .entry((r.user_id.clone(), r.model.clone())) - .or_insert_with(|| UserUsageRow { - user_id: r.user_id.clone(), - model: r.model.clone(), - requests: 0, - prompt_tokens: 0, - completion_tokens: 0, - total_tokens: 0, - cost_micros: 0, - vendor_cost_micros: 0, - }); - e.requests += 1; - e.prompt_tokens = e.prompt_tokens.saturating_add(r.prompt_tokens); - e.completion_tokens = e.completion_tokens.saturating_add(r.completion_tokens); - e.total_tokens = e.total_tokens.saturating_add(r.total_tokens); - e.cost_micros = e.cost_micros.saturating_add(r.cost_micros); - e.vendor_cost_micros = e.vendor_cost_micros.saturating_add(r.vendor_cost_micros); } - Ok(rollup.into_values().collect()) + let written = fresh.len() as u64; + self.rollup + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .extend(fresh); + Ok(written) } async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { @@ -772,6 +877,13 @@ impl SqliteStore { created_at_epoch_secs INTEGER NOT NULL DEFAULT 0, estimated INTEGER NOT NULL DEFAULT 0)", "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", + "CREATE TABLE IF NOT EXISTS usage_rollup ( + minute_epoch INTEGER NOT NULL, tenant TEXT NOT NULL, user_id TEXT NOT NULL, + model TEXT NOT NULL, requests INTEGER NOT NULL, + prompt_tokens INTEGER NOT NULL, completion_tokens INTEGER NOT NULL, + total_tokens INTEGER NOT NULL, cost_micros INTEGER NOT NULL, + vendor_cost_micros INTEGER NOT NULL, + PRIMARY KEY (minute_epoch, tenant, user_id, model))", "CREATE TABLE IF NOT EXISTS files ( n INTEGER PRIMARY KEY AUTOINCREMENT, id TEXT UNIQUE NOT NULL, tenant TEXT NOT NULL DEFAULT 'default', @@ -938,22 +1050,70 @@ impl Store for SqliteStore { since: i64, until: i64, ) -> GResult> { - let rows = sqlx::query( + let watermark: i64 = + sqlx::query_scalar("SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup") + .fetch_one(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; + let rolled = sqlx::query( + "SELECT user_id, model, SUM(requests), SUM(prompt_tokens), SUM(completion_tokens), + SUM(total_tokens), SUM(cost_micros), SUM(vendor_cost_micros) + FROM usage_rollup + WHERE (?1 IS NULL OR tenant = ?1) AND (?2 IS NULL OR user_id = ?2) + AND minute_epoch BETWEEN (?3/60)*60 AND (?4/60)*60 + GROUP BY user_id, model", + ) + .bind(tenant) + .bind(user) + .bind(since) + .bind(until) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read rolled usage", e))?; + let raw = sqlx::query( "SELECT user_id, model, COUNT(*), SUM(prompt_tokens), SUM(completion_tokens), SUM(total_tokens), SUM(cost_micros), SUM(vendor_cost_micros) FROM billing WHERE (?1 IS NULL OR tenant = ?1) AND (?2 IS NULL OR user_id = ?2) - AND created_at_epoch_secs BETWEEN ?3 AND ?4 - GROUP BY user_id, model ORDER BY user_id, model", + AND created_at_epoch_secs BETWEEN MAX(?3, ?5) AND ?4 + GROUP BY user_id, model", ) .bind(tenant) .bind(user) .bind(since) .bind(until) + .bind(watermark) .fetch_all(&self.pool) .await .map_err(|e| crate::sqlx_err("roll up user usage", e))?; - Ok(rows.iter().map(user_usage_row).collect()) + let mut map = std::collections::BTreeMap::new(); + fold_user_usage(&mut map, rolled.iter().map(user_usage_row)); + fold_user_usage(&mut map, raw.iter().map(user_usage_row)); + Ok(map.into_values().collect()) + } + + async fn usage_rollup_advance(&self, now: i64) -> GResult { + let hi = bucket_floor(now); + let r = sqlx::query( + "INSERT INTO usage_rollup (minute_epoch, tenant, user_id, model, requests, + prompt_tokens, completion_tokens, total_tokens, cost_micros, vendor_cost_micros) + SELECT (created_at_epoch_secs/60)*60, tenant, user_id, model, COUNT(*), + SUM(prompt_tokens), SUM(completion_tokens), SUM(total_tokens), + SUM(cost_micros), SUM(vendor_cost_micros) + FROM billing WHERE created_at_epoch_secs >= ?1 AND created_at_epoch_secs < ?2 + GROUP BY 1, 2, 3, 4 + ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET + requests = excluded.requests, prompt_tokens = excluded.prompt_tokens, + completion_tokens = excluded.completion_tokens, + total_tokens = excluded.total_tokens, cost_micros = excluded.cost_micros, + vendor_cost_micros = excluded.vendor_cost_micros", + ) + .bind(hi - ROLLUP_BACKFILL_SECS) + .bind(hi) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("advance usage rollup", e))?; + Ok(r.rows_affected()) } async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { @@ -1235,6 +1395,13 @@ impl PostgresStore { created_at_epoch_secs BIGINT NOT NULL DEFAULT 0, estimated BOOLEAN NOT NULL DEFAULT FALSE)", "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", + "CREATE TABLE IF NOT EXISTS usage_rollup ( + minute_epoch BIGINT NOT NULL, tenant TEXT NOT NULL, user_id TEXT NOT NULL, + model TEXT NOT NULL, requests BIGINT NOT NULL, + prompt_tokens BIGINT NOT NULL, completion_tokens BIGINT NOT NULL, + total_tokens BIGINT NOT NULL, cost_micros BIGINT NOT NULL, + vendor_cost_micros BIGINT NOT NULL, + PRIMARY KEY (minute_epoch, tenant, user_id, model))", "CREATE TABLE IF NOT EXISTS security_events ( n BIGSERIAL PRIMARY KEY, created_at_epoch_secs BIGINT NOT NULL, request_id TEXT NOT NULL DEFAULT '', ak TEXT NOT NULL DEFAULT '', @@ -1412,23 +1579,73 @@ impl Store for PostgresStore { since: i64, until: i64, ) -> GResult> { - let rows = sqlx::query( + let watermark: i64 = + sqlx::query_scalar("SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup") + .fetch_one(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; + let rolled = sqlx::query( + "SELECT user_id, model, SUM(requests)::BIGINT, + SUM(prompt_tokens)::BIGINT, SUM(completion_tokens)::BIGINT, + SUM(total_tokens)::BIGINT, SUM(cost_micros)::BIGINT, SUM(vendor_cost_micros)::BIGINT + FROM usage_rollup + WHERE ($1::text IS NULL OR tenant = $1) AND ($2::text IS NULL OR user_id = $2) + AND minute_epoch BETWEEN ($3/60)*60 AND ($4/60)*60 + GROUP BY user_id, model", + ) + .bind(tenant) + .bind(user) + .bind(since) + .bind(until) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read rolled usage", e))?; + let raw = sqlx::query( "SELECT user_id, model, COUNT(*), SUM(prompt_tokens)::BIGINT, SUM(completion_tokens)::BIGINT, SUM(total_tokens)::BIGINT, SUM(cost_micros)::BIGINT, SUM(vendor_cost_micros)::BIGINT FROM billing WHERE ($1::text IS NULL OR tenant = $1) AND ($2::text IS NULL OR user_id = $2) - AND created_at_epoch_secs BETWEEN $3 AND $4 - GROUP BY user_id, model ORDER BY user_id, model", + AND created_at_epoch_secs BETWEEN GREATEST($3, $5) AND $4 + GROUP BY user_id, model", ) .bind(tenant) .bind(user) .bind(since) .bind(until) + .bind(watermark) .fetch_all(&self.pool) .await .map_err(|e| crate::sqlx_err("roll up user usage", e))?; - Ok(rows.iter().map(user_usage_row).collect()) + let mut map = std::collections::BTreeMap::new(); + fold_user_usage(&mut map, rolled.iter().map(user_usage_row)); + fold_user_usage(&mut map, raw.iter().map(user_usage_row)); + Ok(map.into_values().collect()) + } + + async fn usage_rollup_advance(&self, now: i64) -> GResult { + let hi = bucket_floor(now); + let r = sqlx::query( + "INSERT INTO usage_rollup (minute_epoch, tenant, user_id, model, requests, + prompt_tokens, completion_tokens, total_tokens, cost_micros, vendor_cost_micros) + SELECT (created_at_epoch_secs/60)*60, tenant, user_id, model, COUNT(*), + SUM(prompt_tokens)::BIGINT, SUM(completion_tokens)::BIGINT, + SUM(total_tokens)::BIGINT, SUM(cost_micros)::BIGINT, + SUM(vendor_cost_micros)::BIGINT + FROM billing WHERE created_at_epoch_secs >= $1 AND created_at_epoch_secs < $2 + GROUP BY 1, 2, 3, 4 + ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET + requests = excluded.requests, prompt_tokens = excluded.prompt_tokens, + completion_tokens = excluded.completion_tokens, + total_tokens = excluded.total_tokens, cost_micros = excluded.cost_micros, + vendor_cost_micros = excluded.vendor_cost_micros", + ) + .bind(hi - ROLLUP_BACKFILL_SECS) + .bind(hi) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("advance usage rollup", e))?; + Ok(r.rows_affected()) } async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { @@ -2083,6 +2300,90 @@ mod tests { exercise_audit(&store).await; } + async fn exercise_rollup(store: &dyn Store) { + let mut a = record("m1"); + a.user_id = "alice".into(); + a.created_at_epoch_secs = 400; + let mut b = record("m1"); + b.user_id = "bob".into(); + b.created_at_epoch_secs = 460; + let mut tail = record("m1"); + tail.user_id = "carol".into(); + tail.created_at_epoch_secs = 1_520; + for r in [&a, &b, &tail] { + store.ledger_add(r).await.unwrap(); + } + let baseline = store.usage_by_user(None, None, 0, i64::MAX).await.unwrap(); + assert_eq!(baseline.len(), 3, "raw-only result before any rollup"); + + store.usage_rollup_advance(1_500).await.unwrap(); + store.usage_rollup_advance(1_500).await.unwrap(); + let after = store.usage_by_user(None, None, 0, i64::MAX).await.unwrap(); + assert_eq!( + format!("{baseline:?}"), + format!("{after:?}"), + "rollup + tail union matches raw and advancing twice is idempotent" + ); + + let windowed = store + .usage_by_user(None, None, 450, i64::MAX) + .await + .unwrap(); + assert!( + windowed.iter().all(|r| r.user_id != "alice"), + "window excludes alice's earlier bucket" + ); + assert!( + windowed.iter().any(|r| r.user_id == "bob") + && windowed.iter().any(|r| r.user_id == "carol"), + "window keeps the rolled bucket and the raw tail" + ); + } + + #[tokio::test] + async fn memory_rollup_roundtrip() { + exercise_rollup(&MemoryStore::default()).await; + } + + #[tokio::test] + async fn sqlite_rollup_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let store = SqliteStore::open(dir.path().join("rollup.db").to_str().unwrap()) + .await + .unwrap(); + exercise_rollup(&store).await; + } + + #[tokio::test] + async fn rollup_survives_ledger_prune() { + let store = MemoryStore::with_ledger_cap(1); + let mut a = record("m1"); + a.user_id = "alice".into(); + a.created_at_epoch_secs = 9_000; + store.ledger_add(&a).await.unwrap(); + store.usage_rollup_advance(10_000).await.unwrap(); + + let mut b = record("m1"); + b.user_id = "bob".into(); + b.created_at_epoch_secs = 9_970; + store.ledger_add(&b).await.unwrap(); + let (total, _) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(total, 1, "the cap pruned alice's raw row"); + + store.usage_rollup_advance(10_000).await.unwrap(); + let usage = store.usage_by_user(None, None, 0, i64::MAX).await.unwrap(); + let alice = usage.iter().find(|r| r.user_id == "alice"); + assert_eq!( + alice.map(|r| r.total_tokens), + Some(8), + "pruned row still billed via its bucket (re-advance never deletes)" + ); + assert!( + usage.iter().any(|r| r.user_id == "bob"), + "unrolled tail served from the raw ledger" + ); + } + #[tokio::test] async fn batch_result_is_first_writer_wins() { let store = MemoryStore::default(); diff --git a/crates/task/src/lib.rs b/crates/task/src/lib.rs index 4ae0d1a..46a06f2 100644 --- a/crates/task/src/lib.rs +++ b/crates/task/src/lib.rs @@ -11,6 +11,8 @@ use gw_state::GatewayState; pub const DAILY: Duration = Duration::from_secs(24 * 60 * 60); /// Retained content is swept for expiry this often. pub const PURGE_PERIOD: Duration = Duration::from_secs(60 * 60); +/// Ledger minutes are rolled into the durable usage buckets this often. +pub const ROLLUP_PERIOD: Duration = Duration::from_secs(60); /// Spawn the daily quota reset loop. Returns the join handle (abort to stop). /// `period` is configurable so tests don't wait 24h. @@ -51,6 +53,28 @@ pub fn spawn_content_purge( }) } +/// Spawn the usage-rollup loop: folds completed ledger minutes into the +/// durable per-user buckets. Returns the join handle (abort to stop). +pub fn spawn_usage_rollup( + state: Arc, + period: Duration, +) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + let mut tick = tokio::time::interval(period); + tick.tick().await; + loop { + tick.tick().await; + if let Err(e) = state + .store + .usage_rollup_advance(gw_state::epoch_secs()) + .await + { + tracing::warn!(error = %e, "usage rollup failed"); + } + } + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/docs/governance.md b/docs/governance.md index 3fc1d79..6fed0bf 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -110,6 +110,14 @@ so a fleet drainer still attributes and budgets it). `GET /admin/usage/users?use returns per-(user, model) cost over a billing period (add `format=csv` for export). `TenantConf.user_daily_token_quota` sets a soft per-user daily cap. +A background task folds completed ledger minutes into durable per-(minute, +tenant, user, model) rollup buckets (recomputing a trailing 20-minute window +each pass, so late rows and missed ticks self-heal). Usage queries are served +from those buckets plus the raw ledger tail, so per-user cost stays correct +after `storage.ledger_max_rows` prunes old billing rows — size the cap to hold +at least the backfill window of traffic. Once a period is served from buckets, +its `since`/`until` bounds are minute-aligned. + ## Enterprise content policy `security:` is global by default; a tenant may override it whole with diff --git a/docs/observability.md b/docs/observability.md index c11ce30..84ade90 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -36,6 +36,9 @@ Records persist when a SQLite store is configured and can be capped with product, `user_id` (effective end user), model, protocol, account, token counts, cost, `created_at_epoch_secs`, the PTU-spillover flag, and an `estimated` flag (set when counts came from an aborted stream rather than a vendor usage payload). +Per-user usage additionally rolls into durable minute buckets every minute, so +`GET /admin/usage/users` stays correct after `ledger_max_rows` pruning (see +[Governance](governance.md#per-user-attribution-and-billing)). ## Audit trails From 7706c2cd09dfb3b047fef45c9763b45155f28c7f Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 16 Jul 2026 13:22:14 +0800 Subject: [PATCH 2/9] feat(retention): erase a user's retained content on request (GDPR/PIPL) --- crates/state/src/store.rs | 87 +++++++++++++++++++++++++ crates/views/src/lib.rs | 132 ++++++++++++++++++++++++++++++++++++++ docs/api.md | 1 + docs/governance.md | 4 ++ 4 files changed, 224 insertions(+) diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index aa2536d..e40293f 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -384,6 +384,10 @@ pub trait Store: Send + Sync + std::fmt::Debug { /// Delete content whose `expires_at_epoch_secs` is in `(0, now]`; returns the /// number deleted. Rows with `expires_at = 0` are kept until manual purge. async fn content_purge(&self, now_epoch_secs: i64) -> GResult; + /// Delete every retained content row for `user`, optionally confined to one + /// tenant (a tenant admin's scope); rows deleted. The GDPR/PIPL erasure + /// hook — security events and ledger rows carry no content and are kept. + async fn content_erase_user(&self, tenant: Option<&str>, user: &str) -> GResult; /// The retained content for one request (both prompt and response rows). async fn content_for(&self, request_id: &str) -> GResult>; @@ -686,6 +690,16 @@ impl Store for MemoryStore { Ok((before - content.len()) as u64) } + async fn content_erase_user(&self, tenant: Option<&str>, user: &str) -> GResult { + let mut content = self + .content + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let before = content.len(); + content.retain(|r| r.user_id != user || tenant.is_some_and(|t| t != r.tenant)); + Ok((before - content.len()) as u64) + } + async fn content_for(&self, request_id: &str) -> GResult> { let content = self .content @@ -1217,6 +1231,18 @@ impl Store for SqliteStore { Ok(r.rows_affected()) } + async fn content_erase_user(&self, tenant: Option<&str>, user: &str) -> GResult { + let r = sqlx::query( + "DELETE FROM request_content WHERE user_id = ?1 AND (?2 IS NULL OR tenant = ?2)", + ) + .bind(user) + .bind(tenant) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("erase user content", e))?; + Ok(r.rows_affected()) + } + async fn content_for(&self, request_id: &str) -> GResult> { let rows = sqlx::query( "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, kind, content, @@ -1750,6 +1776,19 @@ impl Store for PostgresStore { Ok(r.rows_affected()) } + async fn content_erase_user(&self, tenant: Option<&str>, user: &str) -> GResult { + let r = sqlx::query( + "DELETE FROM request_content + WHERE user_id = $1 AND ($2::text IS NULL OR tenant = $2)", + ) + .bind(user) + .bind(tenant) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("erase user content", e))?; + Ok(r.rows_affected()) + } + async fn content_for(&self, request_id: &str) -> GResult> { let rows = sqlx::query( "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, kind, content, @@ -2384,6 +2423,54 @@ mod tests { ); } + async fn exercise_erase(store: &dyn Store) { + let rec = |req: &str, user: &str, tenant: &str| crate::ContentRecord { + created_at_epoch_secs: 100, + request_id: req.into(), + ak: "ak".into(), + user_id: user.into(), + tenant: tenant.into(), + kind: "prompt".into(), + content: "hello".into(), + sealed: false, + expires_at_epoch_secs: 0, + }; + store.content_add(&rec("r1", "u1", "t1")).await.unwrap(); + store.content_add(&rec("r2", "u1", "t2")).await.unwrap(); + store.content_add(&rec("r3", "u2", "t1")).await.unwrap(); + + assert_eq!( + store.content_erase_user(Some("t1"), "u1").await.unwrap(), + 1, + "tenant-scoped erase touches only that tenant's rows" + ); + assert_eq!(store.content_for("r2").await.unwrap().len(), 1); + assert_eq!( + store.content_erase_user(None, "u1").await.unwrap(), + 1, + "global erase removes the remaining tenant's row" + ); + assert_eq!( + store.content_for("r3").await.unwrap().len(), + 1, + "other users' content is untouched" + ); + } + + #[tokio::test] + async fn memory_erase_roundtrip() { + exercise_erase(&MemoryStore::default()).await; + } + + #[tokio::test] + async fn sqlite_erase_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let store = SqliteStore::open(dir.path().join("erase.db").to_str().unwrap()) + .await + .unwrap(); + exercise_erase(&store).await; + } + #[tokio::test] async fn batch_result_is_first_writer_wins() { let store = MemoryStore::default(); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index a93cac4..886d319 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -126,6 +126,10 @@ pub fn app(state: AppState) -> Router { .route("/admin/audit/events", get(admin_security_events)) .route("/admin/audit/ops", get(admin_audit_ops)) .route("/admin/audit/content/{request_id}", get(admin_content_get)) + .route( + "/admin/audit/content", + axum::routing::delete(admin_content_erase), + ) .route( "/admin/keys/{ak}", axum::routing::patch(admin_key_patch).delete(admin_key_delete), @@ -1695,6 +1699,47 @@ async fn admin_content_get( Json(json!({ "request_id": request_id, "entries": entries })).into_response() } +/// DELETE /admin/audit/content?user= — erase every retained content row for +/// one end user (the GDPR/PIPL right-to-erasure hook). Tenant-scoped like the +/// other content reads; recorded in the admin-op trail as `content_erase`. +/// Ledger rows and security events carry no content and are kept. +async fn admin_content_erase( + State(s): State, + headers: HeaderMap, + AuditSourceIp(source): AuditSourceIp, + axum::extract::Query(q): axum::extract::Query>, +) -> Response { + let scope = match admin_auth(&s, &headers) { + Ok(scope) => scope, + Err(r) => return r, + }; + let Some(user) = q.get("user").filter(|u| !u.is_empty()) else { + return error_response(400, "user is required"); + }; + let tenant = scope.tenant_filter(&q); + match s + .handler + .state() + .store + .content_erase_user(tenant.as_deref(), user) + .await + { + Ok(deleted) => { + audit_admin( + &s, + &scope, + source, + "content_erase", + user, + format!("rows={deleted}"), + ) + .await; + Json(json!({ "user": user, "deleted": deleted })).into_response() + } + Err(e) => gateway_error(e), + } +} + fn gateway_error(e: GatewayError) -> Response { let code = StatusCode::from_u16(e.http_status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); // OpenAI's error schema types `code` as string-or-null, never a number. @@ -3232,6 +3277,93 @@ mod tests { ); } + #[tokio::test] + async fn content_erase_is_tenant_scoped_and_audited() { + let yaml = "listen: {host: h, port: 1}\nadmin: {token_env: GW_TEST_ERASE_ADMIN}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\ntenants: [{name: t1}, {name: t2, admin_token_env: GW_TEST_ERASE_T2}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 10, daily_token_quota: 1000}]"; + // SAFETY: unique var names for this test; no concurrent reader of them. + unsafe { + std::env::set_var("GW_TEST_ERASE_ADMIN", "root-tok"); + std::env::set_var("GW_TEST_ERASE_T2", "t2-tok"); + } + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let app_state = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + let store = app_state.handler.state().store.clone(); + let rec = |req: &str, tenant: &str| gw_state::ContentRecord { + created_at_epoch_secs: 100, + request_id: req.into(), + ak: "k1".into(), + user_id: "u1".into(), + tenant: tenant.into(), + kind: "prompt".into(), + content: "hello".into(), + sealed: false, + expires_at_epoch_secs: 0, + }; + store.content_add(&rec("r1", "t1")).await.unwrap(); + store.content_add(&rec("r2", "t2")).await.unwrap(); + let router = app(app_state); + + let erase = |token: &'static str| { + Request::builder() + .method("DELETE") + .uri("/admin/audit/content?user=u1") + .header("authorization", format!("Bearer {token}")) + .body(Body::empty()) + .unwrap() + }; + let resp = router.clone().oneshot(erase("t2-tok")).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let j = body_json(resp).await; + assert_eq!(j["deleted"], 1, "tenant admin erases only its own tenant"); + assert_eq!( + store.content_for("r1").await.unwrap().len(), + 1, + "the other tenant's row is untouched" + ); + + let resp = router + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri("/admin/audit/content") + .header("authorization", "Bearer root-tok") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST, "user is required"); + + let resp = router.clone().oneshot(erase("root-tok")).await.unwrap(); + assert_eq!( + body_json(resp).await["deleted"], + 1, + "global erase gets the rest" + ); + + let ops = router + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/audit/ops") + .header("authorization", "Bearer root-tok") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let j = body_json(ops).await; + let erases: Vec<_> = j["entries"] + .as_array() + .unwrap() + .iter() + .filter(|e| e["action"] == "content_erase" && e["target"] == "u1") + .collect(); + assert_eq!(erases.len(), 2, "both erasures audited"); + } + #[tokio::test] async fn full_retention_without_key_never_stores_raw_even_with_dlp_off() { let yaml = "listen: {host: h, port: 1}\nadmin: {token_env: GW_TEST_CONTENT_ADMIN}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, retention: {content: full, days: 1}, security: {dlp_redact: false, detect_secrets: false}}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; diff --git a/docs/api.md b/docs/api.md index 2bce4a7..d00dfd0 100644 --- a/docs/api.md +++ b/docs/api.md @@ -138,6 +138,7 @@ surface on a private network regardless. | GET | `/admin/audit/events` | content-safety hits (blocklist / regex / DLP / moderation) recorded without prompt text; `?limit=`; tenant-scoped | | GET | `/admin/audit/ops` | admin-operation trail (key CRUD, config publish, reload) with actor, target, and source IP; `?limit=`; global token only | | GET | `/admin/audit/content/{request_id}` | retained prompt/response for one request, unsealed when `GW_CONTENT_KEY` is set (sealed rows without it return `content: null`); tenant-scoped | +| DELETE | `/admin/audit/content?user=` | erase all retained content for one end user (GDPR/PIPL); tenant-scoped, audited as `content_erase` | Two token tiers: the global token (`admin.token_env`) manages everything; a tenant's `admin_token_env` token manages only that tenant's keys, usage, and diff --git a/docs/governance.md b/docs/governance.md index 6fed0bf..696908d 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -158,3 +158,7 @@ only its audit is aggregated (a store write per token would be too hot). expiry; an hourly purge deletes elapsed content. Read back with `GET /admin/audit/content/{request_id}` (tenant-scoped; sealed rows are unsealed when the key is present, else returned as `content: null`). + `DELETE /admin/audit/content?user=` erases every retained row for one end + user (the GDPR/PIPL right-to-erasure hook), tenant-scoped and itself + recorded in the admin-op trail; ledger rows and security events carry no + content and are kept (billing/audit legal basis). From f239ac78cddd9fefbeed66c8c7eb52139af50bd4 Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 16 Jul 2026 13:25:41 +0800 Subject: [PATCH 3/9] test: exercise rollup and erasure against real Postgres --- crates/state/src/store.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index e40293f..9e51e48 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -2462,6 +2462,20 @@ mod tests { exercise_erase(&MemoryStore::default()).await; } + #[tokio::test] + async fn pg_rollup_and_erase_roundtrip() { + let Ok(url) = std::env::var("GW_TEST_PG_URL") else { + return; + }; + let store = PostgresStore::connect(&url).await.expect("pg connect"); + sqlx::query("TRUNCATE billing, usage_rollup, request_content") + .execute(&store.pool) + .await + .expect("truncate"); + exercise_rollup(&store).await; + exercise_erase(&store).await; + } + #[tokio::test] async fn sqlite_erase_roundtrip() { let dir = tempfile::tempdir().unwrap(); From e7a4bd330c83d64d99d573367c852f10c929e7a8 Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 16 Jul 2026 14:35:09 +0800 Subject: [PATCH 4/9] fix(review): close rollup correctness and erasure-reach findings Rollup: the advance window now reaches back to the rollup watermark when it trails the 20-minute backfill (first run rolls a pre-existing ledger whole; a stalled task catches up), and bucket upserts take the per-column max so a recompute over a partially pruned ledger can never shrink a bucket. Postgres elects one replica per pass via an advisory xact lock; periodic tasks skip missed ticks. usage_by_user bounds are minute-aligned on entry so repeated queries can't drift as the watermark advances. Erasure: batch inputs are deleted at terminal status; batch results carry the effective user and their messages are erased with retained content; the content_erase audit entry commits in the same transaction as the deletion. DELETE /v1/files/{id} (OpenAI-compatible) lets tenants remove uploaded files, which stay tenant-owned assets outside per-user erasure. --- crates/handler/src/lib.rs | 23 ++ crates/handler/src/offline.rs | 3 + crates/state/src/store.rs | 478 +++++++++++++++++++++++++++++----- crates/task/src/lib.rs | 3 + crates/views/src/lib.rs | 99 +++++-- docs/api.md | 3 +- docs/governance.md | 32 ++- 7 files changed, 548 insertions(+), 93 deletions(-) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index f53595e..4c4f461 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -1016,6 +1016,20 @@ mod tests { users.contains("alice") && users.contains("bob"), "each shared-key batch item bills to its own user: {users:?}" ); + let results = h + .state() + .store + .batch_get(&job.id) + .await + .unwrap() + .unwrap() + .results; + let owners: std::collections::HashSet<&str> = + results.iter().map(|r| r.user.as_str()).collect(); + assert!( + owners.contains("alice") && owners.contains("bob"), + "each generated result carries its owner: {owners:?}" + ); } #[tokio::test] @@ -1086,5 +1100,14 @@ mod tests { users.contains("alice") && users.contains("bob"), "distributed batch preserved per-item user attribution: {users:?}" ); + assert!( + state + .store + .batch_load_items(&job.id) + .await + .unwrap() + .is_empty(), + "a terminal batch's input rows are pruned" + ); } } diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 81e204b..d1afc37 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -111,6 +111,7 @@ impl OfflineHandler { lost.store(true, Relaxed); break; } + let user = ak.attributed_user(&item.user).to_owned(); let request = GatewayRequest { is_online: false, ak: ak.ak.clone(), @@ -132,6 +133,7 @@ impl OfflineHandler { ok: false, message, total_tokens: 0, + user: user.clone(), }; let result = match ran { Ok(Ok(ctx)) => match ctx.outcome { @@ -140,6 +142,7 @@ impl OfflineHandler { ok: true, message: out.response.message, total_tokens: out.response.total_tokens, + user: user.clone(), }, None => fail("pipeline produced no outcome".into()), }, diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 9e51e48..d8ce96c 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -30,11 +30,16 @@ const LEDGER_PRUNE_EVERY: usize = 64; /// Usage-rollup bucket width. const ROLLUP_BUCKET_SECS: i64 = 60; -/// Each rollup advance recomputes this trailing window whole, so late rows and -/// missed ticks self-heal. Size `ledger_max_rows` to hold at least this much -/// traffic, or a recomputed bucket could undercount already-pruned rows. +/// Each rollup advance recomputes at least this trailing window (more when the +/// watermark trails it — first run rolls the whole ledger, a stalled task +/// catches up whole), so late rows and missed ticks self-heal. const ROLLUP_BACKFILL_SECS: i64 = 20 * 60; +/// Postgres advisory-lock key serializing the fleet's rollup: one replica +/// advances per tick, the rest skip (the upsert is idempotent either way — +/// the lock only avoids N replicas repeating the same scan). +const ROLLUP_LOCK_KEY: i64 = 0x6777_726f_6c6c; + /// One billing entry (recorded locally only; no reporting upstream). #[derive(Debug, Clone, serde::Serialize)] pub struct BillingRecord { @@ -111,6 +116,10 @@ pub struct BatchItemResult { pub ok: bool, pub message: String, pub total_tokens: i64, + /// Effective end user the item billed to; ties the generated `message` + /// to an owner so user erasure can reach it. + #[serde(skip_serializing_if = "String::is_empty")] + pub user: String, } /// One offline batch job. @@ -261,6 +270,19 @@ impl UserUsageRow { self.cost_micros = self.cost_micros.saturating_add(o.cost_micros); self.vendor_cost_micros = self.vendor_cost_micros.saturating_add(o.vendor_cost_micros); } + + /// Keep the larger of each counter. A bucket's source rows are append-only + /// and only ever shrink via ledger pruning, so per-column max always picks + /// the more complete snapshot — a recompute from a partially pruned ledger + /// can never shrink a bucket. + fn keep_max(&mut self, o: &UserUsageRow) { + self.requests = self.requests.max(o.requests); + self.prompt_tokens = self.prompt_tokens.max(o.prompt_tokens); + self.completion_tokens = self.completion_tokens.max(o.completion_tokens); + self.total_tokens = self.total_tokens.max(o.total_tokens); + self.cost_micros = self.cost_micros.max(o.cost_micros); + self.vendor_cost_micros = self.vendor_cost_micros.max(o.vendor_cost_micros); + } } /// One raw ledger row as a single-request usage line. @@ -293,6 +315,16 @@ fn bucket_floor(ts: i64) -> i64 { ts - ts.rem_euclid(ROLLUP_BUCKET_SECS) } +/// Minute-align a query window: identical semantics whether a minute is served +/// from its bucket or still raw, so a repeated query can't drift as the +/// watermark advances past its bounds. +fn align_bounds(since: i64, until: i64) -> (i64, i64) { + ( + bucket_floor(since), + bucket_floor(until).saturating_add(ROLLUP_BUCKET_SECS - 1), + ) +} + /// A content-safety outcome, recorded WITHOUT the offending text — only which /// key/user/rule fired and what the gateway did, so hits are queryable per /// ak/tenant without retaining prompt content. @@ -384,16 +416,27 @@ pub trait Store: Send + Sync + std::fmt::Debug { /// Delete content whose `expires_at_epoch_secs` is in `(0, now]`; returns the /// number deleted. Rows with `expires_at = 0` are kept until manual purge. async fn content_purge(&self, now_epoch_secs: i64) -> GResult; - /// Delete every retained content row for `user`, optionally confined to one - /// tenant (a tenant admin's scope); rows deleted. The GDPR/PIPL erasure - /// hook — security events and ledger rows carry no content and are kept. - async fn content_erase_user(&self, tenant: Option<&str>, user: &str) -> GResult; + /// Erase every retained trace of `user`'s content, optionally confined to + /// one tenant (a tenant admin's scope): retained prompt/response rows, + /// generated batch-result messages, and leftover terminal batch inputs. + /// `audit` (its `summary` set to the erased-row count here) is written in + /// the same transaction on the SQL backends, so a recorded success can't + /// separate from the deletion. Returns rows erased. Ledger rows and + /// security events carry no content and are kept. + async fn content_erase_user( + &self, + tenant: Option<&str>, + user: &str, + audit: AdminAudit, + ) -> GResult; /// The retained content for one request (both prompt and response rows). async fn content_for(&self, request_id: &str) -> GResult>; /// Store `content` under a fresh id owned by `tenant`; returns the metadata. async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult; async fn file_get(&self, id: &str) -> GResult>; + /// Delete an uploaded file outright; whether it existed. + async fn file_delete(&self, id: &str) -> GResult; async fn batch_create( &self, @@ -564,6 +607,7 @@ impl Store for MemoryStore { since: i64, until: i64, ) -> GResult> { + let (since, until) = align_bounds(since, until); let mut map = std::collections::BTreeMap::new(); let watermark = { let rollup = self @@ -600,7 +644,15 @@ impl Store for MemoryStore { async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now); - let lo = hi - ROLLUP_BACKFILL_SECS; + let mut rollup = self + .rollup + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let watermark = rollup + .keys() + .next_back() + .map_or(0, |k| k.0 + ROLLUP_BUCKET_SECS); + let lo = (hi - ROLLUP_BACKFILL_SECS).min(watermark); let mut fresh = std::collections::BTreeMap::new(); { let records = self @@ -623,10 +675,9 @@ impl Store for MemoryStore { } } let written = fresh.len() as u64; - self.rollup - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .extend(fresh); + for (k, v) in fresh { + rollup.entry(k).and_modify(|e| e.keep_max(&v)).or_insert(v); + } Ok(written) } @@ -690,14 +741,38 @@ impl Store for MemoryStore { Ok((before - content.len()) as u64) } - async fn content_erase_user(&self, tenant: Option<&str>, user: &str) -> GResult { - let mut content = self - .content + async fn content_erase_user( + &self, + tenant: Option<&str>, + user: &str, + mut audit: AdminAudit, + ) -> GResult { + let mut erased = { + let mut content = self + .content + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let before = content.len(); + content.retain(|r| r.user_id != user || tenant.is_some_and(|t| t != r.tenant)); + (before - content.len()) as u64 + }; + for mut job in self.jobs.iter_mut() { + if tenant.is_some_and(|t| t != job.tenant) { + continue; + } + for r in job.results.iter_mut() { + if r.user == user && !r.message.is_empty() { + r.message = String::new(); + erased += 1; + } + } + } + audit.summary = format!("rows={erased}"); + self.audit .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let before = content.len(); - content.retain(|r| r.user_id != user || tenant.is_some_and(|t| t != r.tenant)); - Ok((before - content.len()) as u64) + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(audit); + Ok(erased) } async fn content_for(&self, request_id: &str) -> GResult> { @@ -732,6 +807,10 @@ impl Store for MemoryStore { Ok(self.files.get(id).map(|f| f.value().clone())) } + async fn file_delete(&self, id: &str) -> GResult { + Ok(self.files.remove(id).is_some()) + } + async fn batch_create( &self, ak: &str, @@ -848,6 +927,7 @@ where ok: row.get(1), message: row.get(2), total_tokens: row.get(3), + user: row.get(4), } } @@ -908,7 +988,8 @@ impl SqliteStore { status TEXT NOT NULL, total INTEGER NOT NULL)", "CREATE TABLE IF NOT EXISTS batch_results ( batch_id TEXT NOT NULL, idx INTEGER NOT NULL, ok INTEGER NOT NULL, - message TEXT NOT NULL, total_tokens INTEGER NOT NULL)", + message TEXT NOT NULL, total_tokens INTEGER NOT NULL, + user_id TEXT NOT NULL DEFAULT '')", "CREATE TABLE IF NOT EXISTS security_events ( n INTEGER PRIMARY KEY AUTOINCREMENT, created_at_epoch_secs INTEGER NOT NULL, request_id TEXT NOT NULL DEFAULT '', ak TEXT NOT NULL DEFAULT '', @@ -946,6 +1027,7 @@ impl SqliteStore { // back-fill pre-tenant rows to an unmatchable '' tenant (fail closed) "ALTER TABLE files ADD COLUMN tenant TEXT NOT NULL DEFAULT ''", "ALTER TABLE batches ADD COLUMN tenant TEXT NOT NULL DEFAULT ''", + "ALTER TABLE batch_results ADD COLUMN user_id TEXT NOT NULL DEFAULT ''", ] { if let Err(e) = sqlx::query(ddl).execute(&pool).await && !e.to_string().contains("duplicate column name") @@ -1064,6 +1146,7 @@ impl Store for SqliteStore { since: i64, until: i64, ) -> GResult> { + let (since, until) = align_bounds(since, until); let watermark: i64 = sqlx::query_scalar("SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup") .fetch_one(&self.pool) @@ -1074,7 +1157,7 @@ impl Store for SqliteStore { SUM(total_tokens), SUM(cost_micros), SUM(vendor_cost_micros) FROM usage_rollup WHERE (?1 IS NULL OR tenant = ?1) AND (?2 IS NULL OR user_id = ?2) - AND minute_epoch BETWEEN (?3/60)*60 AND (?4/60)*60 + AND minute_epoch BETWEEN ?3 AND ?4 GROUP BY user_id, model", ) .bind(tenant) @@ -1108,6 +1191,11 @@ impl Store for SqliteStore { async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now); + let watermark: i64 = + sqlx::query_scalar("SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup") + .fetch_one(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; let r = sqlx::query( "INSERT INTO usage_rollup (minute_epoch, tenant, user_id, model, requests, prompt_tokens, completion_tokens, total_tokens, cost_micros, vendor_cost_micros) @@ -1117,12 +1205,14 @@ impl Store for SqliteStore { FROM billing WHERE created_at_epoch_secs >= ?1 AND created_at_epoch_secs < ?2 GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET - requests = excluded.requests, prompt_tokens = excluded.prompt_tokens, - completion_tokens = excluded.completion_tokens, - total_tokens = excluded.total_tokens, cost_micros = excluded.cost_micros, - vendor_cost_micros = excluded.vendor_cost_micros", + requests = MAX(usage_rollup.requests, excluded.requests), + prompt_tokens = MAX(usage_rollup.prompt_tokens, excluded.prompt_tokens), + completion_tokens = MAX(usage_rollup.completion_tokens, excluded.completion_tokens), + total_tokens = MAX(usage_rollup.total_tokens, excluded.total_tokens), + cost_micros = MAX(usage_rollup.cost_micros, excluded.cost_micros), + vendor_cost_micros = MAX(usage_rollup.vendor_cost_micros, excluded.vendor_cost_micros)", ) - .bind(hi - ROLLUP_BACKFILL_SECS) + .bind((hi - ROLLUP_BACKFILL_SECS).min(watermark)) .bind(hi) .execute(&self.pool) .await @@ -1231,16 +1321,53 @@ impl Store for SqliteStore { Ok(r.rows_affected()) } - async fn content_erase_user(&self, tenant: Option<&str>, user: &str) -> GResult { - let r = sqlx::query( + async fn content_erase_user( + &self, + tenant: Option<&str>, + user: &str, + audit: AdminAudit, + ) -> GResult { + let mut tx = self + .pool + .begin() + .await + .map_err(|e| crate::sqlx_err("begin erase", e))?; + let c = sqlx::query( "DELETE FROM request_content WHERE user_id = ?1 AND (?2 IS NULL OR tenant = ?2)", ) .bind(user) .bind(tenant) - .execute(&self.pool) + .execute(&mut *tx) .await .map_err(|e| crate::sqlx_err("erase user content", e))?; - Ok(r.rows_affected()) + let m = sqlx::query( + "UPDATE batch_results SET message = '' WHERE user_id = ?1 AND message <> '' + AND batch_id IN (SELECT id FROM batches WHERE ?2 IS NULL OR tenant = ?2)", + ) + .bind(user) + .bind(tenant) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("erase batch results", e))?; + let erased = c.rows_affected() + m.rows_affected(); + sqlx::query( + "INSERT INTO admin_audit (created_at_epoch_secs, actor, scope, action, target, + summary, source_ip) VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .bind(audit.created_at_epoch_secs) + .bind(&audit.actor) + .bind(&audit.scope) + .bind(&audit.action) + .bind(&audit.target) + .bind(format!("rows={erased}")) + .bind(&audit.source_ip) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("audit erase", e))?; + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit erase", e))?; + Ok(erased) } async fn content_for(&self, request_id: &str) -> GResult> { @@ -1294,6 +1421,15 @@ impl Store for SqliteStore { })) } + async fn file_delete(&self, id: &str) -> GResult { + let r = sqlx::query("DELETE FROM files WHERE id = ?") + .bind(id) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("delete file", e))?; + Ok(r.rows_affected() > 0) + } + async fn batch_create( &self, ak: &str, @@ -1334,7 +1470,7 @@ impl Store for SqliteStore { .map_err(|e| crate::sqlx_err("read batch", e))?; let Some(row) = row else { return Ok(None) }; let results = sqlx::query( - "SELECT idx, ok, message, total_tokens FROM batch_results + "SELECT idx, ok, message, total_tokens, user_id FROM batch_results WHERE batch_id = ? ORDER BY idx", ) .bind(id) @@ -1366,8 +1502,8 @@ impl Store for SqliteStore { async fn batch_push_result(&self, id: &str, result: BatchItemResult) -> GResult<()> { // reject inserts into a terminal batch (single-node, so no writer race) sqlx::query( - "INSERT INTO batch_results (batch_id, idx, ok, message, total_tokens) - SELECT ?, ?, ?, ?, ? + "INSERT INTO batch_results (batch_id, idx, ok, message, total_tokens, user_id) + SELECT ?, ?, ?, ?, ?, ? WHERE EXISTS (SELECT 1 FROM batches WHERE id = ? AND status NOT IN ('completed', 'failed'))", ) @@ -1376,6 +1512,7 @@ impl Store for SqliteStore { .bind(result.ok) .bind(&result.message) .bind(result.total_tokens) + .bind(&result.user) .bind(id) .execute(&self.pool) .await @@ -1399,6 +1536,20 @@ impl PostgresStore { Self::connect_with_cap(url, 0).await } + /// A terminal batch's input rows have served their purpose — delete them so + /// submitted prompt text does not outlive the run. + async fn prune_terminal_items(&self, id: &str, status: BatchStatus) -> GResult<()> { + if !matches!(status, BatchStatus::Completed | BatchStatus::Failed) { + return Ok(()); + } + sqlx::query("DELETE FROM batch_items WHERE batch_id = $1") + .bind(id) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("prune batch items", e))?; + Ok(()) + } + /// `ledger_max_rows` > 0 prunes the oldest billing rows past the cap on write. pub async fn connect_with_cap(url: &str, ledger_max_rows: u64) -> GResult { let pool = sqlx::postgres::PgPoolOptions::new() @@ -1457,12 +1608,14 @@ impl PostgresStore { status TEXT NOT NULL, total BIGINT NOT NULL)", "CREATE TABLE IF NOT EXISTS batch_results ( batch_id TEXT NOT NULL, idx BIGINT NOT NULL, ok BOOLEAN NOT NULL, - message TEXT NOT NULL, total_tokens BIGINT NOT NULL)", + message TEXT NOT NULL, total_tokens BIGINT NOT NULL, + user_id TEXT NOT NULL DEFAULT '')", "CREATE TABLE IF NOT EXISTS batch_items ( batch_id TEXT NOT NULL, idx BIGINT NOT NULL, messages TEXT NOT NULL, PRIMARY KEY (batch_id, idx))", // per-item end-user attribution so a fleet drainer still bills/budgets it "ALTER TABLE batch_items ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''", + "ALTER TABLE batch_results ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''", "ALTER TABLE batches ADD COLUMN IF NOT EXISTS claimed_at TIMESTAMPTZ", // fence token: bumped on every claim so a reclaimed executor's fenced writes no-op "ALTER TABLE batches ADD COLUMN IF NOT EXISTS claim_seq BIGINT NOT NULL DEFAULT 0", @@ -1605,6 +1758,7 @@ impl Store for PostgresStore { since: i64, until: i64, ) -> GResult> { + let (since, until) = align_bounds(since, until); let watermark: i64 = sqlx::query_scalar("SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup") .fetch_one(&self.pool) @@ -1616,7 +1770,7 @@ impl Store for PostgresStore { SUM(total_tokens)::BIGINT, SUM(cost_micros)::BIGINT, SUM(vendor_cost_micros)::BIGINT FROM usage_rollup WHERE ($1::text IS NULL OR tenant = $1) AND ($2::text IS NULL OR user_id = $2) - AND minute_epoch BETWEEN ($3/60)*60 AND ($4/60)*60 + AND minute_epoch BETWEEN $3 AND $4 GROUP BY user_id, model", ) .bind(tenant) @@ -1651,6 +1805,24 @@ impl Store for PostgresStore { async fn usage_rollup_advance(&self, now: i64) -> GResult { let hi = bucket_floor(now); + let mut tx = self + .pool + .begin() + .await + .map_err(|e| crate::sqlx_err("begin rollup", e))?; + let locked: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock($1)") + .bind(ROLLUP_LOCK_KEY) + .fetch_one(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("acquire rollup lock", e))?; + if !locked { + return Ok(0); + } + let watermark: i64 = + sqlx::query_scalar("SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup") + .fetch_one(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("read rollup watermark", e))?; let r = sqlx::query( "INSERT INTO usage_rollup (minute_epoch, tenant, user_id, model, requests, prompt_tokens, completion_tokens, total_tokens, cost_micros, vendor_cost_micros) @@ -1661,16 +1833,23 @@ impl Store for PostgresStore { FROM billing WHERE created_at_epoch_secs >= $1 AND created_at_epoch_secs < $2 GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET - requests = excluded.requests, prompt_tokens = excluded.prompt_tokens, - completion_tokens = excluded.completion_tokens, - total_tokens = excluded.total_tokens, cost_micros = excluded.cost_micros, - vendor_cost_micros = excluded.vendor_cost_micros", + requests = GREATEST(usage_rollup.requests, excluded.requests), + prompt_tokens = GREATEST(usage_rollup.prompt_tokens, excluded.prompt_tokens), + completion_tokens = + GREATEST(usage_rollup.completion_tokens, excluded.completion_tokens), + total_tokens = GREATEST(usage_rollup.total_tokens, excluded.total_tokens), + cost_micros = GREATEST(usage_rollup.cost_micros, excluded.cost_micros), + vendor_cost_micros = + GREATEST(usage_rollup.vendor_cost_micros, excluded.vendor_cost_micros)", ) - .bind(hi - ROLLUP_BACKFILL_SECS) + .bind((hi - ROLLUP_BACKFILL_SECS).min(watermark)) .bind(hi) - .execute(&self.pool) + .execute(&mut *tx) .await .map_err(|e| crate::sqlx_err("advance usage rollup", e))?; + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit rollup", e))?; Ok(r.rows_affected()) } @@ -1776,17 +1955,67 @@ impl Store for PostgresStore { Ok(r.rows_affected()) } - async fn content_erase_user(&self, tenant: Option<&str>, user: &str) -> GResult { - let r = sqlx::query( + async fn content_erase_user( + &self, + tenant: Option<&str>, + user: &str, + audit: AdminAudit, + ) -> GResult { + let mut tx = self + .pool + .begin() + .await + .map_err(|e| crate::sqlx_err("begin erase", e))?; + let c = sqlx::query( "DELETE FROM request_content WHERE user_id = $1 AND ($2::text IS NULL OR tenant = $2)", ) .bind(user) .bind(tenant) - .execute(&self.pool) + .execute(&mut *tx) .await .map_err(|e| crate::sqlx_err("erase user content", e))?; - Ok(r.rows_affected()) + let m = sqlx::query( + "UPDATE batch_results r SET message = '' FROM batches b + WHERE r.batch_id = b.id AND r.user_id = $1 AND r.message <> '' + AND ($2::text IS NULL OR b.tenant = $2)", + ) + .bind(user) + .bind(tenant) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("erase batch results", e))?; + // leftovers from batches that went terminal before item pruning shipped + let i = sqlx::query( + "DELETE FROM batch_items i USING batches b + WHERE i.batch_id = b.id AND i.user_id = $1 + AND ($2::text IS NULL OR b.tenant = $2) + AND b.status IN ('completed', 'failed')", + ) + .bind(user) + .bind(tenant) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("erase batch items", e))?; + let erased = c.rows_affected() + m.rows_affected() + i.rows_affected(); + sqlx::query( + "INSERT INTO admin_audit (created_at_epoch_secs, actor, scope, action, target, + summary, source_ip) VALUES ($1, $2, $3, $4, $5, $6, $7)", + ) + .bind(audit.created_at_epoch_secs) + .bind(&audit.actor) + .bind(&audit.scope) + .bind(&audit.action) + .bind(&audit.target) + .bind(format!("rows={erased}")) + .bind(&audit.source_ip) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("audit erase", e))?; + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit erase", e))?; + Ok(erased) } async fn content_for(&self, request_id: &str) -> GResult> { @@ -1842,6 +2071,15 @@ impl Store for PostgresStore { })) } + async fn file_delete(&self, id: &str) -> GResult { + let r = sqlx::query("DELETE FROM files WHERE id = $1") + .bind(id) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("delete file", e))?; + Ok(r.rows_affected() > 0) + } + async fn batch_create( &self, ak: &str, @@ -1878,7 +2116,7 @@ impl Store for PostgresStore { .map_err(|e| crate::sqlx_err("read batch", e))?; let Some(row) = row else { return Ok(None) }; let results = sqlx::query( - "SELECT idx, ok, message, total_tokens FROM batch_results + "SELECT idx, ok, message, total_tokens, user_id FROM batch_results WHERE batch_id = $1 ORDER BY idx", ) .bind(id) @@ -1904,7 +2142,7 @@ impl Store for PostgresStore { .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("update batch status", e))?; - Ok(()) + self.prune_terminal_items(id, status).await } async fn batch_set_status_owned( @@ -1920,15 +2158,19 @@ impl Store for PostgresStore { .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("update batch status (fenced)", e))?; - Ok(r.rows_affected() > 0) + let applied = r.rows_affected() > 0; + if applied { + self.prune_terminal_items(id, status).await?; + } + Ok(applied) } async fn batch_push_result(&self, id: &str, result: BatchItemResult) -> GResult<()> { // DO NOTHING (first-writer-wins) + non-terminal guard; the FOR UPDATE row // lock serializes with batch_finalize so no result lands after finalize. sqlx::query( - "INSERT INTO batch_results (batch_id, idx, ok, message, total_tokens) - SELECT $1, $2, $3, $4, $5 + "INSERT INTO batch_results (batch_id, idx, ok, message, total_tokens, user_id) + SELECT $1, $2, $3, $4, $5, $6 WHERE EXISTS (SELECT 1 FROM batches WHERE id = $1 AND status NOT IN ('completed', 'failed') FOR UPDATE) ON CONFLICT (batch_id, idx) DO NOTHING", @@ -1938,6 +2180,7 @@ impl Store for PostgresStore { .bind(result.ok) .bind(&result.message) .bind(result.total_tokens) + .bind(&result.user) .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("insert batch result", e))?; @@ -1986,6 +2229,11 @@ impl Store for PostgresStore { .execute(&mut *tx) .await .map_err(|e| crate::sqlx_err("finalize write", e))?; + sqlx::query("DELETE FROM batch_items WHERE batch_id = $1") + .bind(id) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("finalize prune items", e))?; tx.commit() .await .map_err(|e| crate::sqlx_err("finalize commit", e))?; @@ -2206,6 +2454,7 @@ mod tests { ok: true, message: "ok".into(), total_tokens: 8, + user: String::new(), }, ) .await @@ -2423,6 +2672,54 @@ mod tests { ); } + #[tokio::test] + async fn rollup_backfills_pre_existing_ledger_on_first_run() { + let store = MemoryStore::default(); + let mut a = record("m1"); + a.user_id = "alice".into(); + a.created_at_epoch_secs = 400; + store.ledger_add(&a).await.unwrap(); + // first advance long after the row: the window trails back to the + // empty-rollup watermark (0), not just 20 minutes + store.usage_rollup_advance(400 + 3 * 20 * 60).await.unwrap(); + let usage = store.usage_by_user(None, None, 0, i64::MAX).await.unwrap(); + assert_eq!( + usage + .iter() + .find(|r| r.user_id == "alice") + .map(|r| r.total_tokens), + Some(8), + "pre-existing ledger rows are rolled on the first run, not orphaned below the watermark" + ); + } + + #[tokio::test] + async fn recompute_from_pruned_ledger_never_shrinks_a_bucket() { + let store = MemoryStore::with_ledger_cap(2); + for ts in [100, 110] { + let mut r = record("m1"); + r.user_id = "alice".into(); + r.created_at_epoch_secs = ts; + store.ledger_add(&r).await.unwrap(); + } + store.usage_rollup_advance(1_260).await.unwrap(); + + let mut b = record("m1"); + b.user_id = "bob".into(); + b.created_at_epoch_secs = 1_250; + store.ledger_add(&b).await.unwrap(); // cap=2 prunes one alice row + store.usage_rollup_advance(1_260).await.unwrap(); // recomputes alice's minute from a pruned ledger + let usage = store.usage_by_user(None, None, 0, i64::MAX).await.unwrap(); + assert_eq!( + usage + .iter() + .find(|r| r.user_id == "alice") + .map(|r| r.total_tokens), + Some(16), + "a recompute over the pruned ledger keeps the more complete bucket" + ); + } + async fn exercise_erase(store: &dyn Store) { let rec = |req: &str, user: &str, tenant: &str| crate::ContentRecord { created_at_epoch_secs: 100, @@ -2438,15 +2735,44 @@ mod tests { store.content_add(&rec("r1", "u1", "t1")).await.unwrap(); store.content_add(&rec("r2", "u1", "t2")).await.unwrap(); store.content_add(&rec("r3", "u2", "t1")).await.unwrap(); + let job = store.batch_create("ak", "t1", "m", 1).await.unwrap(); + store + .batch_push_result( + &job.id, + BatchItemResult { + index: 0, + ok: true, + message: "generated for u1".into(), + total_tokens: 3, + user: "u1".into(), + }, + ) + .await + .unwrap(); + let audit = || AdminAudit { + created_at_epoch_secs: 20, + actor: "global".into(), + scope: "global".into(), + action: "content_erase".into(), + target: "u1".into(), + summary: String::new(), + source_ip: "10.0.0.1".into(), + }; assert_eq!( - store.content_erase_user(Some("t1"), "u1").await.unwrap(), - 1, - "tenant-scoped erase touches only that tenant's rows" + store + .content_erase_user(Some("t1"), "u1", audit()) + .await + .unwrap(), + 2, + "tenant-scoped erase: the content row plus the batch-result message" ); assert_eq!(store.content_for("r2").await.unwrap().len(), 1); + let results = store.batch_get(&job.id).await.unwrap().unwrap().results; + assert_eq!(results[0].message, "", "generated output erased"); + assert_eq!(results[0].user, "u1", "attribution survives for billing"); assert_eq!( - store.content_erase_user(None, "u1").await.unwrap(), + store.content_erase_user(None, "u1", audit()).await.unwrap(), 1, "global erase removes the remaining tenant's row" ); @@ -2455,6 +2781,18 @@ mod tests { 1, "other users' content is untouched" ); + let trail = store.admin_audit_list(10).await.unwrap(); + let erases: Vec<_> = trail + .iter() + .filter(|e| e.action == "content_erase" && e.target == "u1") + .collect(); + assert_eq!(erases.len(), 2, "each erasure committed its audit entry"); + let summaries: Vec<&str> = erases.iter().map(|e| e.summary.as_str()).collect(); + assert_eq!( + summaries, + ["rows=1", "rows=2"], + "newest first; each summary carries its erased count" + ); } #[tokio::test] @@ -2468,10 +2806,13 @@ mod tests { return; }; let store = PostgresStore::connect(&url).await.expect("pg connect"); - sqlx::query("TRUNCATE billing, usage_rollup, request_content") - .execute(&store.pool) - .await - .expect("truncate"); + sqlx::query( + "TRUNCATE billing, usage_rollup, request_content, admin_audit, + batches, batch_items, batch_results", + ) + .execute(&store.pool) + .await + .expect("truncate"); exercise_rollup(&store).await; exercise_erase(&store).await; } @@ -2497,6 +2838,7 @@ mod tests { ok, message: msg.into(), total_tokens: 1, + user: String::new(), }, ) }; @@ -2516,6 +2858,7 @@ mod tests { ok, message: String::new(), total_tokens: 0, + user: String::new(), }; let job = store.batch_create("ak", "default", "m", 2).await.unwrap(); store @@ -2673,6 +3016,7 @@ mod tests { ok: true, message: "ok".into(), total_tokens: 5, + user: String::new(), }, ) .await @@ -2688,6 +3032,7 @@ mod tests { ok: false, message: "stale".into(), total_tokens: 0, + user: String::new(), }, ) .await @@ -2710,6 +3055,7 @@ mod tests { ok: true, message: "late".into(), total_tokens: 0, + user: String::new(), }, ) .await @@ -2738,6 +3084,13 @@ mod tests { .await .unwrap(); assert_eq!(qjob.total, 2); + let loaded = store.batch_load_items(&qjob.id).await.unwrap(); + assert_eq!(loaded.len(), 2); + assert_eq!(loaded[1].messages[0].content, "two"); + assert_eq!( + loaded[1].user, "u-two", + "per-item user round-trips through pg" + ); loop { let (c, _claim) = store .batch_claim_pending(120) @@ -2756,12 +3109,9 @@ mod tests { break; } } - let loaded = store.batch_load_items(&qjob.id).await.unwrap(); - assert_eq!(loaded.len(), 2); - assert_eq!(loaded[1].messages[0].content, "two"); - assert_eq!( - loaded[1].user, "u-two", - "per-item user round-trips through pg" + assert!( + store.batch_load_items(&qjob.id).await.unwrap().is_empty(), + "terminal status pruned the batch inputs" ); let fjob = store diff --git a/crates/task/src/lib.rs b/crates/task/src/lib.rs index 46a06f2..eea4989 100644 --- a/crates/task/src/lib.rs +++ b/crates/task/src/lib.rs @@ -22,6 +22,7 @@ pub fn spawn_quota_reset( ) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { let mut tick = tokio::time::interval(period); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); tick.tick().await; // first tick fires immediately; skip it loop { tick.tick().await; @@ -39,6 +40,7 @@ pub fn spawn_content_purge( ) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { let mut tick = tokio::time::interval(period); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); tick.tick().await; loop { tick.tick().await; @@ -61,6 +63,7 @@ pub fn spawn_usage_rollup( ) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { let mut tick = tokio::time::interval(period); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); tick.tick().await; loop { tick.tick().await; diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 886d319..4d138ec 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -113,7 +113,7 @@ pub fn app(state: AppState) -> Router { .route("/v1/batches", post(batches_submit)) .route("/v1/batches/{id}", get(batches_get)) .route("/v1/files", post(files_upload)) - .route("/v1/files/{id}", get(files_get)) + .route("/v1/files/{id}", get(files_get).delete(files_delete)) .route("/v1/files/{id}/content", get(files_content)) .route("/v1/realtime", get(realtime_ws)) .route("/internal/ledger", get(ledger)) @@ -1699,10 +1699,12 @@ async fn admin_content_get( Json(json!({ "request_id": request_id, "entries": entries })).into_response() } -/// DELETE /admin/audit/content?user= — erase every retained content row for -/// one end user (the GDPR/PIPL right-to-erasure hook). Tenant-scoped like the -/// other content reads; recorded in the admin-op trail as `content_erase`. -/// Ledger rows and security events carry no content and are kept. +/// DELETE /admin/audit/content?user= — erase every retained trace of one end +/// user's content (the GDPR/PIPL right-to-erasure hook): retained rows, batch +/// result messages, leftover terminal batch inputs. Tenant-scoped; the +/// `content_erase` audit entry commits with the deletion, so a recorded +/// success can't separate from it. Ledger rows and security events carry no +/// content and are kept. async fn admin_content_erase( State(s): State, headers: HeaderMap, @@ -1717,25 +1719,24 @@ async fn admin_content_erase( return error_response(400, "user is required"); }; let tenant = scope.tenant_filter(&q); + let (actor, scope_kind) = scope.audit_identity(); + let audit = gw_state::AdminAudit { + created_at_epoch_secs: gw_state::epoch_secs(), + actor: actor.to_owned(), + scope: scope_kind.to_owned(), + action: "content_erase".to_owned(), + target: user.clone(), + summary: String::new(), + source_ip: source, + }; match s .handler .state() .store - .content_erase_user(tenant.as_deref(), user) + .content_erase_user(tenant.as_deref(), user, audit) .await { - Ok(deleted) => { - audit_admin( - &s, - &scope, - source, - "content_erase", - user, - format!("rows={deleted}"), - ) - .await; - Json(json!({ "user": user, "deleted": deleted })).into_response() - } + Ok(deleted) => Json(json!({ "user": user, "deleted": deleted })).into_response(), Err(e) => gateway_error(e), } } @@ -3098,6 +3099,28 @@ async fn files_get( } } +/// DELETE /v1/files/{id} — remove an uploaded file (OpenAI-compatible). Files +/// are tenant-owned assets; erasing one end user's rows inside an uploaded +/// JSONL is the tenant's call — delete the file and re-upload if needed. +async fn files_delete( + State(s): State, + headers: HeaderMap, + Path(id): Path, +) -> Response { + let ak = match authenticate(&s, &headers).await { + Ok(ak) => ak, + Err((st, msg)) => return error_response(st, msg), + }; + let found = s.handler.state().store.file_get(&id).await; + if let Err(resp) = tenant_owned(found, |f| &f.tenant, &ak.tenant, "file", &id) { + return resp; + } + match s.handler.state().store.file_delete(&id).await { + Ok(_) => Json(json!({"id": id, "object": "file", "deleted": true})).into_response(), + Err(e) => gateway_error(e), + } +} + /// GET /v1/files/{id}/content (download raw content: batch output, etc). async fn files_content( State(s): State, @@ -3277,6 +3300,46 @@ mod tests { ); } + #[tokio::test] + async fn file_delete_is_tenant_scoped() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\ntenants: [{name: t1}, {name: t2}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 10, daily_token_quota: 1000}, {ak: k2, tenant: t2, product: p, qps: 10, daily_token_quota: 1000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let app_state = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + let store = app_state.handler.state().store.clone(); + let f = store.file_put("t1", "batch", "line".into()).await.unwrap(); + let router = app(app_state); + + let del = |token: &'static str, id: String| { + Request::builder() + .method("DELETE") + .uri(format!("/v1/files/{id}")) + .header("authorization", format!("Bearer {token}")) + .body(Body::empty()) + .unwrap() + }; + let resp = router + .clone() + .oneshot(del("k2", f.id.clone())) + .await + .unwrap(); + assert_eq!( + resp.status(), + StatusCode::NOT_FOUND, + "cross-tenant delete answers 404 and removes nothing" + ); + assert!(store.file_get(&f.id).await.unwrap().is_some()); + + let resp = router + .clone() + .oneshot(del("k1", f.id.clone())) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(body_json(resp).await["deleted"], true); + assert!(store.file_get(&f.id).await.unwrap().is_none()); + } + #[tokio::test] async fn content_erase_is_tenant_scoped_and_audited() { let yaml = "listen: {host: h, port: 1}\nadmin: {token_env: GW_TEST_ERASE_ADMIN}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\ntenants: [{name: t1}, {name: t2, admin_token_env: GW_TEST_ERASE_T2}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 10, daily_token_quota: 1000}]"; diff --git a/docs/api.md b/docs/api.md index d00dfd0..86614a5 100644 --- a/docs/api.md +++ b/docs/api.md @@ -70,6 +70,7 @@ sequence (`message_start` → `content_block_*` → `message_delta` → | POST | `/v1/files` | upload JSONL: `{"purpose":"batch","file":""}` | | GET | `/v1/files/{id}` | file metadata | | GET | `/v1/files/{id}/content` | raw content | +| DELETE | `/v1/files/{id}` | delete an uploaded file (tenant-owned) | | POST | `/v1/batches` | `{"input_file_id":"..."}` or inline `{"items":[...]}` | | GET | `/v1/batches/{id}` | status (`pending`/`running`/`completed`/`failed`) + results | @@ -138,7 +139,7 @@ surface on a private network regardless. | GET | `/admin/audit/events` | content-safety hits (blocklist / regex / DLP / moderation) recorded without prompt text; `?limit=`; tenant-scoped | | GET | `/admin/audit/ops` | admin-operation trail (key CRUD, config publish, reload) with actor, target, and source IP; `?limit=`; global token only | | GET | `/admin/audit/content/{request_id}` | retained prompt/response for one request, unsealed when `GW_CONTENT_KEY` is set (sealed rows without it return `content: null`); tenant-scoped | -| DELETE | `/admin/audit/content?user=` | erase all retained content for one end user (GDPR/PIPL); tenant-scoped, audited as `content_erase` | +| DELETE | `/admin/audit/content?user=` | erase all retained content for one end user — retained rows, batch result messages, leftover batch inputs (GDPR/PIPL); tenant-scoped, audited atomically as `content_erase` | Two token tiers: the global token (`admin.token_env`) manages everything; a tenant's `admin_token_env` token manages only that tenant's keys, usage, and diff --git a/docs/governance.md b/docs/governance.md index 696908d..c01421b 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -111,12 +111,16 @@ returns per-(user, model) cost over a billing period (add `format=csv` for export). `TenantConf.user_daily_token_quota` sets a soft per-user daily cap. A background task folds completed ledger minutes into durable per-(minute, -tenant, user, model) rollup buckets (recomputing a trailing 20-minute window -each pass, so late rows and missed ticks self-heal). Usage queries are served -from those buckets plus the raw ledger tail, so per-user cost stays correct -after `storage.ledger_max_rows` prunes old billing rows — size the cap to hold -at least the backfill window of traffic. Once a period is served from buckets, -its `since`/`until` bounds are minute-aligned. +tenant, user, model) rollup buckets. Each pass recomputes from the rollup +watermark or the trailing 20-minute window, whichever reaches further back — +the first run rolls a pre-existing ledger whole, and a stalled task catches up +on its own. Buckets only ever grow (a recompute over a partially pruned ledger +keeps the more complete aggregate), and on Postgres a fleet elects one replica +per pass via an advisory lock. Usage queries are served from those buckets plus +the raw ledger tail, so per-user cost stays correct after +`storage.ledger_max_rows` prunes old billing rows. `since`/`until` bounds are +minute-aligned, so a repeated query returns the same result whether a minute is +served from its bucket or still raw. ## Enterprise content policy @@ -158,7 +162,15 @@ only its audit is aggregated (a store write per token would be too hot). expiry; an hourly purge deletes elapsed content. Read back with `GET /admin/audit/content/{request_id}` (tenant-scoped; sealed rows are unsealed when the key is present, else returned as `content: null`). - `DELETE /admin/audit/content?user=` erases every retained row for one end - user (the GDPR/PIPL right-to-erasure hook), tenant-scoped and itself - recorded in the admin-op trail; ledger rows and security events carry no - content and are kept (billing/audit legal basis). + `DELETE /admin/audit/content?user=` erases every retained trace of one end + user's content (the GDPR/PIPL right-to-erasure hook): retained rows, batch + result messages, and any leftover terminal batch inputs — tenant-scoped, and + the `content_erase` audit entry commits with the deletion so a recorded + success can't separate from it. Erasure is a point-in-time operation: a + request already in flight may persist content just after it, so quiesce the + user's traffic first or simply repeat the call. Ledger rows and security + events carry no content and are kept (billing/audit legal basis). Uploaded + files are tenant-owned assets with no per-user dimension — the tenant deletes + them via `DELETE /v1/files/{id}`. A batch's input rows are deleted as soon as + the batch reaches a terminal status, so submitted prompt text does not + outlive the run. From 03110d12134b2bb2d3eeecd3c6eec95decbb320c Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 16 Jul 2026 14:55:48 +0800 Subject: [PATCH 5/9] fix(review): settle minute, active-item erasure, atomic file delete, tx pruning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A minute is rolled only after a one-minute settle window, so no in-flight billing write can land in an already-rolled minute — the precondition that makes the max-upsert sound. Erasure now blanks a user's queued batch items in place (they fail at execution instead of running erased content), reaching pending/running batches. File deletion is one guarded DELETE (id + tenant) and SQLite file ids derive from the non-rewinding AUTOINCREMENT sequence, closing an id-reuse/check-then-delete cross-tenant window. Fenced/unfenced terminal status writes prune batch inputs in the same transaction. --- crates/handler/src/lib.rs | 23 +++++ crates/handler/src/offline.rs | 15 +++ crates/state/src/store.rs | 171 +++++++++++++++++++++++++++------- crates/views/src/lib.rs | 11 +-- docs/governance.md | 24 +++-- 5 files changed, 192 insertions(+), 52 deletions(-) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 4c4f461..e15a58a 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -980,6 +980,29 @@ mod tests { ); } + #[tokio::test] + async fn erased_batch_item_fails_instead_of_running() { + let h = handler(); + let off = OfflineHandler::new(h.clone()); + let key = ak(&h).await; + let job = off + .submit( + key, + "gpt-4o".into(), + vec![BatchItem { + messages: Vec::new(), + user: "u1".into(), + }], + ) + .await + .unwrap(); + wait_terminal(&h, &job.id).await; + let done = h.state().store.batch_get(&job.id).await.unwrap().unwrap(); + assert!(!done.results[0].ok, "an erased item must not execute"); + assert_eq!(done.results[0].message, "item content erased"); + assert_eq!(done.results[0].total_tokens, 0, "nothing billed"); + } + #[tokio::test] async fn batch_attributes_each_item_to_its_user() { let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]"; diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index d1afc37..2ce661d 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -112,6 +112,21 @@ impl OfflineHandler { break; } let user = ak.attributed_user(&item.user).to_owned(); + // an erased item (messages blanked by content_erase_user) must not + // run: fail it instead of sending an empty prompt upstream + if item.messages.is_empty() { + let result = BatchItemResult { + index, + ok: false, + message: "item content erased".into(), + total_tokens: 0, + user, + }; + if let Err(e) = store.batch_push_result(id, result).await { + tracing::error!(error = %e, batch = %id, "batch result write failed"); + } + continue; + } let request = GatewayRequest { is_online: false, ak: ak.ak.clone(), diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index d8ce96c..3249bc5 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -35,6 +35,13 @@ const ROLLUP_BUCKET_SECS: i64 = 60; /// catches up whole), so late rows and missed ticks self-heal. const ROLLUP_BACKFILL_SECS: i64 = 20 * 60; +/// A minute is rolled only once it has been closed for at least this long, so +/// an in-flight billing write (or a replica whose clock trails by less than a +/// minute) can never land a row in an already-rolled minute — which is what +/// makes the max-upsert sound: a rolled minute's source set can only shrink +/// (pruning), never grow. +const ROLLUP_SETTLE_SECS: i64 = ROLLUP_BUCKET_SECS; + /// Postgres advisory-lock key serializing the fleet's rollup: one replica /// advances per tick, the rest skip (the upsert is idempotent either way — /// the lock only avoids N replicas repeating the same scan). @@ -435,8 +442,9 @@ pub trait Store: Send + Sync + std::fmt::Debug { /// Store `content` under a fresh id owned by `tenant`; returns the metadata. async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult; async fn file_get(&self, id: &str) -> GResult>; - /// Delete an uploaded file outright; whether it existed. - async fn file_delete(&self, id: &str) -> GResult; + /// Delete `tenant`'s uploaded file in one guarded statement (no + /// check-then-delete window); whether it existed under that tenant. + async fn file_delete(&self, id: &str, tenant: &str) -> GResult; async fn batch_create( &self, @@ -643,7 +651,7 @@ impl Store for MemoryStore { } async fn usage_rollup_advance(&self, now: i64) -> GResult { - let hi = bucket_floor(now); + let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut rollup = self .rollup .lock() @@ -807,8 +815,11 @@ impl Store for MemoryStore { Ok(self.files.get(id).map(|f| f.value().clone())) } - async fn file_delete(&self, id: &str) -> GResult { - Ok(self.files.remove(id).is_some()) + async fn file_delete(&self, id: &str, tenant: &str) -> GResult { + Ok(self + .files + .remove_if(id, |_, f| f.tenant == tenant) + .is_some()) } async fn batch_create( @@ -931,6 +942,25 @@ where } } +/// A terminal batch's input rows have served their purpose — delete them in +/// the same transaction as the status write, so submitted prompt text cannot +/// outlive the run even across a crash between statements. +async fn prune_terminal_items( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + id: &str, + status: BatchStatus, +) -> GResult<()> { + if !matches!(status, BatchStatus::Completed | BatchStatus::Failed) { + return Ok(()); + } + sqlx::query("DELETE FROM batch_items WHERE batch_id = $1") + .bind(id) + .execute(&mut **tx) + .await + .map_err(|e| crate::sqlx_err("prune batch items", e))?; + Ok(()) +} + /// SQLite-backed store (WAL): ledger, files, and batch jobs in one database /// file; ids derive from rowids so they stay unique across restarts. #[derive(Debug)] @@ -1190,7 +1220,7 @@ impl Store for SqliteStore { } async fn usage_rollup_advance(&self, now: i64) -> GResult { - let hi = bucket_floor(now); + let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let watermark: i64 = sqlx::query_scalar("SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup") .fetch_one(&self.pool) @@ -1384,10 +1414,13 @@ impl Store for SqliteStore { async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult { let bytes = content.len(); - // SQLite serializes writers, so the MAX(n)+1 subselect is atomic with the insert + // ids derive from the AUTOINCREMENT sequence (sqlite_sequence), which a + // delete never rewinds — MAX(n)+1 would recycle a deleted file's id let id: String = sqlx::query_scalar( "INSERT INTO files (id, tenant, purpose, bytes, content) - VALUES ('file-' || (SELECT COALESCE(MAX(n), 0) + 1 FROM files), ?, ?, ?, ?) + VALUES ('file-' || (SELECT COALESCE( + (SELECT seq FROM sqlite_sequence WHERE name = 'files'), 0) + 1), + ?, ?, ?, ?) RETURNING id", ) .bind(tenant) @@ -1421,9 +1454,10 @@ impl Store for SqliteStore { })) } - async fn file_delete(&self, id: &str) -> GResult { - let r = sqlx::query("DELETE FROM files WHERE id = ?") + async fn file_delete(&self, id: &str, tenant: &str) -> GResult { + let r = sqlx::query("DELETE FROM files WHERE id = ? AND tenant = ?") .bind(id) + .bind(tenant) .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("delete file", e))?; @@ -1536,20 +1570,6 @@ impl PostgresStore { Self::connect_with_cap(url, 0).await } - /// A terminal batch's input rows have served their purpose — delete them so - /// submitted prompt text does not outlive the run. - async fn prune_terminal_items(&self, id: &str, status: BatchStatus) -> GResult<()> { - if !matches!(status, BatchStatus::Completed | BatchStatus::Failed) { - return Ok(()); - } - sqlx::query("DELETE FROM batch_items WHERE batch_id = $1") - .bind(id) - .execute(&self.pool) - .await - .map_err(|e| crate::sqlx_err("prune batch items", e))?; - Ok(()) - } - /// `ledger_max_rows` > 0 prunes the oldest billing rows past the cap on write. pub async fn connect_with_cap(url: &str, ledger_max_rows: u64) -> GResult { let pool = sqlx::postgres::PgPoolOptions::new() @@ -1804,7 +1824,7 @@ impl Store for PostgresStore { } async fn usage_rollup_advance(&self, now: i64) -> GResult { - let hi = bucket_floor(now); + let hi = bucket_floor(now - ROLLUP_SETTLE_SECS); let mut tx = self .pool .begin() @@ -1985,12 +2005,12 @@ impl Store for PostgresStore { .execute(&mut *tx) .await .map_err(|e| crate::sqlx_err("erase batch results", e))?; - // leftovers from batches that went terminal before item pruning shipped + // pending/running batches included: the emptied item fails at execution + // instead of running erased content (terminal batches already pruned) let i = sqlx::query( - "DELETE FROM batch_items i USING batches b - WHERE i.batch_id = b.id AND i.user_id = $1 - AND ($2::text IS NULL OR b.tenant = $2) - AND b.status IN ('completed', 'failed')", + "UPDATE batch_items i SET messages = '[]' FROM batches b + WHERE i.batch_id = b.id AND i.user_id = $1 AND i.messages <> '[]' + AND ($2::text IS NULL OR b.tenant = $2)", ) .bind(user) .bind(tenant) @@ -2071,9 +2091,10 @@ impl Store for PostgresStore { })) } - async fn file_delete(&self, id: &str) -> GResult { - let r = sqlx::query("DELETE FROM files WHERE id = $1") + async fn file_delete(&self, id: &str, tenant: &str) -> GResult { + let r = sqlx::query("DELETE FROM files WHERE id = $1 AND tenant = $2") .bind(id) + .bind(tenant) .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("delete file", e))?; @@ -2136,13 +2157,21 @@ impl Store for PostgresStore { } async fn batch_set_status(&self, id: &str, status: BatchStatus) -> GResult<()> { + let mut tx = self + .pool + .begin() + .await + .map_err(|e| crate::sqlx_err("begin status", e))?; sqlx::query("UPDATE batches SET status = $1 WHERE id = $2") .bind(status.as_str()) .bind(id) - .execute(&self.pool) + .execute(&mut *tx) .await .map_err(|e| crate::sqlx_err("update batch status", e))?; - self.prune_terminal_items(id, status).await + prune_terminal_items(&mut tx, id, status).await?; + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit status", e)) } async fn batch_set_status_owned( @@ -2151,17 +2180,25 @@ impl Store for PostgresStore { status: BatchStatus, claim: i64, ) -> GResult { + let mut tx = self + .pool + .begin() + .await + .map_err(|e| crate::sqlx_err("begin status", e))?; let r = sqlx::query("UPDATE batches SET status = $1 WHERE id = $2 AND claim_seq = $3") .bind(status.as_str()) .bind(id) .bind(claim) - .execute(&self.pool) + .execute(&mut *tx) .await .map_err(|e| crate::sqlx_err("update batch status (fenced)", e))?; let applied = r.rows_affected() > 0; if applied { - self.prune_terminal_items(id, status).await?; + prune_terminal_items(&mut tx, id, status).await?; } + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit status", e))?; Ok(applied) } @@ -2815,6 +2852,48 @@ mod tests { .expect("truncate"); exercise_rollup(&store).await; exercise_erase(&store).await; + + // erasure reaches a still-pending batch's inputs (PG-only store) + let items = vec![ + gw_models::BatchItem { + messages: vec![gw_models::ChatMsg::text("user", "erase me")], + user: "u1".into(), + }, + gw_models::BatchItem { + messages: vec![gw_models::ChatMsg::text("user", "keep me")], + user: "u2".into(), + }, + ]; + let pending = store + .batch_enqueue("ak", "t1", "gpt-4o", &items) + .await + .unwrap(); + let audit2 = AdminAudit { + created_at_epoch_secs: 21, + actor: "global".into(), + scope: "global".into(), + action: "content_erase".into(), + target: "u1".into(), + summary: String::new(), + source_ip: "10.0.0.1".into(), + }; + assert_eq!( + store + .content_erase_user(Some("t1"), "u1", audit2) + .await + .unwrap(), + 1, + "the pending item's prompt is blanked" + ); + let loaded = store.batch_load_items(&pending.id).await.unwrap(); + assert!( + loaded[0].messages.is_empty(), + "u1's queued prompt erased before any drainer ran" + ); + assert_eq!( + loaded[1].messages[0].content, "keep me", + "other users' queued items untouched" + ); } #[tokio::test] @@ -2826,6 +2905,26 @@ mod tests { exercise_erase(&store).await; } + #[tokio::test] + async fn sqlite_file_ids_never_recycle_after_delete() { + let dir = tempfile::tempdir().unwrap(); + let store = SqliteStore::open(dir.path().join("files.db").to_str().unwrap()) + .await + .unwrap(); + let a = store.file_put("t1", "batch", "one".into()).await.unwrap(); + assert!(store.file_delete(&a.id, "t1").await.unwrap()); + let b = store.file_put("t1", "batch", "two".into()).await.unwrap(); + assert_ne!( + a.id, b.id, + "a deleted file's id must not be handed to a new upload" + ); + assert!( + !store.file_delete(&b.id, "t2").await.unwrap(), + "a foreign tenant's delete is a guarded no-op" + ); + assert!(store.file_get(&b.id).await.unwrap().is_some()); + } + #[tokio::test] async fn batch_result_is_first_writer_wins() { let store = MemoryStore::default(); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 4d138ec..c07a0e5 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -3111,12 +3111,11 @@ async fn files_delete( Ok(ak) => ak, Err((st, msg)) => return error_response(st, msg), }; - let found = s.handler.state().store.file_get(&id).await; - if let Err(resp) = tenant_owned(found, |f| &f.tenant, &ak.tenant, "file", &id) { - return resp; - } - match s.handler.state().store.file_delete(&id).await { - Ok(_) => Json(json!({"id": id, "object": "file", "deleted": true})).into_response(), + // one guarded delete — a check-then-delete pair would race a concurrent + // delete + id reuse into removing another tenant's file + match s.handler.state().store.file_delete(&id, &ak.tenant).await { + Ok(true) => Json(json!({"id": id, "object": "file", "deleted": true})).into_response(), + Ok(false) => error_response(404, format!("file {id} not found")), Err(e) => gateway_error(e), } } diff --git a/docs/governance.md b/docs/governance.md index c01421b..9d2d5d7 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -164,13 +164,17 @@ only its audit is aggregated (a store write per token would be too hot). unsealed when the key is present, else returned as `content: null`). `DELETE /admin/audit/content?user=` erases every retained trace of one end user's content (the GDPR/PIPL right-to-erasure hook): retained rows, batch - result messages, and any leftover terminal batch inputs — tenant-scoped, and - the `content_erase` audit entry commits with the deletion so a recorded - success can't separate from it. Erasure is a point-in-time operation: a - request already in flight may persist content just after it, so quiesce the - user's traffic first or simply repeat the call. Ledger rows and security - events carry no content and are kept (billing/audit legal basis). Uploaded - files are tenant-owned assets with no per-user dimension — the tenant deletes - them via `DELETE /v1/files/{id}`. A batch's input rows are deleted as soon as - the batch reaches a terminal status, so submitted prompt text does not - outlive the run. + result messages, and queued batch inputs — a pending/running item belonging + to the user is blanked in place and fails at execution instead of running + erased content. Tenant-scoped, and the `content_erase` audit entry commits + with the deletion so a recorded success can't separate from it. Erasure + reaches only attributed content: a batch item submitted without any + `user`/`x-gw-user` has no erasure subject (shared-key anonymous submission + opts out of per-user erasure by construction). Erasure is a point-in-time + operation: a request already in flight may persist content just after it, so + quiesce the user's traffic first or simply repeat the call. Ledger rows and + security events carry no content and are kept (billing/audit legal basis). + Uploaded files are tenant-owned assets with no per-user dimension — the + tenant deletes them via `DELETE /v1/files/{id}`. A batch's input rows are + deleted as soon as the batch reaches a terminal status, so submitted prompt + text does not outlive the run. From 8f51e6f2df09635144f4a8473b97b2385898de06 Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 16 Jul 2026 15:08:53 +0800 Subject: [PATCH 6/9] fix(review): pre-dispatch item snapshot and negative-price rejection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A batch executor re-reads each item's stored copy immediately before dispatch (backends that persist items), so an erasure landing while the batch sat queued stops the item instead of running the pre-load snapshot. Config validation rejects negative prices — cost accounting (and the rollup's max-upsert) relies on per-column monotone sums. --- crates/config/src/lib.rs | 40 +++++++++++++++++++++++++++++++++++ crates/handler/src/offline.rs | 7 +++++- crates/state/src/store.rs | 39 ++++++++++++++++++++++++++++++++++ 3 files changed, 85 insertions(+), 1 deletion(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index d6a7b5e..df4d94a 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -46,6 +46,8 @@ pub enum ConfigError { UnknownQuotaModel { owner: String, model: String }, #[error("tenant `{tenant}` fallback model `{model}` is unknown or not entitled")] BadFallbackModel { tenant: String, model: String }, + #[error("`{owner}` sets a negative price")] + NegativePrice { owner: String }, #[error("storage.shared_cache needs storage.redis_url")] SharedCacheNeedsRedis, } @@ -592,6 +594,35 @@ impl GatewayConfig { if self.storage.shared_cache && self.storage.redis_url.is_empty() { return Err(ConfigError::SharedCacheNeedsRedis); } + // negative prices would make cost accounting non-monotonic (the usage + // rollup's max-upsert relies on per-column monotone sums) + let neg = |i: i64, o: i64| i < 0 || o < 0; + for m in &self.models { + if neg(m.input_price_per_1k_micros, m.output_price_per_1k_micros) { + return Err(ConfigError::NegativePrice { + owner: format!("model {}", m.name), + }); + } + } + for a in &self.accounts { + if neg( + a.cost_input_price_per_1k_micros, + a.cost_output_price_per_1k_micros, + ) { + return Err(ConfigError::NegativePrice { + owner: format!("account {}", a.name), + }); + } + } + for t in &self.tenants { + for (model, p) in &t.model_prices { + if neg(p.input_price_per_1k_micros, p.output_price_per_1k_micros) { + return Err(ConfigError::NegativePrice { + owner: format!("tenant {} price for {model}", t.name), + }); + } + } + } for m in &self.models { if m.protocol().is_none() { return Err(ConfigError::UnknownModelMapping { @@ -1080,6 +1111,15 @@ tenants: [{name: t1}, {name: t1}] Err(ConfigError::DuplicateName { kind: "tenant", .. }) )); + let neg_price = "listen: {host: h, port: 1}\nmodels: [{name: m1, protocol: openai-chat, input_price_per_1k_micros: -1}]"; + assert!( + matches!( + GatewayConfig::from_yaml(neg_price), + Err(ConfigError::NegativePrice { .. }) + ), + "negative prices are rejected at load" + ); + let colon = "listen: {host: h, port: 1}\ntenants: [{name: 'a:b'}]"; assert!( GatewayConfig::from_yaml(colon).is_err(), diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 2ce661d..0035f80 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -96,7 +96,7 @@ impl OfflineHandler { } }) }; - for (index, item) in items.into_iter().enumerate() { + for (index, mut item) in items.into_iter().enumerate() { if lost.load(Relaxed) { break; // reclaimed by another instance; stop running new items } @@ -111,6 +111,11 @@ impl OfflineHandler { lost.store(true, Relaxed); break; } + // re-read the stored copy just before dispatch: an erasure that + // landed while this batch sat queued blanks the persisted item + if let Ok(Some(fresh)) = store.batch_item_snapshot(id, index).await { + item = fresh; + } let user = ak.attributed_user(&item.user).to_owned(); // an erased item (messages blanked by content_erase_user) must not // run: fail it instead of sending an empty prompt upstream diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 3249bc5..bf5382d 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -507,6 +507,17 @@ pub trait Store: Send + Sync + std::fmt::Debug { async fn batch_load_items(&self, _id: &str) -> GResult> { Ok(Vec::new()) } + /// The current stored copy of one queued item, on backends that persist + /// items (`None` otherwise). Executors re-read this immediately before + /// dispatch, so an erasure landing while the batch sat queued stops the + /// item instead of letting the pre-load snapshot run. + async fn batch_item_snapshot( + &self, + _id: &str, + _idx: usize, + ) -> GResult> { + Ok(None) + } /// Claim one pending batch (requeuing stale running ones first); `None` = /// nothing to run. The returned fence token (>= 1, bumped per claim) rides /// [`Store::batch_touch`] / [`Store::batch_set_status_owned`] so a reclaimed @@ -2330,6 +2341,25 @@ impl Store for PostgresStore { }) } + async fn batch_item_snapshot( + &self, + id: &str, + idx: usize, + ) -> GResult> { + let row = sqlx::query( + "SELECT messages, user_id FROM batch_items WHERE batch_id = $1 AND idx = $2", + ) + .bind(id) + .bind(idx as i64) + .fetch_optional(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read batch item", e))?; + Ok(row.map(|r| gw_models::BatchItem { + messages: serde_json::from_str(r.get::<&str, _>(0)).unwrap_or_default(), + user: r.get(1), + })) + } + async fn batch_load_items(&self, id: &str) -> GResult> { let rows = sqlx::query( "SELECT messages, user_id FROM batch_items WHERE batch_id = $1 ORDER BY idx", @@ -2890,6 +2920,15 @@ mod tests { loaded[0].messages.is_empty(), "u1's queued prompt erased before any drainer ran" ); + let fresh = store + .batch_item_snapshot(&pending.id, 0) + .await + .unwrap() + .expect("item row still present"); + assert!( + fresh.messages.is_empty(), + "the pre-dispatch snapshot sees the erasure" + ); assert_eq!( loaded[1].messages[0].content, "keep me", "other users' queued items untouched" From ecbf8fff88d3e156e5a8d1963c78a9552fd2c327 Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 16 Jul 2026 15:17:41 +0800 Subject: [PATCH 7/9] fix(review): fail-closed dispatch snapshot; erasure markers stop local mid-batch items MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A distributed executor now fails an item closed when the pre-dispatch re-read errors or finds no row — a stale pre-load copy never dispatches. Local backends (which don't persist items) record an erasure marker in the erase transaction; the executor checks it per item against the batch's start time, so an erasure landing mid-batch stops the user's remaining items while batches submitted after the erasure run normally. --- crates/handler/src/lib.rs | 41 ++++++++++++++++++ crates/handler/src/offline.rs | 37 ++++++++++++++--- crates/state/src/store.rs | 78 ++++++++++++++++++++++++++++++++++- 3 files changed, 149 insertions(+), 7 deletions(-) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index e15a58a..30e808c 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -980,6 +980,47 @@ mod tests { ); } + #[tokio::test] + async fn batch_submitted_after_an_erasure_still_runs() { + let h = handler(); + let off = OfflineHandler::new(h.clone()); + let key = ak(&h).await; + let audit = gw_state::AdminAudit { + created_at_epoch_secs: 1, + actor: "global".into(), + scope: "global".into(), + action: "content_erase".into(), + target: "user-42".into(), + summary: String::new(), + source_ip: String::new(), + }; + h.state() + .store + .content_erase_user(None, "user-42", audit) + .await + .unwrap(); + // cross the second boundary so the batch's start postdates the marker + tokio::time::sleep(std::time::Duration::from_millis(1_100)).await; + let job = off + .submit( + key, + "gpt-4o".into(), + vec![BatchItem { + messages: vec![ChatMsg::text("user", "new content after erasure")], + user: "user-42".into(), + }], + ) + .await + .unwrap(); + wait_terminal(&h, &job.id).await; + let done = h.state().store.batch_get(&job.id).await.unwrap().unwrap(); + assert!( + done.results[0].ok, + "a past erasure must not fail the user's future batches: {}", + done.results[0].message + ); + } + #[tokio::test] async fn erased_batch_item_fails_instead_of_running() { let h = handler(); diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 0035f80..4487770 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -51,6 +51,7 @@ impl OfflineHandler { /// reclaimed, this executor stops rather than double-running items. async fn execute(&self, id: &str, ak: &AkInfo, model: &str, items: Vec, claim: i64) { let store = self.online.state().store.clone(); + let started = gw_state::epoch_secs(); // the distributed claim already set status=running with the fence bump; // only the in-process path needs this write — unfenced on the distributed // path it could resurrect a batch a stale worker no longer owns @@ -112,14 +113,38 @@ impl OfflineHandler { break; } // re-read the stored copy just before dispatch: an erasure that - // landed while this batch sat queued blanks the persisted item - if let Ok(Some(fresh)) = store.batch_item_snapshot(id, index).await { - item = fresh; + // landed while this batch sat queued blanks the persisted item. + // Fail CLOSED — a read error or vanished row can't prove the item + // wasn't erased, so the stale pre-load copy never dispatches + if store.distributed_batches() { + match store.batch_item_snapshot(id, index).await { + Ok(Some(fresh)) => item = fresh, + Ok(None) | Err(_) => { + let result = BatchItemResult { + index, + ok: false, + message: "item unavailable at dispatch".into(), + total_tokens: 0, + user: ak.attributed_user(&item.user).to_owned(), + }; + if let Err(e) = store.batch_push_result(id, result).await { + tracing::error!(error = %e, batch = %id, "batch result write failed"); + } + continue; + } + } } let user = ak.attributed_user(&item.user).to_owned(); - // an erased item (messages blanked by content_erase_user) must not - // run: fail it instead of sending an empty prompt upstream - if item.messages.is_empty() { + // local backends don't persist items, so a mid-batch erasure can't + // blank them — the erasure marker stops the user's remaining items + // (fail closed on a marker read error). An erased item must not + // run: fail it instead of sending an erased prompt upstream + let erased_mid_batch = !store.distributed_batches() + && store + .user_erased_since(&ak.tenant, &user, started) + .await + .unwrap_or(true); + if erased_mid_batch || item.messages.is_empty() { let result = BatchItemResult { index, ok: false, diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index bf5382d..ff8728d 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -518,6 +518,15 @@ pub trait Store: Send + Sync + std::fmt::Debug { ) -> GResult> { Ok(None) } + /// Whether `user`'s content was erased at or after `since` (unix secs). + /// Local executors check this before dispatching an item captured before + /// the erasure, so a long sequential batch can't keep running an erased + /// user's prompts; a batch submitted after the erasure passes (its start + /// postdates `since`). Item-persisting backends re-read rows at dispatch + /// instead and keep the default `false`. + async fn user_erased_since(&self, _tenant: &str, _user: &str, _since: i64) -> GResult { + Ok(false) + } /// Claim one pending batch (requeuing stale running ones first); `None` = /// nothing to run. The returned fence token (>= 1, bumped per claim) rides /// [`Store::batch_touch`] / [`Store::batch_set_status_owned`] so a reclaimed @@ -542,6 +551,8 @@ pub struct MemoryStore { sec_events: Mutex>, audit: Mutex>, content: Mutex>, + /// (tenant, user, erased_at) markers backing [`Store::user_erased_since`]. + erasures: Mutex>, files: DashMap, jobs: DashMap, seq: AtomicUsize, @@ -786,6 +797,14 @@ impl Store for MemoryStore { } } } + self.erasures + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(( + tenant.unwrap_or_default().to_owned(), + user.to_owned(), + crate::epoch_secs(), + )); audit.summary = format!("rows={erased}"); self.audit .lock() @@ -794,6 +813,15 @@ impl Store for MemoryStore { Ok(erased) } + async fn user_erased_since(&self, tenant: &str, user: &str, since: i64) -> GResult { + Ok(self + .erasures + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .any(|(t, u, at)| u == user && *at >= since && (t.is_empty() || t == tenant))) + } + async fn content_for(&self, request_id: &str) -> GResult> { let content = self .content @@ -1012,6 +1040,9 @@ impl SqliteStore { created_at_epoch_secs INTEGER NOT NULL DEFAULT 0, estimated INTEGER NOT NULL DEFAULT 0)", "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", + "CREATE TABLE IF NOT EXISTS erasures ( + tenant TEXT NOT NULL, user_id TEXT NOT NULL, + erased_at_epoch_secs INTEGER NOT NULL)", "CREATE TABLE IF NOT EXISTS usage_rollup ( minute_epoch INTEGER NOT NULL, tenant TEXT NOT NULL, user_id TEXT NOT NULL, model TEXT NOT NULL, requests INTEGER NOT NULL, @@ -1391,6 +1422,15 @@ impl Store for SqliteStore { .await .map_err(|e| crate::sqlx_err("erase batch results", e))?; let erased = c.rows_affected() + m.rows_affected(); + sqlx::query( + "INSERT INTO erasures (tenant, user_id, erased_at_epoch_secs) VALUES (?, ?, ?)", + ) + .bind(tenant.unwrap_or_default()) + .bind(user) + .bind(crate::epoch_secs()) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("record erasure", e))?; sqlx::query( "INSERT INTO admin_audit (created_at_epoch_secs, actor, scope, action, target, summary, source_ip) VALUES (?, ?, ?, ?, ?, ?, ?)", @@ -1465,6 +1505,21 @@ impl Store for SqliteStore { })) } + async fn user_erased_since(&self, tenant: &str, user: &str, since: i64) -> GResult { + let hit: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM erasures + WHERE user_id = ?1 AND erased_at_epoch_secs >= ?2 + AND (tenant = '' OR tenant = ?3))", + ) + .bind(user) + .bind(since) + .bind(tenant) + .fetch_one(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read erasure marker", e))?; + Ok(hit) + } + async fn file_delete(&self, id: &str, tenant: &str) -> GResult { let r = sqlx::query("DELETE FROM files WHERE id = ? AND tenant = ?") .bind(id) @@ -2862,9 +2917,29 @@ mod tests { ); } + /// Local backends only: item-persisting stores (PG) re-read rows at + /// dispatch and keep the trait's `false` default. + async fn exercise_erasure_markers(store: &dyn Store) { + let now = crate::epoch_secs(); + assert!( + store.user_erased_since("t1", "u1", now - 5).await.unwrap(), + "marker visible to a batch started before the erasure" + ); + assert!( + !store.user_erased_since("t1", "u1", now + 5).await.unwrap(), + "a batch started after the erasure is unaffected" + ); + assert!( + !store.user_erased_since("t1", "u2", now - 5).await.unwrap(), + "other users unaffected" + ); + } + #[tokio::test] async fn memory_erase_roundtrip() { - exercise_erase(&MemoryStore::default()).await; + let store = MemoryStore::default(); + exercise_erase(&store).await; + exercise_erasure_markers(&store).await; } #[tokio::test] @@ -2942,6 +3017,7 @@ mod tests { .await .unwrap(); exercise_erase(&store).await; + exercise_erasure_markers(&store).await; } #[tokio::test] From b92d5ad1a3c72e5e7ac1a4ad95126d488ef4c68f Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 16 Jul 2026 15:27:27 +0800 Subject: [PATCH 8/9] fix(review): erasure baseline at item capture; bounded erasure markers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The marker comparison point is now the submit instant (when items are captured), not the spawned executor's first poll, closing the pre-poll erasure window. Markers keep only the latest instant per (tenant, user) — an upserted unique row on SQLite, a map in memory — so repeated erasures can't grow unbounded and the per-item check is a keyed lookup. --- crates/handler/src/offline.rs | 31 ++++++++++++++++++++++++++----- crates/state/src/store.rs | 31 +++++++++++++++++++------------ 2 files changed, 45 insertions(+), 17 deletions(-) diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 4487770..b05d43e 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -39,8 +39,14 @@ impl OfflineHandler { .await?; let this = self.clone(); let (id, model) = (job.id.clone(), model.clone()); + // items are captured HERE: an erasure landing after this instant + // must stop them, so the marker comparison point is submission, + // not the spawned executor's first poll + let captured_at = gw_state::epoch_secs(); // claim 0: non-distributed store — no fence, the heartbeat is a no-op - tokio::spawn(async move { this.execute(&id, &ak, &model, items, 0).await }); + tokio::spawn( + async move { this.execute(&id, &ak, &model, items, 0, captured_at).await }, + ); Ok(job) } } @@ -49,9 +55,16 @@ impl OfflineHandler { /// the terminal status, heartbeating between items. `claim` is the fence /// token (0 for the in-process path); once a heartbeat reports the batch /// reclaimed, this executor stops rather than double-running items. - async fn execute(&self, id: &str, ak: &AkInfo, model: &str, items: Vec, claim: i64) { + async fn execute( + &self, + id: &str, + ak: &AkInfo, + model: &str, + items: Vec, + claim: i64, + captured_at: i64, + ) { let store = self.online.state().store.clone(); - let started = gw_state::epoch_secs(); // the distributed claim already set status=running with the fence bump; // only the in-process path needs this write — unfenced on the distributed // path it could resurrect a batch a stale worker no longer owns @@ -141,7 +154,7 @@ impl OfflineHandler { // run: fail it instead of sending an erased prompt upstream let erased_mid_batch = !store.distributed_batches() && store - .user_erased_since(&ak.tenant, &user, started) + .user_erased_since(&ak.tenant, &user, captured_at) .await .unwrap_or(true); if erased_mid_batch || item.messages.is_empty() { @@ -247,7 +260,15 @@ impl OfflineHandler { continue; } }; - self.execute(&job.id, &ak, &job.model, items, claim).await; + self.execute( + &job.id, + &ak, + &job.model, + items, + claim, + gw_state::epoch_secs(), + ) + .await; } Ok(None) => tokio::time::sleep(poll).await, Err(e) => { diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index ff8728d..8bcb125 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -551,8 +551,9 @@ pub struct MemoryStore { sec_events: Mutex>, audit: Mutex>, content: Mutex>, - /// (tenant, user, erased_at) markers backing [`Store::user_erased_since`]. - erasures: Mutex>, + /// Latest erasure instant per (tenant, user), backing + /// [`Store::user_erased_since`] — one entry per pair, not a log. + erasures: Mutex>, files: DashMap, jobs: DashMap, seq: AtomicUsize, @@ -800,11 +801,10 @@ impl Store for MemoryStore { self.erasures .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .push(( - tenant.unwrap_or_default().to_owned(), - user.to_owned(), + .insert( + (tenant.unwrap_or_default().to_owned(), user.to_owned()), crate::epoch_secs(), - )); + ); audit.summary = format!("rows={erased}"); self.audit .lock() @@ -814,12 +814,16 @@ impl Store for MemoryStore { } async fn user_erased_since(&self, tenant: &str, user: &str, since: i64) -> GResult { - Ok(self + let erasures = self .erasures .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .iter() - .any(|(t, u, at)| u == user && *at >= since && (t.is_empty() || t == tenant))) + .unwrap_or_else(std::sync::PoisonError::into_inner); + let hit = |t: &str| { + erasures + .get(&(t.to_owned(), user.to_owned())) + .is_some_and(|at| *at >= since) + }; + Ok(hit("") || hit(tenant)) } async fn content_for(&self, request_id: &str) -> GResult> { @@ -1042,7 +1046,8 @@ impl SqliteStore { "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", "CREATE TABLE IF NOT EXISTS erasures ( tenant TEXT NOT NULL, user_id TEXT NOT NULL, - erased_at_epoch_secs INTEGER NOT NULL)", + erased_at_epoch_secs INTEGER NOT NULL, + PRIMARY KEY (tenant, user_id))", "CREATE TABLE IF NOT EXISTS usage_rollup ( minute_epoch INTEGER NOT NULL, tenant TEXT NOT NULL, user_id TEXT NOT NULL, model TEXT NOT NULL, requests INTEGER NOT NULL, @@ -1423,7 +1428,9 @@ impl Store for SqliteStore { .map_err(|e| crate::sqlx_err("erase batch results", e))?; let erased = c.rows_affected() + m.rows_affected(); sqlx::query( - "INSERT INTO erasures (tenant, user_id, erased_at_epoch_secs) VALUES (?, ?, ?)", + "INSERT INTO erasures (tenant, user_id, erased_at_epoch_secs) VALUES (?, ?, ?) + ON CONFLICT (tenant, user_id) DO UPDATE SET + erased_at_epoch_secs = excluded.erased_at_epoch_secs", ) .bind(tenant.unwrap_or_default()) .bind(user) From fb86246acbe3717d32b206acc87eba4d0e1f9eab Mon Sep 17 00:00:00 2001 From: CMGS Date: Thu, 16 Jul 2026 16:21:13 +0800 Subject: [PATCH 9/9] =?UTF-8?q?fix(review):=20close=20post-PR=20findings?= =?UTF-8?q?=20=E2=80=94=20prune-vs-rollup=20loss,=20erasure=20identity,=20?= =?UTF-8?q?CI=20isolation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ledger cap now spares rows the rollup hasn't folded (created at or past the watermark) on all backends, so a burst can briefly exceed ledger_max_rows instead of losing usage forever. Distributed enqueue persists the EFFECTIVE user (owner overrides the hint), so execution, billing, and erasure key on one identity; the Postgres migration backfills pre-upgrade batch_results ownership from their surviving item rows. Erasure markers and the submit baseline moved to millisecond precision, so erase-then-resubmit within one second is no longer misjudged. The PG store test runs in a private database (the rollup watermark is global per database and the shared one serves parallel tests) instead of truncating shared tables — the CI parallel-run failure. --- crates/handler/src/lib.rs | 4 +- crates/handler/src/offline.rs | 13 +- crates/state/src/lib.rs | 9 ++ crates/state/src/store.rs | 259 +++++++++++++++++++++++----------- 4 files changed, 197 insertions(+), 88 deletions(-) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 30e808c..0f39519 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -999,8 +999,8 @@ mod tests { .content_erase_user(None, "user-42", audit) .await .unwrap(); - // cross the second boundary so the batch's start postdates the marker - tokio::time::sleep(std::time::Duration::from_millis(1_100)).await; + // millisecond markers: a resubmit moments later must already pass + tokio::time::sleep(std::time::Duration::from_millis(5)).await; let job = off .submit( key, diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index b05d43e..9b39cfc 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -29,6 +29,15 @@ impl OfflineHandler { ) -> gw_models::GResult { let store = self.online.state().store.clone(); if store.distributed_batches() { + // persist the EFFECTIVE user (owner overrides the hint): execution, + // billing, and erasure must all key on the same identity + let items: Vec = items + .into_iter() + .map(|mut i| { + i.user = ak.attributed_user(&i.user).to_owned(); + i + }) + .collect(); // atomic: the job becomes claimable only once all items are saved store .batch_enqueue(&ak.ak, &ak.tenant, &model, &items) @@ -42,7 +51,7 @@ impl OfflineHandler { // items are captured HERE: an erasure landing after this instant // must stop them, so the marker comparison point is submission, // not the spawned executor's first poll - let captured_at = gw_state::epoch_secs(); + let captured_at = gw_state::epoch_millis(); // claim 0: non-distributed store — no fence, the heartbeat is a no-op tokio::spawn( async move { this.execute(&id, &ak, &model, items, 0, captured_at).await }, @@ -266,7 +275,7 @@ impl OfflineHandler { &job.model, items, claim, - gw_state::epoch_secs(), + gw_state::epoch_millis(), ) .await; } diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index b736a35..0c78019 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -863,6 +863,15 @@ pub fn epoch_secs() -> i64 { .unwrap_or(0) } +/// Unix milliseconds; erasure markers use this so an erase-then-resubmit in +/// the same second isn't misjudged as pre-erasure content. +pub fn epoch_millis() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 8bcb125..a899860 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -25,6 +25,9 @@ pub const MAX_METERED_TOKENS: i64 = 1_000_000_000; /// Prune the SQL ledger every Nth insert instead of per write (the cap becomes /// approximate by at most this many rows, saving a round-trip per billing). +/// Pruning spares rows the rollup hasn't folded yet (created at or past the +/// watermark), so a burst can exceed `ledger_max_rows` briefly rather than +/// lose usage — billing integrity outranks a strict cap. const LEDGER_PRUNE_EVERY: usize = 64; /// Usage-rollup bucket width. @@ -518,12 +521,13 @@ pub trait Store: Send + Sync + std::fmt::Debug { ) -> GResult> { Ok(None) } - /// Whether `user`'s content was erased at or after `since` (unix secs). - /// Local executors check this before dispatching an item captured before - /// the erasure, so a long sequential batch can't keep running an erased - /// user's prompts; a batch submitted after the erasure passes (its start - /// postdates `since`). Item-persisting backends re-read rows at dispatch - /// instead and keep the default `false`. + /// Whether `user`'s content was erased at or after `since` (unix MILLIS — + /// second granularity would misjudge an erase-then-resubmit in the same + /// second). Local executors check this before dispatching an item captured + /// before the erasure, so a long sequential batch can't keep running an + /// erased user's prompts; a batch submitted after the erasure passes. + /// Item-persisting backends re-read rows at dispatch instead and keep the + /// default `false`. async fn user_erased_since(&self, _tenant: &str, _user: &str, _since: i64) -> GResult { Ok(false) } @@ -573,6 +577,14 @@ impl MemoryStore { #[async_trait::async_trait] impl Store for MemoryStore { async fn ledger_add(&self, r: &BillingRecord) -> GResult<()> { + // watermark first: rollup-then-records is the lock order advance uses + let watermark = self + .rollup + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .keys() + .next_back() + .map_or(0, |k| k.0 + ROLLUP_BUCKET_SECS); let mut records = self .records .lock() @@ -580,7 +592,14 @@ impl Store for MemoryStore { records.push(r.clone()); if self.ledger_max_rows > 0 && records.len() > self.ledger_max_rows { let excess = records.len() - self.ledger_max_rows; - records.drain(..excess); + // spare rows the rollup hasn't folded yet — lost here means lost + // from usage forever; the cap yields to billing integrity + let cut = records + .iter() + .take(excess) + .take_while(|r| r.created_at_epoch_secs < watermark) + .count(); + records.drain(..cut); } Ok(()) } @@ -803,7 +822,7 @@ impl Store for MemoryStore { .unwrap_or_else(std::sync::PoisonError::into_inner) .insert( (tenant.unwrap_or_default().to_owned(), user.to_owned()), - crate::epoch_secs(), + crate::epoch_millis(), ); audit.summary = format!("rows={erased}"); self.audit @@ -1046,7 +1065,7 @@ impl SqliteStore { "CREATE INDEX IF NOT EXISTS billing_created_idx ON billing (created_at_epoch_secs)", "CREATE TABLE IF NOT EXISTS erasures ( tenant TEXT NOT NULL, user_id TEXT NOT NULL, - erased_at_epoch_secs INTEGER NOT NULL, + erased_at_epoch_millis INTEGER NOT NULL, PRIMARY KEY (tenant, user_id))", "CREATE TABLE IF NOT EXISTS usage_rollup ( minute_epoch INTEGER NOT NULL, tenant TEXT NOT NULL, user_id TEXT NOT NULL, @@ -1161,11 +1180,15 @@ impl Store for SqliteStore { .fetch_add(1, Ordering::Relaxed) .is_multiple_of(LEDGER_PRUNE_EVERY) { - sqlx::query("DELETE FROM billing WHERE n <= (SELECT MAX(n) FROM billing) - ?") - .bind(self.ledger_max_rows as i64) - .execute(&self.pool) - .await - .map_err(|e| crate::sqlx_err("prune billing records", e))?; + sqlx::query( + "DELETE FROM billing WHERE n <= (SELECT MAX(n) FROM billing) - ? + AND created_at_epoch_secs < + (SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup)", + ) + .bind(self.ledger_max_rows as i64) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("prune billing records", e))?; } Ok(()) } @@ -1428,13 +1451,13 @@ impl Store for SqliteStore { .map_err(|e| crate::sqlx_err("erase batch results", e))?; let erased = c.rows_affected() + m.rows_affected(); sqlx::query( - "INSERT INTO erasures (tenant, user_id, erased_at_epoch_secs) VALUES (?, ?, ?) + "INSERT INTO erasures (tenant, user_id, erased_at_epoch_millis) VALUES (?, ?, ?) ON CONFLICT (tenant, user_id) DO UPDATE SET - erased_at_epoch_secs = excluded.erased_at_epoch_secs", + erased_at_epoch_millis = excluded.erased_at_epoch_millis", ) .bind(tenant.unwrap_or_default()) .bind(user) - .bind(crate::epoch_secs()) + .bind(crate::epoch_millis()) .execute(&mut *tx) .await .map_err(|e| crate::sqlx_err("record erasure", e))?; @@ -1515,7 +1538,7 @@ impl Store for SqliteStore { async fn user_erased_since(&self, tenant: &str, user: &str, since: i64) -> GResult { let hit: bool = sqlx::query_scalar( "SELECT EXISTS (SELECT 1 FROM erasures - WHERE user_id = ?1 AND erased_at_epoch_secs >= ?2 + WHERE user_id = ?1 AND erased_at_epoch_millis >= ?2 AND (tenant = '' OR tenant = ?3))", ) .bind(user) @@ -1709,6 +1732,12 @@ impl PostgresStore { // per-item end-user attribution so a fleet drainer still bills/budgets it "ALTER TABLE batch_items ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''", "ALTER TABLE batch_results ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''", + // pre-upgrade results predate the user_id column; their items rows + // (pruned only at terminal status from this version on) still carry + // the owner — backfill so erasure reaches history. Idempotent. + "UPDATE batch_results r SET user_id = i.user_id FROM batch_items i + WHERE r.batch_id = i.batch_id AND r.idx = i.idx + AND r.user_id = '' AND i.user_id <> ''", "ALTER TABLE batches ADD COLUMN IF NOT EXISTS claimed_at TIMESTAMPTZ", // fence token: bumped on every claim so a reclaimed executor's fenced writes no-op "ALTER TABLE batches ADD COLUMN IF NOT EXISTS claim_seq BIGINT NOT NULL DEFAULT 0", @@ -1782,11 +1811,15 @@ impl Store for PostgresStore { .fetch_add(1, Ordering::Relaxed) .is_multiple_of(LEDGER_PRUNE_EVERY) { - sqlx::query("DELETE FROM billing WHERE n <= (SELECT MAX(n) FROM billing) - $1") - .bind(self.ledger_max_rows as i64) - .execute(&self.pool) - .await - .map_err(|e| crate::sqlx_err("prune billing records", e))?; + sqlx::query( + "DELETE FROM billing WHERE n <= (SELECT MAX(n) FROM billing) - $1 + AND created_at_epoch_secs < + (SELECT COALESCE(MAX(minute_epoch), -60) + 60 FROM usage_rollup)", + ) + .bind(self.ledger_max_rows as i64) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("prune billing records", e))?; } Ok(()) } @@ -2717,25 +2750,37 @@ mod tests { exercise_audit(&store).await; } - async fn exercise_rollup(store: &dyn Store) { + /// `ns` namespaces tenant/user ids so shared-database (PG) runs stay + /// isolated from parallel tests without truncating anything. + async fn exercise_rollup(store: &dyn Store, ns: &str) { + let tenant = format!("tr{ns}"); let mut a = record("m1"); + a.tenant = tenant.clone(); a.user_id = "alice".into(); a.created_at_epoch_secs = 400; let mut b = record("m1"); + b.tenant = tenant.clone(); b.user_id = "bob".into(); b.created_at_epoch_secs = 460; let mut tail = record("m1"); + tail.tenant = tenant.clone(); tail.user_id = "carol".into(); tail.created_at_epoch_secs = 1_520; for r in [&a, &b, &tail] { store.ledger_add(r).await.unwrap(); } - let baseline = store.usage_by_user(None, None, 0, i64::MAX).await.unwrap(); + let baseline = store + .usage_by_user(Some(&tenant), None, 0, i64::MAX) + .await + .unwrap(); assert_eq!(baseline.len(), 3, "raw-only result before any rollup"); store.usage_rollup_advance(1_500).await.unwrap(); store.usage_rollup_advance(1_500).await.unwrap(); - let after = store.usage_by_user(None, None, 0, i64::MAX).await.unwrap(); + let after = store + .usage_by_user(Some(&tenant), None, 0, i64::MAX) + .await + .unwrap(); assert_eq!( format!("{baseline:?}"), format!("{after:?}"), @@ -2743,7 +2788,7 @@ mod tests { ); let windowed = store - .usage_by_user(None, None, 450, i64::MAX) + .usage_by_user(Some(&tenant), None, 450, i64::MAX) .await .unwrap(); assert!( @@ -2759,7 +2804,7 @@ mod tests { #[tokio::test] async fn memory_rollup_roundtrip() { - exercise_rollup(&MemoryStore::default()).await; + exercise_rollup(&MemoryStore::default(), "").await; } #[tokio::test] @@ -2768,7 +2813,7 @@ mod tests { let store = SqliteStore::open(dir.path().join("rollup.db").to_str().unwrap()) .await .unwrap(); - exercise_rollup(&store).await; + exercise_rollup(&store, "").await; } #[tokio::test] @@ -2849,7 +2894,10 @@ mod tests { ); } - async fn exercise_erase(store: &dyn Store) { + async fn exercise_erase(store: &dyn Store, ns: &str) { + let (t1, t2) = (format!("t1{ns}"), format!("t2{ns}")); + let (u1, u2) = (format!("u1{ns}"), format!("u2{ns}")); + let (t1, t2, u1, u2) = (t1.as_str(), t2.as_str(), u1.as_str(), u2.as_str()); let rec = |req: &str, user: &str, tenant: &str| crate::ContentRecord { created_at_epoch_secs: 100, request_id: req.into(), @@ -2861,10 +2909,11 @@ mod tests { sealed: false, expires_at_epoch_secs: 0, }; - store.content_add(&rec("r1", "u1", "t1")).await.unwrap(); - store.content_add(&rec("r2", "u1", "t2")).await.unwrap(); - store.content_add(&rec("r3", "u2", "t1")).await.unwrap(); - let job = store.batch_create("ak", "t1", "m", 1).await.unwrap(); + let (r1, r2, r3) = (format!("r1{ns}"), format!("r2{ns}"), format!("r3{ns}")); + store.content_add(&rec(&r1, u1, t1)).await.unwrap(); + store.content_add(&rec(&r2, u1, t2)).await.unwrap(); + store.content_add(&rec(&r3, u2, t1)).await.unwrap(); + let job = store.batch_create("ak", t1, "m", 1).await.unwrap(); store .batch_push_result( &job.id, @@ -2873,7 +2922,7 @@ mod tests { ok: true, message: "generated for u1".into(), total_tokens: 3, - user: "u1".into(), + user: u1.into(), }, ) .await @@ -2884,36 +2933,36 @@ mod tests { actor: "global".into(), scope: "global".into(), action: "content_erase".into(), - target: "u1".into(), + target: u1.into(), summary: String::new(), source_ip: "10.0.0.1".into(), }; assert_eq!( store - .content_erase_user(Some("t1"), "u1", audit()) + .content_erase_user(Some(t1), u1, audit()) .await .unwrap(), 2, "tenant-scoped erase: the content row plus the batch-result message" ); - assert_eq!(store.content_for("r2").await.unwrap().len(), 1); + assert_eq!(store.content_for(&r2).await.unwrap().len(), 1); let results = store.batch_get(&job.id).await.unwrap().unwrap().results; assert_eq!(results[0].message, "", "generated output erased"); - assert_eq!(results[0].user, "u1", "attribution survives for billing"); + assert_eq!(results[0].user, u1, "attribution survives for billing"); assert_eq!( - store.content_erase_user(None, "u1", audit()).await.unwrap(), + store.content_erase_user(None, u1, audit()).await.unwrap(), 1, "global erase removes the remaining tenant's row" ); assert_eq!( - store.content_for("r3").await.unwrap().len(), + store.content_for(&r3).await.unwrap().len(), 1, "other users' content is untouched" ); - let trail = store.admin_audit_list(10).await.unwrap(); + let trail = store.admin_audit_list(1_000).await.unwrap(); let erases: Vec<_> = trail .iter() - .filter(|e| e.action == "content_erase" && e.target == "u1") + .filter(|e| e.action == "content_erase" && e.target == u1) .collect(); assert_eq!(erases.len(), 2, "each erasure committed its audit entry"); let summaries: Vec<&str> = erases.iter().map(|e| e.summary.as_str()).collect(); @@ -2926,18 +2975,19 @@ mod tests { /// Local backends only: item-persisting stores (PG) re-read rows at /// dispatch and keep the trait's `false` default. - async fn exercise_erasure_markers(store: &dyn Store) { - let now = crate::epoch_secs(); + async fn exercise_erasure_markers(store: &dyn Store, ns: &str) { + let (t1, u1, u2) = (format!("t1{ns}"), format!("u1{ns}"), format!("u2{ns}")); + let now = crate::epoch_millis(); assert!( - store.user_erased_since("t1", "u1", now - 5).await.unwrap(), + store.user_erased_since(&t1, &u1, now - 5).await.unwrap(), "marker visible to a batch started before the erasure" ); assert!( - !store.user_erased_since("t1", "u1", now + 5).await.unwrap(), + !store.user_erased_since(&t1, &u1, now + 5).await.unwrap(), "a batch started after the erasure is unaffected" ); assert!( - !store.user_erased_since("t1", "u2", now - 5).await.unwrap(), + !store.user_erased_since(&t1, &u2, now - 5).await.unwrap(), "other users unaffected" ); } @@ -2945,8 +2995,8 @@ mod tests { #[tokio::test] async fn memory_erase_roundtrip() { let store = MemoryStore::default(); - exercise_erase(&store).await; - exercise_erasure_markers(&store).await; + exercise_erase(&store, "").await; + exercise_erasure_markers(&store, "").await; } #[tokio::test] @@ -2954,30 +3004,43 @@ mod tests { let Ok(url) = std::env::var("GW_TEST_PG_URL") else { return; }; - let store = PostgresStore::connect(&url).await.expect("pg connect"); - sqlx::query( - "TRUNCATE billing, usage_rollup, request_content, admin_audit, - batches, batch_items, batch_results", - ) - .execute(&store.pool) - .await - .expect("truncate"); - exercise_rollup(&store).await; - exercise_erase(&store).await; - - // erasure reaches a still-pending batch's inputs (PG-only store) + // a private database, not truncation: the shared one serves parallel + // tests, and the rollup watermark is global per database — leftovers + // from any earlier run would shadow this test's raw fixtures + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let db = format!("gwtest_{nonce}"); + let admin = sqlx::PgPool::connect(&url).await.expect("pg admin"); + sqlx::query(sqlx::AssertSqlSafe(format!("CREATE DATABASE {db}"))) + .execute(&admin) + .await + .expect("create test db"); + let own_url = match url.rfind('/') { + Some(i) => format!("{}/{db}", &url[..i]), + None => url.clone(), + }; + let store = PostgresStore::connect(&own_url).await.expect("pg connect"); + let ns = format!("-{nonce}"); + exercise_rollup(&store, &ns).await; + exercise_erase(&store, &ns).await; + + // erasure reaches a still-pending batch's inputs, keyed by the + // EFFECTIVE user the enqueue persisted + let (t1, erika) = (format!("t1{ns}"), format!("erika{ns}")); let items = vec![ gw_models::BatchItem { messages: vec![gw_models::ChatMsg::text("user", "erase me")], - user: "u1".into(), + user: erika.clone(), }, gw_models::BatchItem { messages: vec![gw_models::ChatMsg::text("user", "keep me")], - user: "u2".into(), + user: format!("other{ns}"), }, ]; let pending = store - .batch_enqueue("ak", "t1", "gpt-4o", &items) + .batch_enqueue("ak", &t1, "gpt-4o", &items) .await .unwrap(); let audit2 = AdminAudit { @@ -2985,13 +3048,13 @@ mod tests { actor: "global".into(), scope: "global".into(), action: "content_erase".into(), - target: "u1".into(), + target: erika.clone(), summary: String::new(), source_ip: "10.0.0.1".into(), }; assert_eq!( store - .content_erase_user(Some("t1"), "u1", audit2) + .content_erase_user(Some(&t1), &erika, audit2) .await .unwrap(), 1, @@ -3000,21 +3063,18 @@ mod tests { let loaded = store.batch_load_items(&pending.id).await.unwrap(); assert!( loaded[0].messages.is_empty(), - "u1's queued prompt erased before any drainer ran" + "queued prompt erased before any drainer ran" + ); + assert_eq!( + loaded[1].messages[0].content, "keep me", + "other users' queued items untouched" ); let fresh = store .batch_item_snapshot(&pending.id, 0) .await .unwrap() .expect("item row still present"); - assert!( - fresh.messages.is_empty(), - "the pre-dispatch snapshot sees the erasure" - ); - assert_eq!( - loaded[1].messages[0].content, "keep me", - "other users' queued items untouched" - ); + assert!(fresh.messages.is_empty(), "pre-dispatch snapshot sees it"); } #[tokio::test] @@ -3023,8 +3083,8 @@ mod tests { let store = SqliteStore::open(dir.path().join("erase.db").to_str().unwrap()) .await .unwrap(); - exercise_erase(&store).await; - exercise_erasure_markers(&store).await; + exercise_erase(&store, "").await; + exercise_erasure_markers(&store, "").await; } #[tokio::test] @@ -3110,13 +3170,32 @@ mod tests { #[tokio::test] async fn ledger_retention_caps_both_stores() { + // rows only prune once rolled: record() stamps created_at=1000, so an + // advance past the settle window folds them and re-arms the cap let mem = MemoryStore::with_ledger_cap(2); - for m in ["a", "b", "c"] { + for m in ["a", "b"] { mem.ledger_add(&record(m)).await.unwrap(); } + mem.ledger_add(&record("unrolled")).await.unwrap(); + let (total, _) = mem.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(total, 3, "unrolled rows are spared from the cap"); + mem.usage_rollup_advance(1_000 + 3 * 60).await.unwrap(); + // a live write always postdates the watermark (settle window) + let mut c = record("c"); + c.created_at_epoch_secs = 1_200; + mem.ledger_add(&c).await.unwrap(); let (total, page) = mem.ledger_snapshot(usize::MAX).await.unwrap(); - assert_eq!(total, 2); - assert_eq!(page[0].model, "b"); + assert_eq!(total, 2, "rolled rows prune to the cap"); + assert_eq!(page[0].model, "unrolled"); + let usage = mem + .usage_by_user(Some("default"), None, 0, i64::MAX) + .await + .unwrap(); + let total_requests: i64 = usage.iter().map(|r| r.requests).sum(); + assert_eq!( + total_requests, 4, + "pruned rows stay billed via their buckets" + ); let dir = tempfile::tempdir().unwrap(); let store = SqliteStore::open_with_cap(dir.path().join("r.db").to_str().unwrap(), 2) @@ -3125,9 +3204,21 @@ mod tests { for i in 0..=LEDGER_PRUNE_EVERY { store.ledger_add(&record(&format!("m{i}"))).await.unwrap(); } - let (total, page) = store.ledger_snapshot(usize::MAX).await.unwrap(); - assert_eq!(total, 2, "prune cycle enforces the cap"); - assert_eq!(page[0].model, format!("m{}", LEDGER_PRUNE_EVERY - 1)); + let (total, _) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!( + total, + LEDGER_PRUNE_EVERY + 1, + "nothing prunes before the rollup folds the rows" + ); + store.usage_rollup_advance(1_000 + 3 * 60).await.unwrap(); + for i in 0..LEDGER_PRUNE_EVERY { + store.ledger_add(&record(&format!("n{i}"))).await.unwrap(); + } + let (total, _) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert!( + total <= 2 + LEDGER_PRUNE_EVERY, + "prune cycle enforces the cap on rolled rows (got {total})" + ); } #[tokio::test]