use deadpool_postgres::{Manager, ManagerConfig, Pool, RecyclingMethod};
use tokio_postgres::NoTls;
use tokio_postgres::config::SslMode;
use crate::core::config::BaselineConfig;
use crate::core::gain::model_pricing::ModelPricing;
use crate::core::ocla::OclaRegistry;
use crate::proxy::usage::RealUsage;
pub const WRITER_QUEUE: usize = 4096;
pub const POOL_MAX_SIZE_ENV: &str = "LEAN_CTX_PG_POOL_MAX_SIZE";
const POOL_MAX_SIZE_DEFAULT: usize = 8;
fn pool_max_size() -> usize {
std::env::var(POOL_MAX_SIZE_ENV)
.ok()
.and_then(|v| v.trim().parse::<usize>().ok())
.map_or(POOL_MAX_SIZE_DEFAULT, |n| n.clamp(2, 64))
}
pub fn pool_from_database_url(database_url: &str) -> anyhow::Result<Pool> {
let pg_cfg: tokio_postgres::Config = database_url.parse()?;
let mgr_cfg = ManagerConfig {
recycling_method: RecyclingMethod::Fast,
};
let mgr = if wants_tls(&pg_cfg) {
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
let roots = rustls::RootCertStore {
roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
};
let tls_cfg = rustls::ClientConfig::builder()
.with_root_certificates(roots)
.with_no_client_auth();
Manager::from_config(
pg_cfg,
tokio_postgres_rustls::MakeRustlsConnect::new(tls_cfg),
mgr_cfg,
)
} else {
Manager::from_config(pg_cfg, NoTls, mgr_cfg)
};
Ok(Pool::builder(mgr).max_size(pool_max_size()).build()?)
}
fn wants_tls(cfg: &tokio_postgres::Config) -> bool {
matches!(cfg.get_ssl_mode(), SslMode::Require)
}
const USAGE_EVENTS_DDL: &str = r"
CREATE TABLE IF NOT EXISTS usage_events (
id BIGSERIAL PRIMARY KEY,
ts TIMESTAMPTZ NOT NULL DEFAULT now(),
person TEXT NOT NULL,
team TEXT,
project TEXT NOT NULL,
tool TEXT,
provider TEXT NOT NULL,
model TEXT NOT NULL,
routed_from TEXT,
input_tokens BIGINT NOT NULL,
output_tokens BIGINT NOT NULL,
cache_read_tokens BIGINT NOT NULL DEFAULT 0,
cache_write_tokens BIGINT NOT NULL DEFAULT 0,
reasoning_tokens BIGINT NOT NULL DEFAULT 0,
cost_usd DOUBLE PRECISION NOT NULL,
saved_tokens BIGINT NOT NULL DEFAULT 0,
saved_usd DOUBLE PRECISION NOT NULL DEFAULT 0,
-- Avoided-cost baseline for the success fee (enterprise#18, Doc 04 §6):
uncompressed_input_tokens BIGINT NOT NULL DEFAULT 0,
reference_model TEXT,
reference_cost_usd DOUBLE PRECISION NOT NULL DEFAULT 0,
is_local BOOLEAN NOT NULL DEFAULT false,
-- Cost provenance (#1179): provider | shadow | list | live | heuristic.
cost_source TEXT NOT NULL DEFAULT 'list'
);
ALTER TABLE usage_events ADD COLUMN IF NOT EXISTS cost_source TEXT NOT NULL DEFAULT 'list';
CREATE INDEX IF NOT EXISTS idx_usage_events_person_ts ON usage_events (person, ts);
CREATE INDEX IF NOT EXISTS idx_usage_events_project_ts ON usage_events (project, ts);
CREATE INDEX IF NOT EXISTS idx_usage_events_model_ts ON usage_events (model, ts);
";
pub async fn init_schema(pool: &Pool) -> anyhow::Result<()> {
let client = pool.get().await?;
client.batch_execute(USAGE_EVENTS_DDL).await?;
Ok(())
}
#[derive(Debug, Clone, PartialEq)]
pub struct UsageEvent {
pub person: String,
pub team: Option<String>,
pub project: String,
pub provider: String,
pub model: String,
pub routed_from: Option<String>,
pub input_tokens: i64,
pub output_tokens: i64,
pub cache_read_tokens: i64,
pub cache_write_tokens: i64,
pub reasoning_tokens: i64,
pub cost_usd: f64,
pub saved_tokens: i64,
pub saved_usd: f64,
pub uncompressed_input_tokens: i64,
pub reference_model: Option<String>,
pub reference_cost_usd: f64,
pub is_local: bool,
pub cost_source: &'static str,
}
fn cost_source_of(kind: crate::core::gain::model_pricing::PricingMatchKind) -> &'static str {
use crate::core::gain::model_pricing::PricingMatchKind as K;
match kind {
K::Exact => "list",
K::Live => "live",
K::Alias | K::Heuristic | K::Fallback => "heuristic",
}
}
const ANONYMOUS_PERSON: &str = "anonymous";
const DEFAULT_PROJECT: &str = "default";
impl UsageEvent {
#[must_use]
pub fn from_usage(
usage: &RealUsage,
pricing: &ModelPricing,
baseline: &BaselineConfig,
) -> Self {
let wire = usage.wire.as_deref();
let quote = pricing.quote(Some(&usage.model));
let is_local = wire.is_some_and(|w| w.is_local);
#[allow(clippy::cast_precision_loss)]
let (cost_usd, cost_source) = if is_local {
let billable = usage.input_tokens
+ usage.output_tokens
+ usage.cache_read_tokens
+ usage.cache_write_tokens;
(
baseline.effective_local_shadow_rate() / 1_000_000.0 * billable as f64,
"shadow",
)
} else if let Some(measured) = usage.provider_cost_usd {
(measured, "provider")
} else {
let estimated = quote.cost.estimate_usd(
usage.input_tokens,
usage.output_tokens,
usage.cache_write_tokens,
usage.cache_read_tokens,
);
(estimated, cost_source_of(quote.match_kind))
};
let saved_tokens = wire.map_or(0, |w| w.saved_tokens);
#[allow(clippy::cast_precision_loss)]
let saved_usd = quote.cost.input_per_m / 1_000_000.0 * saved_tokens as f64;
let uncompressed_input_tokens = wire.map_or(0, |w| w.uncompressed_input_tokens);
let reference_model = baseline
.reference_model
.as_deref()
.map(str::trim)
.filter(|m| !m.is_empty())
.map(str::to_string);
#[allow(clippy::cast_precision_loss)]
let reference_cost_usd = reference_model.as_deref().map_or(0.0, |reference| {
pricing.quote(Some(reference)).cost.input_per_m / 1_000_000.0
* uncompressed_input_tokens as f64
});
Self {
person: wire
.and_then(|w| w.person.clone())
.unwrap_or_else(|| ANONYMOUS_PERSON.to_string()),
team: wire.and_then(|w| w.team.clone()),
project: wire
.and_then(|w| w.project.clone())
.unwrap_or_else(|| DEFAULT_PROJECT.to_string()),
provider: wire.map_or_else(String::new, |w| w.provider.clone()),
model: usage.model.clone(),
routed_from: wire.and_then(|w| w.routed_from.clone()),
input_tokens: to_i64(usage.input_tokens),
output_tokens: to_i64(usage.output_tokens),
cache_read_tokens: to_i64(usage.cache_read_tokens),
cache_write_tokens: to_i64(usage.cache_write_tokens),
reasoning_tokens: to_i64(usage.reasoning_tokens),
cost_usd,
saved_tokens: to_i64(saved_tokens),
saved_usd,
uncompressed_input_tokens: to_i64(uncompressed_input_tokens),
reference_model,
reference_cost_usd,
is_local,
cost_source,
}
}
}
fn to_i64(v: u64) -> i64 {
i64::try_from(v).unwrap_or(i64::MAX)
}
pub async fn insert_event(
client: &deadpool_postgres::Client,
e: &UsageEvent,
) -> anyhow::Result<()> {
client
.execute(
"INSERT INTO usage_events \
(person, team, project, provider, model, routed_from, \
input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, \
reasoning_tokens, cost_usd, saved_tokens, saved_usd, \
uncompressed_input_tokens, reference_model, reference_cost_usd, is_local, \
cost_source) \
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19)",
&[
&e.person,
&e.team,
&e.project,
&e.provider,
&e.model,
&e.routed_from,
&e.input_tokens,
&e.output_tokens,
&e.cache_read_tokens,
&e.cache_write_tokens,
&e.reasoning_tokens,
&e.cost_usd,
&e.saved_tokens,
&e.saved_usd,
&e.uncompressed_input_tokens,
&e.reference_model,
&e.reference_cost_usd,
&e.is_local,
&e.cost_source,
],
)
.await?;
Ok(())
}
pub async fn budget_window_sums(
pool: &Pool,
) -> anyhow::Result<(
std::collections::HashMap<String, f64>,
std::collections::HashMap<String, f64>,
)> {
let client = pool.get().await?;
let mut person_day = std::collections::HashMap::new();
for row in client
.query(
"SELECT person, SUM(cost_usd) FROM usage_events \
WHERE ts >= date_trunc('day', now() AT TIME ZONE 'utc') AT TIME ZONE 'utc' \
GROUP BY person",
&[],
)
.await?
{
person_day.insert(row.get::<_, String>(0), row.get::<_, f64>(1));
}
let mut project_month = std::collections::HashMap::new();
for row in client
.query(
"SELECT project, SUM(cost_usd) FROM usage_events \
WHERE ts >= date_trunc('month', now() AT TIME ZONE 'utc') AT TIME ZONE 'utc' \
GROUP BY project",
&[],
)
.await?
{
project_month.insert(row.get::<_, String>(0), row.get::<_, f64>(1));
}
Ok((person_day, project_month))
}
pub async fn purge_events_older_than(pool: &Pool, days: u32) -> anyhow::Result<u64> {
let client = pool.get().await?;
let purged = client
.execute(
"DELETE FROM usage_events WHERE ts < now() - make_interval(days => $1)",
&[&i32::try_from(days).unwrap_or(i32::MAX)],
)
.await?;
Ok(purged)
}
pub async fn person_events(
pool: &Pool,
person_keys: &[String],
) -> anyhow::Result<Vec<serde_json::Value>> {
let client = pool.get().await?;
let rows = client
.query(
"SELECT to_jsonb(usage_events) FROM usage_events \
WHERE person = ANY($1) ORDER BY ts",
&[&person_keys],
)
.await?;
Ok(rows
.into_iter()
.map(|r| r.get::<_, serde_json::Value>(0))
.collect())
}
pub async fn delete_person_events(pool: &Pool, person_keys: &[String]) -> anyhow::Result<u64> {
let client = pool.get().await?;
let deleted = client
.execute(
"DELETE FROM usage_events WHERE person = ANY($1)",
&[&person_keys],
)
.await?;
Ok(deleted)
}
pub async fn evidence_rows(
pool: &Pool,
from: chrono::DateTime<chrono::Utc>,
to: chrono::DateTime<chrono::Utc>,
) -> anyhow::Result<Vec<serde_json::Value>> {
let client = pool.get().await?;
let rows = client
.query(
"SELECT jsonb_build_object(
'date', to_char(date_trunc('day', ts AT TIME ZONE 'utc'), 'YYYY-MM-DD'),
'person', person,
'project', project,
'model', model,
'provider', provider,
'requests', count(*),
'input_tokens', sum(input_tokens)::BIGINT,
'output_tokens', sum(output_tokens)::BIGINT,
'cache_read_tokens', sum(cache_read_tokens)::BIGINT,
'cost_usd', round(sum(cost_usd)::numeric, 6),
'saved_usd', round(sum(saved_usd)::numeric, 6),
'reference_cost_usd', round(sum(reference_cost_usd)::numeric, 6),
'local_requests', count(*) FILTER (WHERE is_local),
'measured_requests', count(*) FILTER (WHERE cost_source = 'provider'),
'estimated_requests', count(*) FILTER (WHERE cost_source = 'heuristic')
)
FROM usage_events WHERE ts >= $1 AND ts <= $2
GROUP BY
date_trunc('day', ts AT TIME ZONE 'utc'), person, project, model, provider
ORDER BY
date_trunc('day', ts AT TIME ZONE 'utc'), person, project, model, provider",
&[&from, &to],
)
.await?;
Ok(rows
.into_iter()
.map(|r| r.get::<_, serde_json::Value>(0))
.collect())
}
pub fn spawn_writer(pool: Pool) -> bool {
let (tx, mut rx) = tokio::sync::mpsc::channel::<RealUsage>(WRITER_QUEUE);
if !crate::proxy::usage_sink::install(tx) {
return false;
}
tokio::spawn(async move {
let pricing = ModelPricing::load();
let baseline = crate::core::config::Config::load().proxy.baseline.clone();
while let Some(usage) = rx.recv().await {
let event = UsageEvent::from_usage(&usage, &pricing, &baseline);
match pool.get().await {
Ok(client) => {
match insert_event(&client, &event).await {
Ok(()) => {
if let Ok(Some(ocla_record)) = usage.to_ocla_usage_record()
&& let Err(e) =
OclaRegistry::global().usage_sink.record_usage(ocla_record)
{
tracing::debug!("OCLA usage projection: {e}");
}
}
Err(e) => {
tracing::warn!("usage_events insert failed (fail-open): {e:#}");
}
}
}
Err(e) => {
tracing::warn!("usage_events pool unavailable (fail-open): {e:#}");
}
}
}
});
true
}
#[cfg(test)]
mod tests {
use super::*;
use crate::proxy::usage::WireContext;
fn usage_with_wire(wire: Option<Box<WireContext>>) -> RealUsage {
RealUsage {
model: "claude-sonnet-4-5".into(),
input_tokens: 1000,
output_tokens: 500,
cache_read_tokens: 200,
cache_write_tokens: 100,
reasoning_tokens: 50,
provider_cost_usd: None,
cohort: None,
wire,
}
}
#[test]
fn event_carries_identity_and_baseline_fields() {
let usage = usage_with_wire(Some(Box::new(WireContext {
provider: "Anthropic".into(),
person: Some("yves".into()),
team: Some("platform".into()),
project: Some("ai-gateway".into()),
saved_tokens: 4000,
uncompressed_input_tokens: 5000,
is_local: false,
routed_from: Some("claude-opus-4-5".into()),
counterfactual: None,
lineage: None,
})));
let event = UsageEvent::from_usage(
&usage,
&ModelPricing::load(),
&BaselineConfig {
reference_model: Some("claude-opus-4.5".into()),
local_shadow_rate_per_mtok: None,
},
);
assert_eq!(event.person, "yves");
assert_eq!(event.team.as_deref(), Some("platform"));
assert_eq!(event.project, "ai-gateway");
assert_eq!(event.provider, "Anthropic");
assert_eq!(event.model, "claude-sonnet-4-5");
assert_eq!(event.routed_from.as_deref(), Some("claude-opus-4-5"));
assert_eq!(event.input_tokens, 1000);
assert_eq!(event.saved_tokens, 4000);
assert_eq!(event.uncompressed_input_tokens, 5000);
assert!(!event.is_local);
assert!(event.cost_usd > 0.0, "known model must be priced");
assert_eq!(
event.cost_source, "list",
"exact table match books list price"
);
assert!(
event.saved_usd > 0.0,
"saved tokens on a priced model must yield saved USD"
);
assert_eq!(event.reference_model.as_deref(), Some("claude-opus-4.5"));
assert!((event.reference_cost_usd - 0.025).abs() < 1e-9);
}
#[test]
fn provider_reported_cost_beats_the_price_table() {
let mut usage = usage_with_wire(Some(Box::new(WireContext {
provider: "openrouter".into(),
person: Some("nicolas".into()),
team: None,
project: Some("bot".into()),
saved_tokens: 0,
uncompressed_input_tokens: 1000,
is_local: false,
routed_from: None,
counterfactual: None,
lineage: None,
})));
usage.provider_cost_usd = Some(0.0123);
let event =
UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
assert!((event.cost_usd - 0.0123).abs() < 1e-12);
assert_eq!(event.cost_source, "provider");
usage.provider_cost_usd = Some(0.0);
let event =
UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
assert_eq!(event.cost_usd, 0.0);
assert_eq!(event.cost_source, "provider");
}
#[test]
fn unknown_model_is_marked_heuristic_and_local_shadow_beats_measured() {
let mut usage = usage_with_wire(None);
usage.model = "vendor/brand-new-model-20990101".into();
let event =
UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
assert_eq!(event.cost_source, "heuristic");
let mut usage = usage_with_wire(Some(Box::new(WireContext {
provider: "ollama".into(),
person: None,
team: None,
project: None,
saved_tokens: 0,
uncompressed_input_tokens: 0,
is_local: true,
routed_from: None,
counterfactual: None,
lineage: None,
})));
usage.provider_cost_usd = Some(9.99);
let event =
UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
assert_eq!(event.cost_source, "shadow");
assert!(
event.cost_usd < 1.0,
"shadow rate, not the stray measured figure"
);
}
#[test]
fn event_without_wire_context_uses_honest_fallbacks() {
let event = UsageEvent::from_usage(
&usage_with_wire(None),
&ModelPricing::load(),
&BaselineConfig::default(),
);
assert_eq!(event.person, ANONYMOUS_PERSON);
assert_eq!(event.project, DEFAULT_PROJECT);
assert_eq!(event.team, None);
assert_eq!(event.saved_tokens, 0);
assert_eq!(event.saved_usd, 0.0);
assert_eq!(event.uncompressed_input_tokens, 0);
assert!(!event.is_local);
assert_eq!(event.reference_model, None);
assert_eq!(event.reference_cost_usd, 0.0);
}
#[test]
fn local_usage_books_shadow_rate_never_zero() {
let usage = usage_with_wire(Some(Box::new(WireContext {
provider: "ollama".into(),
person: Some("yves".into()),
team: None,
project: None,
saved_tokens: 0,
uncompressed_input_tokens: 2000,
is_local: true,
routed_from: None,
counterfactual: None,
lineage: None,
})));
let event =
UsageEvent::from_usage(&usage, &ModelPricing::load(), &BaselineConfig::default());
assert!(event.is_local);
assert!((event.cost_usd - 0.25 / 1_000_000.0 * 1800.0).abs() < 1e-12);
assert!(event.cost_usd > 0.0, "local cost must never be zero");
let cfg = BaselineConfig {
reference_model: None,
local_shadow_rate_per_mtok: Some(1.0),
};
let event = UsageEvent::from_usage(&usage, &ModelPricing::load(), &cfg);
assert!((event.cost_usd - 1.0 / 1_000_000.0 * 1800.0).abs() < 1e-12);
let zero = BaselineConfig {
reference_model: None,
local_shadow_rate_per_mtok: Some(0.0),
};
assert!(zero.effective_local_shadow_rate() > 0.0);
}
#[test]
fn sslmode_selects_tls_and_pool_builds_for_both() {
let tls: tokio_postgres::Config = "postgres://u:p@db.example.com:5432/app?sslmode=require"
.parse()
.unwrap();
assert!(wants_tls(&tls));
let plain: tokio_postgres::Config = "postgres://u:p@localhost:5432/app".parse().unwrap();
assert!(!wants_tls(&plain));
let disabled: tokio_postgres::Config = "postgres://u:p@localhost:5432/app?sslmode=disable"
.parse()
.unwrap();
assert!(!wants_tls(&disabled));
assert!(
pool_from_database_url("postgres://u:p@db.example.com:5432/app?sslmode=require")
.is_ok()
);
assert!(pool_from_database_url("postgres://u:p@localhost:5432/app").is_ok());
}
#[test]
fn pool_size_env_is_clamped_and_falls_back() {
let _guard = crate::core::data_dir::test_env_lock();
crate::test_env::remove_var(POOL_MAX_SIZE_ENV);
assert_eq!(pool_max_size(), 8, "unset -> default");
crate::test_env::set_var(POOL_MAX_SIZE_ENV, "24");
assert_eq!(pool_max_size(), 24, "explicit value wins");
crate::test_env::set_var(POOL_MAX_SIZE_ENV, "0");
assert_eq!(pool_max_size(), 2, "clamped low");
crate::test_env::set_var(POOL_MAX_SIZE_ENV, "9999");
assert_eq!(pool_max_size(), 64, "clamped high");
crate::test_env::set_var(POOL_MAX_SIZE_ENV, "not-a-number");
assert_eq!(pool_max_size(), 8, "garbage -> default, never panic");
crate::test_env::remove_var(POOL_MAX_SIZE_ENV);
}
#[test]
fn schema_ddl_is_idempotent_by_construction() {
for stmt in ["CREATE TABLE", "CREATE INDEX"] {
for (i, _) in USAGE_EVENTS_DDL.match_indices(stmt) {
let tail = &USAGE_EVENTS_DDL[i..(i + stmt.len() + 14).min(USAGE_EVENTS_DDL.len())];
assert!(
tail.contains("IF NOT EXISTS"),
"non-idempotent DDL statement: {tail}"
);
}
}
for (i, _) in USAGE_EVENTS_DDL.match_indices("ADD COLUMN") {
let tail = &USAGE_EVENTS_DDL[i..(i + 25).min(USAGE_EVENTS_DDL.len())];
assert!(
tail.contains("IF NOT EXISTS"),
"non-idempotent ALTER statement: {tail}"
);
}
for col in [
"uncompressed_input_tokens",
"reference_model",
"reference_cost_usd",
"is_local",
"cost_source",
] {
assert!(
USAGE_EVENTS_DDL.contains(col),
"baseline column {col} missing from schema"
);
}
}
}