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/lib.rs b/crates/handler/src/lib.rs index f53595e..0f39519 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -980,6 +980,70 @@ 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(); + // millisecond markers: a resubmit moments later must already pass + tokio::time::sleep(std::time::Duration::from_millis(5)).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(); + 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}]"; @@ -1016,6 +1080,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 +1164,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..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) @@ -39,8 +48,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_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).await }); + tokio::spawn( + async move { this.execute(&id, &ak, &model, items, 0, captured_at).await }, + ); Ok(job) } } @@ -49,7 +64,15 @@ 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(); // the distributed claim already set status=running with the fence bump; // only the in-process path needs this write — unfenced on the distributed @@ -96,7 +119,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 +134,51 @@ 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. + // 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(); + // 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, captured_at) + .await + .unwrap_or(true); + if erased_mid_batch || 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(), @@ -132,6 +200,7 @@ impl OfflineHandler { ok: false, message, total_tokens: 0, + user: user.clone(), }; let result = match ran { Ok(Ok(ctx)) => match ctx.outcome { @@ -140,6 +209,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()), }, @@ -199,7 +269,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_millis(), + ) + .await; } Ok(None) => tokio::time::sleep(poll).await, Err(e) => { 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/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 08abe30..a899860 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -25,8 +25,31 @@ 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. +const ROLLUP_BUCKET_SECS: i64 = 60; + +/// 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; + +/// 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). +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 { @@ -103,6 +126,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. @@ -230,6 +257,84 @@ 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); + } + + /// 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. +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) +} + +/// 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. @@ -267,7 +372,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 +389,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 +401,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<()>; @@ -312,12 +426,28 @@ 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; + /// 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 `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, @@ -380,6 +510,27 @@ 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) + } + /// 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) + } /// 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 @@ -398,9 +549,15 @@ 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>, + /// 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, @@ -420,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() @@ -427,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(()) } @@ -485,40 +657,78 @@ 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 + .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 - ROLLUP_SETTLE_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 + .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; + for (k, v) in fresh { + rollup.entry(k).and_modify(|e| e.keep_max(&v)).or_insert(v); + } + Ok(written) } async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { @@ -581,6 +791,60 @@ impl Store for MemoryStore { Ok((before - content.len()) as u64) } + 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; + } + } + } + self.erasures + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert( + (tenant.unwrap_or_default().to_owned(), user.to_owned()), + crate::epoch_millis(), + ); + audit.summary = format!("rows={erased}"); + self.audit + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(audit); + Ok(erased) + } + + async fn user_erased_since(&self, tenant: &str, user: &str, since: i64) -> GResult { + let erasures = self + .erasures + .lock() + .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> { let content = self .content @@ -613,6 +877,13 @@ impl Store for MemoryStore { Ok(self.files.get(id).map(|f| f.value().clone())) } + 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( &self, ak: &str, @@ -729,9 +1000,29 @@ where ok: row.get(1), message: row.get(2), total_tokens: row.get(3), + user: row.get(4), } } +/// 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)] @@ -772,6 +1063,17 @@ 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_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, + 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', @@ -782,7 +1084,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 '', @@ -820,6 +1123,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") @@ -876,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(()) } @@ -938,22 +1246,78 @@ impl Store for SqliteStore { since: i64, until: i64, ) -> GResult> { - let rows = sqlx::query( + 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) + .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 AND ?4 + 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 - ROLLUP_SETTLE_SECS); + 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) + 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 = 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).min(watermark)) + .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<()> { @@ -1057,6 +1421,66 @@ impl Store for SqliteStore { Ok(r.rows_affected()) } + 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(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("erase user content", e))?; + 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 erasures (tenant, user_id, erased_at_epoch_millis) VALUES (?, ?, ?) + ON CONFLICT (tenant, user_id) DO UPDATE SET + erased_at_epoch_millis = excluded.erased_at_epoch_millis", + ) + .bind(tenant.unwrap_or_default()) + .bind(user) + .bind(crate::epoch_millis()) + .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 (?, ?, ?, ?, ?, ?, ?)", + ) + .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> { let rows = sqlx::query( "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, kind, content, @@ -1071,10 +1495,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) @@ -1108,6 +1535,31 @@ 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_millis >= ?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) + .bind(tenant) + .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, @@ -1148,7 +1600,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) @@ -1180,8 +1632,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'))", ) @@ -1190,6 +1642,7 @@ impl Store for SqliteStore { .bind(result.ok) .bind(&result.message) .bind(result.total_tokens) + .bind(&result.user) .bind(id) .execute(&self.pool) .await @@ -1235,6 +1688,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 '', @@ -1264,12 +1724,20 @@ 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 ''", + // 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", @@ -1343,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(()) } @@ -1412,23 +1884,99 @@ impl Store for PostgresStore { since: i64, until: i64, ) -> GResult> { - let rows = sqlx::query( + 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) + .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 AND $4 + 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 - ROLLUP_SETTLE_SECS); + 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) + 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 = 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).min(watermark)) + .bind(hi) + .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()) } async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { @@ -1533,6 +2081,69 @@ impl Store for PostgresStore { Ok(r.rows_affected()) } + 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(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("erase user content", e))?; + 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))?; + // pending/running batches included: the emptied item fails at execution + // instead of running erased content (terminal batches already pruned) + let i = sqlx::query( + "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) + .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> { let rows = sqlx::query( "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, kind, content, @@ -1586,6 +2197,16 @@ impl Store for PostgresStore { })) } + 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))?; + Ok(r.rows_affected() > 0) + } + async fn batch_create( &self, ak: &str, @@ -1622,7 +2243,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) @@ -1642,13 +2263,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))?; - Ok(()) + 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( @@ -1657,22 +2286,34 @@ 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))?; - Ok(r.rows_affected() > 0) + let applied = r.rows_affected() > 0; + if applied { + prune_terminal_items(&mut tx, id, status).await?; + } + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit status", e))?; + 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", @@ -1682,6 +2323,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))?; @@ -1730,6 +2372,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))?; @@ -1789,6 +2436,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", @@ -1950,6 +2616,7 @@ mod tests { ok: true, message: "ok".into(), total_tokens: 8, + user: String::new(), }, ) .await @@ -2083,6 +2750,363 @@ mod tests { exercise_audit(&store).await; } + /// `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(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(Some(&tenant), 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(Some(&tenant), 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 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, 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(), + ak: "ak".into(), + user_id: user.into(), + tenant: tenant.into(), + kind: "prompt".into(), + content: "hello".into(), + sealed: false, + expires_at_epoch_secs: 0, + }; + 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, + 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, 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, audit()).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" + ); + let trail = store.admin_audit_list(1_000).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" + ); + } + + /// 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, 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(), + "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() { + let store = MemoryStore::default(); + exercise_erase(&store, "").await; + exercise_erasure_markers(&store, "").await; + } + + #[tokio::test] + async fn pg_rollup_and_erase_roundtrip() { + let Ok(url) = std::env::var("GW_TEST_PG_URL") else { + return; + }; + // 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: erika.clone(), + }, + gw_models::BatchItem { + messages: vec![gw_models::ChatMsg::text("user", "keep me")], + user: format!("other{ns}"), + }, + ]; + 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: erika.clone(), + summary: String::new(), + source_ip: "10.0.0.1".into(), + }; + assert_eq!( + store + .content_erase_user(Some(&t1), &erika, 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(), + "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(), "pre-dispatch snapshot sees it"); + } + + #[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; + exercise_erasure_markers(&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(); @@ -2095,6 +3119,7 @@ mod tests { ok, message: msg.into(), total_tokens: 1, + user: String::new(), }, ) }; @@ -2114,6 +3139,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 @@ -2144,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) @@ -2159,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] @@ -2271,6 +3328,7 @@ mod tests { ok: true, message: "ok".into(), total_tokens: 5, + user: String::new(), }, ) .await @@ -2286,6 +3344,7 @@ mod tests { ok: false, message: "stale".into(), total_tokens: 0, + user: String::new(), }, ) .await @@ -2308,6 +3367,7 @@ mod tests { ok: true, message: "late".into(), total_tokens: 0, + user: String::new(), }, ) .await @@ -2336,6 +3396,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) @@ -2354,12 +3421,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 4ae0d1a..eea4989 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. @@ -20,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; @@ -37,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; @@ -51,6 +55,29 @@ 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.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + 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/crates/views/src/lib.rs b/crates/views/src/lib.rs index a93cac4..c07a0e5 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)) @@ -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,48 @@ async fn admin_content_get( Json(json!({ "request_id": request_id, "entries": entries })).into_response() } +/// 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, + 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); + 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, audit) + .await + { + Ok(deleted) => 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. @@ -3053,6 +3099,27 @@ 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), + }; + // 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), + } +} + /// GET /v1/files/{id}/content (download raw content: batch output, etc). async fn files_content( State(s): State, @@ -3232,6 +3299,133 @@ 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}]"; + // 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..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,6 +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 — 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 3fc1d79..9d2d5d7 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -110,6 +110,18 @@ 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. 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 `security:` is global by default; a tenant may override it whole with @@ -150,3 +162,19 @@ 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 trace of one end + user's content (the GDPR/PIPL right-to-erasure hook): retained rows, batch + 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. 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