1use 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
24pub const WRITER_QUEUE: usize = 4096;
27
28pub const POOL_MAX_SIZE_ENV: &str = "LEAN_CTX_PG_POOL_MAX_SIZE";
30
31const POOL_MAX_SIZE_DEFAULT: usize = 8;
36
37fn 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
47pub 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 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
82fn wants_tls(cfg: &tokio_postgres::Config) -> bool {
85 matches!(cfg.get_ssl_mode(), SslMode::Require)
86}
87
88const 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
124pub 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#[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 pub cost_source: &'static str,
156}
157
158fn 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
169const ANONYMOUS_PERSON: &str = "anonymous";
174const DEFAULT_PROJECT: &str = "default";
175
176impl UsageEvent {
177 #[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 #[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
271pub 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
311pub 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
348pub 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
362pub 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
382pub 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
394pub 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
436pub 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}