diff --git a/Cargo.lock b/Cargo.lock index 56b3fdf..d2b92d5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,16 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common 0.1.7", + "generic-array", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -241,6 +251,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures 0.2.17", +] + [[package]] name = "chacha20" version = "0.10.1" @@ -252,6 +273,30 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "chacha20poly1305" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" +dependencies = [ + "aead", + "chacha20 0.9.1", + "cipher", + "poly1305", + "zeroize", +] + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common 0.1.7", + "inout", + "zeroize", +] + [[package]] name = "cmov" version = "0.5.4" @@ -354,6 +399,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", + "rand_core 0.6.4", "typenum", ] @@ -749,9 +795,11 @@ name = "gw-config" version = "0.1.0" dependencies = [ "gw-consts", + "regex", "serde", "serde_yaml", "thiserror", + "tracing", ] [[package]] @@ -813,6 +861,7 @@ dependencies = [ "gw-engines", "gw-models", "gw-state", + "regex", "serde_json", "tokio", "tracing", @@ -869,12 +918,15 @@ version = "0.1.0" dependencies = [ "arc-swap", "async-trait", + "base64", + "chacha20poly1305", "dashmap", "futures", "governor", "gw-config", "gw-consts", "gw-models", + "hex", "metrics", "moka", "redis", @@ -899,6 +951,7 @@ dependencies = [ name = "gw-views" version = "0.1.0" dependencies = [ + "async-trait", "axum", "base64", "bytes", @@ -1223,6 +1276,15 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -1463,6 +1525,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "opaque-debug" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" + [[package]] name = "parking" version = "2.2.1" @@ -1510,6 +1578,17 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +[[package]] +name = "poly1305" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" +dependencies = [ + "cpufeatures 0.2.17", + "opaque-debug", + "universal-hash", +] + [[package]] name = "portable-atomic" version = "1.13.1" @@ -1651,7 +1730,7 @@ version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" dependencies = [ - "chacha20", + "chacha20 0.10.1", "getrandom 0.4.3", "rand_core 0.10.1", ] @@ -1666,6 +1745,15 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + [[package]] name = "rand_core" version = "0.9.5" @@ -2663,6 +2751,16 @@ version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common 0.1.7", + "subtle", +] + [[package]] name = "unsafe-libyaml" version = "0.2.11" diff --git a/Cargo.toml b/Cargo.toml index e295295..368451a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -76,10 +76,12 @@ gw-task = { path = "crates/task" } gw-views = { path = "crates/views" } base64 = "0.22" tiktoken-rs = "0.12" -# pure-compute crypto (SigV4 signing, cache keys) — no network capability +# pure-compute crypto (SigV4 signing, cache keys, content-at-rest AEAD) — no network capability sha2 = "0.10" hmac = "0.12" hex = "0.4" +chacha20poly1305 = "0.10" +regex = "1" # Faster iterative builds; keep release optimized. [profile.dev] diff --git a/README.md b/README.md index c80c512..5a64c12 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ key-based auth, quotas, rate limits, failover, and a billing ledger. - **Staged request pipeline** — a 4-layer DAG per request: model resolve / quota / cache lookup → account selection (priority, PTU-first, failover) → rate limits + engine call (retry on upstream 5xx) → usage extraction, billing, cache store - **Governance built in** — access-key auth, daily token quotas, QPS / QPM / TPM limits at key, product, and model level, request-level TTL cache, account cooldown and recovery, DLP redaction and blocklist plugins. Admission reserves then settles, so concurrent requests can't overshoot a quota - **Multi-tenant** — keys carry a tenant; tenants get a pooled QPS bucket, a model entitlement allowlist, per-(key, model) quota defaults with an optional fallback-model degrade, key lifecycle (expiry/ban), and tenant-scoped admin tokens. Billing records charged cost and (optionally) vendor cost per row, so margin is queryable per tenant × model +- **Per-user billing & enterprise audit** — every ledger row attributes to an effective end user (the key's `owner`, else the request's `x-gw-user` / `user` hint) with a `request_id`, so cost rolls up per user (`/admin/usage/users`) and a soft per-user daily budget applies on every surface. Per-tenant content policy adds blocklist action tiers (block / flag / shadow), regex recognizers, secret masking, and an external-moderation seam; every hit is recorded without prompt text. An admin-operation trail (key CRUD / config / reload, with source IP) and optional at-rest content retention complete the audit surfaces (`/admin/audit/*`) - **Fleet-ready** — run N instances behind a load balancer: Postgres shares config (versioned + a change feed), the access-key table, the ledger/files/batches store, and a distributed batch queue any instance drains; Redis shares rate/quota/TPM counters, account health, and optionally the response cache. Single-node stays zero-dependency - **Providers behind traits** — engines talk to upstreams through a `Transport` seam; accounts with a real endpoint go over HTTP (reqwest + rustls), accounts without one are served by a deterministic in-process mock; AWS SigV4 signing included - **Observability built in** — Prometheus `/metrics` (per-route request/status counters, per-pipeline-stage latency, token counters), structured access logs diff --git a/conf/gateway.yaml b/conf/gateway.yaml index acdfd49..ede32b6 100644 --- a/conf/gateway.yaml +++ b/conf/gateway.yaml @@ -86,6 +86,17 @@ tenants: fallback_model: gpt-4o-mini model_prices: # e2e: tenant pays a premium over the list price gpt-4o: {input_price_per_1k_micros: 5000, output_price_per_1k_micros: 20000} + # Optional per-tenant knobs (enterprise audit/compliance): + # user_daily_token_quota: 100000 # soft per-end-user daily cap + # retention: {content: redacted, days: 30} # none | redacted | full (full needs GW_CONTENT_KEY) + # security: # overrides the global `security:` for this tenant + # blocklist: ["forbidden"] + # blocklist_action: flag # block | flag | shadow + # dlp_redact: true + # detect_secrets: true # mask API keys / credentials + # moderate: false # route to the wired external moderator + # regex_rules: + # - {name: ssn, pattern: '\d{3}-\d{2}-\d{4}', action: block} # Public model name → wire protocol, optional provider binding (only that # provider's accounts serve the model), plus demo prices (micros per 1k tokens). diff --git a/crates/config/Cargo.toml b/crates/config/Cargo.toml index b716bd0..6d11285 100644 --- a/crates/config/Cargo.toml +++ b/crates/config/Cargo.toml @@ -11,6 +11,8 @@ gw-consts = { workspace = true } serde = { workspace = true } serde_yaml = { workspace = true } thiserror = { workspace = true } +regex = { workspace = true } +tracing = { workspace = true } [lints] workspace = true diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 5cf88ab..d6a7b5e 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -64,6 +64,9 @@ pub struct AkConf { /// Tenant this key belongs to; empty = the implicit `default` tenant. #[serde(default)] pub tenant: String, + /// End user this key is issued to (one key = one user); `None` = shared key. + #[serde(default)] + pub owner: Option, pub qps: f64, pub daily_token_quota: i64, /// tokens-per-minute window limit; None = unlimited. @@ -149,15 +152,77 @@ fn default_priority() -> i32 { 1 } -/// Local security policy (rule-based; no cloud security service). +/// What a fired content rule does. `block` denies the request; `flag` lets it +/// through but records the hit; `shadow` is `flag` for a rule under evaluation — +/// same recording, and the caller can tell trial rules apart when auditing. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Action { + #[default] + Block, + Flag, + Shadow, +} + +impl Action { + pub fn as_str(self) -> &'static str { + match self { + Action::Block => "block", + Action::Flag => "flag", + Action::Shadow => "shadow", + } + } +} + +/// A named regex recognizer applied to the same inbound text the blocklist +/// scans. `pattern` is compiled once at config load into [`SecurityConf::regexes`]. +#[derive(Debug, Clone, Deserialize)] +pub struct RegexRule { + pub name: String, + pub pattern: String, + #[serde(default)] + pub action: Action, +} + +/// Local security policy (rule-based; no cloud security service). Lives globally +/// (`security:`) and per-tenant ([`TenantConf::security`]); the tenant's wins +/// whole when present. #[derive(Debug, Clone, Default, Deserialize)] pub struct SecurityConf { /// Blocklist terms; normalized to lower-case (empties dropped) at load. #[serde(default)] pub blocklist: Vec, + /// What a blocklist hit does (default: block). + #[serde(default)] + pub blocklist_action: Action, /// Whether to DLP-redact inbound/outbound content (emails/phone numbers). #[serde(default)] pub dlp_redact: bool, + /// Detect API keys / credentials in inbound text and redact them. + #[serde(default)] + pub detect_secrets: bool, + /// Route inbound text through the wired external moderator; needs a + /// moderator plugged into the handler (the default one allows everything). + #[serde(default)] + pub moderate: bool, + /// On a moderator error, admit the request (`true`) or deny it (`false`). + #[serde(default)] + pub moderation_fail_open: bool, + /// Named regex recognizers. + #[serde(default)] + pub regex_rules: Vec, + /// Compiled `regex_rules`, built at config load (rules that fail to compile + /// are dropped with a warning). Never deserialized. + #[serde(skip)] + pub regexes: Vec, +} + +/// A compiled [`RegexRule`], ready to match on the hot path. +#[derive(Debug, Clone)] +pub struct CompiledRule { + pub name: String, + pub action: Action, + pub re: regex::Regex, } /// Account stability policy (in-memory). @@ -267,6 +332,16 @@ pub struct TenantConf { /// Per-model charged-price overrides (else the model's list price applies). #[serde(default)] pub model_prices: std::collections::HashMap, + /// Per-user daily token budget (a soft cap keyed by end user); `None` = + /// unlimited. Enforced only when the request carries a user attribution. + #[serde(default)] + pub user_daily_token_quota: Option, + /// Content-safety policy for this tenant; `None` = use the global `security:`. + #[serde(default)] + pub security: Option, + /// Prompt/response retention for this tenant; `None` = retain nothing. + #[serde(default)] + pub retention: Option, } impl TenantConf { @@ -276,6 +351,26 @@ impl TenantConf { } } +/// How much request/response content a tenant retains, and for how long. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ContentLevel { + /// Store nothing (the default posture). + None, + /// Store the post-DLP redacted text (privacy-preserving audit). + Redacted, + /// Store the raw text (needs `GW_CONTENT_KEY` for at-rest encryption). + Full, +} + +#[derive(Debug, Clone, Copy, Deserialize)] +pub struct RetentionConf { + pub content: ContentLevel, + /// Days after which stored content is purged; 0 = keep until manually purged. + #[serde(default)] + pub days: u32, +} + /// First-class provider preset: `kind` fixes the endpoint, auth style, and /// served wire types, so going live is `kind` + `api_key_env`. #[derive(Debug, Clone, Deserialize)] @@ -370,6 +465,11 @@ pub struct GatewayConfig { /// Admin surface gate (dynamic config reload / key management). #[serde(default)] pub admin: AdminConf, + /// Trust `x-real-ip` / `x-forwarded-for` for the audit source IP. Off by + /// default: the audit records the real TCP peer, which a client can't forge. + /// Enable only when a trusted proxy fronts the gateway and sets those headers. + #[serde(default)] + pub trust_proxy_headers: bool, /// Stable hash of the source document; see [`Self::generation`]. #[serde(skip)] generation: u64, @@ -464,14 +564,13 @@ impl GatewayConfig { protocols: preset.wires.iter().map(|w| (*w).to_owned()).collect(), }); } - // lower-case once here so security_check needn't rebuild the list per request - self.security.blocklist = self - .security - .blocklist - .iter() - .filter(|w| !w.is_empty()) - .map(|w| w.to_lowercase()) - .collect(); + // normalize the global policy and every tenant override once at load + compile_security(&mut self.security); + for t in &mut self.tenants { + if let Some(sec) = t.security.as_mut() { + compile_security(sec); + } + } Ok(()) } @@ -516,6 +615,15 @@ impl GatewayConfig { check_unique("product", self.products.iter().map(|p| p.name.as_str()))?; check_unique("provider", self.providers.iter().map(|p| p.name.as_str()))?; check_unique("tenant", self.tenants.iter().map(|t| t.name.as_str()))?; + // a colon in a tenant name would alias another tenant's `ub:{tenant}:{user}` budget key + for t in &self.tenants { + if t.name.contains(':') { + return Err(ConfigError::DuplicateName { + kind: "tenant (':' not allowed in name)", + name: t.name.clone(), + }); + } + } // a typo'd tenant would silently fall back to the unrestricted default — reject at load for k in &self.access_keys { if !self.is_known_tenant(&k.tenant) { @@ -599,6 +707,19 @@ impl GatewayConfig { Ok(()) } + /// The effective content-safety policy for `tenant`: the tenant's own + /// override when present, else the global `security:`. + pub fn security_for(&self, tenant: &str) -> &SecurityConf { + self.find_tenant(tenant) + .and_then(|t| t.security.as_ref()) + .unwrap_or(&self.security) + } + + /// The retention policy for `tenant`, if it configured one. + pub fn retention_for(&self, tenant: &str) -> Option<&RetentionConf> { + self.find_tenant(tenant).and_then(|t| t.retention.as_ref()) + } + /// Whether `tenant` may call `model`: a declared tenant without an /// allowlist (and the implicit default) allows every model. An undeclared /// non-default tenant fails closed — a runtime key outliving the reload @@ -641,6 +762,33 @@ impl GatewayConfig { } } +/// Normalize a security policy at load: lower-case the blocklist (so scans +/// don't rebuild it per request) and compile the regex rules (dropping any that +/// fail to compile, loudly). +fn compile_security(sec: &mut SecurityConf) { + sec.blocklist = sec + .blocklist + .iter() + .filter(|w| !w.is_empty()) + .map(|w| w.to_lowercase()) + .collect(); + sec.regexes = sec + .regex_rules + .iter() + .filter_map(|r| match regex::Regex::new(&r.pattern) { + Ok(re) => Some(CompiledRule { + name: r.name.clone(), + action: r.action, + re, + }), + Err(e) => { + tracing::warn!(rule = %r.name, error = %e, "dropping uncompilable regex rule"); + None + } + }) + .collect(); +} + /// A token read from the named env var at call time; `None` when the name is /// empty or the var is unset/empty. fn token_from_env(var: &str) -> Option { @@ -932,6 +1080,12 @@ tenants: [{name: t1}, {name: t1}] Err(ConfigError::DuplicateName { kind: "tenant", .. }) )); + let colon = "listen: {host: h, port: 1}\ntenants: [{name: 'a:b'}]"; + assert!( + GatewayConfig::from_yaml(colon).is_err(), + "a colon in a tenant name is rejected (budget-key aliasing)" + ); + let bad_quota = r#" listen: {host: h, port: 1} models: [{name: m1, protocol: openai-chat}] diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index 96be7dc..25dcf76 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -71,6 +71,14 @@ impl DagContext { self.decisions.push((node, what.into())); } + /// The effective end user: the key's `owner` (authoritative) else request + /// metadata; `""` when neither is present. Resolution lives on [`AkInfo`] so + /// REST and realtime can't diverge on an empty owner. + pub fn effective_user_id(&self) -> &str { + self.ak + .attributed_user(self.request.user_id.as_deref().unwrap_or_default()) + } + /// The decision trail as `"stage: detail"` lines. pub fn decision_lines(&self) -> impl Iterator + '_ { self.decisions.iter().map(|(n, w)| format!("{n}: {w}")) diff --git a/crates/dag/src/executor.rs b/crates/dag/src/executor.rs index 7f9b3fc..e6e0412 100644 --- a/crates/dag/src/executor.rs +++ b/crates/dag/src/executor.rs @@ -137,6 +137,7 @@ mod tests { ak: "t".into(), product: "demo".into(), tenant: "default".into(), + owner: None, qps: 10.0, daily_token_quota: 1000, tokens_per_minute: None, diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 996cc46..bf0341a 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -384,6 +384,30 @@ impl DagNode for AkTpmLimit { } } +/// model_access/user_budget: per-user daily token budget (soft cap), keyed by +/// the effective end user. No-op without a tenant limit or a user attribution. +pub struct UserBudgetGate; + +#[async_trait::async_trait] +impl DagNode for UserBudgetGate { + fn name(&self) -> &'static str { + "user_budget" + } + fn deps(&self) -> &'static [&'static str] { + &["ak_tpm"] + } + async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { + admission::check_user_budget( + ctx.state.governance.as_ref(), + &ctx.cfg, + &ctx.ak.tenant, + ctx.effective_user_id(), + ) + .await + .map_err(limit_denied) + } +} + /// model_access/call_engine: factory dispatch + engine execution + failover. /// On an upstream 5xx the failed account is excluded and reselected once (a /// PTU → paygo spill sets `ptu_spillover`); a second failure propagates as-is. @@ -395,7 +419,7 @@ impl DagNode for CallEngine { "call_engine" } fn deps(&self) -> &'static [&'static str] { - &["ak_tpm"] + &["user_budget"] } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { let threshold = ctx.cfg.stability.failure_threshold; @@ -536,7 +560,7 @@ impl DagNode for CostCalc { }; let ct = enc.encode_len(&resp.message) as i64; ctx.decide("cost_calc", format!("aborted stream, billed {pt}+{ct}")); - return bill(ctx, pt, ct, pt.saturating_add(ct)).await; + return bill(ctx, pt, ct, pt.saturating_add(ct), true).await; } // default rate is 1:1; the formula carries future weighted rates. // saturating sums so a malformed usage subtree can't overflow the totals @@ -564,13 +588,20 @@ impl DagNode for CostCalc { resp.prompt_tokens.saturating_add(resp.completion_tokens), ), }; - bill(ctx, prompt, completion, total).await + bill(ctx, prompt, completion, total, false).await } } /// Settle reserves and write the ledger for one served request via the shared -/// [`admission::settle_and_bill`] orchestration. -async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> GResult<()> { +/// [`admission::settle_and_bill`] orchestration. `estimated` marks a bill from +/// an aborted stream's estimated counts rather than a vendor usage payload. +async fn bill( + ctx: &mut DagContext, + prompt: i64, + completion: i64, + total: i64, + estimated: bool, +) -> GResult<()> { let ptu_spillover = ctx .outcome .as_ref() @@ -583,6 +614,12 @@ async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> let requested = param .and_then(|p| p.fallback_from.as_deref()) .unwrap_or(served); + // locals first: the whole-ctx borrow `effective_user_id` needs can't overlap these takes + let (reserved, tpm_reserved, model_quota_key) = ( + ctx.quota_reserved.take().unwrap_or(0), + ctx.tpm_reserved.take(), + ctx.model_quota_key.take(), + ); let record = admission::settle_and_bill( ctx.state.governance.as_ref(), ctx.state.store.as_ref(), @@ -592,6 +629,8 @@ async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> ak: &ctx.ak.ak, product: &ctx.ak.product, tenant: &ctx.ak.tenant, + user_id: ctx.effective_user_id(), + request_id: &ctx.request.request_id, requested_model: requested, served_model: served, protocol: param.map(|p| p.protocol.as_str()).unwrap_or_default(), @@ -600,14 +639,23 @@ async fn bill(ctx: &mut DagContext, prompt: i64, completion: i64, total: i64) -> completion, total, ptu_spillover, + estimated, }, - reserved: ctx.quota_reserved.take().unwrap_or(0), - tpm_reserved: ctx.tpm_reserved.take(), + reserved, + tpm_reserved, reserved_at: ctx.quota_at, - model_quota_key: ctx.model_quota_key.take(), + model_quota_key, }, ) .await; + admission::consume_user_budget( + ctx.state.governance.as_ref(), + &ctx.cfg, + &ctx.ak.tenant, + ctx.effective_user_id(), + record.total_tokens, + ) + .await; ctx.decide( "cost_calc", format!( @@ -690,6 +738,7 @@ pub fn default_layers() -> Vec { Box::new(ProductQpmLimit), Box::new(ModelQpmLimit), Box::new(AkTpmLimit), + Box::new(UserBudgetGate), Box::new(CallEngine), ], }, diff --git a/crates/handler/Cargo.toml b/crates/handler/Cargo.toml index 7561dca..9778db7 100644 --- a/crates/handler/Cargo.toml +++ b/crates/handler/Cargo.toml @@ -17,10 +17,11 @@ gw-dag = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } futures = { workspace = true } +regex = { workspace = true } +async-trait = { workspace = true } [dev-dependencies] tokio = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } [lints] diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 7cab625..f53595e 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -3,21 +3,28 @@ //! DAG layers, then the plugin post-stage; `OfflineHandler` reuses the same //! chain for batches. +pub mod moderation; pub mod offline; pub mod plugins; use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use futures::FutureExt; use gw_config::GatewayConfig; use gw_dag::DagContext; use gw_engines::http_transport::UpstreamPolicy; use gw_engines::{EngineOutcome, SharedTransport}; -use gw_models::{GResult, GatewayError, GatewayRequest, GatewayResponse}; +use gw_models::{Block, GResult, GatewayError, GatewayRequest, GatewayResponse}; use gw_state::{AkInfo, GatewayState, SharedConfig}; -pub use offline::{BatchItem, OfflineHandler}; +pub use gw_models::BatchItem; +pub use offline::OfflineHandler; + +const MODERATION_UNAVAILABLE: &str = "content moderation is unavailable"; + +static REQ_SEQ: AtomicU64 = AtomicU64::new(0); /// Runs one request through the plugin pre-stage, the DAG, and the plugin post-stage. #[derive(Clone)] @@ -25,6 +32,7 @@ pub struct OnlineHandler { pub config: SharedConfig, pub transport: SharedTransport, plan: Arc, + moderator: Arc, } impl OnlineHandler { @@ -38,11 +46,19 @@ impl OnlineHandler { config, transport, plan, + moderator: moderation::default_moderator(), }; handler.push_policies(&handler.cfg()); handler } + /// Plug an external content moderator into the pre-stage (enable per tenant + /// via `security.moderate`). + pub fn with_moderator(mut self, moderator: Arc) -> Self { + self.moderator = moderator; + self + } + /// The live config snapshot (cheap atomic load). Introspection surfaces read /// through this so a runtime reload takes effect immediately. pub fn cfg(&self) -> Arc { @@ -67,43 +83,21 @@ impl OnlineHandler { /// Run one request: plugin pre → DAG (4 layers) → plugin post. /// The returned context carries the outcome, decision log, billing effects. pub async fn run(&self, mut request: GatewayRequest, ak: AkInfo) -> GResult { + if request.request_id.is_empty() { + request.request_id = new_request_id(); + } // one consistent snapshot for the whole request let snap = self.config.load(); + let sec = snap.cfg.security_for(&ak.tenant); // outbound DLP is a response-buffering boundary: a masked span can // straddle deltas, so no engine may stream raw ones — enforced here so // no caller can opt out - let dlp = snap.cfg.security.dlp_redact; + let dlp = sec.dlp_redact; if dlp { request.stream_tx = None; } - // blocklist on the ORIGINAL content, before DLP — else a blocklisted term - // inside a redacted span (a domain in an email) is masked out and slips - if let Some(block) = plugins::security_check(&snap.cfg.security, &mut request) { - let mut ctx = DagContext::new( - snap.cfg.clone(), - snap.state.clone(), - self.transport.clone(), - request, - ak, - ); - ctx.decide( - "security_check", - format!("blocked (code {})", block.err_code), - ); - let response = GatewayResponse { - message: block.message.clone(), - finish_reason: "content_filter".to_owned(), - ..Default::default() - }; - ctx.outcome = Some(EngineOutcome { - response, - http_code: 200, - block, - ..Default::default() - }); - return Ok(ctx); - } - let redacted = plugins::dlp_redact_request(&snap.cfg.security, &mut request); + // scan the ORIGINAL content pre-DLP: a blocklisted term inside a redacted span would slip + let scan = plugins::security_check(sec, &mut request); let mut ctx = DagContext::new( snap.cfg.clone(), @@ -112,8 +106,42 @@ impl OnlineHandler { request, ak, ); + // every fired rule is recorded (block/flag/shadow alike); only a block-action hit denies + for hit in &scan.hits { + emit_security_event(&ctx, &hit.rule, hit.action.as_str(), hit.count).await; + } + if let Some(block) = scan.block { + ctx.decide( + "security_check", + format!("blocked (code {})", block.err_code), + ); + ctx.outcome = Some(content_filter_outcome(block)); + return Ok(ctx); + } + + // pre-DLP text, computed once for moderation and the retained prompt + let retention = snap + .cfg + .retention_for(&ctx.ak.tenant) + .copied() + .filter(|r| r.content != gw_config::ContentLevel::None); + let inbound = + (sec.moderate || retention.is_some()).then(|| plugins::inbound_text(&mut ctx.request)); + + if sec.moderate + && let Some(block) = self + .moderate(&ctx, sec, inbound.as_deref().unwrap_or_default()) + .await + { + ctx.decide("moderation", "denied"); + ctx.outcome = Some(content_filter_outcome(block)); + return Ok(ctx); + } + + let redacted = plugins::dlp_redact_request(sec, &mut ctx.request); if redacted > 0 { ctx.decide("dlp", format!("redacted {redacted} span(s) inbound")); + emit_security_event(&ctx, "dlp", "redact", redacted as i64).await; } // a panicking node must refund too, not leak the reserves; unwind-safe — @@ -148,8 +176,15 @@ impl OnlineHandler { outcome.response.model = requested; } + // raw response pre-outbound-DLP, only when full retention can store it (key present) + let capture_raw = matches!(retention, Some(r) if r.content == gw_config::ContentLevel::Full) + && gw_state::sealing_available(); + let raw_response = capture_raw + .then(|| ctx.outcome.as_ref().map(|o| o.response.message.clone())) + .flatten(); + let redacted_out = if let Some(outcome) = ctx.outcome.as_mut() { - let n = plugins::dlp_redact_response(&snap.cfg.security, &mut outcome.response); + let n = plugins::dlp_redact_response(sec, &mut outcome.response); // raw decoded deltas are pre-redaction; drop them so no downstream // reconstruction can replay unmasked text past the boundary if dlp { @@ -159,12 +194,75 @@ impl OnlineHandler { } else { 0 }; + if let Some(r) = retention { + persist_content( + &ctx, + r, + capture_raw, + inbound.unwrap_or_default(), + raw_response, + ) + .await; + } if redacted_out > 0 { ctx.decide("dlp", format!("redacted {redacted_out} span(s) outbound")); + emit_security_event(&ctx, "dlp", "redact_out", redacted_out as i64).await; } Ok(ctx) } + /// Run the wired moderator over raw text; `Some(reason)` to deny, `None` to + /// allow. The seam the realtime surface uses (it has no `DagContext`); the + /// caller records the security event on its own surface. + pub async fn moderate_text(&self, sec: &gw_config::SecurityConf, text: &str) -> Option { + match self.moderation(sec, text).await { + Moderation::Allow => None, + Moderation::Deny(reason) => Some(reason), + Moderation::Unavailable => Some(MODERATION_UNAVAILABLE.to_owned()), + } + } + + /// Run the wired moderator over the request's pre-DLP inbound `text`; + /// `Some(Block)` to deny. Records a security event on a moderator deny. + async fn moderate( + &self, + ctx: &DagContext, + sec: &gw_config::SecurityConf, + text: &str, + ) -> Option { + match self.moderation(sec, text).await { + Moderation::Allow => None, + Moderation::Deny(reason) => { + emit_security_event(ctx, "moderation", "block", 1).await; + Some(Block::blocked( + reason, + gw_consts::ErrCode::EMPTY_RESP.value() as i32, + )) + } + Moderation::Unavailable => Some(Block::blocked( + MODERATION_UNAVAILABLE, + gw_consts::ErrCode::SYSTEM_ERROR.value() as i32, + )), + } + } + + /// The one moderator-verdict resolution every surface shares, so the + /// fail-open posture can't drift between REST and realtime. + async fn moderation(&self, sec: &gw_config::SecurityConf, text: &str) -> Moderation { + match self.moderator.review(text).await { + Ok(moderation::Verdict::Allow) => Moderation::Allow, + Ok(moderation::Verdict::Deny(reason)) => Moderation::Deny(reason), + Err(e) => { + tracing::warn!(error = %e, fail_open = sec.moderation_fail_open, "moderator error"); + if sec.moderation_fail_open { + Moderation::Allow + } else { + Moderation::Unavailable + } + } + } + } + /// Derive the upstream policies (timeouts/connect-retries) from `cfg` and /// apply them to the transport live. fn push_policies(&self, cfg: &GatewayConfig) { @@ -190,6 +288,137 @@ impl OnlineHandler { } } +/// One resolved moderator verdict: `Unavailable` is a moderator error under a +/// fail-closed posture (fail-open resolves to `Allow`). +enum Moderation { + Allow, + Deny(String), + Unavailable, +} + +/// A per-request correlation id: `req--`, time-sortable and +/// unique within the process (the seq disambiguates same-millisecond requests). +pub fn new_request_id() -> String { + let ms = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis()) + .unwrap_or(0); + format!("req-{ms}-{}", REQ_SEQ.fetch_add(1, Ordering::Relaxed)) +} + +/// The 200-with-`content_filter` outcome every pre-stage denial returns. +fn content_filter_outcome(block: Block) -> EngineOutcome { + EngineOutcome { + response: GatewayResponse { + message: block.message.clone(), + finish_reason: "content_filter".to_owned(), + ..Default::default() + }, + http_code: 200, + block, + ..Default::default() + } +} + +/// Record a content-safety outcome (no prompt text) against this request's +/// key/user/tenant. Best-effort. Surface = the request protocol, or "batch" +/// for offline items. +async fn emit_security_event(ctx: &DagContext, rule: &str, action: &str, hits: i64) { + let surface = if ctx.request.is_online { + ctx.request + .model_param_v2 + .as_ref() + .map(|p| p.protocol.as_str()) + .unwrap_or_default() + .to_owned() + } else { + "batch".to_owned() + }; + gw_state::SecurityEvent { + created_at_epoch_secs: gw_state::epoch_secs(), + request_id: ctx.request.request_id.clone(), + ak: ctx.ak.ak.clone(), + user_id: ctx.effective_user_id().to_owned(), + tenant: ctx.ak.tenant.clone(), + surface, + rule: rule.to_owned(), + action: action.to_owned(), + hits, + } + .record(ctx.state.store.as_ref()) + .await; +} + +/// Persist this request's prompt and response per the tenant's retention policy. +/// `inbound` is the pre-DLP prompt text; `store_full` (resolved once by the +/// caller: full level AND a content key) stores it sealed, else it is stored +/// PII/secret-stripped — retention owns that redaction, so a row can't hold +/// raw content even with DLP off. Best-effort — a store failure is logged. +async fn persist_content( + ctx: &DagContext, + retention: gw_config::RetentionConf, + store_full: bool, + inbound: String, + raw_response: Option, +) { + if retention.content == gw_config::ContentLevel::Full && !store_full { + tracing::warn!("full retention configured but GW_CONTENT_KEY unset; storing redacted text"); + } + + let now = gw_state::epoch_secs(); + let expires = if retention.days > 0 { + now + retention.days as i64 * 86_400 + } else { + 0 + }; + let redacted_response = || { + plugins::redact_retained( + ctx.outcome + .as_ref() + .map(|o| o.response.message.as_str()) + .unwrap_or_default(), + ) + }; + let prompt = if store_full { + inbound + } else { + plugins::redact_retained(&inbound) + }; + let response = if store_full { + raw_response.unwrap_or_else(redacted_response) + } else { + redacted_response() + }; + + let writes = [("prompt", prompt), ("response", response)] + .into_iter() + .filter(|(_, text)| !text.is_empty()) + .map(|(kind, text)| { + // seal whenever a key exists (defense in depth even for redacted text) + let (content, sealed) = match gw_state::content::seal(&text) { + Some(ct) => (ct, true), + None => (text, false), + }; + let record = gw_state::ContentRecord { + created_at_epoch_secs: now, + request_id: ctx.request.request_id.clone(), + ak: ctx.ak.ak.clone(), + user_id: ctx.effective_user_id().to_owned(), + tenant: ctx.ak.tenant.clone(), + kind: kind.to_owned(), + content, + sealed, + expires_at_epoch_secs: expires, + }; + async move { + if let Err(e) = ctx.state.store.content_add(&record).await { + tracing::warn!(error = %e, kind, "content retention write failed"); + } + } + }); + futures::future::join_all(writes).await; +} + #[cfg(test)] mod tests { use super::*; @@ -232,6 +461,110 @@ mod tests { } } + #[derive(Debug)] + struct DenyModerator; + + #[async_trait::async_trait] + impl moderation::Moderator for DenyModerator { + async fn review(&self, _text: &str) -> Result { + Ok(moderation::Verdict::Deny("nope".into())) + } + } + + #[tokio::test] + async fn moderator_denies_when_tenant_enables_it() { + let mut cfg = GatewayConfig::embedded_default().unwrap(); + cfg.security.moderate = true; + let cfg = Arc::new(cfg); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ) + .with_moderator(Arc::new(DenyModerator)); + let ctx = h + .run(chat_req("gpt-4o", "hello"), ak(&h).await) + .await + .unwrap(); + let out = ctx.outcome.expect("outcome"); + assert!(out.block.block); + assert_eq!(out.response.finish_reason, "content_filter"); + assert!( + h.state() + .store + .ledger_snapshot(1) + .await + .unwrap() + .1 + .is_empty(), + "a moderated deny bills nothing" + ); + } + + #[derive(Debug)] + struct ErrModerator; + + #[async_trait::async_trait] + impl moderation::Moderator for ErrModerator { + async fn review(&self, _text: &str) -> Result { + Err("moderator upstream down".into()) + } + } + + #[tokio::test] + async fn moderate_text_allow_deny_and_failure_posture() { + let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let base = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ); + let mut sec = gw_config::SecurityConf::default(); + assert_eq!(base.moderate_text(&sec, "x").await, None, "default allows"); + let deny = base.clone().with_moderator(Arc::new(DenyModerator)); + assert_eq!(deny.moderate_text(&sec, "x").await.as_deref(), Some("nope")); + let err = base.with_moderator(Arc::new(ErrModerator)); + sec.moderation_fail_open = true; + assert_eq!( + err.moderate_text(&sec, "x").await, + None, + "error + fail-open allows" + ); + sec.moderation_fail_open = false; + assert!( + err.moderate_text(&sec, "x").await.is_some(), + "error + fail-closed denies" + ); + } + + #[tokio::test] + async fn per_user_budget_denies_over_the_cap() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, user_daily_token_quota: 5}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let with_user = |content: &str| GatewayRequest { + is_online: true, + message: vec![ChatMsg::text("user", content)], + model_param_v2: Some(ModelParamV2::with_name(Protocol::OpenaiChat, "gpt-4o")), + user_id: Some("u1".into()), + ..Default::default() + }; + h.run(with_user("first burns the budget"), key.clone()) + .await + .unwrap(); + let err = h + .run(with_user("second is over"), key) + .await + .err() + .expect("second denied by the per-user budget"); + assert_eq!(err.http_status, 429); + } + #[tokio::test] async fn full_pipeline_openai() { let h = handler(); @@ -558,6 +891,7 @@ mod tests { ak: "ak-t1".into(), product: "p".into(), tenant: "t1".into(), + owner: None, qps: 10.0, daily_token_quota: 100_000, tokens_per_minute: None, @@ -598,9 +932,11 @@ mod tests { let items = vec![ BatchItem { messages: vec![ChatMsg::text("user", "same prompt")], + user: String::new(), }, BatchItem { messages: vec![ChatMsg::text("user", "same prompt")], + user: String::new(), }, ]; let job = off @@ -623,9 +959,11 @@ mod tests { vec![ BatchItem { messages: vec![ChatMsg::text("user", "one")], + user: String::new(), }, BatchItem { messages: vec![ChatMsg::text("user", "two")], + user: String::new(), }, ], ) @@ -642,6 +980,44 @@ mod tests { ); } + #[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}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ); + let off = OfflineHandler::new(h.clone()); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let job = off + .submit( + key, + "gpt-4o".into(), + vec![ + BatchItem { + messages: vec![ChatMsg::text("user", "for alice")], + user: "alice".into(), + }, + BatchItem { + messages: vec![ChatMsg::text("user", "for bob")], + user: "bob".into(), + }, + ], + ) + .await + .unwrap(); + wait_terminal(&h, &job.id).await; + let (_, ledger) = h.state().store.ledger_snapshot(usize::MAX).await.unwrap(); + let users: std::collections::HashSet<&str> = + ledger.iter().map(|r| r.user_id.as_str()).collect(); + assert!( + users.contains("alice") && users.contains("bob"), + "each shared-key batch item bills to its own user: {users:?}" + ); + } + #[tokio::test] async fn distributed_batch_drained_by_a_separate_handler() { let Ok(url) = std::env::var("GW_TEST_PG_URL") else { @@ -670,9 +1046,11 @@ mod tests { vec![ BatchItem { messages: vec![ChatMsg::text("user", "alpha")], + user: "alice".into(), }, BatchItem { messages: vec![ChatMsg::text("user", "beta")], + user: "bob".into(), }, ], ) @@ -701,5 +1079,12 @@ mod tests { let j = completed.expect("drain completed the batch"); assert_eq!(j.results.len(), 2, "both items executed exactly once"); assert!(j.results.iter().all(|r| r.ok && r.total_tokens > 0)); + let (_, ledger) = state.store.ledger_snapshot(usize::MAX).await.unwrap(); + let users: std::collections::HashSet<&str> = + ledger.iter().map(|r| r.user_id.as_str()).collect(); + assert!( + users.contains("alice") && users.contains("bob"), + "distributed batch preserved per-item user attribution: {users:?}" + ); } } diff --git a/crates/handler/src/moderation.rs b/crates/handler/src/moderation.rs new file mode 100644 index 0000000..d3edd3e --- /dev/null +++ b/crates/handler/src/moderation.rs @@ -0,0 +1,40 @@ +//! The moderation seam: an optional external content-review pass, plugged into +//! the pre-stage. The default [`AllowModerator`] is a deterministic no-op, so +//! the hot path pays nothing until a real moderator is wired in and a tenant +//! turns `moderate` on. External review is latency-bearing and pull-driven, so +//! this ships as a trait, not a built-in integration. + +use std::sync::Arc; + +/// A moderator's decision on one request's text. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Verdict { + Allow, + /// Deny with a user-facing reason. + Deny(String), +} + +/// A pluggable content moderator. `review` sees the request's concatenated +/// inbound text; a real impl calls an external service. An `Err` is a moderator +/// failure — the caller resolves it per its fail-open/closed posture. +#[async_trait::async_trait] +pub trait Moderator: Send + Sync + std::fmt::Debug { + async fn review(&self, text: &str) -> Result; +} + +/// The default: allow everything (a real moderator replaces it via +/// [`crate::OnlineHandler::with_moderator`]). +#[derive(Debug, Default)] +pub struct AllowModerator; + +#[async_trait::async_trait] +impl Moderator for AllowModerator { + async fn review(&self, _text: &str) -> Result { + Ok(Verdict::Allow) + } +} + +/// The default moderator handle. +pub fn default_moderator() -> Arc { + Arc::new(AllowModerator) +} diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 5636468..81e204b 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -2,17 +2,11 @@ //! instance's background task; a distributed store (Postgres) only persists the //! items and a fleet drain loop on any instance claims and runs them. -use gw_models::{ChatMsg, GatewayRequest, ModelParamV2}; +use gw_models::{BatchItem, GatewayRequest, ModelParamV2}; use gw_state::{AkInfo, BatchItemResult, BatchJob, BatchStatus}; use crate::OnlineHandler; -/// One batch item: a self-contained message list. -#[derive(Debug, Clone)] -pub struct BatchItem { - pub messages: Vec, -} - /// Batch orchestration built on top of the online handler. #[derive(Clone)] pub struct OfflineHandler { @@ -34,18 +28,19 @@ impl OfflineHandler { items: Vec, ) -> gw_models::GResult { let store = self.online.state().store.clone(); - let msgs: Vec> = items.into_iter().map(|i| i.messages).collect(); if store.distributed_batches() { // atomic: the job becomes claimable only once all items are saved - store.batch_enqueue(&ak.ak, &ak.tenant, &model, &msgs).await + store + .batch_enqueue(&ak.ak, &ak.tenant, &model, &items) + .await } else { let job = store - .batch_create(&ak.ak, &ak.tenant, &model, msgs.len()) + .batch_create(&ak.ak, &ak.tenant, &model, items.len()) .await?; let this = self.clone(); let (id, model) = (job.id.clone(), model.clone()); // claim 0: non-distributed store — no fence, the heartbeat is a no-op - tokio::spawn(async move { this.execute(&id, &ak, &model, msgs, 0).await }); + tokio::spawn(async move { this.execute(&id, &ak, &model, items, 0).await }); Ok(job) } } @@ -54,14 +49,7 @@ 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) { 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 @@ -108,7 +96,7 @@ impl OfflineHandler { } }) }; - for (index, messages) in items.into_iter().enumerate() { + for (index, item) in items.into_iter().enumerate() { if lost.load(Relaxed) { break; // reclaimed by another instance; stop running new items } @@ -126,7 +114,8 @@ impl OfflineHandler { let request = GatewayRequest { is_online: false, ak: ak.ak.clone(), - message: messages, + message: item.messages, + user_id: (!item.user.is_empty()).then_some(item.user), model_param_v2: Some(ModelParamV2::with_name( gw_consts::Protocol::OpenaiChat, model.to_owned(), diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 712cf4a..795dfd5 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -4,40 +4,137 @@ //! typed params. The post-stage redacts the outbound message; streaming //! surfaces buffer and replay the redacted text when outbound DLP is on. -use gw_config::SecurityConf; +use gw_config::{Action, SecurityConf}; use gw_models::{Block, GatewayRequest, GatewayResponse}; const BLOCKED_MSG: &str = "this content cannot be answered, please try a different request"; -/// Blocklist check. Returns Block on a hit (block=true implies hit=true). -/// Scans the same inbound text DLP covers, so no surface is a bypass; takes -/// `&mut` only to share the one traversal with the redaction pass. -pub fn security_check(sec: &SecurityConf, request: &mut GatewayRequest) -> Option { - if sec.blocklist.is_empty() { - return None; +/// One content rule that fired on a request, for the security-event stream. +pub struct RuleHit { + pub rule: String, + pub action: Action, + pub count: i64, +} + +/// The result of scanning one request: whether to deny it, plus every rule that +/// fired (for the audit stream). A `block`-action hit fills `block`; `flag`/ +/// `shadow` hits only populate `hits`. +#[derive(Default)] +pub struct ScanOutcome { + pub block: Option, + pub hits: Vec, +} + +/// Scan inbound text against the blocklist and the regex recognizers. One +/// traversal of the same fields DLP covers, so no surface is a bypass. `&mut` +/// only to share the traversal; nothing is rewritten here. +pub fn security_check(sec: &SecurityConf, request: &mut GatewayRequest) -> ScanOutcome { + if sec.blocklist.is_empty() && sec.regexes.is_empty() { + return ScanOutcome::default(); } - // the shared traversal counts hits; a String is never rewritten here - let mut scan = |s: &mut String| usize::from(blocklist_hit(sec, s)); - let mut blocked = 0; + let mut counts = ScanCounts::new(sec); + let mut visit = |s: &mut String| counts.visit(s); for msg in &mut request.message { - blocked += usize::from(blocklist_hit(sec, &msg.content)); + visit(&mut msg.content); if let Some(parts) = &mut msg.parts { - blocked += walk_json_strings(parts, &mut scan); + walk_json_strings(parts, &mut visit); + } + } + if let Some(param) = request.model_param_v2.as_mut() { + walk_json_strings(&mut param.raw, &mut visit); + if let Some(typed) = param.typed.as_mut() { + for_each_typed_text(typed, &mut visit); + } + } + counts.outcome() +} + +/// Per-rule hit counters for one scan, shared by the REST request scan and the +/// realtime frame scan so the two apply identical action semantics. +struct ScanCounts<'a> { + sec: &'a SecurityConf, + blocklist: i64, + regex: Vec, +} + +impl<'a> ScanCounts<'a> { + fn new(sec: &'a SecurityConf) -> Self { + Self { + sec, + blocklist: 0, + regex: vec![0; sec.regexes.len()], + } + } + + fn visit(&mut self, s: &str) -> usize { + self.blocklist += i64::from(blocklist_hit(self.sec, s)); + for (i, r) in self.sec.regexes.iter().enumerate() { + self.regex[i] += r.re.find_iter(s).count() as i64; + } + 0 + } + + /// Fold the counts into a [`ScanOutcome`]: a rule with count > 0 is a hit + /// at its action; any block-action hit denies. + fn outcome(self) -> ScanOutcome { + let mut hits = Vec::new(); + if self.blocklist > 0 { + hits.push(RuleHit { + rule: "blocklist".to_owned(), + action: self.sec.blocklist_action, + count: self.blocklist, + }); + } + for (r, count) in self.sec.regexes.iter().zip(self.regex) { + if count > 0 { + hits.push(RuleHit { + rule: r.name.clone(), + action: r.action, + count, + }); + } + } + let block = hits + .iter() + .any(|h| h.action == Action::Block) + .then(|| Block::blocked(BLOCKED_MSG, gw_consts::ErrCode::EMPTY_RESP.value() as i32)); + ScanOutcome { block, hits } + } +} + +/// All inbound text a request carries — message content, multimodal text +/// parts, the Responses raw body, and the family typed params — collected via +/// the same traversals the blocklist scan and DLP run, so the field lists +/// cannot drift apart. The one text view moderation and content retention +/// operate on. `&mut` only to share those traversals; nothing is rewritten. +pub fn inbound_text(request: &mut GatewayRequest) -> String { + let mut out = String::new(); + let mut collect = |s: &mut String| { + push_text(&mut out, s); + 0 + }; + for m in &mut request.message { + collect(&mut m.content); + if let Some(parts) = &mut m.parts { + for_each_part_text(parts, &mut collect); } } if let Some(param) = request.model_param_v2.as_mut() { - blocked += walk_json_strings(&mut param.raw, &mut scan); + walk_json_strings(&mut param.raw, &mut collect); if let Some(typed) = param.typed.as_mut() { - blocked += for_each_typed_text(typed, &mut scan); + for_each_typed_text(typed, &mut collect); } } - if blocked > 0 { - return Some(Block::blocked( - BLOCKED_MSG, - gw_consts::ErrCode::EMPTY_RESP.value() as i32, - )); + out +} + +fn push_text(out: &mut String, s: &str) { + if !s.is_empty() { + if !out.is_empty() { + out.push('\n'); + } + out.push_str(s); } - None } /// Case-insensitive blocklist test; terms are pre-lowercased at config load. @@ -101,83 +198,145 @@ fn for_each_typed_text( } } -/// Blocklist scan over a realtime frame's text-bearing fields — the same terms -/// every REST surface blocks, so the WebSocket surface is not a bypass. -pub fn realtime_frame_blocked(sec: &SecurityConf, frame: &mut serde_json::Value) -> Option { - if sec.blocklist.is_empty() { - return None; +/// Visit a multimodal `parts` array's text blocks; returns summed hits. +/// Deliberately narrower than the blocklist scan: non-text parts (image URLs, +/// base64 data) are never visited, so a rewrite can't corrupt them. +fn for_each_part_text( + parts: &mut serde_json::Value, + f: &mut impl FnMut(&mut String) -> usize, +) -> usize { + let Some(arr) = parts.as_array_mut() else { + return 0; + }; + let mut hits = 0; + for p in arr { + if p["type"] == "text" + && let Some(serde_json::Value::String(s)) = p.get_mut("text") + { + hits += f(s); + } + } + hits +} + +/// Scan a realtime frame's text-bearing fields against the blocklist AND the +/// regex recognizers, honoring their actions — the same policy every REST +/// surface runs, so the WebSocket surface is not a bypass. A block-action hit +/// denies the frame. `collect_text` also gathers the frame's text (for +/// moderation) in the same traversal, so a moderated tenant pays one walk per +/// frame, not two. +pub fn realtime_frame_scan( + sec: &SecurityConf, + frame: &mut serde_json::Value, + collect_text: bool, +) -> (ScanOutcome, String) { + let scan_rules = !sec.blocklist.is_empty() || !sec.regexes.is_empty(); + let mut counts = ScanCounts::new(sec); + let mut text = String::new(); + if scan_rules || collect_text { + gw_engines::realtime::visit_frame_text(frame, &mut |s| { + if scan_rules { + counts.visit(s); + } + if collect_text { + push_text(&mut text, s); + } + 0 + }); } - let hits = - gw_engines::realtime::visit_frame_text(frame, &mut |s| usize::from(blocklist_hit(sec, s))); - (hits > 0).then(|| Block::blocked(BLOCKED_MSG, gw_consts::ErrCode::EMPTY_RESP.value() as i32)) + (counts.outcome(), text) } /// DLP-redact a realtime frame's text-bearing fields in place; the hit count. +/// Honors both `dlp_redact` (PII) and `detect_secrets` (credentials), the same +/// as the REST path — a realtime frame must not be a secret-redaction bypass. /// Per-frame best effort: a PII span straddling two deltas is not caught — a /// realtime relay cannot buffer the way the REST stream surfaces do. pub fn dlp_redact_realtime_frame(sec: &SecurityConf, frame: &mut serde_json::Value) -> usize { - if !sec.dlp_redact { + if !sec.dlp_redact && !sec.detect_secrets { return 0; } - gw_engines::realtime::visit_frame_text(frame, &mut redact_str) + let (pii, secrets) = (sec.dlp_redact, sec.detect_secrets); + gw_engines::realtime::visit_frame_text(frame, &mut |s| redact_in_place(s, pii, secrets)) } -/// DLP inbound redaction: emails and 11-digit phone numbers. +/// DLP inbound redaction: emails, 11-digit phone numbers, and — when +/// `detect_secrets` is on — API keys / credentials. pub fn dlp_redact_request(sec: &SecurityConf, request: &mut GatewayRequest) -> usize { - if !sec.dlp_redact { + if !sec.dlp_redact && !sec.detect_secrets { return 0; } + let secrets = sec.detect_secrets; + let pii = sec.dlp_redact; + let mut redact_field = |s: &mut String| redact_in_place(s, pii, secrets); let mut hits = 0; for msg in &mut request.message { - hits += redact_str(&mut msg.content); + hits += redact_field(&mut msg.content); // engines forward `parts` (not `content`) when present, so PII must be // scrubbed inside the parts' text blocks too if let Some(parts) = &mut msg.parts { - hits += redact_parts_text(parts); + hits += for_each_part_text(parts, &mut redact_field); } } // non-chat surfaces carry user text outside `message` (Responses raw body, // family typed params) — scrub those too or they reach the vendor unredacted if let Some(param) = request.model_param_v2.as_mut() { - hits += walk_json_strings(&mut param.raw, &mut redact_str); + hits += walk_json_strings(&mut param.raw, &mut redact_field); if let Some(typed) = param.typed.as_mut() { - hits += for_each_typed_text(typed, &mut redact_str); + hits += for_each_typed_text(typed, &mut redact_field); } } hits } -/// Redact one string in place; the hit count. -fn redact_str(s: &mut String) -> usize { - match redact(s) { - Some((redacted, n)) => { - *s = redacted; - n - } - None => 0, - } +/// A privacy-safe copy of `text` for content retention: PII and secrets are +/// ALWAYS stripped, independent of the tenant's forwarding DLP flags. Retention +/// owns its own redaction so a `redacted` row — or a keyless `full` downgrade — +/// can never persist raw secrets/PII even when inline DLP is disabled. +pub fn redact_retained(text: &str) -> String { + let mut s = text.to_owned(); + redact_in_place(&mut s, true, true); + s } -/// Redact PII inside a multimodal `parts` array's text blocks, in place. -/// Deliberately narrower than the blocklist scan: non-text parts (image URLs, -/// base64 data) are never rewritten, so a redaction can't corrupt them. -fn redact_parts_text(parts: &mut serde_json::Value) -> usize { - let Some(arr) = parts.as_array_mut() else { - return 0; - }; +/// Redact one string in place (email/phone via `pii`, secrets via `secrets`); +/// the hit count. +fn redact_in_place(s: &mut String, pii: bool, secrets: bool) -> usize { let mut hits = 0; - for p in arr { - if p["type"] == "text" - && let Some(t) = p["text"].as_str() - && let Some((redacted, n)) = redact(t) - { - p["text"] = serde_json::Value::String(redacted); - hits += n; - } + if pii && let Some((redacted, n)) = redact(s) { + *s = redacted; + hits += n; + } + if secrets && let Some((redacted, n)) = redact_secrets(s) { + *s = redacted; + hits += n; } hits } +/// Redact one string in place; email/phone only (the outbound response path). +fn redact_str(s: &mut String) -> usize { + redact_in_place(s, true, false) +} + +/// Mask credential shapes (API keys, tokens, private-key headers) with +/// `[REDACTED_SECRET]`. High-precision patterns to avoid mauling normal text. +fn redact_secrets(text: &str) -> Option<(String, usize)> { + static SECRETS: std::sync::LazyLock = std::sync::LazyLock::new(|| { + #[allow(clippy::expect_used)] // a compile-time literal, covered by tests + regex::Regex::new( + r"sk-[A-Za-z0-9_-]{20,}|AKIA[0-9A-Z]{16}|gh[pousr]_[A-Za-z0-9]{36,}|xox[baprs]-[A-Za-z0-9-]{10,}|-----BEGIN [A-Z ]*PRIVATE KEY-----", + ) + .expect("static secret patterns compile") + }); + let mut count = 0; + let out = SECRETS.replace_all(text, |_: ®ex::Captures| { + count += 1; + "[REDACTED_SECRET]" + }); + (count > 0).then(|| (out.into_owned(), count)) +} + /// DLP outbound redaction: the flat `message` plus the structured payloads the /// non-chat surfaces return (`response_v2`, `tool_calls`), so vendor-introduced /// PII can't leak through a field the surface serializes verbatim. @@ -297,6 +456,7 @@ mod tests { SecurityConf { blocklist: vec!["forbiddenword".into()], dlp_redact: true, + ..Default::default() } } @@ -306,10 +466,14 @@ mod tests { message: vec![ChatMsg::text("user", "say ForbiddenWord now")], ..Default::default() }; - let block = security_check(&sec(), &mut req).unwrap(); + let block = security_check(&sec(), &mut req).block.unwrap(); assert!(block.block && block.hit); assert_eq!(block.err_code, 4003); - assert!(security_check(&sec(), &mut GatewayRequest::default()).is_none()); + assert!( + security_check(&sec(), &mut GatewayRequest::default()) + .block + .is_none() + ); } #[test] @@ -317,17 +481,156 @@ mod tests { let s = SecurityConf { blocklist: vec!["forbiddenword".into(), "禁词".into()], dlp_redact: false, + ..Default::default() }; let mut req = GatewayRequest { message: vec![ChatMsg::text("user", "前文 FORBIDDENWORD 后文")], ..Default::default() }; - assert!(security_check(&s, &mut req).is_some()); + assert!(security_check(&s, &mut req).block.is_some()); let mut req = GatewayRequest { message: vec![ChatMsg::text("user", "包含 禁词 的内容")], ..Default::default() }; - assert!(security_check(&s, &mut req).is_some()); + assert!(security_check(&s, &mut req).block.is_some()); + } + + #[test] + fn inbound_text_covers_non_chat_raw_and_typed() { + use gw_models::{ModelParamV2, TypedParams}; + let mut param = ModelParamV2::with_name(gw_consts::Protocol::Responses, "m"); + param.raw = serde_json::json!({"input": "secret responses text", "model": "m"}); + let mut req = GatewayRequest { + model_param_v2: Some(param), + ..Default::default() + }; + assert!( + inbound_text(&mut req).contains("secret responses text"), + "Responses input rides in raw with an empty message" + ); + + let mut param = ModelParamV2::with_name(gw_consts::Protocol::Embeddings, "m"); + param.typed = Some(TypedParams::Embeddings(gw_models::EmbeddingParams { + input: vec!["typed embed text".into()], + dimensions: None, + })); + let mut req = GatewayRequest { + model_param_v2: Some(param), + ..Default::default() + }; + assert!( + inbound_text(&mut req).contains("typed embed text"), + "embeddings input rides in typed with an empty message" + ); + } + + fn ssn_block() -> SecurityConf { + SecurityConf { + regexes: vec![gw_config::CompiledRule { + name: "ssn".into(), + action: Action::Block, + re: regex::Regex::new(r"\d{3}-\d{2}-\d{4}").unwrap(), + }], + ..Default::default() + } + } + + #[test] + fn realtime_frame_scan_honors_regex_and_actions() { + let s = ssn_block(); + let mut frame = serde_json::json!({"type":"input_text","text":"my ssn is 123-45-6789"}); + let (out, text) = realtime_frame_scan(&s, &mut frame, true); + assert!(out.block.is_some(), "regex Block denies on realtime too"); + assert!( + text.contains("123-45-6789"), + "collect_text gathers the frame text in the same walk" + ); + + let s2 = SecurityConf { + blocklist: vec!["watch".into()], + blocklist_action: Action::Flag, + ..Default::default() + }; + let mut frame = serde_json::json!({"type":"input_text","text":"please watch this"}); + let (out, text) = realtime_frame_scan(&s2, &mut frame, false); + assert!(out.block.is_none(), "flag does not block realtime"); + assert_eq!(out.hits.len(), 1); + assert!(text.is_empty(), "no text collected unless asked"); + } + + #[test] + fn realtime_frame_redacts_secrets_when_detect_secrets_on() { + let s = SecurityConf { + detect_secrets: true, + dlp_redact: false, + ..Default::default() + }; + let mut frame = serde_json::json!({ + "type":"input_text","text":"here is sk-abcdefghijklmnopqrstuvwxyz012345" + }); + let n = dlp_redact_realtime_frame(&s, &mut frame); + assert_eq!(n, 1, "secret masked even with dlp_redact off"); + let text = frame["text"].as_str().unwrap(); + assert!( + text.contains("[REDACTED_SECRET]") && !text.contains("sk-abc"), + "{text}" + ); + + let none = SecurityConf::default(); + let mut frame = + serde_json::json!({"type":"input_text","text":"sk-abcdefghijklmnopqrstuvwxyz012345"}); + assert_eq!( + dlp_redact_realtime_frame(&none, &mut frame), + 0, + "both flags off leaves the frame untouched" + ); + } + + #[test] + fn redact_retained_strips_pii_and_secrets_unconditionally() { + let out = + redact_retained("mail john.doe@example.com key sk-abcdefghijklmnopqrstuvwxyz012345"); + assert!(out.contains("[REDACTED_EMAIL]"), "{out}"); + assert!(out.contains("[REDACTED_SECRET]"), "{out}"); + assert!( + !out.contains("example.com") && !out.contains("sk-abc"), + "{out}" + ); + assert_eq!(redact_retained("clean text"), "clean text"); + } + + #[test] + fn flag_action_records_a_hit_without_blocking() { + let s = SecurityConf { + blocklist: vec!["watchword".into()], + blocklist_action: Action::Flag, + ..Default::default() + }; + let mut req = GatewayRequest { + message: vec![ChatMsg::text("user", "contains watchword here")], + ..Default::default() + }; + let out = security_check(&s, &mut req); + assert!(out.block.is_none(), "flag does not deny"); + assert_eq!(out.hits.len(), 1); + assert_eq!(out.hits[0].action, Action::Flag); + } + + #[test] + fn regex_rule_blocks_and_redact_secrets_masks() { + let s = ssn_block(); + let mut req = GatewayRequest { + message: vec![ChatMsg::text("user", "my ssn is 123-45-6789")], + ..Default::default() + }; + assert!( + security_check(&s, &mut req).block.is_some(), + "regex Block denies" + ); + + let (masked, n) = redact_secrets("key sk-abcdefghijklmnopqrstuvwxyz012345 end").unwrap(); + assert_eq!(n, 1); + assert!(masked.contains("[REDACTED_SECRET]") && !masked.contains("sk-abc")); } #[test] diff --git a/crates/models/src/lib.rs b/crates/models/src/lib.rs index f58975b..208e61f 100644 --- a/crates/models/src/lib.rs +++ b/crates/models/src/lib.rs @@ -20,7 +20,7 @@ pub use params::{ VideoParams, }; pub use request::domain::{Account, ChatMsg}; -pub use request::{GatewayRequest, ModelParamV2}; +pub use request::{BatchItem, GatewayRequest, ModelParamV2}; pub use response::GatewayResponse; pub use response::StreamChunk; pub use token_estimate::{HeuristicEncoder, TokenEncoder, estimate_prompt_tokens}; diff --git a/crates/models/src/request.rs b/crates/models/src/request.rs index b9525e8..e260ffa 100644 --- a/crates/models/src/request.rs +++ b/crates/models/src/request.rs @@ -12,6 +12,13 @@ pub struct GatewayRequest { pub model_param_v2: Option, pub ak: String, pub is_online: bool, + /// End-user attribution from request metadata (OpenAI `user` / Anthropic + /// `metadata.user_id` / `x-gw-user`); only trusted when the key is shared — + /// a key's own `owner` overrides it at billing. + pub user_id: Option, + /// Correlation id assigned at ingress; joins access log, ledger, and audit + /// events for one request. Empty until the handler stamps it. + pub request_id: String, /// When set, a streaming-capable engine forwards chunks here as they arrive /// instead of buffering; the bounded channel is the backpressure seam. pub stream_tx: Option>, @@ -29,6 +36,16 @@ impl GatewayRequest { } } +/// One queued batch item: a message list plus the client-supplied end-user +/// attribution. `user` is persisted with the item so a distributed drainer on +/// another instance still attributes and budgets it (owner still overrides at +/// billing time); empty when the submitter gave no `user`/`x-gw-user`. +#[derive(Debug, Clone)] +pub struct BatchItem { + pub messages: Vec, + pub user: String, +} + /// Domain types referenced by `GatewayRequest`. Per-vendor long-tail fields /// ride in `raw`/`extra` passthroughs rather than being individually modeled. pub mod domain { diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 5a883db..77f5106 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -83,6 +83,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 distributed_batches = state.store.distributed_batches(); let transport = select_transport()?; @@ -185,11 +186,16 @@ async fn main() -> anyhow::Result<()> { let listener = tokio::net::TcpListener::bind(&addr).await?; tracing::info!("gw listening on http://{addr}"); - axum::serve(listener, router) - .with_graceful_shutdown(shutdown_signal()) - .await?; + // connect-info so the audit trail can root the source IP at the TCP peer + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(shutdown_signal()) + .await?; quota_task.abort(); + purge_task.abort(); tracing::info!("gw drained and exiting"); Ok(()) } diff --git a/crates/server/tests/e2e.rs b/crates/server/tests/e2e.rs index f51e815..3ce7511 100644 --- a/crates/server/tests/e2e.rs +++ b/crates/server/tests/e2e.rs @@ -27,6 +27,72 @@ fn app() -> Router { )) } +#[tokio::test] +async fn admin_audit_freezes_source_ip_before_a_trust_flip() { + const V1: &str = r#" +listen: {host: 127.0.0.1, port: 0} +admin: {token_env: GW_TEST_ADMIN_TOKEN_TRUST} +trust_proxy_headers: false +access_keys: [{ak: ak-x, product: demo, qps: 100, daily_token_quota: 1000000}] +models: [{name: gpt-4o, protocol: openai-chat}] +accounts: [{name: mock-openai-1, provider: openai, protocols: ["openai-chat"]}] +"#; + // SAFETY: unique var name for this test; no concurrent reader of it. + unsafe { std::env::set_var("GW_TEST_ADMIN_TOKEN_TRUST", "s3cret") }; + let v1 = GatewayConfig::from_yaml(V1).unwrap(); + let v2_yaml = V1.replace("trust_proxy_headers: false", "trust_proxy_headers: true"); + let loader: gw_views::ConfigLoader = Arc::new(move || { + let yaml = v2_yaml.clone(); + Box::pin(async move { GatewayConfig::from_yaml(&yaml).map_err(|e| e.to_string()) }) + as gw_views::ConfigFuture + }); + let state = Arc::new(GatewayState::from_config(&v1)); + let shared = gw_state::SharedConfig::new(Arc::new(v1), state); + let app = gw_views::app(gw_views::AppState::with_config( + shared, + Arc::new(gw_engines::MockTransport), + Some(loader), + )); + + let r = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/admin/reload") + .header("authorization", "Bearer s3cret") + .header("x-real-ip", "9.9.9.9") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(r.status(), StatusCode::OK); + + let ops = app + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/audit/ops") + .header("authorization", "Bearer s3cret") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let j = body_json(ops).await; + let reload = j["entries"] + .as_array() + .unwrap() + .iter() + .find(|e| e["action"] == "reload") + .expect("reload audited"); + assert_ne!( + reload["source_ip"], "9.9.9.9", + "the trust-enabling op must not trust its own forged header" + ); +} + #[tokio::test] async fn admin_reload_is_gated_and_swaps_keys_live() { let r = app() diff --git a/crates/state/Cargo.toml b/crates/state/Cargo.toml index 3ad744c..0fe9db1 100644 --- a/crates/state/Cargo.toml +++ b/crates/state/Cargo.toml @@ -23,6 +23,9 @@ serde = { workspace = true } tokio = { workspace = true } futures = { workspace = true } metrics = { workspace = true } +chacha20poly1305 = { workspace = true } +hex = { workspace = true } +base64 = { workspace = true } [dev-dependencies] tempfile = { workspace = true } diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index bcc6962..1dcbf05 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -25,6 +25,58 @@ pub fn model_quota_key(ak: &str, model: &str) -> String { format!("{ak}|{model}") } +/// The per-user daily-budget counter key, namespaced by tenant so the same +/// user id under two tenants meters separately. +pub fn user_budget_key(tenant: &str, user: &str) -> String { + format!("ub:{tenant}:{user}") +} + +/// The tenant's per-user daily token budget, if it configured one. +pub fn user_budget_limit(cfg: &GatewayConfig, tenant: &str) -> Option { + cfg.find_tenant(tenant) + .and_then(|t| t.user_daily_token_quota) +} + +/// Per-user daily token budget (a soft cap): admit while the user is under the +/// tenant's limit. Skipped when the tenant sets no limit or the request carries +/// no user attribution. Consumed at billing via [`consume_user_budget`]. +pub async fn check_user_budget( + gov: &dyn Governance, + cfg: &GatewayConfig, + tenant: &str, + user: &str, +) -> Result<(), String> { + if user.is_empty() { + return Ok(()); + } + let Some(limit) = user_budget_limit(cfg, tenant) else { + return Ok(()); + }; + if gov.quota_check(&user_budget_key(tenant, user), limit).await { + Ok(()) + } else { + Err(format!("daily token budget exhausted for user `{user}`")) + } +} + +/// Accrue actual usage to the per-user daily budget (no-op without a limit or +/// user). Soft cap: check-then-consume, so a burst can overshoot by one turn. +pub async fn consume_user_budget( + gov: &dyn Governance, + cfg: &GatewayConfig, + tenant: &str, + user: &str, + total: i64, +) { + if user.is_empty() || total <= 0 { + return; + } + if user_budget_limit(cfg, tenant).is_some() { + gov.quota_consume(&user_budget_key(tenant, user), total) + .await; + } +} + /// The per-(AK, model) daily cap: AK override, else tenant default, else none. pub fn model_quota_limit(cfg: &GatewayConfig, ak: &AkInfo, model: &str) -> Option { ak.model_quotas.get(model).copied().or_else(|| { diff --git a/crates/state/src/content.rs b/crates/state/src/content.rs new file mode 100644 index 0000000..721193f --- /dev/null +++ b/crates/state/src/content.rs @@ -0,0 +1,99 @@ +//! Content-at-rest sealing for the retention store: authenticated encryption +//! (XChaCha20-Poly1305) under a process-wide deployment key from +//! `GW_CONTENT_KEY` (64 hex chars = 32 bytes). No key configured → sealing is +//! unavailable, and `full` retention refuses to store raw content. + +use std::sync::LazyLock; + +use base64::Engine as _; +use chacha20poly1305::aead::{Aead, KeyInit, OsRng}; +use chacha20poly1305::{AeadCore, XChaCha20Poly1305}; + +/// The process-wide content key, loaded once from `GW_CONTENT_KEY`. +static CIPHER: LazyLock> = LazyLock::new(|| { + let raw = std::env::var("GW_CONTENT_KEY").ok()?; + match hex::decode(raw.trim()) { + Ok(bytes) if bytes.len() == 32 => Some(XChaCha20Poly1305::new(bytes.as_slice().into())), + _ => { + tracing::error!("GW_CONTENT_KEY is not 64 hex chars (32 bytes); content sealing off"); + None + } + } +}); + +/// One stored prompt or response. +#[derive(Debug, Clone, serde::Serialize)] +pub struct ContentRecord { + pub created_at_epoch_secs: i64, + pub request_id: String, + pub ak: String, + pub user_id: String, + pub tenant: String, + /// "prompt" | "response". + pub kind: String, + /// Sealed (base64 nonce||ciphertext) or plaintext, per [`ContentRecord::sealed`]. + pub content: String, + /// Whether `content` is sealed ciphertext. + pub sealed: bool, + /// Unix seconds after which the purge job deletes this row; 0 = keep. + pub expires_at_epoch_secs: i64, +} + +/// Whether a deployment key is configured (so `full` retention may store raw). +pub fn sealing_available() -> bool { + CIPHER.is_some() +} + +/// Seal `plaintext` into base64(nonce ‖ ciphertext); `None` when no key is set. +pub fn seal(plaintext: &str) -> Option { + seal_with(CIPHER.as_ref()?, plaintext) +} + +/// Reverse [`seal`]; `None` when no key is set or the input is malformed/forged. +pub fn open(sealed: &str) -> Option { + open_with(CIPHER.as_ref()?, sealed) +} + +fn seal_with(cipher: &XChaCha20Poly1305, plaintext: &str) -> Option { + let nonce = XChaCha20Poly1305::generate_nonce(&mut OsRng); + let ct = cipher.encrypt(&nonce, plaintext.as_bytes()).ok()?; + let mut out = nonce.to_vec(); + out.extend_from_slice(&ct); + Some(base64::engine::general_purpose::STANDARD.encode(out)) +} + +fn open_with(cipher: &XChaCha20Poly1305, sealed: &str) -> Option { + let raw = base64::engine::general_purpose::STANDARD + .decode(sealed) + .ok()?; + if raw.len() < 24 { + return None; + } + let (nonce, ct) = raw.split_at(24); + let pt = cipher.decrypt(nonce.into(), ct).ok()?; + String::from_utf8(pt).ok() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn seal_open_roundtrips_and_rejects_tamper() { + let cipher = XChaCha20Poly1305::new((&[7u8; 32]).into()); + let sealed = seal_with(&cipher, "credit card 4111 1111 1111 1111").unwrap(); + assert!(!sealed.contains("4111"), "ciphertext hides the plaintext"); + assert_eq!( + open_with(&cipher, &sealed).unwrap(), + "credit card 4111 1111 1111 1111" + ); + let other = XChaCha20Poly1305::new((&[8u8; 32]).into()); + assert!(open_with(&other, &sealed).is_none(), "wrong key can't open"); + let mut tampered = sealed.clone(); + tampered.push('x'); + assert!( + open_with(&cipher, &tampered).is_none(), + "tamper is rejected" + ); + } +} diff --git a/crates/state/src/keystore.rs b/crates/state/src/keystore.rs index 477be82..6277d93 100644 --- a/crates/state/src/keystore.rs +++ b/crates/state/src/keystore.rs @@ -31,8 +31,11 @@ pub trait KeyStore: Send + Sync + std::fmt::Debug { async fn patch(&self, ak: &str, patch: &KeyPatch) -> GResult>; /// Remove a key regardless of source; whether it existed. async fn revoke(&self, ak: &str) -> GResult; - /// Every key, sorted by ak for stable listings. - async fn list(&self) -> GResult>; + /// A page of keys, sorted by ak, optionally confined to one tenant. The + /// tenant filter applies before paging, so a scoped page is never emptied + /// by a later filter; `offset`/`limit` bound the scan. + async fn list(&self, tenant: Option<&str>, offset: usize, limit: usize) + -> GResult>; /// Re-apply the config file's key set, leaving admin-created keys untouched. async fn reload_config_keys(&self, keys: &[gw_config::AkConf]) -> GResult<()>; } @@ -67,11 +70,16 @@ impl PostgresKeyStore { expires_at_epoch_secs BIGINT, banned BOOLEAN NOT NULL DEFAULT FALSE, model_quotas TEXT NOT NULL DEFAULT '{}', + owner TEXT, source TEXT NOT NULL DEFAULT 'admin')", ) .execute(&pool) .await .map_err(|e| crate::sqlx_err("create access_keys schema", e))?; + sqlx::query("ALTER TABLE access_keys ADD COLUMN IF NOT EXISTS owner TEXT") + .execute(&pool) + .await + .map_err(|e| crate::sqlx_err("migrate access_keys owner", e))?; Ok(Self { pool, cache: moka::sync::Cache::builder() @@ -91,7 +99,7 @@ impl PostgresKeyStore { async fn fetch(&self, ak: &str) -> Result, sqlx::Error> { let row = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas FROM access_keys WHERE ak = $1", + expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys WHERE ak = $1", ) .bind(ak) .fetch_optional(&self.pool) @@ -139,7 +147,7 @@ impl KeyStore for PostgresKeyStore { .map_err(|e| crate::sqlx_err("begin patch", e))?; let row = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas FROM access_keys + expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys WHERE ak = $1 FOR UPDATE", ) .bind(ak) @@ -181,11 +189,20 @@ impl KeyStore for PostgresKeyStore { Ok(n > 0) } - async fn list(&self) -> GResult> { + async fn list( + &self, + tenant: Option<&str>, + offset: usize, + limit: usize, + ) -> GResult> { let rows = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas FROM access_keys ORDER BY ak", + expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys + WHERE ($1::text IS NULL OR tenant = $1) ORDER BY ak LIMIT $2 OFFSET $3", ) + .bind(tenant) + .bind(limit.min(i64::MAX as usize) as i64) + .bind(offset.min(i64::MAX as usize) as i64) .fetch_all(&self.pool) .await .map_err(|e| crate::sqlx_err("list keys", e))?; @@ -229,14 +246,15 @@ async fn upsert( let quotas = serde_json::to_string(&*info.model_quotas).unwrap_or_else(|_| "{}".into()); sqlx::query( "INSERT INTO access_keys (ak, product, tenant, qps, daily_token_quota, - tokens_per_minute, expires_at_epoch_secs, banned, model_quotas, source) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) + tokens_per_minute, expires_at_epoch_secs, banned, model_quotas, owner, source) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) ON CONFLICT (ak) DO UPDATE SET product = EXCLUDED.product, tenant = EXCLUDED.tenant, qps = EXCLUDED.qps, daily_token_quota = EXCLUDED.daily_token_quota, tokens_per_minute = EXCLUDED.tokens_per_minute, expires_at_epoch_secs = EXCLUDED.expires_at_epoch_secs, banned = EXCLUDED.banned, model_quotas = EXCLUDED.model_quotas, + owner = EXCLUDED.owner, source = CASE WHEN access_keys.source = 'config' AND EXCLUDED.source = 'admin' THEN 'config' ELSE EXCLUDED.source END", ) @@ -249,6 +267,7 @@ async fn upsert( .bind(info.expires_at_epoch_secs) .bind(info.banned) .bind("as) + .bind(&info.owner) .bind(source_str(source)) .execute(exec) .await @@ -268,6 +287,7 @@ fn row_to_info(row: &sqlx::postgres::PgRow) -> AkInfo { model_quotas: std::sync::Arc::new( serde_json::from_str(row.get::<&str, _>(8)).unwrap_or_default(), ), + owner: row.get(9), } } @@ -287,6 +307,7 @@ mod tests { ak: ak.into(), product: "p".into(), tenant: "default".into(), + owner: None, qps, daily_token_quota: 10, tokens_per_minute: None, diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 734b4a6..b736a35 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -15,12 +15,14 @@ use gw_models::Account; pub mod admission; pub mod configstore; +pub mod content; pub mod governance; pub mod health; pub mod keystore; pub mod store; pub use configstore::{CONFIG_CHANNEL, PostgresConfigStore}; +pub use content::{ContentRecord, sealing_available}; pub use governance::{Governance, MemoryGovernance, RedisGovernance}; pub use health::{HealthStore, RedisHealth}; pub use keystore::{KeyStore, PostgresKeyStore}; @@ -35,6 +37,8 @@ pub struct AkInfo { pub product: String, /// Tenant this key belongs to (`gw_config::DEFAULT_TENANT` when undeclared). pub tenant: String, + /// End user this key is issued to (one key = one user); `None` = shared key. + pub owner: Option, pub qps: f64, pub daily_token_quota: i64, /// tokens-per-minute window cap; None = unlimited. @@ -60,6 +64,16 @@ impl AkInfo { } } + /// The attributed end user: this key's non-empty `owner` (authoritative), + /// else the caller-supplied `fallback` (request metadata / the realtime + /// `x-gw-user` hint). The one resolution every surface shares. + pub fn attributed_user<'a>(&'a self, fallback: &'a str) -> &'a str { + self.owner + .as_deref() + .filter(|s| !s.is_empty()) + .unwrap_or(fallback) + } + /// Apply a partial quota/lifecycle patch. pub fn apply_patch(&mut self, patch: &KeyPatch) { if let Some(v) = patch.qps { @@ -97,6 +111,7 @@ impl From<&gw_config::AkConf> for AkInfo { ak: k.ak.clone(), product: k.product.clone(), tenant: k.tenant.clone(), + owner: k.owner.clone(), qps: k.qps, daily_token_quota: k.daily_token_quota, tokens_per_minute: k.tokens_per_minute, @@ -165,11 +180,17 @@ impl AkAuth { Some(e.0.clone()) } - /// Every key in the table, sorted by ak for stable listings. - pub fn list(&self) -> Vec { - let mut keys: Vec = self.keys.iter().map(|e| e.value().0.clone()).collect(); + /// A page of keys, sorted by ak (stable), optionally confined to `tenant`, + /// `offset..offset+limit` — the filter applies before paging. + pub fn list(&self, tenant: Option<&str>, offset: usize, limit: usize) -> Vec { + let mut keys: Vec = self + .keys + .iter() + .map(|e| e.value().0.clone()) + .filter(|k| tenant.is_none_or(|t| t == k.tenant)) + .collect(); keys.sort_by(|a, b| a.ak.cmp(&b.ak)); - keys + keys.into_iter().skip(offset).take(limit).collect() } /// Remove a key regardless of source; returns whether it existed. @@ -205,8 +226,13 @@ impl KeyStore for AkAuth { async fn revoke(&self, ak: &str) -> gw_models::GResult { Ok(AkAuth::revoke(self, ak)) } - async fn list(&self) -> gw_models::GResult> { - Ok(AkAuth::list(self)) + async fn list( + &self, + tenant: Option<&str>, + offset: usize, + limit: usize, + ) -> gw_models::GResult> { + Ok(AkAuth::list(self, tenant, offset, limit)) } async fn reload_config_keys(&self, keys: &[gw_config::AkConf]) -> gw_models::GResult<()> { AkAuth::reload_config_keys(self, keys); @@ -850,6 +876,7 @@ mod tests { ak: ak.into(), product: "p".into(), tenant: "default".into(), + owner: None, qps: 1.0, daily_token_quota: 10, tokens_per_minute: None, @@ -859,6 +886,25 @@ mod tests { } } + #[test] + fn attributed_user_prefers_nonempty_owner_else_fallback() { + let mut ak = ak_info("k"); + assert_eq!(ak.attributed_user("hint"), "hint", "no owner → fallback"); + ak.owner = Some(String::new()); + assert_eq!( + ak.attributed_user("hint"), + "hint", + "empty owner is not an identity → fallback" + ); + ak.owner = Some("alice".into()); + assert_eq!( + ak.attributed_user("hint"), + "alice", + "owner wins over fallback" + ); + assert_eq!(ak.attributed_user(""), "alice"); + } + #[test] fn rate_limit_qps_change_takes_effect_immediately() { let rl = RateLimiter::default(); diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 4f83d15..08abe30 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -33,6 +33,16 @@ pub struct BillingRecord { pub ak: String, pub product: String, pub tenant: String, + /// Effective end user: the key's `owner` if set, else request metadata; empty + /// when neither is present. The precise per-user billing dimension. + #[serde(default)] + pub user_id: String, + /// Ingress correlation id, joins this row to the access log and audit events. + #[serde(default)] + pub request_id: String, + /// Unix seconds the call was billed — the billing-period axis. + #[serde(default)] + pub created_at_epoch_secs: i64, /// Public model the caller requested. pub model: String, /// Model that actually served (differs from `model` after a quota fallback). @@ -49,6 +59,10 @@ pub struct BillingRecord { /// PTU spilled over to a paygo account (a failover occurred). #[serde(default)] pub ptu_spillover: bool, + /// Token counts were estimated (an aborted stream billed from delivered + /// text), not read from a vendor usage payload. + #[serde(default)] + pub estimated: bool, } /// Offline batch job status. @@ -125,6 +139,10 @@ pub struct BillingInput<'a> { pub ak: &'a str, pub product: &'a str, pub tenant: &'a str, + /// Effective end user (key owner else request metadata); empty when neither. + pub user_id: &'a str, + /// Ingress correlation id for this request. + pub request_id: &'a str, /// Public model the caller requested (accrues the per-(AK, model) counter). pub requested_model: &'a str, /// Model that actually served — charged at its price (may differ on fallback). @@ -135,6 +153,8 @@ pub struct BillingInput<'a> { pub completion: i64, pub total: i64, pub ptu_spillover: bool, + /// Counts are estimated (aborted stream), not vendor-reported. + pub estimated: bool, } /// Clamp a metered token count into `[0, MAX_METERED_TOKENS]`. @@ -167,6 +187,9 @@ pub fn billing_record(cfg: &gw_config::GatewayConfig, b: &BillingInput) -> Billi ak: b.ak.to_owned(), product: b.product.to_owned(), tenant: b.tenant.to_owned(), + user_id: b.user_id.to_owned(), + request_id: b.request_id.to_owned(), + created_at_epoch_secs: crate::epoch_secs(), model: b.requested_model.to_owned(), served_model: b.served_model.to_owned(), protocol: b.protocol.to_owned(), @@ -177,6 +200,7 @@ pub fn billing_record(cfg: &gw_config::GatewayConfig, b: &BillingInput) -> Billi cost_micros: gw_models::cost_micros(prompt, completion, charged), vendor_cost_micros: gw_models::cost_micros(prompt, completion, vendor), ptu_spillover: b.ptu_spillover, + estimated: b.estimated, } } @@ -193,6 +217,64 @@ pub struct UsageRow { pub vendor_cost_micros: i64, } +/// One row of the per-(user, model) usage rollup over a billing period. +#[derive(Debug, Clone, serde::Serialize)] +pub struct UserUsageRow { + pub user_id: String, + pub model: String, + pub requests: i64, + pub prompt_tokens: i64, + pub completion_tokens: i64, + pub total_tokens: i64, + pub cost_micros: i64, + pub vendor_cost_micros: i64, +} + +/// 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. +#[derive(Debug, Clone, serde::Serialize)] +pub struct SecurityEvent { + pub created_at_epoch_secs: i64, + pub request_id: String, + pub ak: String, + pub user_id: String, + pub tenant: String, + /// Which surface: chat/messages/responses/realtime/… + pub surface: String, + /// The rule family that fired: "blocklist" | "dlp" | a recognizer name. + pub rule: String, + /// What the gateway did: "block" | "redact" | "flag". + pub action: String, + pub hits: i64, +} + +impl SecurityEvent { + /// Append to `store`. Best-effort: a write failure is logged, never fails + /// the request being audited. + pub async fn record(&self, store: &dyn Store) { + if let Err(e) = store.security_event_add(self).await { + tracing::warn!(error = %e, rule = %self.rule, "security event write failed"); + } + } +} + +/// One admin-plane mutation, recorded with who/what/when for compliance. +#[derive(Debug, Clone, serde::Serialize)] +pub struct AdminAudit { + pub created_at_epoch_secs: i64, + /// The admin identity: "global" or the tenant name whose token was used. + pub actor: String, + /// The presented scope: "global" | "tenant". + pub scope: String, + /// The mutation: "key_create" | "key_patch" | "key_delete" | "config_publish" | "reload". + pub action: String, + /// The object acted on (an ak, a config version, …). + pub target: String, + pub summary: String, + pub source_ip: String, +} + #[async_trait::async_trait] pub trait Store: Send + Sync + std::fmt::Debug { async fn ledger_add(&self, r: &BillingRecord) -> GResult<()>; @@ -201,6 +283,37 @@ 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. + async fn usage_by_user( + &self, + tenant: Option<&str>, + user: Option<&str>, + since: i64, + until: i64, + ) -> GResult>; + + /// Append a content-safety event (no prompt text retained). + async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()>; + /// The most recent `limit` security events, newest first; optional tenant filter. + async fn security_events( + &self, + tenant: Option<&str>, + limit: usize, + ) -> GResult>; + + async fn admin_audit_add(&self, e: &AdminAudit) -> GResult<()>; + /// The most recent `limit` admin-audit entries, newest first. + async fn admin_audit_list(&self, limit: usize) -> GResult>; + + /// Store one retained prompt/response record (per-tenant retention policy). + async fn content_add(&self, r: &crate::ContentRecord) -> GResult<()>; + /// 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; + /// 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; @@ -259,12 +372,12 @@ pub trait Store: Send + Sync + std::fmt::Debug { ak: &str, tenant: &str, model: &str, - items: &[Vec], + items: &[gw_models::BatchItem], ) -> GResult { self.batch_create(ak, tenant, model, items.len()).await } /// Load a batch's input items for execution. - async fn batch_load_items(&self, _id: &str) -> GResult>> { + async fn batch_load_items(&self, _id: &str) -> GResult> { Ok(Vec::new()) } /// Claim one pending batch (requeuing stale running ones first); `None` = @@ -285,6 +398,9 @@ pub trait Store: Send + Sync + std::fmt::Debug { #[derive(Debug, Default)] pub struct MemoryStore { records: Mutex>, + sec_events: Mutex>, + audit: Mutex>, + content: Mutex>, files: DashMap, jobs: DashMap, seq: AtomicUsize, @@ -362,6 +478,121 @@ impl Store for MemoryStore { Ok(rollup.into_values().collect()) } + async fn usage_by_user( + &self, + tenant: Option<&str>, + user: Option<&str>, + since: i64, + until: i64, + ) -> GResult> { + 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 + { + continue; + } + 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()) + } + + async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { + self.sec_events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(e.clone()); + Ok(()) + } + + async fn security_events( + &self, + tenant: Option<&str>, + limit: usize, + ) -> GResult> { + let events = self + .sec_events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Ok(events + .iter() + .rev() + .filter(|e| tenant.is_none_or(|t| t == e.tenant)) + .take(limit) + .cloned() + .collect()) + } + + async fn admin_audit_add(&self, e: &AdminAudit) -> GResult<()> { + self.audit + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(e.clone()); + Ok(()) + } + + async fn admin_audit_list(&self, limit: usize) -> GResult> { + let audit = self + .audit + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Ok(audit.iter().rev().take(limit).cloned().collect()) + } + + async fn content_add(&self, r: &crate::ContentRecord) -> GResult<()> { + self.content + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(r.clone()); + Ok(()) + } + + async fn content_purge(&self, now: i64) -> GResult { + let mut content = self + .content + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let before = content.len(); + content.retain(|r| r.expires_at_epoch_secs == 0 || r.expires_at_epoch_secs > now); + Ok((before - content.len()) as u64) + } + + async fn content_for(&self, request_id: &str) -> GResult> { + let content = self + .content + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Ok(content + .iter() + .filter(|r| r.request_id == request_id) + .cloned() + .collect()) + } + async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult { let id = format!( "file-local-{}", @@ -428,51 +659,63 @@ impl Store for MemoryStore { } } -/// Positional row → record shared by the SQL backends (identical SELECT order). -fn row_to_billing<'r, R>(row: &'r R) -> BillingRecord -where - R: sqlx::Row, - usize: sqlx::ColumnIndex, - String: sqlx::Decode<'r, R::Database> + sqlx::Type, - i64: sqlx::Decode<'r, R::Database> + sqlx::Type, - bool: sqlx::Decode<'r, R::Database> + sqlx::Type, -{ - BillingRecord { - ak: row.get(0), - product: row.get(1), - tenant: row.get(2), - model: row.get(3), - served_model: row.get(4), - protocol: row.get(5), - account: row.get(6), - prompt_tokens: row.get(7), - completion_tokens: row.get(8), - total_tokens: row.get(9), - cost_micros: row.get(10), - vendor_cost_micros: row.get(11), - ptu_spillover: row.get(12), - } +/// Positional row → record mappers shared by the SQL backends: fields decode +/// in the SELECT's column order. One macro so the sqlx trait-bound boilerplate +/// lives once. +macro_rules! row_mapper { + ($name:ident -> $ty:path { $($field:ident),+ $(,)? }) => { + fn $name<'r, R>(row: &'r R) -> $ty + where + R: sqlx::Row, + usize: sqlx::ColumnIndex, + String: sqlx::Decode<'r, R::Database> + sqlx::Type, + i64: sqlx::Decode<'r, R::Database> + sqlx::Type, + bool: sqlx::Decode<'r, R::Database> + sqlx::Type, + { + let mut col = 0usize; + $(let $field = row.get(next_col(&mut col));)+ + $ty { $($field),+ } + } + }; } -fn usage_row<'r, R>(row: &'r R) -> UsageRow -where - R: sqlx::Row, - usize: sqlx::ColumnIndex, - String: sqlx::Decode<'r, R::Database> + sqlx::Type, - i64: sqlx::Decode<'r, R::Database> + sqlx::Type, -{ - UsageRow { - tenant: row.get(0), - model: row.get(1), - requests: row.get(2), - prompt_tokens: row.get(3), - completion_tokens: row.get(4), - total_tokens: row.get(5), - cost_micros: row.get(6), - vendor_cost_micros: row.get(7), - } +fn next_col(col: &mut usize) -> usize { + let i = *col; + *col += 1; + i } +row_mapper!(row_to_billing -> BillingRecord { + ak, product, tenant, model, served_model, protocol, account, + prompt_tokens, completion_tokens, total_tokens, cost_micros, + vendor_cost_micros, ptu_spillover, user_id, request_id, + created_at_epoch_secs, estimated, +}); + +row_mapper!(usage_row -> UsageRow { + tenant, model, requests, prompt_tokens, completion_tokens, total_tokens, + cost_micros, vendor_cost_micros, +}); + +row_mapper!(user_usage_row -> UserUsageRow { + user_id, model, requests, prompt_tokens, completion_tokens, total_tokens, + cost_micros, vendor_cost_micros, +}); + +row_mapper!(security_event_row -> SecurityEvent { + created_at_epoch_secs, request_id, ak, user_id, tenant, surface, rule, + action, hits, +}); + +row_mapper!(admin_audit_row -> AdminAudit { + created_at_epoch_secs, actor, scope, action, target, summary, source_ip, +}); + +row_mapper!(content_row -> crate::ContentRecord { + created_at_epoch_secs, request_id, ak, user_id, tenant, kind, content, + sealed, expires_at_epoch_secs, +}); + fn batch_item_row<'r, R>(row: &'r R) -> BatchItemResult where R: sqlx::Row, @@ -524,7 +767,11 @@ impl SqliteStore { 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 DEFAULT 0, - ptu_spillover INTEGER NOT NULL DEFAULT 0)", + ptu_spillover INTEGER NOT NULL DEFAULT 0, + user_id TEXT NOT NULL DEFAULT '', request_id TEXT NOT NULL DEFAULT '', + 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 files ( n INTEGER PRIMARY KEY AUTOINCREMENT, id TEXT UNIQUE NOT NULL, tenant TEXT NOT NULL DEFAULT 'default', @@ -536,6 +783,25 @@ impl SqliteStore { "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)", + "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 '', + user_id TEXT NOT NULL DEFAULT '', tenant TEXT NOT NULL DEFAULT '', + surface TEXT NOT NULL DEFAULT '', rule TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL DEFAULT '', hits INTEGER NOT NULL DEFAULT 0)", + "CREATE TABLE IF NOT EXISTS admin_audit ( + n INTEGER PRIMARY KEY AUTOINCREMENT, created_at_epoch_secs INTEGER NOT NULL, + actor TEXT NOT NULL DEFAULT '', scope TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL DEFAULT '', target TEXT NOT NULL DEFAULT '', + summary TEXT NOT NULL DEFAULT '', source_ip TEXT NOT NULL DEFAULT '')", + "CREATE TABLE IF NOT EXISTS request_content ( + n INTEGER PRIMARY KEY AUTOINCREMENT, created_at_epoch_secs INTEGER NOT NULL, + request_id TEXT NOT NULL DEFAULT '', ak TEXT NOT NULL DEFAULT '', + user_id TEXT NOT NULL DEFAULT '', tenant TEXT NOT NULL DEFAULT '', + kind TEXT NOT NULL DEFAULT '', content TEXT NOT NULL DEFAULT '', + sealed INTEGER NOT NULL DEFAULT 0, expires_at_epoch_secs INTEGER NOT NULL DEFAULT 0)", + "CREATE INDEX IF NOT EXISTS content_expiry_idx ON request_content (expires_at_epoch_secs)", + "CREATE INDEX IF NOT EXISTS content_request_idx ON request_content (request_id)", ] { sqlx::query(ddl) .execute(&pool) @@ -547,6 +813,10 @@ impl SqliteStore { "ALTER TABLE billing ADD COLUMN tenant TEXT NOT NULL DEFAULT 'default'", "ALTER TABLE billing ADD COLUMN served_model TEXT NOT NULL DEFAULT ''", "ALTER TABLE billing ADD COLUMN vendor_cost_micros INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE billing ADD COLUMN user_id TEXT NOT NULL DEFAULT ''", + "ALTER TABLE billing ADD COLUMN request_id TEXT NOT NULL DEFAULT ''", + "ALTER TABLE billing ADD COLUMN created_at_epoch_secs INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE billing ADD COLUMN estimated INTEGER NOT NULL DEFAULT 0", // 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 ''", @@ -576,8 +846,9 @@ impl Store for SqliteStore { sqlx::query( "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs, + estimated) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", ) .bind(&r.ak) .bind(&r.product) @@ -592,6 +863,10 @@ impl Store for SqliteStore { .bind(r.cost_micros) .bind(r.vendor_cost_micros) .bind(r.ptu_spillover) + .bind(&r.user_id) + .bind(&r.request_id) + .bind(r.created_at_epoch_secs) + .bind(r.estimated) .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("insert billing record", e))?; @@ -618,7 +893,8 @@ impl Store for SqliteStore { let mut rows = sqlx::query( "SELECT ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs, + estimated FROM billing ORDER BY n DESC LIMIT ?", ) .bind(limit.min(i64::MAX as usize) as i64) @@ -655,6 +931,144 @@ impl Store for SqliteStore { Ok(rows.iter().map(usage_row).collect()) } + async fn usage_by_user( + &self, + tenant: Option<&str>, + user: Option<&str>, + since: i64, + until: i64, + ) -> GResult> { + let rows = 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", + ) + .bind(tenant) + .bind(user) + .bind(since) + .bind(until) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("roll up user usage", e))?; + Ok(rows.iter().map(user_usage_row).collect()) + } + + async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { + sqlx::query( + "INSERT INTO security_events (created_at_epoch_secs, request_id, ak, user_id, + tenant, surface, rule, action, hits) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(e.created_at_epoch_secs) + .bind(&e.request_id) + .bind(&e.ak) + .bind(&e.user_id) + .bind(&e.tenant) + .bind(&e.surface) + .bind(&e.rule) + .bind(&e.action) + .bind(e.hits) + .execute(&self.pool) + .await + .map_err(|err| crate::sqlx_err("insert security event", err))?; + Ok(()) + } + + async fn security_events( + &self, + tenant: Option<&str>, + limit: usize, + ) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, surface, rule, + action, hits FROM security_events + WHERE (?1 IS NULL OR tenant = ?1) ORDER BY n DESC LIMIT ?2", + ) + .bind(tenant) + .bind(limit.min(i64::MAX as usize) as i64) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read security events", e))?; + Ok(rows.iter().map(security_event_row).collect()) + } + + async fn admin_audit_add(&self, e: &AdminAudit) -> GResult<()> { + sqlx::query( + "INSERT INTO admin_audit (created_at_epoch_secs, actor, scope, action, target, + summary, source_ip) VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .bind(e.created_at_epoch_secs) + .bind(&e.actor) + .bind(&e.scope) + .bind(&e.action) + .bind(&e.target) + .bind(&e.summary) + .bind(&e.source_ip) + .execute(&self.pool) + .await + .map_err(|err| crate::sqlx_err("insert admin audit", err))?; + Ok(()) + } + + async fn admin_audit_list(&self, limit: usize) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, actor, scope, action, target, summary, source_ip + FROM admin_audit ORDER BY n DESC LIMIT ?", + ) + .bind(limit.min(i64::MAX as usize) as i64) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read admin audit", e))?; + Ok(rows.iter().map(admin_audit_row).collect()) + } + + async fn content_add(&self, r: &crate::ContentRecord) -> GResult<()> { + sqlx::query( + "INSERT INTO request_content (created_at_epoch_secs, request_id, ak, user_id, + tenant, kind, content, sealed, expires_at_epoch_secs) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(r.created_at_epoch_secs) + .bind(&r.request_id) + .bind(&r.ak) + .bind(&r.user_id) + .bind(&r.tenant) + .bind(&r.kind) + .bind(&r.content) + .bind(r.sealed) + .bind(r.expires_at_epoch_secs) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("insert content", e))?; + Ok(()) + } + + async fn content_purge(&self, now: i64) -> GResult { + let r = sqlx::query( + "DELETE FROM request_content + WHERE expires_at_epoch_secs > 0 AND expires_at_epoch_secs <= ?", + ) + .bind(now) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("purge 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, + sealed, expires_at_epoch_secs FROM request_content WHERE request_id = ? ORDER BY n", + ) + .bind(request_id) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read content", e))?; + Ok(rows.iter().map(content_row).collect()) + } + 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 @@ -816,7 +1230,30 @@ impl PostgresStore { 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 DEFAULT 0, - ptu_spillover BOOLEAN NOT NULL DEFAULT FALSE)", + ptu_spillover BOOLEAN NOT NULL DEFAULT FALSE, + user_id TEXT NOT NULL DEFAULT '', request_id TEXT NOT NULL DEFAULT '', + 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 security_events ( + n BIGSERIAL PRIMARY KEY, created_at_epoch_secs BIGINT NOT NULL, + request_id TEXT NOT NULL DEFAULT '', ak TEXT NOT NULL DEFAULT '', + user_id TEXT NOT NULL DEFAULT '', tenant TEXT NOT NULL DEFAULT '', + surface TEXT NOT NULL DEFAULT '', rule TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL DEFAULT '', hits BIGINT NOT NULL DEFAULT 0)", + "CREATE TABLE IF NOT EXISTS admin_audit ( + n BIGSERIAL PRIMARY KEY, created_at_epoch_secs BIGINT NOT NULL, + actor TEXT NOT NULL DEFAULT '', scope TEXT NOT NULL DEFAULT '', + action TEXT NOT NULL DEFAULT '', target TEXT NOT NULL DEFAULT '', + summary TEXT NOT NULL DEFAULT '', source_ip TEXT NOT NULL DEFAULT '')", + "CREATE TABLE IF NOT EXISTS request_content ( + n BIGSERIAL PRIMARY KEY, created_at_epoch_secs BIGINT NOT NULL, + request_id TEXT NOT NULL DEFAULT '', ak TEXT NOT NULL DEFAULT '', + user_id TEXT NOT NULL DEFAULT '', tenant TEXT NOT NULL DEFAULT '', + kind TEXT NOT NULL DEFAULT '', content TEXT NOT NULL DEFAULT '', + sealed BOOLEAN NOT NULL DEFAULT FALSE, expires_at_epoch_secs BIGINT NOT NULL DEFAULT 0)", + "CREATE INDEX IF NOT EXISTS content_expiry_idx ON request_content (expires_at_epoch_secs)", + "CREATE INDEX IF NOT EXISTS content_request_idx ON request_content (request_id)", "CREATE TABLE IF NOT EXISTS files ( n BIGSERIAL PRIMARY KEY, id TEXT UNIQUE NOT NULL, tenant TEXT NOT NULL DEFAULT 'default', @@ -831,6 +1268,8 @@ impl PostgresStore { "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 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", @@ -848,13 +1287,18 @@ impl PostgresStore { .await .map_err(|e| crate::sqlx_err("create postgres schema", e))?; } - sqlx::query( - "ALTER TABLE billing ADD COLUMN IF NOT EXISTS - vendor_cost_micros BIGINT NOT NULL DEFAULT 0", - ) - .execute(&pool) - .await - .map_err(|e| crate::sqlx_err("migrate postgres billing schema", e))?; + for ddl in [ + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS vendor_cost_micros BIGINT NOT NULL DEFAULT 0", + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''", + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS request_id TEXT NOT NULL DEFAULT ''", + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS created_at_epoch_secs BIGINT NOT NULL DEFAULT 0", + "ALTER TABLE billing ADD COLUMN IF NOT EXISTS estimated BOOLEAN NOT NULL DEFAULT FALSE", + ] { + sqlx::query(ddl) + .execute(&pool) + .await + .map_err(|e| crate::sqlx_err("migrate postgres billing schema", e))?; + } Ok(Self { pool, ledger_max_rows, @@ -869,8 +1313,9 @@ impl Store for PostgresStore { sqlx::query( "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)", + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs, + estimated) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)", ) .bind(&r.ak) .bind(&r.product) @@ -885,6 +1330,10 @@ impl Store for PostgresStore { .bind(r.cost_micros) .bind(r.vendor_cost_micros) .bind(r.ptu_spillover) + .bind(&r.user_id) + .bind(&r.request_id) + .bind(r.created_at_epoch_secs) + .bind(r.estimated) .execute(&self.pool) .await .map_err(|e| crate::sqlx_err("insert billing record", e))?; @@ -911,7 +1360,8 @@ impl Store for PostgresStore { let mut rows = sqlx::query( "SELECT ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover + vendor_cost_micros, ptu_spillover, user_id, request_id, created_at_epoch_secs, + estimated FROM billing ORDER BY n DESC LIMIT $1", ) .bind(limit.min(i64::MAX as usize) as i64) @@ -955,6 +1405,146 @@ impl Store for PostgresStore { Ok(rows.iter().map(usage_row).collect()) } + async fn usage_by_user( + &self, + tenant: Option<&str>, + user: Option<&str>, + since: i64, + until: i64, + ) -> GResult> { + let rows = 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", + ) + .bind(tenant) + .bind(user) + .bind(since) + .bind(until) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("roll up user usage", e))?; + Ok(rows.iter().map(user_usage_row).collect()) + } + + async fn security_event_add(&self, e: &SecurityEvent) -> GResult<()> { + sqlx::query( + "INSERT INTO security_events (created_at_epoch_secs, request_id, ak, user_id, + tenant, surface, rule, action, hits) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)", + ) + .bind(e.created_at_epoch_secs) + .bind(&e.request_id) + .bind(&e.ak) + .bind(&e.user_id) + .bind(&e.tenant) + .bind(&e.surface) + .bind(&e.rule) + .bind(&e.action) + .bind(e.hits) + .execute(&self.pool) + .await + .map_err(|err| crate::sqlx_err("insert security event", err))?; + Ok(()) + } + + async fn security_events( + &self, + tenant: Option<&str>, + limit: usize, + ) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, request_id, ak, user_id, tenant, surface, rule, + action, hits FROM security_events + WHERE ($1::text IS NULL OR tenant = $1) ORDER BY n DESC LIMIT $2", + ) + .bind(tenant) + .bind(limit.min(i64::MAX as usize) as i64) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read security events", e))?; + Ok(rows.iter().map(security_event_row).collect()) + } + + async fn admin_audit_add(&self, e: &AdminAudit) -> GResult<()> { + 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(e.created_at_epoch_secs) + .bind(&e.actor) + .bind(&e.scope) + .bind(&e.action) + .bind(&e.target) + .bind(&e.summary) + .bind(&e.source_ip) + .execute(&self.pool) + .await + .map_err(|err| crate::sqlx_err("insert admin audit", err))?; + Ok(()) + } + + async fn admin_audit_list(&self, limit: usize) -> GResult> { + let rows = sqlx::query( + "SELECT created_at_epoch_secs, actor, scope, action, target, summary, source_ip + FROM admin_audit ORDER BY n DESC LIMIT $1", + ) + .bind(limit.min(i64::MAX as usize) as i64) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read admin audit", e))?; + Ok(rows.iter().map(admin_audit_row).collect()) + } + + async fn content_add(&self, r: &crate::ContentRecord) -> GResult<()> { + sqlx::query( + "INSERT INTO request_content (created_at_epoch_secs, request_id, ak, user_id, + tenant, kind, content, sealed, expires_at_epoch_secs) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)", + ) + .bind(r.created_at_epoch_secs) + .bind(&r.request_id) + .bind(&r.ak) + .bind(&r.user_id) + .bind(&r.tenant) + .bind(&r.kind) + .bind(&r.content) + .bind(r.sealed) + .bind(r.expires_at_epoch_secs) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("insert content", e))?; + Ok(()) + } + + async fn content_purge(&self, now: i64) -> GResult { + let r = sqlx::query( + "DELETE FROM request_content + WHERE expires_at_epoch_secs > 0 AND expires_at_epoch_secs <= $1", + ) + .bind(now) + .execute(&self.pool) + .await + .map_err(|e| crate::sqlx_err("purge 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, + sealed, expires_at_epoch_secs FROM request_content WHERE request_id = $1 ORDER BY n", + ) + .bind(request_id) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("read content", e))?; + Ok(rows.iter().map(content_row).collect()) + } + async fn file_put(&self, tenant: &str, purpose: &str, content: String) -> GResult { let bytes = content.len(); // consume the sequence explicitly — concurrent writers race a MAX(n)+1 subselect @@ -1155,7 +1745,7 @@ impl Store for PostgresStore { ak: &str, tenant: &str, model: &str, - items: &[Vec], + items: &[gw_models::BatchItem], ) -> GResult { // the batch becomes claimable only once all its items are committed let mut tx = self @@ -1172,15 +1762,18 @@ impl Store for PostgresStore { .fetch_one(&mut *tx) .await .map_err(|e| crate::sqlx_err("insert batch", e))?; - for (idx, msgs) in items.iter().enumerate() { - let json = serde_json::to_string(msgs).unwrap_or_else(|_| "[]".into()); - sqlx::query("INSERT INTO batch_items (batch_id, idx, messages) VALUES ($1, $2, $3)") - .bind(&id) - .bind(idx as i64) - .bind(json) - .execute(&mut *tx) - .await - .map_err(|e| crate::sqlx_err("save batch item", e))?; + for (idx, item) in items.iter().enumerate() { + let json = serde_json::to_string(&item.messages).unwrap_or_else(|_| "[]".into()); + sqlx::query( + "INSERT INTO batch_items (batch_id, idx, messages, user_id) VALUES ($1, $2, $3, $4)", + ) + .bind(&id) + .bind(idx as i64) + .bind(json) + .bind(&item.user) + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("save batch item", e))?; } tx.commit() .await @@ -1196,15 +1789,20 @@ impl Store for PostgresStore { }) } - async fn batch_load_items(&self, id: &str) -> GResult>> { - let rows = sqlx::query("SELECT messages FROM batch_items WHERE batch_id = $1 ORDER BY idx") - .bind(id) - .fetch_all(&self.pool) - .await - .map_err(|e| crate::sqlx_err("load batch items", e))?; + 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", + ) + .bind(id) + .fetch_all(&self.pool) + .await + .map_err(|e| crate::sqlx_err("load batch items", e))?; Ok(rows .iter() - .map(|r| serde_json::from_str(r.get::<&str, _>(0)).unwrap_or_default()) + .map(|r| gw_models::BatchItem { + messages: serde_json::from_str(r.get::<&str, _>(0)).unwrap_or_default(), + user: r.get::(1), + }) .collect()) } @@ -1271,6 +1869,8 @@ mod tests { ak: "k", product: "demo", tenant: "default", + user_id: "u1", + request_id: "req-1", requested_model: "gpt-4o", served_model: "gpt-4o", protocol: "openai-chat", @@ -1279,6 +1879,7 @@ mod tests { completion: i64::MAX, total: i64::MAX, ptu_spillover: false, + estimated: false, }, ); assert_eq!(rec.prompt_tokens, MAX_METERED_TOKENS); @@ -1292,6 +1893,9 @@ mod tests { ak: "ak-t".into(), product: "p".into(), tenant: "default".into(), + user_id: "u1".into(), + request_id: "req-1".into(), + created_at_epoch_secs: 1_000, model: model.into(), served_model: model.into(), protocol: "openai-chat".into(), @@ -1302,6 +1906,7 @@ mod tests { cost_micros: 42, vendor_cost_micros: 7, ptu_spillover: false, + estimated: false, } } @@ -1364,6 +1969,120 @@ mod tests { exercise(&MemoryStore::default()).await; } + async fn exercise_audit(store: &dyn Store) { + let mut a = record("m1"); + a.user_id = "alice".into(); + a.created_at_epoch_secs = 500; + let mut b = record("m1"); + b.user_id = "bob".into(); + b.created_at_epoch_secs = 1_500; + store.ledger_add(&a).await.unwrap(); + store.ledger_add(&b).await.unwrap(); + + let all = store.usage_by_user(None, None, 0, i64::MAX).await.unwrap(); + assert_eq!(all.len(), 2, "two users"); + let alice = store + .usage_by_user(Some("default"), Some("alice"), 0, i64::MAX) + .await + .unwrap(); + assert_eq!(alice.len(), 1); + assert_eq!(alice[0].user_id, "alice"); + assert_eq!(alice[0].total_tokens, 8); + let windowed = store + .usage_by_user(None, None, 1_000, i64::MAX) + .await + .unwrap(); + assert_eq!(windowed.len(), 1, "only bob is in the window"); + assert_eq!(windowed[0].user_id, "bob"); + + store + .security_event_add(&SecurityEvent { + created_at_epoch_secs: 10, + request_id: "req-9".into(), + ak: "ak-t".into(), + user_id: "alice".into(), + tenant: "default".into(), + surface: "openai-chat".into(), + rule: "blocklist".into(), + action: "block".into(), + hits: 1, + }) + .await + .unwrap(); + let events = store.security_events(Some("default"), 10).await.unwrap(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].rule, "blocklist"); + assert!( + store + .security_events(Some("ghost"), 10) + .await + .unwrap() + .is_empty(), + "tenant filter excludes others" + ); + + store + .admin_audit_add(&AdminAudit { + created_at_epoch_secs: 20, + actor: "global".into(), + scope: "global".into(), + action: "key_create".into(), + target: "ak-new".into(), + summary: "tenant=t1".into(), + source_ip: "10.0.0.1".into(), + }) + .await + .unwrap(); + let audit = store.admin_audit_list(10).await.unwrap(); + assert_eq!(audit.len(), 1); + assert_eq!(audit[0].action, "key_create"); + assert_eq!(audit[0].target, "ak-new"); + } + + #[tokio::test] + async fn memory_audit_roundtrip() { + exercise_audit(&MemoryStore::default()).await; + } + + #[tokio::test] + async fn content_retention_stores_and_purges() { + let store = MemoryStore::default(); + let rec = |kind: &str, expires: i64| crate::ContentRecord { + created_at_epoch_secs: 100, + request_id: "req-1".into(), + ak: "ak".into(), + user_id: "u".into(), + tenant: "default".into(), + kind: kind.into(), + content: "hello".into(), + sealed: false, + expires_at_epoch_secs: expires, + }; + store.content_add(&rec("prompt", 200)).await.unwrap(); + store.content_add(&rec("response", 0)).await.unwrap(); + let got = store.content_for("req-1").await.unwrap(); + assert_eq!(got.len(), 2); + + assert_eq!( + store.content_purge(150).await.unwrap(), + 0, + "not yet expired" + ); + assert_eq!(store.content_purge(250).await.unwrap(), 1, "prompt expired"); + let kept = store.content_for("req-1").await.unwrap(); + assert_eq!(kept.len(), 1, "the keep-forever response survives"); + assert_eq!(kept[0].kind, "response"); + } + + #[tokio::test] + async fn sqlite_audit_roundtrip() { + let dir = tempfile::tempdir().unwrap(); + let store = SqliteStore::open(dir.path().join("audit.db").to_str().unwrap()) + .await + .unwrap(); + exercise_audit(&store).await; + } + #[tokio::test] async fn batch_result_is_first_writer_wins() { let store = MemoryStore::default(); @@ -1603,8 +2322,14 @@ mod tests { assert!(store.distributed_batches()); let qmsgs = vec![ - vec![gw_models::ChatMsg::text("user", "one")], - vec![gw_models::ChatMsg::text("user", "two")], + gw_models::BatchItem { + messages: vec![gw_models::ChatMsg::text("user", "one")], + user: "u-one".into(), + }, + gw_models::BatchItem { + messages: vec![gw_models::ChatMsg::text("user", "two")], + user: "u-two".into(), + }, ]; let qjob = store .batch_enqueue("ak-b", "default", "gpt-4o", &qmsgs) @@ -1631,7 +2356,11 @@ mod tests { } let loaded = store.batch_load_items(&qjob.id).await.unwrap(); assert_eq!(loaded.len(), 2); - assert_eq!(loaded[1][0].content, "two"); + assert_eq!(loaded[1].messages[0].content, "two"); + assert_eq!( + loaded[1].user, "u-two", + "per-item user round-trips through pg" + ); let fjob = store .batch_enqueue("ak-f", "default", "gpt-4o", &qmsgs) diff --git a/crates/task/src/lib.rs b/crates/task/src/lib.rs index 42df45b..4ae0d1a 100644 --- a/crates/task/src/lib.rs +++ b/crates/task/src/lib.rs @@ -1,8 +1,6 @@ -//! Local background tasks. -//! -//! Currently just one pure in-process task: periodic AK daily quota reset. -//! Batch job execution lives in gw-handler::offline (spawned on submit) and -//! needs no separate poller. +//! Local background tasks: periodic AK daily quota reset and retained-content +//! purge. Batch job execution lives in gw-handler::offline (spawned on submit) +//! and needs no separate poller. use std::sync::Arc; use std::time::Duration; @@ -11,6 +9,8 @@ use gw_state::GatewayState; /// The production period: once a day. 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); /// Spawn the daily quota reset loop. Returns the join handle (abort to stop). /// `period` is configurable so tests don't wait 24h. @@ -29,6 +29,28 @@ pub fn spawn_quota_reset( }) } +/// Spawn the retained-content purge loop: deletes content rows whose retention +/// window has elapsed. Returns the join handle (abort to stop). +pub fn spawn_content_purge( + 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; + match state.store.content_purge(gw_state::epoch_secs()).await { + Ok(n) if n > 0 => { + tracing::info!(target: "task", purged = n, "content purge") + } + Ok(_) => {} + Err(e) => tracing::warn!(error = %e, "content purge failed"), + } + } + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/views/Cargo.toml b/crates/views/Cargo.toml index 3dac2cc..b72f0c2 100644 --- a/crates/views/Cargo.toml +++ b/crates/views/Cargo.toml @@ -27,6 +27,7 @@ base64 = { workspace = true } [dev-dependencies] tower = { workspace = true } +async-trait = { workspace = true } [lints] workspace = true diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 44fe90f..a93cac4 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -3,6 +3,7 @@ //! structured access-log line per request. use std::convert::Infallible; +use std::fmt::Write as _; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Instant; @@ -32,6 +33,7 @@ use gw_state::{AkInfo, GatewayState}; use serde_json::{Value, json}; const LEDGER_PAGE_DEFAULT: usize = 100; +const KEY_PAGE_DEFAULT: usize = 200; const STREAM_CHANNEL_CAP: usize = 64; /// Per-turn token reserve against the AK daily quota; settled to actuals at billing. const REALTIME_TURN_RESERVE: i64 = 1_000; @@ -120,6 +122,10 @@ pub fn app(state: AppState) -> Router { .route("/admin/config", axum::routing::put(admin_config_put)) .route("/admin/keys", post(admin_key_create).get(admin_key_list)) .route("/admin/usage", get(admin_usage)) + .route("/admin/usage/users", get(admin_usage_users)) + .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/keys/{ak}", axum::routing::patch(admin_key_patch).delete(admin_key_delete), @@ -222,8 +228,10 @@ async fn realtime_ws( }; // select "realtime" so subprotocol-offering clients get a valid handshake let ws = ws.protocols(["realtime"]); + // client attribution hint captured at connect (no per-turn body user field) + let hint = user_header(&headers).unwrap_or_default(); if account.endpoint.is_empty() { - ws.on_upgrade(move |socket| realtime_session(socket, s, ak, model, mt, account.name)) + ws.on_upgrade(move |socket| realtime_session(socket, s, ak, model, mt, account.name, hint)) } else if gw_engines::realtime::is_gemini_realtime(&account.provider) { // no pre-generation gate signal in this dialect — refuse rather than bill after the fact error_response( @@ -234,7 +242,7 @@ async fn realtime_ws( ), ) } else { - ws.on_upgrade(move |socket| realtime_bridge(socket, s, ak, model, mt, account)) + ws.on_upgrade(move |socket| realtime_bridge(socket, s, ak, model, mt, account, hint)) } } @@ -244,10 +252,16 @@ async fn realtime_ws( /// from the admission config when a reload lands mid-turn). struct RealtimeAdmit { ak: AkInfo, + /// Effective attribution user for this turn: the key's owner if set, else + /// the client's connect-time `x-gw-user` hint; empty for an ownerless key + /// with no hint. Captured at admission so billing and budget agree. + user: String, reserved: i64, /// Tokens reserved in the AK TPM window; `None` when the key has no TPM cap. tpm_reserved: Option, at: i64, + /// Per-turn correlation id for the ledger row. + request_id: String, snap: Arc, } @@ -269,7 +283,12 @@ impl RealtimeAdmit { /// mid-stream), and the reserve is a fixed turn estimate. Reserves are taken /// last so a denial never leaves one behind; a failed TPM reserve rolls back /// the daily reserve just taken. -async fn realtime_gate(s: &AppState, ak: &AkInfo, model: &str) -> Result { +async fn realtime_gate( + s: &AppState, + ak: &AkInfo, + model: &str, + hint: &str, +) -> Result { let snap = s.handler.config.load(); let (cfg, state) = (&snap.cfg, &snap.state); let ak = match state.auth.authenticate(&ak.ak).await { @@ -289,6 +308,7 @@ async fn realtime_gate(s: &AppState, ak: &AkInfo, model: &str) -> Result Result "prompt").increment(it as u64); metrics::counter!("gateway_tokens_total", "kind" => "completion").increment(ot as u64); } @@ -373,6 +407,7 @@ async fn realtime_session( model: String, mt: gw_consts::Protocol, account: String, + hint: String, ) { use axum::extract::ws::Message; let send = |v: Value| Message::Text(v.to_string().into()); @@ -394,18 +429,15 @@ async fn realtime_session( .await; continue; }; - // same blocklist + inbound DLP every REST surface runs - let sec = &s.handler.cfg().security; - if let Some(block) = gw_handler::plugins::realtime_frame_blocked(sec, &mut ev) { + if let Err(reason) = rt_inbound_policy(&s, &ak, &hint, &mut ev).await { let _ = socket - .send(send(json!({"type":"error","message": block.message}))) + .send(send(json!({"type":"error","message": reason}))) .await; continue; } - gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut ev); match ev["type"].as_str().unwrap_or_default() { "input_text" => { - let admit = match realtime_gate(&s, &ak, &model).await { + let admit = match realtime_gate(&s, &ak, &model, &hint).await { Ok(a) => a, Err(denied) => { let _ = socket @@ -489,6 +521,7 @@ async fn realtime_bridge( model: String, mt: gw_consts::Protocol, account: gw_models::Account, + hint: String, ) { use axum::extract::ws::Message as CMsg; use futures::{SinkExt, StreamExt}; @@ -542,6 +575,8 @@ async fn realtime_bridge( let mut pending: std::collections::VecDeque = std::collections::VecDeque::new(); // denied server-VAD turn: swallow its upstream frames until its terminal frame let mut suppress = false; + // outbound DLP redactions summed within a turn, recorded once at its boundary + let mut out_redacted = 0i64; loop { tokio::select! { m = cl_rx.next() => { @@ -555,24 +590,26 @@ async fn realtime_bridge( Some(Ok(_)) => continue, // ping/pong handled by the ws stacks }; if let Some(mut frame) = frame { - // same blocklist + inbound DLP every REST surface runs - let sec = &s.handler.cfg().security; - if let Some(block) = gw_handler::plugins::realtime_frame_blocked(sec, &mut frame) { - if cl_tx - .send(CMsg::Text(send_err(block.message).to_string().into())) - .await - .is_err() - { - break; + match rt_inbound_policy(&s, &ak, &hint, &mut frame).await { + Err(reason) => { + if cl_tx + .send(CMsg::Text(send_err(reason).to_string().into())) + .await + .is_err() + { + break; + } + continue; + } + Ok(redacted) => { + if redacted > 0 { + forward = UMsg::text(frame.to_string()); + } } - continue; - } - if gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut frame) > 0 { - forward = UMsg::text(frame.to_string()); } // gate each generation trigger, not every control frame if is_response_create(&frame) { - match realtime_gate(&s, &ak, &model).await { + match realtime_gate(&s, &ak, &model, &hint).await { Ok(admit) => { pending.push_back(admit); generations += 1; @@ -610,6 +647,7 @@ async fn realtime_bridge( }; let mut relay = true; let mut redacted: Option = None; + let mut turn_ended = false; match frame { Some(mut v) => { if suppress { @@ -621,7 +659,7 @@ async fn realtime_bridge( // server-VAD: OpenAI auto-starts a turn with no client // response.create — gate it here like a manual one else if realtime_turn_started(&account.provider, &v) && pending.is_empty() { - match realtime_gate(&s, &ak, &model).await { + match realtime_gate(&s, &ak, &model, &hint).await { Ok(admit) => pending.push_back(admit), Err(denied) => { let _ = up_tx @@ -651,11 +689,14 @@ async fn realtime_bridge( .authenticate(&ak.ak) .await .unwrap_or_else(|| ak.clone()); + let user = billed.attributed_user(&hint).to_owned(); let unreserved = RealtimeAdmit { ak: billed, + user, reserved: 0, tpm_reserved: None, at: gw_state::epoch_secs(), + request_id: gw_handler::new_request_id(), snap, }; bill_realtime_turn( @@ -671,17 +712,28 @@ async fn realtime_bridge( None => {} } recognized += 1; + turn_ended = true; } // outbound DLP, per frame (a span straddling deltas is // beyond a relay that cannot buffer) - if relay - && gw_handler::plugins::dlp_redact_realtime_frame( - &s.handler.cfg().security, + let cfg = s.handler.cfg(); + let n = if relay { + gw_handler::plugins::dlp_redact_realtime_frame( + cfg.security_for(&ak.tenant), &mut v, - ) > 0 - { + ) + } else { + 0 + }; + if n > 0 { redacted = Some(v.to_string()); } + // per-token events would be too hot: sum the turn, record once at its boundary + out_redacted += n as i64; + if turn_ended { + flush_rt_out_dlp(&s, &ak, &hint, out_redacted).await; + out_redacted = 0; + } } // a denied turn's non-JSON output (e.g. audio deltas) is dropped too None => relay = !suppress, @@ -704,6 +756,9 @@ async fn realtime_bridge( for a in pending { a.refund().await; } + // a turn aborted before its boundary (upstream drop) still applied its + // redactions per frame — flush the pending count so the audit isn't lost + flush_rt_out_dlp(&s, &ak, &hint, out_redacted).await; if generations > 0 && recognized == 0 { tracing::warn!( account = %account.name, @@ -740,13 +795,16 @@ fn log_access(surface: &str, ctx: &DagContext, started: Instant) { }) .unwrap_or_default(); let latency = started.elapsed(); + let user_id = ctx.effective_user_id(); metrics::counter!("gateway_tokens_total", "kind" => "prompt").increment(pt.max(0) as u64); metrics::counter!("gateway_tokens_total", "kind" => "completion").increment(ct.max(0) as u64); tracing::info!( target: "access", surface, + request_id = %ctx.request.request_id, ak = %ctx.ak.ak, product = %ctx.ak.product, + user_id, model = %model, protocol = mt, account, @@ -792,10 +850,7 @@ async fn ledger( State(s): State, axum::extract::Query(q): axum::extract::Query>, ) -> Response { - let limit = q - .get("limit") - .and_then(|v| v.parse::().ok()) - .unwrap_or(LEDGER_PAGE_DEFAULT); + let limit = q_num(&q, "limit", LEDGER_PAGE_DEFAULT); match s.handler.state().store.ledger_snapshot(limit).await { Ok((count, records)) => Json(json!({ "count": count, "records": records })).into_response(), Err(e) => gateway_error(e), @@ -822,6 +877,24 @@ async fn accounts(State(s): State) -> Json { Json(resp) } +/// The `x-gw-user` attribution hint; surfaces fall back to the body's own user +/// field. See [`gw_models::GatewayRequest::user_id`] for the trust model. +fn user_header(headers: &HeaderMap) -> Option { + headers + .get("x-gw-user") + .and_then(|v| v.to_str().ok()) + .map(str::to_owned) + .filter(|s| !s.is_empty()) +} + +/// The one request-metadata attribution precedence the REST surfaces apply: +/// `x-gw-user` header, else the dialect's own user field (OpenAI `user`, +/// Anthropic `metadata.user_id`). Batch items invert it — per-item `user` +/// first — so shared-key batches keep per-item attribution. +fn user_hint(headers: &HeaderMap, field: &Value) -> Option { + user_header(headers).or_else(|| field.as_str().map(str::to_owned)) +} + /// AK auth: `Authorization: Bearer ` or `x-api-key: `. The error is /// `(status, message)` so each surface can shape it to its own wire dialect. async fn authenticate(s: &AppState, headers: &HeaderMap) -> Result { @@ -880,6 +953,203 @@ impl AdminScope { AdminScope::Tenant(t) => t == tenant, } } + + /// (actor, scope) for the audit trail: who the token spoke for. + fn audit_identity(&self) -> (&str, &'static str) { + match self { + AdminScope::Global => ("global", "global"), + AdminScope::Tenant(t) => (t.as_str(), "tenant"), + } + } + + /// The tenant a scoped read is confined to: a tenant admin sees only its + /// own; the global admin may narrow with `?tenant=`. + fn tenant_filter(&self, q: &std::collections::HashMap) -> Option { + match self { + AdminScope::Tenant(t) => Some(t.clone()), + AdminScope::Global => q.get("tenant").cloned(), + } + } +} + +/// A numeric query param, or `default` when absent/unparseable. +fn q_num( + q: &std::collections::HashMap, + key: &str, + default: T, +) -> T { + q.get(key).and_then(|v| v.parse().ok()).unwrap_or(default) +} + +/// Write one realtime security event (`user` already resolved). The shared sink +/// for realtime blocklist/regex hits, moderation denials, and inbound DLP hits. +async fn write_rt_event( + s: &AppState, + ak: &AkInfo, + user: &str, + rule: &str, + action: &str, + hits: i64, +) { + gw_state::SecurityEvent { + created_at_epoch_secs: gw_state::epoch_secs(), + request_id: String::new(), + ak: ak.ak.clone(), + user_id: user.to_owned(), + tenant: ak.tenant.clone(), + surface: "realtime".to_owned(), + rule: rule.to_owned(), + action: action.to_owned(), + hits, + } + .record(s.handler.state().store.as_ref()) + .await; +} + +/// Record a turn's summed outbound DLP redactions as one event; no-op at zero. +async fn flush_rt_out_dlp(s: &AppState, ak: &AkInfo, hint: &str, count: i64) { + if count > 0 { + write_rt_event(s, ak, ak.attributed_user(hint), "dlp", "redact", count).await; + } +} + +/// The full inbound content policy for one realtime frame — the same chain +/// every REST surface runs (scan + hit events, moderation, DLP + event), +/// shared by both WebSocket paths. `Err(reason)` denies the frame; `Ok(n)` is +/// the DLP redaction count (n > 0 means the frame was rewritten). +async fn rt_inbound_policy( + s: &AppState, + ak: &AkInfo, + hint: &str, + frame: &mut Value, +) -> Result { + let cfg = s.handler.cfg(); + let sec = cfg.security_for(&ak.tenant); + let (scan, text) = gw_handler::plugins::realtime_frame_scan(sec, frame, sec.moderate); + emit_rt_hits(s, ak, &scan.hits, hint).await; + if let Some(block) = scan.block { + return Err(block.message); + } + if let Some(reason) = realtime_moderate(s, sec, ak, hint, &text).await { + return Err(reason); + } + let redacted = gw_handler::plugins::dlp_redact_realtime_frame(sec, frame); + if redacted > 0 { + write_rt_event( + s, + ak, + ak.attributed_user(hint), + "dlp", + "redact", + redacted as i64, + ) + .await; + } + Ok(redacted) +} + +/// Record a realtime frame's content-safety hits to the security-event stream +/// (parity with the REST surfaces). +async fn emit_rt_hits( + s: &AppState, + ak: &AkInfo, + hits: &[gw_handler::plugins::RuleHit], + hint: &str, +) { + let user = ak.attributed_user(hint); + for hit in hits { + write_rt_event(s, ak, user, &hit.rule, hit.action.as_str(), hit.count).await; + } +} + +/// Moderate a realtime frame's inbound `text` (collected by the frame scan) +/// via the wired moderator — parity with the REST surface, so a moderated +/// tenant can't be bypassed over the WebSocket. `Some(reason)` denies the +/// frame; records a moderation event. +async fn realtime_moderate( + s: &AppState, + sec: &gw_config::SecurityConf, + ak: &AkInfo, + hint: &str, + text: &str, +) -> Option { + if !sec.moderate || text.is_empty() { + return None; + } + let reason = s.handler.moderate_text(sec, text).await?; + write_rt_event(s, ak, ak.attributed_user(hint), "moderation", "block", 1).await; + Some(reason) +} + +/// The caller IP for the admin audit trail, resolved at request entry — before +/// any config mutation the handler performs — so the op that flips +/// `trust_proxy_headers` is audited under the policy in effect when it +/// arrived, not the one it just installed. Empty when the router is driven +/// without connect info (the test harness). +struct AuditSourceIp(String); + +impl axum::extract::FromRequestParts for AuditSourceIp { + type Rejection = std::convert::Infallible; + + async fn from_request_parts( + parts: &mut axum::http::request::Parts, + s: &AppState, + ) -> Result { + let peer = parts + .extensions + .get::>() + .map(|ci| ci.0); + Ok(AuditSourceIp(source_ip( + peer, + &parts.headers, + s.handler.cfg().trust_proxy_headers, + ))) + } +} + +/// The caller IP for the audit trail. Roots at the real TCP `peer`, which a +/// client cannot forge. Only when `trust_proxy` is set (a trusted proxy fronts +/// the gateway) does it read `x-real-ip`, then the RIGHTMOST `x-forwarded-for` +/// hop (the one that proxy appended) — never the leftmost, which a client forges. +fn source_ip(peer: Option, headers: &HeaderMap, trust_proxy: bool) -> String { + if trust_proxy { + if let Some(ip) = headers.get("x-real-ip").and_then(|v| v.to_str().ok()) { + return ip.trim().to_owned(); + } + if let Some(ip) = headers + .get("x-forwarded-for") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.rsplit(',').next()) + { + return ip.trim().to_owned(); + } + } + peer.map(|p| p.ip().to_string()).unwrap_or_default() +} + +/// Record one admin-plane mutation to the audit trail (who/what/when/where). +/// Best-effort: a store failure is logged, never fails the operation. +async fn audit_admin( + s: &AppState, + scope: &AdminScope, + source: String, + action: &str, + target: &str, + summary: String, +) { + let (actor, scope_kind) = scope.audit_identity(); + let entry = gw_state::AdminAudit { + created_at_epoch_secs: gw_state::epoch_secs(), + actor: actor.to_owned(), + scope: scope_kind.to_owned(), + action: action.to_owned(), + target: target.to_owned(), + summary, + source_ip: source, + }; + if let Err(e) = s.handler.state().store.admin_audit_add(&entry).await { + tracing::warn!(error = %e, action, "admin audit write failed"); + } } /// Admin gate: the global token is checked first (a colliding tenant token @@ -940,7 +1210,7 @@ async fn scoped_key( /// The admin surfaces' public view of a key — one shape for PATCH and GET. fn ak_public_json(k: &AkInfo) -> Value { json!({ - "ak": k.ak, "product": k.product, "tenant": k.tenant, + "ak": k.ak, "product": k.product, "tenant": k.tenant, "owner": k.owner, "qps": k.qps, "daily_token_quota": k.daily_token_quota, "tokens_per_minute": k.tokens_per_minute, "expires_at_epoch_secs": k.expires_at_epoch_secs, @@ -976,7 +1246,11 @@ fn ct_eq(a: &str, b: &str) -> bool { /// POST /admin/reload — re-read config from source and swap it in atomically; /// governance, store, health, and cache are preserved. -async fn admin_reload(State(s): State, headers: HeaderMap) -> Response { +async fn admin_reload( + State(s): State, + headers: HeaderMap, + AuditSourceIp(source): AuditSourceIp, +) -> Response { if let Err(r) = require_global_admin(&s, &headers) { return r; } @@ -989,6 +1263,7 @@ async fn admin_reload(State(s): State, headers: HeaderMap) -> Response accounts = cfg.accounts.len(), "config reloaded" ); + audit_admin(&s, &AdminScope::Global, source, "reload", "", String::new()).await; ( StatusCode::OK, Json(json!({ @@ -1009,6 +1284,7 @@ async fn admin_reload(State(s): State, headers: HeaderMap) -> Response async fn admin_key_create( State(s): State, headers: HeaderMap, + AuditSourceIp(source): AuditSourceIp, Json(body): Json, ) -> Response { let scope = match admin_auth(&s, &headers) { @@ -1040,6 +1316,7 @@ async fn admin_key_create( ak: ak.to_owned(), product: product.to_owned(), tenant: tenant.to_owned(), + owner: body["owner"].as_str().map(str::to_owned), qps: body["qps"].as_f64().unwrap_or(0.0), daily_token_quota: body["daily_token_quota"].as_i64().unwrap_or(0), tokens_per_minute: body["tokens_per_minute"].as_i64(), @@ -1065,6 +1342,15 @@ async fn admin_key_create( { return gateway_error(e); } + audit_admin( + &s, + &scope, + source, + "key_create", + ak, + format!("tenant={tenant}"), + ) + .await; ( StatusCode::CREATED, Json(json!({ "ak": ak, "status": "created" })), @@ -1076,6 +1362,7 @@ async fn admin_key_create( async fn admin_key_patch( State(s): State, headers: HeaderMap, + AuditSourceIp(source): AuditSourceIp, Path(ak): Path, Json(body): Json, ) -> Response { @@ -1102,7 +1389,10 @@ async fn admin_key_patch( let patched = s.handler.state().auth.patch(&ak, &patch).await; match patched { Err(e) => gateway_error(e), - Ok(Some(info)) => (StatusCode::OK, Json(ak_public_json(&info))).into_response(), + Ok(Some(info)) => { + audit_admin(&s, &scope, source, "key_patch", &ak, String::new()).await; + (StatusCode::OK, Json(ak_public_json(&info))).into_response() + } Ok(None) => error_response(404, format!("key {ak} not found")), } } @@ -1111,6 +1401,7 @@ async fn admin_key_patch( async fn admin_key_delete( State(s): State, headers: HeaderMap, + AuditSourceIp(source): AuditSourceIp, Path(ak): Path, ) -> Response { let scope = match admin_auth(&s, &headers) { @@ -1122,18 +1413,26 @@ async fn admin_key_delete( } match s.handler.state().auth.revoke(&ak).await { Err(e) => gateway_error(e), - Ok(true) => ( - StatusCode::OK, - Json(json!({ "ak": ak, "status": "revoked" })), - ) - .into_response(), + Ok(true) => { + audit_admin(&s, &scope, source, "key_delete", &ak, String::new()).await; + ( + StatusCode::OK, + Json(json!({ "ak": ak, "status": "revoked" })), + ) + .into_response() + } Ok(false) => error_response(404, format!("key {ak} not found")), } } /// PUT /admin/config — validate, publish to the fleet config store, and reload /// this instance; peers converge via the store's change feed. Global admin only. -async fn admin_config_put(State(s): State, headers: HeaderMap, body: String) -> Response { +async fn admin_config_put( + State(s): State, + headers: HeaderMap, + AuditSourceIp(source): AuditSourceIp, + body: String, +) -> Response { if let Err(r) = require_global_admin(&s, &headers) { return r; } @@ -1150,7 +1449,22 @@ async fn admin_config_put(State(s): State, headers: HeaderMap, body: S Ok(v) => v, Err(e) => return gateway_error(e), }; - match s.reload().await { + // audit before the local reload can fail — the published version already leads the fleet + let reload = s.reload().await; + let detail = match &reload { + Ok(()) => String::new(), + Err(e) => format!("local reload failed: {e}"), + }; + audit_admin( + &s, + &AdminScope::Global, + source, + "config_publish", + &version.to_string(), + detail, + ) + .await; + match reload { Ok(()) => ( StatusCode::OK, Json(json!({ "status": "published", "version": version })), @@ -1163,13 +1477,28 @@ async fn admin_config_put(State(s): State, headers: HeaderMap, body: S } } -/// GET /admin/keys — the key table, scoped: a tenant admin sees only its own keys. -async fn admin_key_list(State(s): State, headers: HeaderMap) -> Response { +/// GET /admin/keys?offset=&limit= — a page of the key table, scoped: a tenant +/// admin sees only its own keys. Paginated so a fleet key table never loads whole. +async fn admin_key_list( + State(s): State, + headers: HeaderMap, + axum::extract::Query(q): axum::extract::Query>, +) -> Response { let scope = match admin_auth(&s, &headers) { Ok(scope) => scope, Err(r) => return r, }; - let listed = match s.handler.state().auth.list().await { + let offset = q_num(&q, "offset", 0); + let limit = q_num(&q, "limit", KEY_PAGE_DEFAULT); + // the scope filters in the store before paging, or a tenant admin's page could come back empty + let tenant = scope.tenant_filter(&q); + let listed = match s + .handler + .state() + .auth + .list(tenant.as_deref(), offset, limit) + .await + { Ok(v) => v, Err(e) => return gateway_error(e), }; @@ -1178,7 +1507,7 @@ async fn admin_key_list(State(s): State, headers: HeaderMap) -> Respon .filter(|k| scope.covers(&k.tenant)) .map(|k| ak_public_json(&k)) .collect(); - let mut resp = json!({ "count": keys.len() }); + let mut resp = json!({ "count": keys.len(), "offset": offset }); resp["keys"] = Value::Array(keys); Json(resp).into_response() } @@ -1194,10 +1523,7 @@ async fn admin_usage( Ok(scope) => scope, Err(r) => return r, }; - let filter = match &scope { - AdminScope::Tenant(t) => Some(t.clone()), - AdminScope::Global => q.get("tenant").cloned(), - }; + let filter = scope.tenant_filter(&q); let usage = match s .handler .state() @@ -1211,6 +1537,164 @@ async fn admin_usage( Json(json!({ "usage": usage })).into_response() } +/// GET /admin/usage/users?user=&since=&until= — precise per-user cost over a +/// billing period, grouped by (user, requested model). Tenant-scoped like +/// [`admin_usage`]; `since`/`until` are unix seconds (default: all time). +async fn admin_usage_users( + State(s): State, + headers: HeaderMap, + axum::extract::Query(q): axum::extract::Query>, +) -> Response { + let scope = match admin_auth(&s, &headers) { + Ok(scope) => scope, + Err(r) => return r, + }; + let tenant = scope.tenant_filter(&q); + let since = q_num(&q, "since", 0); + let until = q_num(&q, "until", i64::MAX); + let usage = match s + .handler + .state() + .store + .usage_by_user( + tenant.as_deref(), + q.get("user").map(String::as_str), + since, + until, + ) + .await + { + Ok(rows) => rows, + Err(e) => return gateway_error(e), + }; + if q.get("format").map(String::as_str) == Some("csv") { + let mut csv = String::from( + "user_id,model,requests,prompt_tokens,completion_tokens,total_tokens,cost_micros,vendor_cost_micros\n", + ); + for u in &usage { + let _ = writeln!( + csv, + "{},{},{},{},{},{},{},{}", + csv_field(&u.user_id), + csv_field(&u.model), + u.requests, + u.prompt_tokens, + u.completion_tokens, + u.total_tokens, + u.cost_micros, + u.vendor_cost_micros, + ); + } + return ([("content-type", "text/csv")], csv).into_response(); + } + Json(json!({ "usage": usage })).into_response() +} + +/// A CSV field, RFC-4180 quoted AND neutralized against spreadsheet formula +/// injection: a field opening with a formula trigger (`= + - @` / tab / CR) is +/// prefixed with `'` so Excel/Sheets treat it as text (the value is +/// attacker-controlled — it can carry a user id). +fn csv_field(s: &str) -> String { + let needs_prefix = s + .chars() + .next() + .is_some_and(|c| matches!(c, '=' | '+' | '-' | '@' | '\t' | '\r')); + let body = if needs_prefix { + format!("'{s}") + } else { + s.to_owned() + }; + if body.contains([',', '"', '\n', '\r']) { + format!("\"{}\"", body.replace('"', "\"\"")) + } else { + body + } +} + +/// GET /admin/audit/events?limit= — content-safety hits (no prompt text), newest +/// first. Tenant-scoped: a tenant admin sees only its own tenant's events. +async fn admin_security_events( + State(s): State, + headers: HeaderMap, + axum::extract::Query(q): axum::extract::Query>, +) -> Response { + let scope = match admin_auth(&s, &headers) { + Ok(scope) => scope, + Err(r) => return r, + }; + let tenant = scope.tenant_filter(&q); + let limit = q_num(&q, "limit", LEDGER_PAGE_DEFAULT); + match s + .handler + .state() + .store + .security_events(tenant.as_deref(), limit) + .await + { + Ok(events) => Json(json!({ "events": events })).into_response(), + Err(e) => gateway_error(e), + } +} + +/// GET /admin/audit/ops?limit= — admin-operation audit trail, newest first. +/// Global admin only (the trail spans all tenants). +async fn admin_audit_ops( + State(s): State, + headers: HeaderMap, + axum::extract::Query(q): axum::extract::Query>, +) -> Response { + if let Err(r) = require_global_admin(&s, &headers) { + return r; + } + let limit = q_num(&q, "limit", LEDGER_PAGE_DEFAULT); + match s.handler.state().store.admin_audit_list(limit).await { + Ok(entries) => Json(json!({ "entries": entries })).into_response(), + Err(e) => gateway_error(e), + } +} + +/// GET /admin/audit/content/{request_id} — the retained prompt/response rows +/// for one request, unsealed when the content key is present (a sealed row +/// without it returns `content: null`). Tenant-scoped like the other reads. +async fn admin_content_get( + State(s): State, + headers: HeaderMap, + Path(request_id): Path, +) -> Response { + let scope = match admin_auth(&s, &headers) { + Ok(scope) => scope, + Err(r) => return r, + }; + let rows = match s.handler.state().store.content_for(&request_id).await { + Ok(rows) => rows, + Err(e) => return gateway_error(e), + }; + let entries: Vec = rows + .into_iter() + .filter(|r| scope.covers(&r.tenant)) + .map(|r| { + let content = if r.sealed { + gw_state::content::open(&r.content) + .map(Value::String) + .unwrap_or(Value::Null) + } else { + Value::String(r.content) + }; + json!({ + "created_at_epoch_secs": r.created_at_epoch_secs, + "kind": r.kind, + "ak": r.ak, + "user_id": r.user_id, + "tenant": r.tenant, + "sealed": r.sealed, + "expires_at_epoch_secs": r.expires_at_epoch_secs, + "content": content, + }) + }) + .collect(); + Json(json!({ "request_id": request_id, "entries": entries })).into_response() +} + 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. @@ -1345,6 +1829,7 @@ async fn chat_completions( ); param.typed = Some(typed); param.raw = Value::Object(body.extra); + let user_id = user_hint(&headers, ¶m.raw["user"]); let request = GatewayRequest { is_online: true, @@ -1352,6 +1837,7 @@ async fn chat_completions( ak: ak.ak.clone(), message: messages, model_param_v2: Some(param), + user_id, ..Default::default() }; @@ -1409,7 +1895,7 @@ fn spawn_stream_pipeline( started: Instant, ) -> tokio::sync::mpsc::Receiver { let (tx, rx) = tokio::sync::mpsc::channel::(STREAM_CHANNEL_CAP); - let dlp = s.handler.cfg().security.dlp_redact; + let dlp = s.handler.cfg().security_for(&ak.tenant).dlp_redact; if !dlp { request.stream_tx = Some(tx.clone()); } @@ -1675,6 +2161,7 @@ async fn messages( ModelParamV2::with_name(gw_consts::Protocol::AnthropicMessages, body.model.clone()); param.typed = Some(typed); param.raw = Value::Object(body.extra); + let user_id = user_hint(&headers, ¶m.raw["metadata"]["user_id"]); let request = GatewayRequest { is_online: true, @@ -1693,6 +2180,7 @@ async fn messages( }) .collect(), model_param_v2: Some(param), + user_id, ..Default::default() }; @@ -1957,6 +2445,7 @@ async fn run_family( mt: gw_consts::Protocol, typed: TypedParams, messages: Vec, + user_id: Option, ) -> Result { let mut param = ModelParamV2::with_name(mt, model); param.typed = Some(typed); @@ -1965,6 +2454,7 @@ async fn run_family( ak: ak.ak.clone(), message: messages, model_param_v2: Some(param), + user_id, ..Default::default() }; match run_pipeline(s, request, ak).await { @@ -2019,6 +2509,7 @@ async fn completions( gw_consts::Protocol::Completions, typed, vec![ChatMsg::text("user", prompt)], + user_hint(&headers, &body["user"]), ) .await { @@ -2066,6 +2557,7 @@ async fn responses( return error_response(400, "input is required"); } let stream = body["stream"].as_bool().unwrap_or(false); + let user_id = user_hint(&headers, &body["user"]); let mut param = ModelParamV2::with_name(gw_consts::Protocol::Responses, model.clone()); param.raw = body; let request = GatewayRequest { @@ -2073,6 +2565,7 @@ async fn responses( stream, ak: ak.ak.clone(), model_param_v2: Some(param), + user_id, ..Default::default() }; @@ -2216,6 +2709,7 @@ async fn embeddings( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_hint(&headers, &body["user"]), ) .await { @@ -2255,6 +2749,7 @@ async fn images_generations( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_hint(&headers, &body["user"]), ) .await { @@ -2297,6 +2792,7 @@ async fn images_edits( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_hint(&headers, &body["user"]), ) .await { @@ -2336,6 +2832,7 @@ async fn audio_speech( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_hint(&headers, &body["user"]), ) .await { @@ -2391,6 +2888,7 @@ async fn audio_transcriptions( gw_consts::Protocol::OpenaiChat, typed, vec![], + user_hint(&headers, &body["user"]), ) .await { @@ -2432,6 +2930,15 @@ async fn batches_submit( }; let mut model = body["model"].as_str().unwrap_or_default().to_owned(); let mut batch_items = Vec::new(); + // batch-level attribution hint; a per-item body `user` overrides it + let hint = user_header(&headers); + let item_user = |v: &Value| { + v["user"] + .as_str() + .or(hint.as_deref()) + .unwrap_or_default() + .to_owned() + }; if let Some(file_id) = body["input_file_id"].as_str() { let found = s.handler.state().store.file_get(file_id).await; @@ -2453,7 +2960,10 @@ async fn batches_submit( if msgs.is_empty() { return error_response(400, "input file line missing a messages array"); } - batch_items.push(BatchItem { messages: msgs }); + batch_items.push(BatchItem { + messages: msgs, + user: item_user(reqbody), + }); } } else if let Some(items) = body["items"].as_array() { for it in items { @@ -2461,7 +2971,10 @@ async fn batches_submit( if msgs.is_empty() { return error_response(400, "each item needs a non-empty messages array"); } - batch_items.push(BatchItem { messages: msgs }); + batch_items.push(BatchItem { + messages: msgs, + user: item_user(it), + }); } } else { return error_response(400, "either items or input_file_id is required"); @@ -2621,6 +3134,239 @@ mod tests { assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); } + #[tokio::test] + async fn user_attribution_and_request_id_land_in_the_ledger() { + let cfg = Arc::new(GatewayConfig::embedded_default().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 req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .header("authorization", "Bearer ak-demo-123") + .header("x-gw-user", "user-42") + .body(Body::from( + r#"{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}"#, + )) + .unwrap(); + let resp = app(app_state).oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let (_, records) = store.ledger_snapshot(10).await.unwrap(); + let row = records.last().expect("a ledger row"); + assert_eq!(row.user_id, "user-42", "x-gw-user attributed the cost"); + assert!(!row.request_id.is_empty(), "request_id stamped"); + assert!(row.created_at_epoch_secs > 0, "created_at stamped"); + let by_user = store + .usage_by_user(None, Some("user-42"), 0, i64::MAX) + .await + .unwrap(); + assert_eq!(by_user.len(), 1); + assert!(by_user[0].total_tokens > 0); + } + + #[test] + fn source_ip_roots_at_peer_and_ignores_forgeable_headers_untrusted() { + let peer: std::net::SocketAddr = "203.0.113.7:5000".parse().unwrap(); + let mut h = HeaderMap::new(); + h.insert("x-real-ip", "10.0.0.5".parse().unwrap()); + h.insert("x-forwarded-for", "1.2.3.4, 10.0.0.9".parse().unwrap()); + assert_eq!( + source_ip(Some(peer), &h, false), + "203.0.113.7", + "untrusted: forgeable headers ignored, the TCP peer wins" + ); + assert_eq!( + source_ip(None, &h, false), + "", + "no peer, no forgeable header" + ); + assert_eq!( + source_ip(Some(peer), &h, true), + "10.0.0.5", + "trusted proxy: x-real-ip wins" + ); + h.remove("x-real-ip"); + assert_eq!(source_ip(Some(peer), &h, true), "10.0.0.9", "rightmost hop"); + } + + #[test] + fn csv_field_neutralizes_formula_injection() { + assert_eq!(csv_field("alice"), "alice"); + assert_eq!( + csv_field("+cmd"), + "'+cmd", + "formula trigger prefixed with '" + ); + assert_eq!( + csv_field("=SUM(A1,A2)"), + "\"'=SUM(A1,A2)\"", + "prefixed AND quoted (has a comma)" + ); + assert_eq!(csv_field("a,b"), "\"a,b\""); + } + + #[tokio::test] + async fn dlp_hit_is_recorded_as_a_security_event() { + let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); + assert!(cfg.security.dlp_redact, "embedded config has DLP on"); + 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 req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .header("authorization", "Bearer ak-demo-123") + .body(Body::from( + r#"{"model":"gpt-4o","messages":[{"role":"user","content":"mail me at a@b.com"}]}"#, + )) + .unwrap(); + let resp = app(app_state).oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let events = store.security_events(None, 10).await.unwrap(); + assert!( + events.iter().any(|e| e.rule == "dlp"), + "an inbound PII redaction was recorded, no prompt text stored" + ); + } + + #[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}]"; + // SAFETY: unique var name for this test; no concurrent reader of it. + unsafe { std::env::set_var("GW_TEST_CONTENT_ADMIN", "s3cret") }; + assert!( + !gw_state::sealing_available(), + "test env has no content key" + ); + 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 router = app(app_state); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .header("authorization", "Bearer k1") + .body(Body::from( + r#"{"model":"gpt-4o","messages":[{"role":"user","content":"here is sk-abcdefghijklmnopqrstuvwxyz012345"}]}"#, + )) + .unwrap(); + let resp = router.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let (_, rows) = store.ledger_snapshot(1).await.unwrap(); + let stored = store.content_for(&rows[0].request_id).await.unwrap(); + let prompt = stored + .iter() + .find(|c| c.kind == "prompt") + .expect("prompt stored"); + assert!(!prompt.sealed, "no key → unsealed"); + assert!( + prompt.content.contains("[REDACTED_SECRET]"), + "secret masked: {}", + prompt.content + ); + for c in &stored { + assert!( + !c.content.contains("sk-abc"), + "raw secret never persisted: {}", + c.content + ); + } + + let read = router + .oneshot( + Request::builder() + .method("GET") + .uri(format!("/admin/audit/content/{}", rows[0].request_id)) + .header("authorization", "Bearer s3cret") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(read.status(), StatusCode::OK); + let j = body_json(read).await; + let entries = j["entries"].as_array().unwrap(); + assert_eq!(entries.len(), 2, "prompt and response rows read back"); + let prompt_entry = entries + .iter() + .find(|e| e["kind"] == "prompt") + .expect("prompt entry"); + assert!( + prompt_entry["content"] + .as_str() + .unwrap() + .contains("[REDACTED_SECRET]"), + "read-back returns the redacted text" + ); + } + + #[derive(Debug)] + struct DenyModerator; + + #[async_trait::async_trait] + impl gw_handler::moderation::Moderator for DenyModerator { + async fn review(&self, _text: &str) -> Result { + Ok(gw_handler::moderation::Verdict::Deny( + "blocked by moderator".into(), + )) + } + } + + #[tokio::test] + async fn realtime_moderates_and_records_inbound_dlp() { + let yaml = "listen: {host: h, port: 1}\nsecurity: {moderate: true, detect_secrets: true}\nmodels: [{name: rt, protocol: realtime}]\naccess_keys: [{ak: k1, product: p, qps: 10, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let handler = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ) + .with_moderator(Arc::new(DenyModerator)); + let offline = OfflineHandler::new(handler.clone()); + let app = AppState { + handler, + offline, + loader: None, + config_store: None, + }; + let ak = app.handler.state().auth.authenticate("k1").await.unwrap(); + let cfg = app.handler.cfg(); + let sec = cfg.security_for(&ak.tenant); + + assert_eq!( + realtime_moderate(&app, sec, &ak, "", "hello there") + .await + .as_deref(), + Some("blocked by moderator") + ); + let mut secret = json!({"type":"input_text","text":"sk-abcdefghijklmnopqrstuvwxyz012345"}); + let n = gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut secret); + assert!(n > 0); + write_rt_event(&app, &ak, ak.attributed_user(""), "dlp", "redact", n as i64).await; + + let events = app + .handler + .state() + .store + .security_events(None, 10) + .await + .unwrap(); + assert!( + events.iter().any(|e| e.rule == "moderation"), + "moderation event" + ); + assert!( + events.iter().any(|e| e.rule == "dlp"), + "inbound realtime DLP event" + ); + } + #[tokio::test] async fn chat_non_stream_ok() { let resp = test_app() @@ -2669,18 +3415,18 @@ mod tests { let gov = || s.handler.state().governance.clone(); let used = || async { gov().quota_used(&ak.ak).await }; - let a1 = realtime_gate(&s, &ak, "gpt-4o").await.expect("admit"); + let a1 = realtime_gate(&s, &ak, "gpt-4o", "").await.expect("admit"); assert_eq!(used().await, REALTIME_TURN_RESERVE, "reserved up front"); bill_realtime_turn(&a1, "gpt-4o", gw_consts::Protocol::Realtime, "acc", 30, 70).await; assert_eq!(used().await, 100, "settled to actual (30 + 70)"); - let a2 = realtime_gate(&s, &ak, "gpt-4o").await.expect("admit"); + let a2 = realtime_gate(&s, &ak, "gpt-4o", "").await.expect("admit"); assert_eq!(used().await, 100 + REALTIME_TURN_RESERVE); gov().quota_settle(&a2.ak.ak, -a2.reserved, a2.at).await; assert_eq!(used().await, 100, "dropped turn refunded whole"); - let a3 = realtime_gate(&s, &ak, "gpt-4o").await.expect("admit"); + let a3 = realtime_gate(&s, &ak, "gpt-4o", "").await.expect("admit"); assert_eq!(used().await, 100 + REALTIME_TURN_RESERVE); let ledger_before = s.handler.state().store.ledger_snapshot(1).await.unwrap().0; bill_realtime_turn(&a3, "gpt-4o", gw_consts::Protocol::Realtime, "acc", 0, 0).await; @@ -2706,14 +3452,14 @@ mod tests { .unwrap(); let gov = s.handler.state().governance.clone(); - let a1 = realtime_gate(&s, &ak, "gpt-4o") + let a1 = realtime_gate(&s, &ak, "gpt-4o", "") .await .expect("first admits"); assert_eq!(a1.tpm_reserved, Some(REALTIME_TURN_RESERVE)); let daily_before = gov.quota_used(&ak.ak).await; assert!( - realtime_gate(&s, &ak, "gpt-4o").await.is_err(), + realtime_gate(&s, &ak, "gpt-4o", "").await.is_err(), "second turn denied by the TPM reserve" ); assert_eq!( @@ -2735,7 +3481,7 @@ mod tests { let s = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); let ak = s.handler.state().auth.authenticate("k-rt").await.unwrap(); - let admit = realtime_gate(&s, &ak, "rt").await.expect("admit"); + let admit = realtime_gate(&s, &ak, "rt", "").await.expect("admit"); s.handler .reload(GatewayConfig::from_yaml(&price(2_000_000)).unwrap()) .await @@ -2749,6 +3495,56 @@ mod tests { ); } + #[tokio::test] + async fn realtime_attributes_user_from_owner_then_header_hint() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: rt, protocol: realtime, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 1000}]\naccess_keys: [{ak: k-shared, product: p, qps: 10, daily_token_quota: 100000}, {ak: k-owned, product: p, qps: 10, daily_token_quota: 100000, owner: bob}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let s = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + + let shared = s + .handler + .state() + .auth + .authenticate("k-shared") + .await + .unwrap(); + let admit = realtime_gate(&s, &shared, "rt", "alice") + .await + .expect("admit"); + assert_eq!(admit.user, "alice", "ownerless key attributes to the hint"); + bill_realtime_turn(&admit, "rt", gw_consts::Protocol::Realtime, "acc", 40, 60).await; + + let owned = s + .handler + .state() + .auth + .authenticate("k-owned") + .await + .unwrap(); + let admit = realtime_gate(&s, &owned, "rt", "mallory") + .await + .expect("admit"); + assert_eq!( + admit.user, "bob", + "owner is authoritative over a spoofed hint" + ); + bill_realtime_turn(&admit, "rt", gw_consts::Protocol::Realtime, "acc", 10, 20).await; + + let (_, records) = s.handler.state().store.ledger_snapshot(2).await.unwrap(); + let users: std::collections::HashSet<&str> = + records.iter().map(|r| r.user_id.as_str()).collect(); + assert!( + users.contains("alice"), + "shared-key turn billed to header hint" + ); + assert!(users.contains("bob"), "owned-key turn billed to owner"); + assert!( + !users.contains("mallory"), + "spoofed hint never overrides owner" + ); + } + #[test] fn finish_reason_mapping_both_directions() { assert_eq!(finish_openai("end_turn"), "stop"); diff --git a/docs/api.md b/docs/api.md index e9fe2a7..2bce4a7 100644 --- a/docs/api.md +++ b/docs/api.md @@ -20,6 +20,11 @@ error shape, so its SDKs can dispatch on it: {"type": "error", "error": {"type": "invalid_request_error", "message": "..."}} ``` +For per-user attribution on a shared key, send `x-gw-user: ` (it also reads +OpenAI's body `user` field and Anthropic's `metadata.user_id`). A key's own +`owner` overrides the hint, so a key issued to one user always bills to that +user. See [Governance](governance.md#per-user-attribution-and-billing). + ## OpenAI-compatible | Method | Path | Notes | @@ -89,10 +94,13 @@ A realtime model bound to an account with a real `endpoint` bridges the session to that vendor's realtime WebSocket: a transparent relay, with the gateway enforcing the same governance chain as the REST path per generation — tenant and AK QPS, product/model QPM, per-(key, model) and daily-token quota, TPM — plus -billing (shared pricing) from the vendor's usage. Content security also applies: -the blocklist gates inbound frames and DLP redacts text fields in both -directions (per frame — a PII span straddling two deltas is beyond a relay -that cannot buffer). Each generation re-checks the +billing (shared pricing) from the vendor's usage. The full content policy also +applies, so the WebSocket is not a bypass: the blocklist, regex recognizers, and +(when enabled) the external moderator gate inbound frames, and DLP — emails, +phone numbers, and credential masking — redacts text fields in both directions +(per frame — a PII span straddling two deltas is beyond a relay that cannot +buffer). Every hit is audited without prompt text; per-user attribution comes +from the `x-gw-user` hint captured at connect. Each generation re-checks the key, so a key banned, expired, or revoked (or a model de-entitled) mid-session stops generating. An endpoint-less account serves a local mock session (OpenAI Realtime event shape) for offline development. @@ -122,14 +130,20 @@ surface on a private network regardless. | POST | `/admin/reload` | re-read config from source and swap it in atomically (global token only) | | PUT | `/admin/config` | validate + publish a new config document to the fleet config store; every instance reloads via the change feed (global token; needs `storage.postgres_url`) | | GET | `/admin/keys` | list keys (a tenant token sees only its own tenant's) | -| POST | `/admin/keys` | create/replace a key: `{ak, product, tenant?, qps, daily_token_quota, tokens_per_minute?, expires_at_epoch_secs?, banned?, model_quotas?}` | +| POST | `/admin/keys` | create/replace a key: `{ak, product, tenant?, owner?, qps, daily_token_quota, tokens_per_minute?, expires_at_epoch_secs?, banned?, model_quotas?}` (`owner` binds the key to one end user — authoritative for attribution) | | PATCH | `/admin/keys/{ak}` | update any of `qps` / `daily_token_quota` / `tokens_per_minute` / `expires_at_epoch_secs` (null clears) / `banned` | | DELETE | `/admin/keys/{ak}` | revoke a key | | GET | `/admin/usage` | ledger rollup by tenant × model (requests, tokens, charged `cost_micros`, `vendor_cost_micros` for margin); `?tenant=` filter for the global token | +| GET | `/admin/usage/users` | per-user cost rollup (user × model) over a billing period: `?since=&until=` (unix secs), `?user=` filter, `?format=csv` export; tenant-scoped | +| 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 | Two token tiers: the global token (`admin.token_env`) manages everything; a -tenant's `admin_token_env` token manages only that tenant's keys and usage -(cross-tenant keys answer 404, reload/config-publish answer 403). +tenant's `admin_token_env` token manages only that tenant's keys, usage, and +content-safety events, scoped to its own tenant (cross-tenant keys answer 404; +reload, config-publish, and the cross-tenant `/admin/audit/ops` trail answer +403). A reload rebuilds the AK table (config keys), models, providers, tenants, and accounts while preserving the runtime seams — governance counters, the durable diff --git a/docs/configuration.md b/docs/configuration.md index 8ec8d43..4878b38 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -6,7 +6,9 @@ One YAML file configures the gateway. Resolution order: 2. otherwise the embedded default (the repo's `conf/gateway.yaml`) `GW_HOST` / `GW_PORT` override `listen.host` / `listen.port` at runtime -(the container image sets `GW_HOST=0.0.0.0`). +(the container image sets `GW_HOST=0.0.0.0`). `GW_CONTENT_KEY` (64 hex chars = +32 bytes) is the deployment key that seals retained content at rest; without it, +`full` retention refuses to store raw text and falls back to redacted. ## Sections @@ -47,6 +49,9 @@ access_keys: - ak: ak-demo-123 # bearer / x-api-key value clients send product: demo # product group (for product-level QPM) tenant: acme # optional; absent = the unrestricted `default` tenant + owner: alice # optional; binds the key to one end user (authoritative + # for per-user attribution; a shared key omits it and + # falls back to the request's `x-gw-user` / `user`) qps: 100 # per-key request rate daily_token_quota: 1000000 tokens_per_minute: 600 # optional TPM window limit @@ -69,12 +74,26 @@ tenants: admin_token_env: ACME_ADMIN_TOKEN # optional tenant-scoped /admin token model_prices: # optional per-model charged-price override for this tenant gpt-4o: {input_price_per_1k_micros: 5000, output_price_per_1k_micros: 20000} + user_daily_token_quota: 100000 # optional soft per-end-user daily cap + security: # optional; overrides the global `security:` WHOLE for this tenant + blocklist: ["forbidden"] + blocklist_action: flag # block | flag | shadow + detect_secrets: true + regex_rules: + - {name: ssn, pattern: '\d{3}-\d{2}-\d{4}', action: block} + retention: # optional prompt/response retention; absent = retain nothing + content: redacted # none | redacted | full (full needs GW_CONTENT_KEY) + days: 30 # purge after N days; 0 = keep until manually purged ``` Keys without a `tenant` join the implicit `default` tenant (no pooled limits, entitled to every model), so a flat config keeps working unchanged. The model catalog (`GET /v1/models`) filters to the caller's entitlement. +`user_daily_token_quota`, `security`, and `retention` are enterprise controls +detailed in [Governance](governance.md); `security` replaces the global policy +outright when present (it is not merged field-by-field). + ### `models` — public model names and dispatch ```yaml @@ -134,9 +153,15 @@ the charged `cost_micros` and margin is queryable per tenant/model via ### `security`, `stability`, `products` ```yaml -security: - dlp_redact: true # redact emails/phone numbers before egress - blocklist: ["badword"] # reject requests containing listed terms +security: # global default; a tenant may override it whole + dlp_redact: true # redact emails/phone numbers, both directions + detect_secrets: true # also mask API keys / credentials in inbound text + blocklist: ["badword"] # reject/flag requests containing listed terms + blocklist_action: block # block (deny) | flag (record) | shadow (trial a rule) + regex_rules: # named recognizers, each with its own action + - {name: ssn, pattern: '\d{3}-\d{2}-\d{4}', action: block} + moderate: false # route inbound text through the wired external moderator + moderation_fail_open: false # on a moderator error: admit (true) or deny (false) stability: failure_threshold: 3 # consecutive failures before an account cools down @@ -147,6 +172,20 @@ products: qpm: 120 # product-level request rate ``` +Every rule that fires (block / flag / DLP / moderation) is recorded without the +prompt text to the security-event stream (`GET /admin/audit/events`). The same +policy runs on the realtime WebSocket, so it is not a bypass. `moderate` needs a +moderator wired into the handler — the default one allows everything. See +[Governance](governance.md#enterprise-content-policy). + +### Top-level flags + +```yaml +trust_proxy_headers: false # audit source IP: false = the real TCP peer (unforgeable); + # true = trust x-real-ip / rightmost x-forwarded-for hop + # (only behind a proxy that sets them) +``` + ## Observability `GET /metrics` serves the Prometheus registry: `gateway_requests_total` diff --git a/docs/deployment.md b/docs/deployment.md index e95e32c..2269ae6 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -24,6 +24,7 @@ GW_CONFIG=/etc/gateway.yaml ./target/release/gw | `GW_HOST` | override `listen.host` (containers set `0.0.0.0`) | | `GW_PORT` | override `listen.port` | | `GW_TRANSPORT` | `mock` (zero egress) / `http` (no mock) / unset (auto-route) | +| `GW_CONTENT_KEY` | 64 hex chars (32 bytes); seals retained content at rest. Without it, `full` retention stores redacted text instead of raw | | `RUST_LOG` | log level, e.g. `info`, `gw_views=debug` | | provider key vars | named by each account's `api_key_env` | diff --git a/docs/governance.md b/docs/governance.md index f9975a1..3fc1d79 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -96,3 +96,57 @@ redacted text replayed** rather than forwarded token-by-token — DLP trades incremental delivery for a guarantee that no unmasked text reaches the client. Turn `dlp_redact` off to keep incremental delivery; note the embedded demo config ships with it on. + +## Per-user attribution and billing + +Every ledger row carries a `user_id`, `request_id`, and `created_at_epoch_secs`. +The effective user is the key's `owner` (one key = one user, the enterprise +model) if set, else request metadata — the `x-gw-user` header, OpenAI's `user` +field, or Anthropic's `metadata.user_id`. `owner` is authoritative; the +metadata hint is only trusted for shared keys. This holds on every surface — +REST, realtime (the `x-gw-user` hint is captured at WS connect), and batch +(each item's `user`, or the submitter's `x-gw-user`, is persisted with the item +so a fleet drainer still attributes and budgets it). `GET /admin/usage/users?user=&since=&until=` +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. + +## Enterprise content policy + +`security:` is global by default; a tenant may override it whole with +`tenants[].security`. Beyond the blocklist and DLP, a policy can carry +`blocklist_action` (`block` denies, `flag` records, `shadow` trials a rule), +named `regex_rules` (each with its own action), and `detect_secrets` (masks API +keys / credentials). Every rule that fires is recorded — without the offending +text — to the security-event stream, queryable at `GET /admin/audit/events`. +Set `moderate: true` to route inbound text through an external moderator wired +into the handler (`moderation_fail_open` picks the posture on a moderator +error). + +The same policy applies on every surface, realtime included: a `/v1/realtime` +WebSocket runs the blocklist, regex rules, DLP redaction, and the external +moderator on inbound frames, so it is not a bypass. Every hit is audited: inbound +block/flag/moderation and DLP redactions are recorded per frame; outbound DLP +redactions (which stream token-by-token) are summed across the turn and recorded +as one event at the turn boundary — the redaction still applies to every frame, +only its audit is aggregated (a store write per token would be too hot). + +## Audit trails + +- **Content-safety events** (`/admin/audit/events`): who/which-rule/what-action, + no prompt text; tenant-scoped. +- **Admin operations** (`/admin/audit/ops`, global admin only): every key CRUD, + config publish, and reload with actor, action, target, and source IP. A + config publish is recorded even when the local reload fails, since the version + is already the fleet's source of truth. The source IP is the real TCP peer; + set `trust_proxy_headers: true` (top level) to instead trust `x-real-ip` / + `x-forwarded-for` — do that ONLY behind a proxy that sets them, or a direct + client could forge the recorded IP. +- **Content retention** (`tenants[].retention`): `none` (default), `redacted` + (PII/secrets stripped), or `full` (raw). Stored in `request_content`, sealed + at rest with XChaCha20-Poly1305 under `GW_CONTENT_KEY` (64 hex chars). + Retention owns its redaction, so `redacted` — and a keyless `full` that falls + back to it — never persists raw secrets/PII even if the tenant forwards + traffic with DLP off. `full` refuses to store raw without a key. `days` sets + 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`). diff --git a/docs/observability.md b/docs/observability.md index 824c81f..c11ce30 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -22,14 +22,26 @@ status codes, protocol/stage names) — no per-key or per-model cardinality. ## Access log One structured line per request goes to stdout (via `tracing`; control level -with `RUST_LOG`), carrying `ak`, `product`, `model`, `protocol`, `account`, -`status`, `prompt_tokens`, `completion_tokens`, `total_tokens`, and -`latency_ms`. +with `RUST_LOG`), carrying `request_id`, `ak`, `product`, `user_id`, `model`, +`protocol`, `account`, `status`, `prompt_tokens`, `completion_tokens`, +`total_tokens`, and `latency_ms`. `request_id` joins the access log to the +ledger row and the audit events for the same request. ## Billing ledger `GET /internal/ledger?limit=N` returns the most recent `N` billing records (newest first); `count` is always the true total, independent of the page size. Records persist when a SQLite store is configured and can be capped with -`storage.ledger_max_rows`. Each record has the access key, product, model, -protocol, account, token counts, cost, and the PTU-spillover flag. +`storage.ledger_max_rows`. Each record carries `request_id`, the access key, +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). + +## Audit trails + +Three operator-facing audit surfaces, all under the gated `/admin` prefix and +covered in [Governance](governance.md#audit-trails): `GET /admin/audit/events` +(content-safety hits, no prompt text), `GET /admin/audit/ops` (admin-plane +mutations with source IP), and `GET /admin/usage/users` (per-user cost). Content +retention, when a tenant enables it, is read back via +`GET /admin/audit/content/{request_id}`.