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"
90CREATE TABLE IF NOT EXISTS usage_events (
91 id BIGSERIAL PRIMARY KEY,
92 ts TIMESTAMPTZ NOT NULL DEFAULT now(),
93 person TEXT NOT NULL,
94 team TEXT,
95 project TEXT NOT NULL,
96 tool TEXT,
97 provider TEXT NOT NULL,
98 model TEXT NOT NULL,
99 routed_from TEXT,
100 input_tokens BIGINT NOT NULL,
101 output_tokens BIGINT NOT NULL,
102 cache_read_tokens BIGINT NOT NULL DEFAULT 0,
103 cache_write_tokens BIGINT NOT NULL DEFAULT 0,
104 reasoning_tokens BIGINT NOT NULL DEFAULT 0,
105 cost_usd DOUBLE PRECISION NOT NULL,
106 saved_tokens BIGINT NOT NULL DEFAULT 0,
107 saved_usd DOUBLE PRECISION NOT NULL DEFAULT 0,
108 -- Avoided-cost baseline for the success fee (enterprise#18, Doc 04 §6):
109 uncompressed_input_tokens BIGINT NOT NULL DEFAULT 0,
110 reference_model TEXT,
111 reference_cost_usd DOUBLE PRECISION NOT NULL DEFAULT 0,
112 is_local BOOLEAN NOT NULL DEFAULT false
113);
114CREATE INDEX IF NOT EXISTS idx_usage_events_person_ts ON usage_events (person, ts);
115CREATE INDEX IF NOT EXISTS idx_usage_events_project_ts ON usage_events (project, ts);
116CREATE INDEX IF NOT EXISTS idx_usage_events_model_ts ON usage_events (model, ts);
117";
118
119pub async fn init_schema(pool: &Pool) -> anyhow::Result<()> {
121 let client = pool.get().await?;
122 client.batch_execute(USAGE_EVENTS_DDL).await?;
123 Ok(())
124}
125
126#[derive(Debug, Clone, PartialEq)]
128pub struct UsageEvent {
129 pub person: String,
130 pub team: Option<String>,
131 pub project: String,
132 pub provider: String,
133 pub model: String,
134 pub routed_from: Option<String>,
135 pub input_tokens: i64,
136 pub output_tokens: i64,
137 pub cache_read_tokens: i64,
138 pub cache_write_tokens: i64,
139 pub reasoning_tokens: i64,
140 pub cost_usd: f64,
141 pub saved_tokens: i64,
142 pub saved_usd: f64,
143 pub uncompressed_input_tokens: i64,
144 pub reference_model: Option<String>,
145 pub reference_cost_usd: f64,
146 pub is_local: bool,
147}
148
149const ANONYMOUS_PERSON: &str = "anonymous";
154const DEFAULT_PROJECT: &str = "default";
155
156impl UsageEvent {
157 #[must_use]
170 pub fn from_usage(
171 usage: &RealUsage,
172 pricing: &ModelPricing,
173 baseline: &BaselineConfig,
174 ) -> Self {
175 let wire = usage.wire.as_deref();
176 let quote = pricing.quote(Some(&usage.model));
177 let is_local = wire.is_some_and(|w| w.is_local);
178 #[allow(clippy::cast_precision_loss)]
179 let cost_usd = if is_local {
180 let billable = usage.input_tokens
181 + usage.output_tokens
182 + usage.cache_read_tokens
183 + usage.cache_write_tokens;
184 baseline.effective_local_shadow_rate() / 1_000_000.0 * billable as f64
185 } else {
186 quote.cost.estimate_usd(
187 usage.input_tokens,
188 usage.output_tokens,
189 usage.cache_write_tokens,
190 usage.cache_read_tokens,
191 )
192 };
193 let saved_tokens = wire.map_or(0, |w| w.saved_tokens);
194 #[allow(clippy::cast_precision_loss)]
196 let saved_usd = quote.cost.input_per_m / 1_000_000.0 * saved_tokens as f64;
197
198 let uncompressed_input_tokens = wire.map_or(0, |w| w.uncompressed_input_tokens);
199 let reference_model = baseline
200 .reference_model
201 .as_deref()
202 .map(str::trim)
203 .filter(|m| !m.is_empty())
204 .map(str::to_string);
205 #[allow(clippy::cast_precision_loss)]
206 let reference_cost_usd = reference_model.as_deref().map_or(0.0, |reference| {
207 pricing.quote(Some(reference)).cost.input_per_m / 1_000_000.0
208 * uncompressed_input_tokens as f64
209 });
210
211 Self {
212 person: wire
213 .and_then(|w| w.person.clone())
214 .unwrap_or_else(|| ANONYMOUS_PERSON.to_string()),
215 team: wire.and_then(|w| w.team.clone()),
216 project: wire
217 .and_then(|w| w.project.clone())
218 .unwrap_or_else(|| DEFAULT_PROJECT.to_string()),
219 provider: wire.map_or_else(String::new, |w| w.provider.clone()),
220 model: usage.model.clone(),
221 routed_from: wire.and_then(|w| w.routed_from.clone()),
222 input_tokens: to_i64(usage.input_tokens),
223 output_tokens: to_i64(usage.output_tokens),
224 cache_read_tokens: to_i64(usage.cache_read_tokens),
225 cache_write_tokens: to_i64(usage.cache_write_tokens),
226 reasoning_tokens: to_i64(usage.reasoning_tokens),
227 cost_usd,
228 saved_tokens: to_i64(saved_tokens),
229 saved_usd,
230 uncompressed_input_tokens: to_i64(uncompressed_input_tokens),
231 reference_model,
232 reference_cost_usd,
233 is_local,
234 }
235 }
236}
237
238fn to_i64(v: u64) -> i64 {
239 i64::try_from(v).unwrap_or(i64::MAX)
240}
241
242pub async fn insert_event(
244 client: &deadpool_postgres::Client,
245 e: &UsageEvent,
246) -> anyhow::Result<()> {
247 client
248 .execute(
249 "INSERT INTO usage_events \
250 (person, team, project, provider, model, routed_from, \
251 input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, \
252 reasoning_tokens, cost_usd, saved_tokens, saved_usd, \
253 uncompressed_input_tokens, reference_model, reference_cost_usd, is_local) \
254 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18)",
255 &[
256 &e.person,
257 &e.team,
258 &e.project,
259 &e.provider,
260 &e.model,
261 &e.routed_from,
262 &e.input_tokens,
263 &e.output_tokens,
264 &e.cache_read_tokens,
265 &e.cache_write_tokens,
266 &e.reasoning_tokens,
267 &e.cost_usd,
268 &e.saved_tokens,
269 &e.saved_usd,
270 &e.uncompressed_input_tokens,
271 &e.reference_model,
272 &e.reference_cost_usd,
273 &e.is_local,
274 ],
275 )
276 .await?;
277 Ok(())
278}
279
280pub async fn budget_window_sums(
284 pool: &Pool,
285) -> anyhow::Result<(
286 std::collections::HashMap<String, f64>,
287 std::collections::HashMap<String, f64>,
288)> {
289 let client = pool.get().await?;
290 let mut person_day = std::collections::HashMap::new();
291 for row in client
292 .query(
293 "SELECT person, SUM(cost_usd) FROM usage_events \
294 WHERE ts >= date_trunc('day', now() AT TIME ZONE 'utc') AT TIME ZONE 'utc' \
295 GROUP BY person",
296 &[],
297 )
298 .await?
299 {
300 person_day.insert(row.get::<_, String>(0), row.get::<_, f64>(1));
301 }
302 let mut project_month = std::collections::HashMap::new();
303 for row in client
304 .query(
305 "SELECT project, SUM(cost_usd) FROM usage_events \
306 WHERE ts >= date_trunc('month', now() AT TIME ZONE 'utc') AT TIME ZONE 'utc' \
307 GROUP BY project",
308 &[],
309 )
310 .await?
311 {
312 project_month.insert(row.get::<_, String>(0), row.get::<_, f64>(1));
313 }
314 Ok((person_day, project_month))
315}
316
317pub async fn purge_events_older_than(pool: &Pool, days: u32) -> anyhow::Result<u64> {
321 let client = pool.get().await?;
322 let purged = client
323 .execute(
324 "DELETE FROM usage_events WHERE ts < now() - make_interval(days => $1)",
325 &[&i32::try_from(days).unwrap_or(i32::MAX)],
326 )
327 .await?;
328 Ok(purged)
329}
330
331pub async fn person_events(
334 pool: &Pool,
335 person_keys: &[String],
336) -> anyhow::Result<Vec<serde_json::Value>> {
337 let client = pool.get().await?;
338 let rows = client
339 .query(
340 "SELECT to_jsonb(usage_events) FROM usage_events \
341 WHERE person = ANY($1) ORDER BY ts",
342 &[&person_keys],
343 )
344 .await?;
345 Ok(rows
346 .into_iter()
347 .map(|r| r.get::<_, serde_json::Value>(0))
348 .collect())
349}
350
351pub async fn delete_person_events(pool: &Pool, person_keys: &[String]) -> anyhow::Result<u64> {
353 let client = pool.get().await?;
354 let deleted = client
355 .execute(
356 "DELETE FROM usage_events WHERE person = ANY($1)",
357 &[&person_keys],
358 )
359 .await?;
360 Ok(deleted)
361}
362
363pub async fn evidence_rows(
367 pool: &Pool,
368 from: chrono::DateTime<chrono::Utc>,
369 to: chrono::DateTime<chrono::Utc>,
370) -> anyhow::Result<Vec<serde_json::Value>> {
371 let client = pool.get().await?;
372 let rows = client
373 .query(
374 "SELECT jsonb_build_object(
375 'date', to_char(date_trunc('day', ts AT TIME ZONE 'utc'), 'YYYY-MM-DD'),
376 'person', person,
377 'project', project,
378 'model', model,
379 'provider', provider,
380 'requests', count(*),
381 'input_tokens', sum(input_tokens)::BIGINT,
382 'output_tokens', sum(output_tokens)::BIGINT,
383 'cache_read_tokens', sum(cache_read_tokens)::BIGINT,
384 'cost_usd', round(sum(cost_usd)::numeric, 6),
385 'saved_usd', round(sum(saved_usd)::numeric, 6),
386 'reference_cost_usd', round(sum(reference_cost_usd)::numeric, 6),
387 'local_requests', count(*) FILTER (WHERE is_local)
388 )
389 FROM usage_events WHERE ts >= $1 AND ts <= $2
390 GROUP BY
391 date_trunc('day', ts AT TIME ZONE 'utc'), person, project, model, provider
392 ORDER BY
393 date_trunc('day', ts AT TIME ZONE 'utc'), person, project, model, provider",
394 &[&from, &to],
395 )
396 .await?;
397 Ok(rows
398 .into_iter()
399 .map(|r| r.get::<_, serde_json::Value>(0))
400 .collect())
401}
402
403pub fn spawn_writer(pool: Pool) -> bool {
409 let (tx, mut rx) = tokio::sync::mpsc::channel::<RealUsage>(WRITER_QUEUE);
410 if !crate::proxy::usage_sink::install(tx) {
411 return false;
412 }
413 tokio::spawn(async move {
414 let pricing = ModelPricing::load();
418 let baseline = crate::core::config::Config::load().proxy.baseline.clone();
419 while let Some(usage) = rx.recv().await {
420 let event = UsageEvent::from_usage(&usage, &pricing, &baseline);
421 match pool.get().await {
422 Ok(client) => {
423 if let Err(e) = insert_event(&client, &event).await {
424 tracing::warn!("usage_events insert failed (fail-open): {e:#}");
425 }
426 }
427 Err(e) => {
428 tracing::warn!("usage_events pool unavailable (fail-open): {e:#}");
429 }
430 }
431 }
432 });
433 true
434}
435
436#[cfg(test)]
437mod tests {
438 use super::*;
439 use crate::proxy::usage::WireContext;
440
441 fn usage_with_wire(wire: Option<Box<WireContext>>) -> RealUsage {
442 RealUsage {
443 model: "claude-sonnet-4-5".into(),
444 input_tokens: 1000,
445 output_tokens: 500,
446 cache_read_tokens: 200,
447 cache_write_tokens: 100,
448 reasoning_tokens: 50,
449 cohort: None,
450 wire,
451 }
452 }
453
454 #[test]
455 fn event_carries_identity_and_baseline_fields() {
456 let usage = usage_with_wire(Some(Box::new(WireContext {
457 provider: "Anthropic".into(),
458 person: Some("yves".into()),
459 team: Some("platform".into()),
460 project: Some("ai-gateway".into()),
461 saved_tokens: 4000,
462 uncompressed_input_tokens: 5000,
463 is_local: false,
464 routed_from: Some("claude-opus-4-5".into()),
465 counterfactual: None,
466 })));
467 let event = UsageEvent::from_usage(
468 &usage,
469 &ModelPricing::load(),
470 &BaselineConfig {
471 reference_model: Some("claude-opus-4.5".into()),
472 local_shadow_rate_per_mtok: None,
473 },
474 );
475
476 assert_eq!(event.person, "yves");
477 assert_eq!(event.team.as_deref(), Some("platform"));
478 assert_eq!(event.project, "ai-gateway");
479 assert_eq!(event.provider, "Anthropic");
480 assert_eq!(event.model, "claude-sonnet-4-5");
481 assert_eq!(event.routed_from.as_deref(), Some("claude-opus-4-5"));
482 assert_eq!(event.input_tokens, 1000);
483 assert_eq!(event.saved_tokens, 4000);
484 assert_eq!(event.uncompressed_input_tokens, 5000);
485 assert!(!event.is_local);
486 assert!(event.cost_usd > 0.0, "known model must be priced");
487 assert!(
488 event.saved_usd > 0.0,
489 "saved tokens on a priced model must yield saved USD"
490 );
491 assert_eq!(event.reference_model.as_deref(), Some("claude-opus-4.5"));
494 assert!((event.reference_cost_usd - 0.025).abs() < 1e-9);
495 }
496
497 #[test]
498 fn event_without_wire_context_uses_honest_fallbacks() {
499 let event = UsageEvent::from_usage(
500 &usage_with_wire(None),
501 &ModelPricing::load(),
502 &BaselineConfig::default(),
503 );
504 assert_eq!(event.person, ANONYMOUS_PERSON);
505 assert_eq!(event.project, DEFAULT_PROJECT);
506 assert_eq!(event.team, None);
507 assert_eq!(event.saved_tokens, 0);
508 assert_eq!(event.saved_usd, 0.0);
509 assert_eq!(event.uncompressed_input_tokens, 0);
510 assert!(!event.is_local);
511 assert_eq!(event.reference_model, None);
513 assert_eq!(event.reference_cost_usd, 0.0);
514 }
515
516 #[test]
517 fn local_usage_books_shadow_rate_never_zero() {
518 let usage = usage_with_wire(Some(Box::new(WireContext {
521 provider: "ollama".into(),
522 person: Some("yves".into()),
523 team: None,
524 project: None,
525 saved_tokens: 0,
526 uncompressed_input_tokens: 2000,
527 is_local: true,
528 routed_from: None,
529 counterfactual: None,
530 })));
531 let event =
532 UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
533 assert!(event.is_local);
534 assert!((event.cost_usd - 0.25 / 1_000_000.0 * 1800.0).abs() < 1e-12);
537 assert!(event.cost_usd > 0.0, "local cost must never be zero");
538
539 let cfg = BaselineConfig {
541 reference_model: None,
542 local_shadow_rate_per_mtok: Some(1.0),
543 };
544 let event = UsageEvent::from_usage(&usage, &ModelPricing::load(), &cfg);
545 assert!((event.cost_usd - 1.0 / 1_000_000.0 * 1800.0).abs() < 1e-12);
546 let zero = BaselineConfig {
547 reference_model: None,
548 local_shadow_rate_per_mtok: Some(0.0),
549 };
550 assert!(zero.effective_local_shadow_rate() > 0.0);
551 }
552
553 #[test]
554 fn sslmode_selects_tls_and_pool_builds_for_both() {
555 let tls: tokio_postgres::Config = "postgres://u:p@db.example.com:5432/app?sslmode=require"
557 .parse()
558 .unwrap();
559 assert!(wants_tls(&tls));
560 let plain: tokio_postgres::Config = "postgres://u:p@localhost:5432/app".parse().unwrap();
561 assert!(!wants_tls(&plain));
562 let disabled: tokio_postgres::Config = "postgres://u:p@localhost:5432/app?sslmode=disable"
563 .parse()
564 .unwrap();
565 assert!(!wants_tls(&disabled));
566
567 assert!(
570 pool_from_database_url("postgres://u:p@db.example.com:5432/app?sslmode=require")
571 .is_ok()
572 );
573 assert!(pool_from_database_url("postgres://u:p@localhost:5432/app").is_ok());
574 }
575
576 #[test]
577 fn pool_size_env_is_clamped_and_falls_back() {
578 let _guard = crate::core::data_dir::test_env_lock();
580 crate::test_env::remove_var(POOL_MAX_SIZE_ENV);
581 assert_eq!(pool_max_size(), 8, "unset -> default");
582 crate::test_env::set_var(POOL_MAX_SIZE_ENV, "24");
583 assert_eq!(pool_max_size(), 24, "explicit value wins");
584 crate::test_env::set_var(POOL_MAX_SIZE_ENV, "0");
585 assert_eq!(pool_max_size(), 2, "clamped low");
586 crate::test_env::set_var(POOL_MAX_SIZE_ENV, "9999");
587 assert_eq!(pool_max_size(), 64, "clamped high");
588 crate::test_env::set_var(POOL_MAX_SIZE_ENV, "not-a-number");
589 assert_eq!(pool_max_size(), 8, "garbage -> default, never panic");
590 crate::test_env::remove_var(POOL_MAX_SIZE_ENV);
591 }
592
593 #[test]
594 fn schema_ddl_is_idempotent_by_construction() {
595 for stmt in ["CREATE TABLE", "CREATE INDEX"] {
598 for (i, _) in USAGE_EVENTS_DDL.match_indices(stmt) {
599 let tail = &USAGE_EVENTS_DDL[i..(i + stmt.len() + 14).min(USAGE_EVENTS_DDL.len())];
600 assert!(
601 tail.contains("IF NOT EXISTS"),
602 "non-idempotent DDL statement: {tail}"
603 );
604 }
605 }
606 for col in [
608 "uncompressed_input_tokens",
609 "reference_model",
610 "reference_cost_usd",
611 "is_local",
612 ] {
613 assert!(
614 USAGE_EVENTS_DDL.contains(col),
615 "baseline column {col} missing from schema"
616 );
617 }
618 }
619}