Skip to main content

lean_ctx/gateway_server/
store.rs

1//! `usage_events` Postgres store (enterprise#17, baseline fields enterprise#18).
2//!
3//! One row per measured LLM turn: who (person/team/project, enterprise#11),
4//! what (provider/model/tokens), what it cost (priced with the shared
5//! `ModelPricing` table) and the counterfactual-baseline inputs that make the
6//! success fee provable (`uncompressed_input_tokens`, `reference_model`,
7//! `reference_cost_usd`, `is_local` — Doc 08 §2).
8//!
9//! Schema management follows the repo rule: `init_schema` is idempotent
10//! `batch_execute` DDL (`CREATE TABLE IF NOT EXISTS …`), no migration files.
11//!
12//! The writer consumes the `proxy::usage_sink` stream: bounded channel, spawned
13//! task, INSERT per event. Fail-open (enterprise#12): insert errors are logged
14//! and counted, never propagated to the request path.
15
16use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
17use tokio_postgres::NoTls;
18use tokio_postgres::config::SslMode;
19
20use crate::core::config::BaselineConfig;
21use crate::core::gain::model_pricing::ModelPricing;
22use crate::proxy::usage::RealUsage;
23
24/// Buffered events between the proxy choke-point and the Postgres writer.
25/// Sized for bursts (a full channel drops events, counted in `usage_sink`).
26pub const WRITER_QUEUE: usize = 4096;
27
28/// Env var overriding the store pool's `max_size` (chart: `database.poolMaxSize`).
29pub const POOL_MAX_SIZE_ENV: &str = "LEAN_CTX_PG_POOL_MAX_SIZE";
30
31/// Default pool size. The writer is a single sequential task (one connection),
32/// the rest serves the admin API/report queries — 8 is comfortable for one
33/// replica; K8s replicas each get their own pool (load-test: deploy-repo
34/// `docs/ops/load-test.md`).
35const POOL_MAX_SIZE_DEFAULT: usize = 8;
36
37/// Pool size from `LEAN_CTX_PG_POOL_MAX_SIZE`, clamped to a sane band.
38/// Invalid/unset values fall back to the default — a typo can never produce
39/// a 1-connection or 10k-connection pool.
40fn pool_max_size() -> usize {
41    std::env::var(POOL_MAX_SIZE_ENV)
42        .ok()
43        .and_then(|v| v.trim().parse::<usize>().ok())
44        .map_or(POOL_MAX_SIZE_DEFAULT, |n| n.clamp(2, 64))
45}
46
47/// Builds the store pool from a `DATABASE_URL`, honoring `sslmode` (#54/#58).
48///
49/// - `sslmode=disable`/`prefer`/unset: plain TCP (the pilot/in-cluster case;
50///   `prefer`'s opportunistic upgrade would mask misconfiguration, so it stays
51///   plain — deployments that need TLS must say `require`).
52/// - `sslmode=require`: rustls with the webpki root store — the managed-
53///   Postgres case (Azure/AWS/GCP enforce TLS). Unlike libpq's `require`,
54///   certificate and hostname are **always verified** (verify-full rigor);
55///   lean-ctx does not implement an unverified-TLS downgrade.
56pub fn pool_from_database_url(database_url: &str) -> anyhow::Result<Pool> {
57    let pg_cfg: tokio_postgres::Config = database_url.parse()?;
58    let mgr_cfg = ManagerConfig {
59        recycling_method: RecyclingMethod::Fast,
60    };
61    let mgr = if wants_tls(&pg_cfg) {
62        // The pool is built before the proxy installs the process-default
63        // CryptoProvider (#597) — make sure one exists (idempotent).
64        let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
65        let roots = rustls::RootCertStore {
66            roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
67        };
68        let tls_cfg = rustls::ClientConfig::builder()
69            .with_root_certificates(roots)
70            .with_no_client_auth();
71        Manager::from_config(
72            pg_cfg,
73            tokio_postgres_rustls::MakeRustlsConnect::new(tls_cfg),
74            mgr_cfg,
75        )
76    } else {
77        Manager::from_config(pg_cfg, NoTls, mgr_cfg)
78    };
79    Ok(Pool::builder(mgr).max_size(pool_max_size()).build()?)
80}
81
82/// True when the URL's `sslmode` asks for TLS. tokio-postgres 0.7 models
83/// `disable`/`prefer`/`require`; anything else fails URL parsing upstream.
84fn wants_tls(cfg: &tokio_postgres::Config) -> bool {
85    matches!(cfg.get_ssl_mode(), SslMode::Require)
86}
87
88/// Idempotent DDL (Doc 08 §2): `IF NOT EXISTS` only, run on every start.
89/// Columns added after the first release ride along as idempotent
90/// `ALTER TABLE … ADD COLUMN IF NOT EXISTS` — same rule, no migration files.
91const USAGE_EVENTS_DDL: &str = r"
92CREATE TABLE IF NOT EXISTS usage_events (
93  id                 BIGSERIAL PRIMARY KEY,
94  ts                 TIMESTAMPTZ      NOT NULL DEFAULT now(),
95  person             TEXT             NOT NULL,
96  team               TEXT,
97  project            TEXT             NOT NULL,
98  tool               TEXT,
99  provider           TEXT             NOT NULL,
100  model              TEXT             NOT NULL,
101  routed_from        TEXT,
102  input_tokens       BIGINT           NOT NULL,
103  output_tokens      BIGINT           NOT NULL,
104  cache_read_tokens  BIGINT           NOT NULL DEFAULT 0,
105  cache_write_tokens BIGINT           NOT NULL DEFAULT 0,
106  reasoning_tokens   BIGINT           NOT NULL DEFAULT 0,
107  cost_usd           DOUBLE PRECISION NOT NULL,
108  saved_tokens       BIGINT           NOT NULL DEFAULT 0,
109  saved_usd          DOUBLE PRECISION NOT NULL DEFAULT 0,
110  -- Avoided-cost baseline for the success fee (enterprise#18, Doc 04 §6):
111  uncompressed_input_tokens BIGINT    NOT NULL DEFAULT 0,
112  reference_model    TEXT,
113  reference_cost_usd DOUBLE PRECISION NOT NULL DEFAULT 0,
114  is_local           BOOLEAN          NOT NULL DEFAULT false,
115  -- Cost provenance (#1179): provider | shadow | list | live | heuristic.
116  cost_source        TEXT             NOT NULL DEFAULT 'list'
117);
118ALTER TABLE usage_events ADD COLUMN IF NOT EXISTS cost_source TEXT NOT NULL DEFAULT 'list';
119CREATE INDEX IF NOT EXISTS idx_usage_events_person_ts  ON usage_events (person, ts);
120CREATE INDEX IF NOT EXISTS idx_usage_events_project_ts ON usage_events (project, ts);
121CREATE INDEX IF NOT EXISTS idx_usage_events_model_ts   ON usage_events (model, ts);
122";
123
124/// Applies the usage-store DDL. Safe to run on every start (idempotent).
125pub async fn init_schema(pool: &Pool) -> anyhow::Result<()> {
126    let client = pool.get().await?;
127    client.batch_execute(USAGE_EVENTS_DDL).await?;
128    Ok(())
129}
130
131/// One `usage_events` row, fully derived from a finalized [`RealUsage`].
132#[derive(Debug, Clone, PartialEq)]
133pub struct UsageEvent {
134    pub person: String,
135    pub team: Option<String>,
136    pub project: String,
137    pub provider: String,
138    pub model: String,
139    pub routed_from: Option<String>,
140    pub input_tokens: i64,
141    pub output_tokens: i64,
142    pub cache_read_tokens: i64,
143    pub cache_write_tokens: i64,
144    pub reasoning_tokens: i64,
145    pub cost_usd: f64,
146    pub saved_tokens: i64,
147    pub saved_usd: f64,
148    pub uncompressed_input_tokens: i64,
149    pub reference_model: Option<String>,
150    pub reference_cost_usd: f64,
151    pub is_local: bool,
152    /// Where `cost_usd` came from (#1179): `provider` (the response's own
153    /// charge), `shadow` (local shadow rate), `live` (current provider price
154    /// list), `list` (embedded list price) or `heuristic` (estimate).
155    pub cost_source: &'static str,
156}
157
158/// Maps a pricing match onto the stored `cost_source` value for table-priced
159/// rows. `provider`/`shadow` are decided before pricing is consulted.
160fn cost_source_of(kind: crate::core::gain::model_pricing::PricingMatchKind) -> &'static str {
161    use crate::core::gain::model_pricing::PricingMatchKind as K;
162    match kind {
163        K::Exact => "list",
164        K::Live => "live",
165        K::Alias | K::Heuristic | K::Fallback => "heuristic",
166    }
167}
168
169/// Identity fallbacks when a request carried no gateway key/tags: the row must
170/// still be attributable (`NOT NULL`), and "anonymous/default" is honest about
171/// what the gateway knew. Strict deployments make keys mandatory via
172/// `proxy_require_token` + gateway-keys, so these appear only in solo mode.
173const ANONYMOUS_PERSON: &str = "anonymous";
174const DEFAULT_PROJECT: &str = "default";
175
176impl UsageEvent {
177    /// Derives the row from a measured turn, pricing both the actual cost and
178    /// the compression saving with the shared pricing table, and stamping the
179    /// counterfactual baseline (enterprise#15/#18):
180    ///
181    /// - `cost_usd`, in precedence order (#1179): the provider's own reported
182    ///   charge (`usage.cost`, OpenRouter) — except `is_local`, which books
183    ///   the transparent `local_shadow_rate` (never $0; Doc 04 §6) — then the
184    ///   live/list price table. `cost_source` records which one applied.
185    /// - `reference_cost_usd`: the request's **uncompressed** input tokens
186    ///   priced at the contract-frozen `reference_model`'s input rate (Doc 08
187    ///   §2) — the counterfactual the avoided-cost ledger settles against.
188    /// - `saved_usd`: the SEE (compression) component only — saved request
189    ///   tokens at the served model's input rate. Full mechanism attribution
190    ///   (routing/caching) is the signed ledger's job (wave 4, enterprise#19).
191    #[must_use]
192    pub fn from_usage(
193        usage: &RealUsage,
194        pricing: &ModelPricing,
195        baseline: &BaselineConfig,
196    ) -> Self {
197        let wire = usage.wire.as_deref();
198        let quote = pricing.quote(Some(&usage.model));
199        let is_local = wire.is_some_and(|w| w.is_local);
200        #[allow(clippy::cast_precision_loss)]
201        let (cost_usd, cost_source) = if is_local {
202            let billable = usage.input_tokens
203                + usage.output_tokens
204                + usage.cache_read_tokens
205                + usage.cache_write_tokens;
206            (
207                baseline.effective_local_shadow_rate() / 1_000_000.0 * billable as f64,
208                "shadow",
209            )
210        } else if let Some(measured) = usage.provider_cost_usd {
211            (measured, "provider")
212        } else {
213            let estimated = quote.cost.estimate_usd(
214                usage.input_tokens,
215                usage.output_tokens,
216                usage.cache_write_tokens,
217                usage.cache_read_tokens,
218            );
219            (estimated, cost_source_of(quote.match_kind))
220        };
221        let saved_tokens = wire.map_or(0, |w| w.saved_tokens);
222        // Input-side saving: input-rate USD per token × saved request tokens.
223        #[allow(clippy::cast_precision_loss)]
224        let saved_usd = quote.cost.input_per_m / 1_000_000.0 * saved_tokens as f64;
225
226        let uncompressed_input_tokens = wire.map_or(0, |w| w.uncompressed_input_tokens);
227        let reference_model = baseline
228            .reference_model
229            .as_deref()
230            .map(str::trim)
231            .filter(|m| !m.is_empty())
232            .map(str::to_string);
233        #[allow(clippy::cast_precision_loss)]
234        let reference_cost_usd = reference_model.as_deref().map_or(0.0, |reference| {
235            pricing.quote(Some(reference)).cost.input_per_m / 1_000_000.0
236                * uncompressed_input_tokens as f64
237        });
238
239        Self {
240            person: wire
241                .and_then(|w| w.person.clone())
242                .unwrap_or_else(|| ANONYMOUS_PERSON.to_string()),
243            team: wire.and_then(|w| w.team.clone()),
244            project: wire
245                .and_then(|w| w.project.clone())
246                .unwrap_or_else(|| DEFAULT_PROJECT.to_string()),
247            provider: wire.map_or_else(String::new, |w| w.provider.clone()),
248            model: usage.model.clone(),
249            routed_from: wire.and_then(|w| w.routed_from.clone()),
250            input_tokens: to_i64(usage.input_tokens),
251            output_tokens: to_i64(usage.output_tokens),
252            cache_read_tokens: to_i64(usage.cache_read_tokens),
253            cache_write_tokens: to_i64(usage.cache_write_tokens),
254            reasoning_tokens: to_i64(usage.reasoning_tokens),
255            cost_usd,
256            saved_tokens: to_i64(saved_tokens),
257            saved_usd,
258            uncompressed_input_tokens: to_i64(uncompressed_input_tokens),
259            reference_model,
260            reference_cost_usd,
261            is_local,
262            cost_source,
263        }
264    }
265}
266
267fn to_i64(v: u64) -> i64 {
268    i64::try_from(v).unwrap_or(i64::MAX)
269}
270
271/// Inserts one event. Errors bubble to the writer loop, which logs and moves on.
272pub async fn insert_event(
273    client: &deadpool_postgres::Client,
274    e: &UsageEvent,
275) -> anyhow::Result<()> {
276    client
277        .execute(
278            "INSERT INTO usage_events \
279             (person, team, project, provider, model, routed_from, \
280              input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, \
281              reasoning_tokens, cost_usd, saved_tokens, saved_usd, \
282              uncompressed_input_tokens, reference_model, reference_cost_usd, is_local, \
283              cost_source) \
284             VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19)",
285            &[
286                &e.person,
287                &e.team,
288                &e.project,
289                &e.provider,
290                &e.model,
291                &e.routed_from,
292                &e.input_tokens,
293                &e.output_tokens,
294                &e.cache_read_tokens,
295                &e.cache_write_tokens,
296                &e.reasoning_tokens,
297                &e.cost_usd,
298                &e.saved_tokens,
299                &e.saved_usd,
300                &e.uncompressed_input_tokens,
301                &e.reference_model,
302                &e.reference_cost_usd,
303                &e.is_local,
304                &e.cost_source,
305            ],
306        )
307        .await?;
308    Ok(())
309}
310
311/// Current-window spend sums for the budget gate (enterprise#25):
312/// per-person spend for the running UTC day and per-project spend for the
313/// running UTC month, straight from `usage_events`.
314pub async fn budget_window_sums(
315    pool: &Pool,
316) -> anyhow::Result<(
317    std::collections::HashMap<String, f64>,
318    std::collections::HashMap<String, f64>,
319)> {
320    let client = pool.get().await?;
321    let mut person_day = std::collections::HashMap::new();
322    for row in client
323        .query(
324            "SELECT person, SUM(cost_usd) FROM usage_events \
325             WHERE ts >= date_trunc('day', now() AT TIME ZONE 'utc') AT TIME ZONE 'utc' \
326             GROUP BY person",
327            &[],
328        )
329        .await?
330    {
331        person_day.insert(row.get::<_, String>(0), row.get::<_, f64>(1));
332    }
333    let mut project_month = std::collections::HashMap::new();
334    for row in client
335        .query(
336            "SELECT project, SUM(cost_usd) FROM usage_events \
337             WHERE ts >= date_trunc('month', now() AT TIME ZONE 'utc') AT TIME ZONE 'utc' \
338             GROUP BY project",
339            &[],
340        )
341        .await?
342    {
343        project_month.insert(row.get::<_, String>(0), row.get::<_, f64>(1));
344    }
345    Ok((person_day, project_month))
346}
347
348/// Deletes `usage_events` rows older than `days` (enterprise#36). Returns the
349/// number of purged rows. `days == 0` is rejected by the caller (retention
350/// disabled), never here — this function always deletes what it is told.
351pub async fn purge_events_older_than(pool: &Pool, days: u32) -> anyhow::Result<u64> {
352    let client = pool.get().await?;
353    let purged = client
354        .execute(
355            "DELETE FROM usage_events WHERE ts < now() - make_interval(days => $1)",
356            &[&i32::try_from(days).unwrap_or(i32::MAX)],
357        )
358        .await?;
359    Ok(purged)
360}
361
362/// All events attributed to one of `person_keys` (raw + pseudonym, GDPR
363/// Art. 15 export), as self-describing JSON rows.
364pub async fn person_events(
365    pool: &Pool,
366    person_keys: &[String],
367) -> anyhow::Result<Vec<serde_json::Value>> {
368    let client = pool.get().await?;
369    let rows = client
370        .query(
371            "SELECT to_jsonb(usage_events) FROM usage_events \
372             WHERE person = ANY($1) ORDER BY ts",
373            &[&person_keys],
374        )
375        .await?;
376    Ok(rows
377        .into_iter()
378        .map(|r| r.get::<_, serde_json::Value>(0))
379        .collect())
380}
381
382/// Deletes all events of `person_keys` (GDPR Art. 17). Returns rows removed.
383pub async fn delete_person_events(pool: &Pool, person_keys: &[String]) -> anyhow::Result<u64> {
384    let client = pool.get().await?;
385    let deleted = client
386        .execute(
387            "DELETE FROM usage_events WHERE person = ANY($1)",
388            &[&person_keys],
389        )
390        .await?;
391    Ok(deleted)
392}
393
394/// Daily evidence aggregates for the export window (enterprise#36): bounded
395/// output regardless of event volume, yet fine-grained enough for an EU-AI-Act
396/// usage-evidence audit (per day × person × project × model).
397pub async fn evidence_rows(
398    pool: &Pool,
399    from: chrono::DateTime<chrono::Utc>,
400    to: chrono::DateTime<chrono::Utc>,
401) -> anyhow::Result<Vec<serde_json::Value>> {
402    let client = pool.get().await?;
403    let rows = client
404        .query(
405            "SELECT jsonb_build_object(
406               'date', to_char(date_trunc('day', ts AT TIME ZONE 'utc'), 'YYYY-MM-DD'),
407               'person', person,
408               'project', project,
409               'model', model,
410               'provider', provider,
411               'requests', count(*),
412               'input_tokens', sum(input_tokens)::BIGINT,
413               'output_tokens', sum(output_tokens)::BIGINT,
414               'cache_read_tokens', sum(cache_read_tokens)::BIGINT,
415               'cost_usd', round(sum(cost_usd)::numeric, 6),
416               'saved_usd', round(sum(saved_usd)::numeric, 6),
417               'reference_cost_usd', round(sum(reference_cost_usd)::numeric, 6),
418               'local_requests', count(*) FILTER (WHERE is_local),
419               'measured_requests', count(*) FILTER (WHERE cost_source = 'provider'),
420               'estimated_requests', count(*) FILTER (WHERE cost_source = 'heuristic')
421             )
422             FROM usage_events WHERE ts >= $1 AND ts <= $2
423             GROUP BY
424               date_trunc('day', ts AT TIME ZONE 'utc'), person, project, model, provider
425             ORDER BY
426               date_trunc('day', ts AT TIME ZONE 'utc'), person, project, model, provider",
427            &[&from, &to],
428        )
429        .await?;
430    Ok(rows
431        .into_iter()
432        .map(|r| r.get::<_, serde_json::Value>(0))
433        .collect())
434}
435
436/// Wires the usage stream into Postgres: installs the process-wide sink
437/// (`proxy::usage_sink`) and spawns the writer task. Call once at gateway
438/// startup, after `init_schema`.
439///
440/// Returns `false` when a sink was already installed (double start).
441pub fn spawn_writer(pool: Pool) -> bool {
442    let (tx, mut rx) = tokio::sync::mpsc::channel::<RealUsage>(WRITER_QUEUE);
443    if !crate::proxy::usage_sink::install(tx) {
444        return false;
445    }
446    tokio::spawn(async move {
447        // One pricing table + baseline for the writer's lifetime: rows are
448        // priced at insert time (the ledger re-values against frozen
449        // references); the baseline is contract-frozen anyway (#41).
450        let pricing = ModelPricing::load();
451        let baseline = crate::core::config::Config::load().proxy.baseline.clone();
452        while let Some(usage) = rx.recv().await {
453            let event = UsageEvent::from_usage(&usage, &pricing, &baseline);
454            match pool.get().await {
455                Ok(client) => {
456                    if let Err(e) = insert_event(&client, &event).await {
457                        tracing::warn!("usage_events insert failed (fail-open): {e:#}");
458                    }
459                }
460                Err(e) => {
461                    tracing::warn!("usage_events pool unavailable (fail-open): {e:#}");
462                }
463            }
464        }
465    });
466    true
467}
468
469#[cfg(test)]
470mod tests {
471    use super::*;
472    use crate::proxy::usage::WireContext;
473
474    fn usage_with_wire(wire: Option<Box<WireContext>>) -> RealUsage {
475        RealUsage {
476            model: "claude-sonnet-4-5".into(),
477            input_tokens: 1000,
478            output_tokens: 500,
479            cache_read_tokens: 200,
480            cache_write_tokens: 100,
481            reasoning_tokens: 50,
482            provider_cost_usd: None,
483            cohort: None,
484            wire,
485        }
486    }
487
488    #[test]
489    fn event_carries_identity_and_baseline_fields() {
490        let usage = usage_with_wire(Some(Box::new(WireContext {
491            provider: "Anthropic".into(),
492            person: Some("yves".into()),
493            team: Some("platform".into()),
494            project: Some("ai-gateway".into()),
495            saved_tokens: 4000,
496            uncompressed_input_tokens: 5000,
497            is_local: false,
498            routed_from: Some("claude-opus-4-5".into()),
499            counterfactual: None,
500        })));
501        let event = UsageEvent::from_usage(
502            &usage,
503            &ModelPricing::load(),
504            &BaselineConfig {
505                reference_model: Some("claude-opus-4.5".into()),
506                local_shadow_rate_per_mtok: None,
507            },
508        );
509
510        assert_eq!(event.person, "yves");
511        assert_eq!(event.team.as_deref(), Some("platform"));
512        assert_eq!(event.project, "ai-gateway");
513        assert_eq!(event.provider, "Anthropic");
514        assert_eq!(event.model, "claude-sonnet-4-5");
515        assert_eq!(event.routed_from.as_deref(), Some("claude-opus-4-5"));
516        assert_eq!(event.input_tokens, 1000);
517        assert_eq!(event.saved_tokens, 4000);
518        assert_eq!(event.uncompressed_input_tokens, 5000);
519        assert!(!event.is_local);
520        assert!(event.cost_usd > 0.0, "known model must be priced");
521        assert_eq!(
522            event.cost_source, "list",
523            "exact table match books list price"
524        );
525        assert!(
526            event.saved_usd > 0.0,
527            "saved tokens on a priced model must yield saved USD"
528        );
529        // Counterfactual (enterprise#15): 5000 uncompressed input tokens at
530        // claude-opus-4.5's $5/MTok input rate = $0.025.
531        assert_eq!(event.reference_model.as_deref(), Some("claude-opus-4.5"));
532        assert!((event.reference_cost_usd - 0.025).abs() < 1e-9);
533    }
534
535    #[test]
536    fn provider_reported_cost_beats_the_price_table() {
537        // #1179: OpenRouter's `usage.cost` is the bill — the table estimate for
538        // these tokens (claude-sonnet at list price) would be ~50× higher and
539        // must NOT be booked when a measured figure exists.
540        let mut usage = usage_with_wire(Some(Box::new(WireContext {
541            provider: "openrouter".into(),
542            person: Some("nicolas".into()),
543            team: None,
544            project: Some("bot".into()),
545            saved_tokens: 0,
546            uncompressed_input_tokens: 1000,
547            is_local: false,
548            routed_from: None,
549            counterfactual: None,
550        })));
551        usage.provider_cost_usd = Some(0.0123);
552        let event =
553            UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
554        assert!((event.cost_usd - 0.0123).abs() < 1e-12);
555        assert_eq!(event.cost_source, "provider");
556
557        // A measured zero (":free" model) is a real price, not a missing one.
558        usage.provider_cost_usd = Some(0.0);
559        let event =
560            UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
561        assert_eq!(event.cost_usd, 0.0);
562        assert_eq!(event.cost_source, "provider");
563    }
564
565    #[test]
566    fn unknown_model_is_marked_heuristic_and_local_shadow_beats_measured() {
567        // An unpriced model falls into the blended fallback → cost_source
568        // must say so instead of presenting the estimate as exact.
569        let mut usage = usage_with_wire(None);
570        usage.model = "vendor/brand-new-model-20990101".into();
571        let event =
572            UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
573        assert_eq!(event.cost_source, "heuristic");
574
575        // Local turns book the shadow rate even if a cost slipped through —
576        // the shadow rate is the contract for local inference (Doc 04 §6).
577        let mut usage = usage_with_wire(Some(Box::new(WireContext {
578            provider: "ollama".into(),
579            person: None,
580            team: None,
581            project: None,
582            saved_tokens: 0,
583            uncompressed_input_tokens: 0,
584            is_local: true,
585            routed_from: None,
586            counterfactual: None,
587        })));
588        usage.provider_cost_usd = Some(9.99);
589        let event =
590            UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
591        assert_eq!(event.cost_source, "shadow");
592        assert!(
593            event.cost_usd < 1.0,
594            "shadow rate, not the stray measured figure"
595        );
596    }
597
598    #[test]
599    fn event_without_wire_context_uses_honest_fallbacks() {
600        let event = UsageEvent::from_usage(
601            &usage_with_wire(None),
602            &ModelPricing::load(),
603            &BaselineConfig::default(),
604        );
605        assert_eq!(event.person, ANONYMOUS_PERSON);
606        assert_eq!(event.project, DEFAULT_PROJECT);
607        assert_eq!(event.team, None);
608        assert_eq!(event.saved_tokens, 0);
609        assert_eq!(event.saved_usd, 0.0);
610        assert_eq!(event.uncompressed_input_tokens, 0);
611        assert!(!event.is_local);
612        // No reference model configured → no counterfactual claimed.
613        assert_eq!(event.reference_model, None);
614        assert_eq!(event.reference_cost_usd, 0.0);
615    }
616
617    #[test]
618    fn local_usage_books_shadow_rate_never_zero() {
619        // enterprise#15/#18: local inference is billed via the transparent
620        // shadow rate — savings against local models stay honest, not infinite.
621        let usage = usage_with_wire(Some(Box::new(WireContext {
622            provider: "ollama".into(),
623            person: Some("yves".into()),
624            team: None,
625            project: None,
626            saved_tokens: 0,
627            uncompressed_input_tokens: 2000,
628            is_local: true,
629            routed_from: None,
630            counterfactual: None,
631        })));
632        let event =
633            UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
634        assert!(event.is_local);
635        // billable = 1000 in + 500 out + 200 cache-read + 100 cache-write =
636        // 1800 tokens × $0.25/MTok default shadow rate.
637        assert!((event.cost_usd - 0.25 / 1_000_000.0 * 1800.0).abs() < 1e-12);
638        assert!(event.cost_usd > 0.0, "local cost must never be zero");
639
640        // A configured rate wins; a zero/negative config falls back to default.
641        let cfg = BaselineConfig {
642            reference_model: None,
643            local_shadow_rate_per_mtok: Some(1.0),
644        };
645        let event = UsageEvent::from_usage(&usage, &ModelPricing::load(), &cfg);
646        assert!((event.cost_usd - 1.0 / 1_000_000.0 * 1800.0).abs() < 1e-12);
647        let zero = BaselineConfig {
648            reference_model: None,
649            local_shadow_rate_per_mtok: Some(0.0),
650        };
651        assert!(zero.effective_local_shadow_rate() > 0.0);
652    }
653
654    #[test]
655    fn sslmode_selects_tls_and_pool_builds_for_both() {
656        // require → TLS connector; disable/unset → plain (#54/#58).
657        let tls: tokio_postgres::Config = "postgres://u:p@db.example.com:5432/app?sslmode=require"
658            .parse()
659            .unwrap();
660        assert!(wants_tls(&tls));
661        let plain: tokio_postgres::Config = "postgres://u:p@localhost:5432/app".parse().unwrap();
662        assert!(!wants_tls(&plain));
663        let disabled: tokio_postgres::Config = "postgres://u:p@localhost:5432/app?sslmode=disable"
664            .parse()
665            .unwrap();
666        assert!(!wants_tls(&disabled));
667
668        // Pool construction (no connection attempt) must succeed on both paths —
669        // this exercises the rustls config + root store wiring.
670        assert!(
671            pool_from_database_url("postgres://u:p@db.example.com:5432/app?sslmode=require")
672                .is_ok()
673        );
674        assert!(pool_from_database_url("postgres://u:p@localhost:5432/app").is_ok());
675    }
676
677    #[test]
678    fn pool_size_env_is_clamped_and_falls_back() {
679        // Env mutation is serialized process-wide through test_env_lock().
680        let _guard = crate::core::data_dir::test_env_lock();
681        crate::test_env::remove_var(POOL_MAX_SIZE_ENV);
682        assert_eq!(pool_max_size(), 8, "unset -> default");
683        crate::test_env::set_var(POOL_MAX_SIZE_ENV, "24");
684        assert_eq!(pool_max_size(), 24, "explicit value wins");
685        crate::test_env::set_var(POOL_MAX_SIZE_ENV, "0");
686        assert_eq!(pool_max_size(), 2, "clamped low");
687        crate::test_env::set_var(POOL_MAX_SIZE_ENV, "9999");
688        assert_eq!(pool_max_size(), 64, "clamped high");
689        crate::test_env::set_var(POOL_MAX_SIZE_ENV, "not-a-number");
690        assert_eq!(pool_max_size(), 8, "garbage -> default, never panic");
691        crate::test_env::remove_var(POOL_MAX_SIZE_ENV);
692    }
693
694    #[test]
695    fn schema_ddl_is_idempotent_by_construction() {
696        // The gateway runs this DDL on every start against a live database, so
697        // every CREATE must carry IF NOT EXISTS, and post-release columns ride
698        // along as ALTER TABLE … ADD COLUMN IF NOT EXISTS.
699        for stmt in ["CREATE TABLE", "CREATE INDEX"] {
700            for (i, _) in USAGE_EVENTS_DDL.match_indices(stmt) {
701                let tail = &USAGE_EVENTS_DDL[i..(i + stmt.len() + 14).min(USAGE_EVENTS_DDL.len())];
702                assert!(
703                    tail.contains("IF NOT EXISTS"),
704                    "non-idempotent DDL statement: {tail}"
705                );
706            }
707        }
708        for (i, _) in USAGE_EVENTS_DDL.match_indices("ADD COLUMN") {
709            let tail = &USAGE_EVENTS_DDL[i..(i + 25).min(USAGE_EVENTS_DDL.len())];
710            assert!(
711                tail.contains("IF NOT EXISTS"),
712                "non-idempotent ALTER statement: {tail}"
713            );
714        }
715        // And the baseline fields (enterprise#18) are part of the schema.
716        for col in [
717            "uncompressed_input_tokens",
718            "reference_model",
719            "reference_cost_usd",
720            "is_local",
721            "cost_source",
722        ] {
723            assert!(
724                USAGE_EVENTS_DDL.contains(col),
725                "baseline column {col} missing from schema"
726            );
727        }
728    }
729}