use std::fs::{self, File};
use std::io::{BufWriter, Write};
use std::path::Path;
use chrono::{DateTime, Datelike, Duration, TimeZone, Utc};
use sea_orm::sea_query::OnConflict;
use sea_orm::{
ActiveModelTrait, ColumnTrait, Condition, ConnectionTrait, DatabaseConnection, DbBackend,
EntityTrait, PaginatorTrait, QueryFilter, QueryOrder, QuerySelect, Set, Statement,
};
use serde::{Deserialize, Serialize};
use bitrouter_ai::types::{Usage, UsageOrigin};
use bitrouter_sdk::{BitrouterError, Result};
use crate::cloud::settlement::{SettlementReceipt, SettlementState};
use crate::metering::db::{ReconciliationStatus, RequestMetric};
use crate::metering::entities::requests;
use crate::metering::pricing::{
ChargeEvidence, ChargeStatus, ModelPricing, PricingSource, unavailable_charge_evidence,
};
#[derive(Debug, Clone, Copy)]
pub enum TimeWindow {
LastMinute,
LastHour,
Today,
ThisWeek,
ThisMonth,
Custom {
start: DateTime<Utc>,
end: DateTime<Utc>,
},
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TokenUsage {
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
}
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct MeteringUsageRecord {
pub id: Option<String>,
pub request_id: Option<String>,
pub provider_id: String,
pub model_id: String,
#[serde(default)]
pub prompt_tokens: u64,
#[serde(default)]
pub completion_tokens: u64,
#[serde(default)]
pub reasoning_tokens: u64,
#[serde(default)]
pub uncached_input_tokens: u64,
#[serde(default)]
pub cache_read_tokens: u64,
#[serde(default)]
pub cache_write_tokens: u64,
#[serde(default)]
pub output_tokens: u64,
#[serde(default)]
pub usage_origin: UsageOrigin,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub raw_usage: Option<serde_json::Value>,
pub final_charge_micro_usd: Option<u64>,
#[serde(default)]
pub charge_status: ChargeStatus,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub charge_evidence: Option<ChargeEvidence>,
#[serde(default)]
pub reconciliation_status: ReconciliationStatus,
#[serde(default)]
pub reconciliation_attempts: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub authoritative_receipt: Option<serde_json::Value>,
pub status: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error_code: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct UsagePriceOverride {
pub provider_id: String,
pub model_id: String,
pub input_micro_usd_per_token: f64,
pub cache_read_micro_usd_per_token: Option<f64>,
pub cache_write_micro_usd_per_token: Option<f64>,
pub output_micro_usd_per_token: f64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReconciliationRecord {
pub request_id: String,
pub status: ReconciliationStatus,
pub attempts: u32,
}
impl MeteringUsageRecord {
pub fn apply_price_overrides(records: &mut [Self], prices: &[UsagePriceOverride]) {
for record in records {
if record.charge_status == ChargeStatus::Computed
|| record.reconciliation_status != ReconciliationStatus::NotApplicable
|| record.usage_origin == UsageOrigin::Unknown
{
continue;
}
let Some(price) = prices.iter().find(|price| {
price.provider_id == record.provider_id && price.model_id == record.model_id
}) else {
continue;
};
if (record.cache_read_tokens > 0 || record.cache_write_tokens > 0)
&& (price.cache_read_micro_usd_per_token.is_none()
|| price.cache_write_micro_usd_per_token.is_none())
{
continue;
}
let usage = Usage {
prompt_tokens: record.prompt_tokens,
completion_tokens: record.completion_tokens,
reasoning_tokens: record.reasoning_tokens,
cache_read_tokens: record.cache_read_tokens,
cache_write_tokens: record.cache_write_tokens,
origin: record.usage_origin,
raw: record.raw_usage.clone().map(Box::new),
..Default::default()
};
let pricing = ModelPricing::cache_aware(
Some(price.input_micro_usd_per_token),
price.cache_read_micro_usd_per_token,
price.cache_write_micro_usd_per_token,
Some(price.output_micro_usd_per_token),
);
let evidence = super::tariff::override_charge_evidence(
&usage,
&pricing,
record
.charge_evidence
.as_ref()
.and_then(|evidence| evidence.tariff_snapshot.as_ref()),
);
if evidence.status != ChargeStatus::Computed {
record.charge_status = evidence.status;
record.charge_evidence = Some(evidence);
record.final_charge_micro_usd = None;
continue;
}
record.uncached_input_tokens = evidence.normalized_usage.uncached_input_tokens;
record.output_tokens = evidence.normalized_usage.output_tokens;
record.final_charge_micro_usd =
evidence.charge_micro_usd.map(|charge| charge.max(0) as u64);
record.charge_status = evidence.status;
record.charge_evidence = Some(evidence);
}
}
pub fn write_jsonl(path: impl AsRef<Path>, records: &[Self]) -> Result<()> {
if let Some(parent) = path.as_ref().parent()
&& !parent.as_os_str().is_empty()
{
fs::create_dir_all(parent).map_err(|e| {
BitrouterError::internal(format!(
"metering usage jsonl mkdir {}: {e}",
parent.display()
))
})?;
}
let file = File::create(path.as_ref()).map_err(|e| {
BitrouterError::internal(format!(
"metering usage jsonl create {}: {e}",
path.as_ref().display()
))
})?;
let mut writer = BufWriter::new(file);
for record in records {
serde_json::to_writer(&mut writer, record).map_err(|e| {
BitrouterError::internal(format!("metering usage jsonl serialize: {e}"))
})?;
writer.write_all(b"\n").map_err(|e| {
BitrouterError::internal(format!("metering usage jsonl write: {e}"))
})?;
}
writer
.flush()
.map_err(|e| BitrouterError::internal(format!("metering usage jsonl flush: {e}")))
}
}
impl UsagePriceOverride {
pub fn parse(value: &str) -> Result<Self> {
const EXPECTED: &str =
"provider:model=uncached,cache_read,cache_write,output (or legacy input,output)";
let (route, prices) = value.split_once('=').ok_or_else(|| {
BitrouterError::bad_request(format!(
"invalid price override {value:?}; expected {EXPECTED}"
))
})?;
let (provider_id, model_id) = route.split_once(':').ok_or_else(|| {
BitrouterError::bad_request(format!(
"invalid price override {value:?}; expected {EXPECTED}"
))
})?;
let rates = prices
.split(',')
.map(str::trim)
.map(|rate| {
rate.parse::<f64>().map_err(|error| {
BitrouterError::bad_request(format!(
"invalid price in override {value:?}: {error}"
))
})
})
.collect::<Result<Vec<_>>>()?;
if rates.iter().any(|rate| !rate.is_finite() || *rate < 0.0) {
return Err(BitrouterError::bad_request(format!(
"invalid price override {value:?}; rates must be finite and non-negative"
)));
}
let (input, cache_read, cache_write, output) = match rates.as_slice() {
[input, output] => (*input, None, None, *output),
[input, cache_read, cache_write, output] => {
(*input, Some(*cache_read), Some(*cache_write), *output)
}
_ => {
return Err(BitrouterError::bad_request(format!(
"invalid price override {value:?}; expected {EXPECTED}"
)));
}
};
Ok(Self {
provider_id: provider_id.trim().to_string(),
model_id: model_id.trim().to_string(),
input_micro_usd_per_token: input,
cache_read_micro_usd_per_token: cache_read,
cache_write_micro_usd_per_token: cache_write,
output_micro_usd_per_token: output,
})
}
}
#[derive(Debug, Clone, Copy, Default)]
pub struct RateMetrics {
pub requests_per_minute: f64,
pub tokens_per_minute: f64,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct SpendSummary {
pub spend_micro_usd: u64,
pub requests: u64,
pub unpriced: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RequestRow {
pub request_id: String,
pub created_at: String,
pub router_id: Option<String>,
pub binding_digest: Option<String>,
pub original_selector: Option<String>,
pub model_id: String,
pub provider_id: String,
pub prompt_tokens: i64,
pub completion_tokens: i64,
pub cache_read_tokens: i64,
pub cache_write_tokens: i64,
pub estimated_charge_micro_usd: i64,
pub latency_ms: i64,
pub error: Option<String>,
pub charge_status: ChargeStatus,
pub episode_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RequestPage {
pub rows: Vec<RequestRow>,
pub truncated: bool,
}
impl From<requests::Model> for RequestRow {
fn from(m: requests::Model) -> Self {
Self {
request_id: m.request_id,
created_at: m.created_at,
router_id: m.router_id,
binding_digest: m.binding_digest,
original_selector: m.original_selector,
model_id: m.model_id,
provider_id: m.provider_id,
prompt_tokens: m.prompt_tokens,
completion_tokens: m.completion_tokens,
cache_read_tokens: m.cache_read_tokens,
cache_write_tokens: m.cache_write_tokens,
estimated_charge_micro_usd: m.estimated_charge_micro_usd,
latency_ms: m.latency_ms,
error: m.error,
charge_status: ChargeStatus::from_persisted(&m.charge_status),
episode_id: None,
}
}
}
fn summarize(charges: Vec<(i64, String)>) -> SpendSummary {
let mut summary = SpendSummary {
requests: charges.len() as u64,
..Default::default()
};
for (charge, status) in charges {
match ChargeStatus::from_persisted(&status) {
ChargeStatus::Computed | ChargeStatus::NotCharged => {
summary.spend_micro_usd += charge.max(0) as u64;
}
ChargeStatus::Unknown | ChargeStatus::LegacyUnknown => summary.unpriced += 1,
}
}
summary
}
#[derive(Clone)]
pub struct MeteringStore {
db: DatabaseConnection,
}
impl MeteringStore {
pub fn new(db: DatabaseConnection) -> Self {
Self { db }
}
pub async fn get_spend(&self, api_key_id: &str, window: TimeWindow) -> Result<u64> {
let start = window_start(window).to_rfc3339();
let charges: Vec<i64> = requests::Entity::find()
.select_only()
.column(requests::Column::EstimatedChargeMicroUsd)
.filter(requests::Column::ApiKeyId.eq(api_key_id))
.filter(requests::Column::CreatedAt.gte(start))
.into_tuple()
.all(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("get_spend: {e}")))?;
Ok(charges.into_iter().map(|c| c.max(0) as u64).sum())
}
pub async fn get_enforceable_spend(
&self,
api_key_id: &str,
window: TimeWindow,
) -> Result<Option<u64>> {
let start = window_start(window).to_rfc3339();
let charges: Vec<(i64, String)> = requests::Entity::find()
.select_only()
.column(requests::Column::EstimatedChargeMicroUsd)
.column(requests::Column::ChargeStatus)
.filter(requests::Column::ApiKeyId.eq(api_key_id))
.filter(requests::Column::CreatedAt.gte(start))
.into_tuple()
.all(&self.db)
.await
.map_err(|error| BitrouterError::internal(format!("get_enforceable_spend: {error}")))?;
let mut total = 0_u64;
for (charge, status) in charges {
match ChargeStatus::from_persisted(&status) {
ChargeStatus::Computed | ChargeStatus::NotCharged => {
total = total.saturating_add(charge.max(0) as u64);
}
ChargeStatus::Unknown | ChargeStatus::LegacyUnknown => return Ok(None),
}
}
Ok(Some(total))
}
pub async fn get_enforceable_total_spend(&self, window: TimeWindow) -> Result<Option<u64>> {
let start = window_start(window).to_rfc3339();
let mut query = requests::Entity::find()
.select_only()
.column(requests::Column::EstimatedChargeMicroUsd)
.column(requests::Column::ChargeStatus)
.filter(requests::Column::CreatedAt.gte(start));
if let TimeWindow::Custom { end, .. } = window {
query = query.filter(requests::Column::CreatedAt.lt(end.to_rfc3339()));
}
let charges: Vec<(i64, String)> =
query.into_tuple().all(&self.db).await.map_err(|error| {
BitrouterError::internal(format!("get_enforceable_total_spend: {error}"))
})?;
let mut total = 0_u64;
for (charge, status) in charges {
match ChargeStatus::from_persisted(&status) {
ChargeStatus::Computed | ChargeStatus::NotCharged => {
total = total.saturating_add(charge.max(0) as u64);
}
ChargeStatus::Unknown | ChargeStatus::LegacyUnknown => return Ok(None),
}
}
Ok(Some(total))
}
pub async fn get_request_count(&self, api_key_id: &str, window: TimeWindow) -> Result<u64> {
let start = window_start(window).to_rfc3339();
requests::Entity::find()
.filter(requests::Column::ApiKeyId.eq(api_key_id))
.filter(requests::Column::CreatedAt.gte(start))
.count(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("get_request_count: {e}")))
}
pub async fn get_token_usage(
&self,
api_key_id: &str,
model: &str,
window: TimeWindow,
) -> Result<TokenUsage> {
let start = window_start(window).to_rfc3339();
let rows: Vec<(i64, i64)> = requests::Entity::find()
.select_only()
.column(requests::Column::PromptTokens)
.column(requests::Column::CompletionTokens)
.filter(requests::Column::ApiKeyId.eq(api_key_id))
.filter(requests::Column::ModelId.eq(model))
.filter(requests::Column::CreatedAt.gte(start))
.into_tuple()
.all(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("get_token_usage: {e}")))?;
let prompt: u64 = rows.iter().map(|(p, _)| (*p).max(0) as u64).sum();
let completion: u64 = rows.iter().map(|(_, c)| (*c).max(0) as u64).sum();
Ok(TokenUsage {
prompt_tokens: prompt,
completion_tokens: completion,
total_tokens: prompt + completion,
})
}
pub async fn get_rate(&self, api_key_id: &str) -> Result<RateMetrics> {
let start = window_start(TimeWindow::LastMinute).to_rfc3339();
let rows: Vec<(i64, i64)> = requests::Entity::find()
.select_only()
.column(requests::Column::PromptTokens)
.column(requests::Column::CompletionTokens)
.filter(requests::Column::ApiKeyId.eq(api_key_id))
.filter(requests::Column::CreatedAt.gte(start))
.into_tuple()
.all(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("get_rate: {e}")))?;
let tokens: u64 = rows
.iter()
.map(|(p, c)| (*p).max(0) as u64 + (*c).max(0) as u64)
.sum();
Ok(RateMetrics {
requests_per_minute: rows.len() as f64,
tokens_per_minute: tokens as f64,
})
}
pub async fn recent_requests(
&self,
window: TimeWindow,
limit: u64,
launch_id: Option<&str>,
) -> Result<Vec<RequestRow>> {
let start = window_start(window).to_rfc3339();
let mut query = requests::Entity::find()
.filter(requests::Column::CreatedAt.gte(start))
.order_by_desc(requests::Column::CreatedAt)
.order_by_desc(requests::Column::RequestId);
if let Some(launch) = launch_id {
query = query.filter(requests::Column::LaunchId.eq(launch));
}
if let TimeWindow::Custom { end, .. } = window {
query = query.filter(requests::Column::CreatedAt.lt(end.to_rfc3339()));
}
let rows = query
.limit(limit)
.all(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("recent_requests: {e}")))?;
let mut rows: Vec<RequestRow> = rows.into_iter().map(RequestRow::from).collect();
self.attach_episodes(&mut rows).await;
Ok(rows)
}
pub async fn recent_requests_filtered(
&self,
window: TimeWindow,
model: Option<&str>,
provider: Option<&str>,
limit: u64,
) -> Result<RequestPage> {
let start = window_start(window).to_rfc3339();
let mut query = requests::Entity::find()
.filter(requests::Column::CreatedAt.gte(start))
.order_by_desc(requests::Column::CreatedAt)
.order_by_desc(requests::Column::RequestId);
if let Some(model) = model {
query = query.filter(requests::Column::ModelId.eq(model));
}
if let Some(provider) = provider {
query = query.filter(requests::Column::ProviderId.eq(provider));
}
if let TimeWindow::Custom { end, .. } = window {
query = query.filter(requests::Column::CreatedAt.lt(end.to_rfc3339()));
}
let mut rows = query
.limit(limit.saturating_add(1))
.all(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("recent_requests_filtered: {e}")))?;
let truncated = rows.len() > limit as usize;
if truncated {
rows.pop();
}
let mut rows: Vec<RequestRow> = rows.into_iter().map(RequestRow::from).collect();
self.attach_episodes(&mut rows).await;
Ok(RequestPage { rows, truncated })
}
async fn attach_episodes(&self, rows: &mut [RequestRow]) {
if rows.is_empty() {
return;
}
let backend = self.db.get_database_backend();
let placeholders: Vec<String> = (1..=rows.len())
.map(|i| match backend {
DbBackend::Postgres => format!("${i}"),
_ => "?".to_string(),
})
.collect();
let sql = format!(
"SELECT request_id, episode_id FROM trajectory_requests WHERE request_id IN ({})",
placeholders.join(", ")
);
let values: Vec<sea_orm::Value> = rows
.iter()
.map(|r| sea_orm::Value::from(r.request_id.clone()))
.collect();
let Ok(found) = self
.db
.query_all(Statement::from_sql_and_values(backend, &sql, values))
.await
else {
return;
};
let mut episodes: std::collections::HashMap<String, String> =
std::collections::HashMap::new();
for row in found {
if let (Ok(request_id), Ok(episode_id)) = (
row.try_get::<String>("", "request_id"),
row.try_get::<String>("", "episode_id"),
) {
episodes.insert(request_id, episode_id);
}
}
for row in rows.iter_mut() {
row.episode_id = episodes.get(&row.request_id).cloned();
}
}
pub async fn get_total_rate(&self) -> Result<RateMetrics> {
let start = window_start(TimeWindow::LastMinute).to_rfc3339();
let rows: Vec<(i64, i64)> = requests::Entity::find()
.select_only()
.column(requests::Column::PromptTokens)
.column(requests::Column::CompletionTokens)
.filter(requests::Column::CreatedAt.gte(start))
.into_tuple()
.all(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("get_total_rate: {e}")))?;
let tokens: u64 = rows
.iter()
.map(|(p, c)| (*p).max(0) as u64 + (*c).max(0) as u64)
.sum();
Ok(RateMetrics {
requests_per_minute: rows.len() as f64,
tokens_per_minute: tokens as f64,
})
}
pub async fn export_usage(&self, window: TimeWindow) -> Result<Vec<MeteringUsageRecord>> {
let start = window_start(window).to_rfc3339();
let mut query = requests::Entity::find()
.filter(requests::Column::CreatedAt.gte(start))
.order_by_asc(requests::Column::CreatedAt)
.order_by_asc(requests::Column::RequestId);
if let TimeWindow::Custom { end, .. } = window {
query = query.filter(requests::Column::CreatedAt.lt(end.to_rfc3339()));
}
let rows = query
.all(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("export_usage: {e}")))?;
Ok(rows.into_iter().map(MeteringUsageRecord::from).collect())
}
pub async fn spend_summary_for_launch(
&self,
launch_id: &str,
window: TimeWindow,
) -> Result<SpendSummary> {
let start = window_start(window).to_rfc3339();
let mut query = requests::Entity::find()
.select_only()
.column(requests::Column::EstimatedChargeMicroUsd)
.column(requests::Column::ChargeStatus)
.filter(requests::Column::LaunchId.eq(launch_id))
.filter(requests::Column::CreatedAt.gte(start));
if let TimeWindow::Custom { end, .. } = window {
query = query.filter(requests::Column::CreatedAt.lt(end.to_rfc3339()));
}
let charges: Vec<(i64, String)> = query
.into_tuple()
.all(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("spend_summary_for_launch: {e}")))?;
Ok(summarize(charges))
}
pub async fn spend_summary_for_acp_session(
&self,
api_principal: &str,
controller_instance_id: &str,
session_id: &str,
window: TimeWindow,
) -> Result<SpendSummary> {
let start = window_start(window).to_rfc3339();
let mut query = requests::Entity::find()
.select_only()
.column(requests::Column::EstimatedChargeMicroUsd)
.column(requests::Column::ChargeStatus)
.filter(requests::Column::RouteScopeId.eq(api_principal))
.filter(requests::Column::ControllerInstanceId.eq(controller_instance_id))
.filter(
Condition::any()
.add(requests::Column::AcpSessionId.eq(session_id))
.add(requests::Column::NativeRootSessionId.eq(session_id))
.add(requests::Column::NativeAgentThreadId.eq(session_id)),
)
.filter(requests::Column::CreatedAt.gte(start));
if let TimeWindow::Custom { end, .. } = window {
query = query.filter(requests::Column::CreatedAt.lt(end.to_rfc3339()));
}
let charges: Vec<(i64, String)> =
query.into_tuple().all(&self.db).await.map_err(|e| {
BitrouterError::internal(format!("spend_summary_for_acp_session: {e}"))
})?;
Ok(summarize(charges))
}
pub async fn spend_summary(&self, window: TimeWindow) -> Result<SpendSummary> {
self.spend_summary_filtered(window, None, None).await
}
pub async fn spend_summary_filtered(
&self,
window: TimeWindow,
model: Option<&str>,
provider: Option<&str>,
) -> Result<SpendSummary> {
let start = window_start(window).to_rfc3339();
let mut query = requests::Entity::find()
.select_only()
.column(requests::Column::EstimatedChargeMicroUsd)
.column(requests::Column::ChargeStatus)
.filter(requests::Column::CreatedAt.gte(start));
if let Some(model) = model {
query = query.filter(requests::Column::ModelId.eq(model));
}
if let Some(provider) = provider {
query = query.filter(requests::Column::ProviderId.eq(provider));
}
if let TimeWindow::Custom { end, .. } = window {
query = query.filter(requests::Column::CreatedAt.lt(end.to_rfc3339()));
}
let charges: Vec<(i64, String)> = query
.into_tuple()
.all(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("spend_summary_filtered: {e}")))?;
Ok(summarize(charges))
}
pub(super) async fn native_request_rows(
&self,
caller: &bitrouter_sdk::caller::CallerContext,
request_ids: &[String],
) -> Result<Vec<requests::Model>> {
if request_ids.is_empty() {
return Ok(Vec::new());
}
if request_ids.len() > 64 || caller.is_anonymous() {
return Err(BitrouterError::bad_request(
"invalid native cost lookup scope",
));
}
requests::Entity::find()
.filter(requests::Column::RequestId.is_in(request_ids.iter().cloned()))
.filter(requests::Column::UserId.eq(caller.user_id()))
.filter(requests::Column::ApiKeyId.eq(caller.api_key_id()))
.all(&self.db)
.await
.map_err(|_| BitrouterError::internal("native cost evidence unavailable"))
}
pub async fn reconciliation_records(
&self,
request_ids: &[String],
) -> Result<Vec<ReconciliationRecord>> {
if request_ids.is_empty() {
return Ok(Vec::new());
}
let rows = requests::Entity::find()
.filter(requests::Column::RequestId.is_in(request_ids.iter().cloned()))
.order_by_asc(requests::Column::RequestId)
.all(&self.db)
.await
.map_err(|error| {
BitrouterError::internal(format!("load reconciliation records: {error}"))
})?;
Ok(rows
.into_iter()
.map(|row| ReconciliationRecord {
request_id: row.request_id,
status: ReconciliationStatus::from_persisted(&row.reconciliation_status),
attempts: row.reconciliation_attempts.max(0) as u32,
})
.collect())
}
pub async fn start_reconciliation_attempt(&self, request_id: &str) -> Result<u32> {
let row = requests::Entity::find_by_id(request_id)
.one(&self.db)
.await
.map_err(|error| {
BitrouterError::internal(format!("load reconciliation attempt: {error}"))
})?
.ok_or_else(|| {
BitrouterError::bad_request(format!(
"reconciliation request not found: {request_id}"
))
})?;
if ReconciliationStatus::from_persisted(&row.reconciliation_status)
!= ReconciliationStatus::Pending
{
return Err(BitrouterError::bad_request(format!(
"reconciliation request {request_id} is not pending"
)));
}
let attempts = row.reconciliation_attempts.max(0).saturating_add(1);
let mut active: requests::ActiveModel = row.into();
active.reconciliation_attempts = Set(attempts);
active.reconciliation_last_attempt_at = Set(Some(Utc::now().to_rfc3339()));
active.reconciliation_last_error = Set(None);
active.update(&self.db).await.map_err(|error| {
BitrouterError::internal(format!("persist reconciliation attempt: {error}"))
})?;
Ok(attempts as u32)
}
pub async fn reopen_unknown_reconciliation(
&self,
request_id: &str,
max_attempts: u32,
) -> Result<bool> {
let row = requests::Entity::find_by_id(request_id)
.one(&self.db)
.await
.map_err(|error| {
BitrouterError::internal(format!("load unknown reconciliation row: {error}"))
})?
.ok_or_else(|| {
BitrouterError::bad_request(format!(
"reconciliation request not found: {request_id}"
))
})?;
if ReconciliationStatus::from_persisted(&row.reconciliation_status)
!= ReconciliationStatus::Unknown
|| row.reconciliation_attempts.max(0) as u32 >= max_attempts
{
return Ok(false);
}
let mut active: requests::ActiveModel = row.into();
active.reconciliation_status = Set(ReconciliationStatus::Pending.as_str().to_string());
active.authoritative_settled_at = Set(None);
active.update(&self.db).await.map_err(|error| {
BitrouterError::internal(format!("reopen unknown reconciliation: {error}"))
})?;
Ok(true)
}
pub async fn record_reconciliation_error(&self, request_id: &str, error: &str) -> Result<()> {
let row = requests::Entity::find_by_id(request_id)
.one(&self.db)
.await
.map_err(|load_error| {
BitrouterError::internal(format!("load reconciliation error row: {load_error}"))
})?
.ok_or_else(|| {
BitrouterError::bad_request(format!(
"reconciliation request not found: {request_id}"
))
})?;
let mut active: requests::ActiveModel = row.into();
active.reconciliation_last_error = Set(Some(truncate_error(error)));
active.update(&self.db).await.map_err(|update_error| {
BitrouterError::internal(format!("persist reconciliation error: {update_error}"))
})?;
Ok(())
}
pub async fn exhaust_reconciliation(&self, request_id: &str, reason: &str) -> Result<()> {
let row = requests::Entity::find_by_id(request_id)
.one(&self.db)
.await
.map_err(|error| {
BitrouterError::internal(format!("load exhausted reconciliation: {error}"))
})?
.ok_or_else(|| {
BitrouterError::bad_request(format!(
"reconciliation request not found: {request_id}"
))
})?;
if ReconciliationStatus::from_persisted(&row.reconciliation_status)
!= ReconciliationStatus::Pending
{
return Ok(());
}
let mut active: requests::ActiveModel = row.into();
active.reconciliation_status = Set(ReconciliationStatus::Unknown.as_str().to_string());
active.charge_status = Set(ChargeStatus::Unknown.as_str().to_string());
active.estimated_charge_micro_usd = Set(0);
active.reconciliation_last_error = Set(Some(truncate_error(reason)));
active.authoritative_settled_at = Set(Some(Utc::now().to_rfc3339()));
active.update(&self.db).await.map_err(|error| {
BitrouterError::internal(format!("persist exhausted reconciliation: {error}"))
})?;
Ok(())
}
pub async fn apply_authoritative_receipt(
&self,
receipt: &SettlementReceipt,
prices: &[UsagePriceOverride],
) -> Result<ReconciliationStatus> {
self.apply_authoritative_receipt_inner(receipt, prices, false)
.await
}
pub async fn apply_authoritative_receipt_charge(
&self,
receipt: &SettlementReceipt,
) -> Result<ReconciliationStatus> {
self.apply_authoritative_receipt_inner(receipt, &[], true)
.await
}
async fn apply_authoritative_receipt_inner(
&self,
receipt: &SettlementReceipt,
prices: &[UsagePriceOverride],
accept_receipt_charge: bool,
) -> Result<ReconciliationStatus> {
let row = requests::Entity::find_by_id(&receipt.request_id)
.one(&self.db)
.await
.map_err(|error| BitrouterError::internal(format!("load reconciliation row: {error}")))?
.ok_or_else(|| {
BitrouterError::bad_request(format!(
"reconciliation request not found: {}",
receipt.request_id
))
})?;
let admitted_tariff = stored_tariff(&row)?;
let current = ReconciliationStatus::from_persisted(&row.reconciliation_status);
if current != ReconciliationStatus::Pending {
return Err(BitrouterError::bad_request(format!(
"reconciliation request {} is already terminal ({})",
receipt.request_id,
current.as_str()
)));
}
let receipt_json = serde_json::to_value(receipt).map_err(|error| {
BitrouterError::internal(format!("serialize authoritative receipt: {error}"))
})?;
let usage = authoritative_usage(receipt, receipt_json.clone())?;
let (status, charge_status, mut evidence, provider_id, model_id) = match receipt.state {
SettlementState::Pending => return Ok(ReconciliationStatus::Pending),
SettlementState::NotCharged => {
if usage.prompt_tokens != 0 || usage.completion_tokens != 0 {
return self
.persist_unknown_reconciliation(
row,
&usage,
receipt_json,
"not_charged_receipt_has_usage",
)
.await;
}
let mut evidence = unavailable_charge_evidence(&usage, "authoritative_not_charged");
evidence.status = ChargeStatus::NotCharged;
(
ReconciliationStatus::NotCharged,
ChargeStatus::NotCharged,
evidence,
receipt
.provider_id
.clone()
.unwrap_or(row.provider_id.clone()),
receipt.model_id.clone().unwrap_or(row.model_id.clone()),
)
}
SettlementState::Unknown => {
return self
.persist_unknown_reconciliation(
row,
&usage,
receipt_json,
"authoritative_settlement_unknown",
)
.await;
}
SettlementState::Computed => {
let Some(provider_id) = receipt.provider_id.as_deref() else {
return self
.persist_unknown_reconciliation(
row,
&usage,
receipt_json,
"authoritative_provider_missing",
)
.await;
};
let Some(model_id) = receipt.model_id.as_deref() else {
return self
.persist_unknown_reconciliation(
row,
&usage,
receipt_json,
"authoritative_model_missing",
)
.await;
};
let evidence = if accept_receipt_charge {
let charge = receipt.final_charge_micro_usd.ok_or_else(|| {
BitrouterError::bad_request(
"computed authoritative receipt has no final charge",
)
})?;
ChargeEvidence {
tariff_snapshot: None,
status: ChargeStatus::Computed,
charge_micro_usd: Some(charge),
normalized_usage: usage.normalized_buckets().map_err(|error| {
BitrouterError::bad_request(format!(
"authoritative usage normalization failed: {error:?}"
))
})?,
effective_rates: Default::default(),
pricing_source: PricingSource::AuthoritativeReceipt,
pricing_version: crate::eval::types::canonical_digest(receipt)
.map_err(|error| BitrouterError::internal(error.to_string()))?,
unknown_reason: None,
}
} else {
let candidates = prices.iter().filter(|price| {
price.provider_id == provider_id && price.model_id == model_id
});
if candidates.clone().next().is_none() {
return self
.persist_unknown_reconciliation(
row,
&usage,
receipt_json,
"authoritative_pricing_not_found",
)
.await;
}
let mut matching = candidates
.map(|price| {
let pricing = ModelPricing::cache_aware(
Some(price.input_micro_usd_per_token),
price.cache_read_micro_usd_per_token,
price.cache_write_micro_usd_per_token,
Some(price.output_micro_usd_per_token),
);
super::tariff::override_charge_evidence(
&usage,
&pricing,
admitted_tariff.as_ref(),
)
})
.filter(|evidence| {
evidence.status == ChargeStatus::Computed
&& evidence.charge_micro_usd == receipt.final_charge_micro_usd
});
let Some(evidence) = matching.next() else {
return self
.persist_unknown_reconciliation(
row,
&usage,
receipt_json,
"authoritative_charge_mismatch",
)
.await;
};
if matching
.any(|candidate| candidate.pricing_version != evidence.pricing_version)
{
return self
.persist_unknown_reconciliation(
row,
&usage,
receipt_json,
"authoritative_pricing_ambiguous",
)
.await;
}
evidence
};
(
ReconciliationStatus::Computed,
ChargeStatus::Computed,
evidence,
provider_id.to_string(),
model_id.to_string(),
)
}
};
evidence.tariff_snapshot = admitted_tariff;
let raw_json = serde_json::to_string(&receipt_json).map_err(|error| {
BitrouterError::internal(format!("serialize authoritative usage: {error}"))
})?;
let evidence_json = serde_json::to_string(&evidence).map_err(|error| {
BitrouterError::internal(format!("serialize reconciled charge evidence: {error}"))
})?;
let settled_at = Utc::now().to_rfc3339();
let mut active: requests::ActiveModel = row.into();
active.provider_id = Set(provider_id);
active.model_id = Set(model_id);
active.prompt_tokens = Set(usage.prompt_tokens as i64);
active.completion_tokens = Set(usage.completion_tokens as i64);
active.reasoning_tokens = Set(usage.reasoning_tokens as i64);
active.cache_read_tokens = Set(usage.cache_read_tokens as i64);
active.cache_write_tokens = Set(usage.cache_write_tokens as i64);
let normalized = &evidence.normalized_usage;
active.uncached_input_tokens = Set(normalized.uncached_input_tokens as i64);
active.output_tokens = Set(normalized.output_tokens as i64);
active.usage_origin = Set(UsageOrigin::AuthoritativeReceipt.as_str().to_string());
active.raw_usage_json = Set(Some(raw_json.clone()));
active.charge_status = Set(charge_status.as_str().to_string());
active.charge_evidence_json = Set(Some(evidence_json));
active.estimated_charge_micro_usd = Set(evidence.charge_micro_usd.unwrap_or(0));
active.reconciliation_status = Set(status.as_str().to_string());
active.reconciliation_last_error = Set(None);
active.authoritative_settled_at = Set(Some(settled_at));
active.authoritative_receipt_json = Set(Some(raw_json));
active.update(&self.db).await.map_err(|error| {
BitrouterError::internal(format!("persist authoritative receipt: {error}"))
})?;
Ok(status)
}
async fn persist_unknown_reconciliation(
&self,
row: requests::Model,
usage: &Usage,
receipt_json: serde_json::Value,
reason: &str,
) -> Result<ReconciliationStatus> {
let mut evidence = unavailable_charge_evidence(usage, reason);
evidence.tariff_snapshot = stored_tariff(&row)?;
let evidence_json = serde_json::to_string(&evidence).map_err(|error| {
BitrouterError::internal(format!("serialize unknown reconciliation: {error}"))
})?;
let receipt_json = serde_json::to_string(&receipt_json).map_err(|error| {
BitrouterError::internal(format!("serialize authoritative receipt: {error}"))
})?;
let mut active: requests::ActiveModel = row.into();
active.prompt_tokens = Set(usage.prompt_tokens as i64);
active.completion_tokens = Set(usage.completion_tokens as i64);
active.reasoning_tokens = Set(usage.reasoning_tokens as i64);
active.cache_read_tokens = Set(usage.cache_read_tokens as i64);
active.cache_write_tokens = Set(usage.cache_write_tokens as i64);
active.uncached_input_tokens = Set(evidence.normalized_usage.uncached_input_tokens as i64);
active.output_tokens = Set(evidence.normalized_usage.output_tokens as i64);
active.usage_origin = Set(UsageOrigin::AuthoritativeReceipt.as_str().to_string());
active.raw_usage_json = Set(Some(receipt_json.clone()));
active.charge_status = Set(ChargeStatus::Unknown.as_str().to_string());
active.charge_evidence_json = Set(Some(evidence_json));
active.estimated_charge_micro_usd = Set(0);
active.reconciliation_status = Set(ReconciliationStatus::Unknown.as_str().to_string());
active.reconciliation_last_error = Set(Some(reason.to_string()));
active.authoritative_settled_at = Set(Some(Utc::now().to_rfc3339()));
active.authoritative_receipt_json = Set(Some(receipt_json));
active.update(&self.db).await.map_err(|error| {
BitrouterError::internal(format!("persist unknown reconciliation: {error}"))
})?;
Ok(ReconciliationStatus::Unknown)
}
pub async fn record_request(&self, record: RequestMetric) -> Result<()> {
let session_identity = record.session_identity.as_ref();
let raw_usage_json = record
.raw_usage
.as_ref()
.map(serde_json::to_string)
.transpose()
.map_err(|error| BitrouterError::internal(format!("serialize raw usage: {error}")))?;
let charge_evidence_json = serde_json::to_string(&record.charge_evidence)
.map(Some)
.map_err(|error| {
BitrouterError::internal(format!("serialize charge evidence: {error}"))
})?;
let row = requests::ActiveModel {
request_id: Set(record.request_id),
user_id: Set(record.user_id),
api_key_id: Set(record.api_key_id),
launch_id: Set(record.launch_id),
route_scope_id: Set(session_identity.map(|identity| identity.route_scope_id.clone())),
agent_harness: Set(session_identity.and_then(|identity| identity.agent_harness.clone())),
controller_instance_id: Set(
session_identity.and_then(|identity| identity.controller_instance_id.clone())
),
acp_session_id: Set(
session_identity.and_then(|identity| identity.acp_session_id.clone())
),
native_root_session_id: Set(
session_identity.and_then(|identity| identity.native_root_session_id.clone())
),
native_agent_thread_id: Set(
session_identity.and_then(|identity| identity.native_agent_thread_id.clone())
),
native_parent_agent_thread_id: Set(session_identity
.and_then(|identity| identity.native_parent_agent_thread_id.clone())),
native_turn_id: Set(
session_identity.and_then(|identity| identity.native_turn_id.clone())
),
route_lease_id: Set(
session_identity.and_then(|identity| identity.route_lease_id.clone())
),
session_identity_json: Set(session_identity.map(|identity| identity.serialized.clone())),
router_id: Set(record.router_id),
binding_digest: Set(record.binding_digest),
original_selector: Set(record.original_selector),
model_id: Set(record.model_id),
provider_id: Set(record.provider_id),
prompt_tokens: Set(record.prompt_tokens as i64),
completion_tokens: Set(record.completion_tokens as i64),
reasoning_tokens: Set(record.reasoning_tokens as i64),
cache_read_tokens: Set(record.cache_read_tokens as i64),
cache_write_tokens: Set(record.cache_write_tokens as i64),
uncached_input_tokens: Set(record.uncached_input_tokens as i64),
output_tokens: Set(record.output_tokens as i64),
usage_origin: Set(record.usage_origin.as_str().to_string()),
raw_usage_json: Set(raw_usage_json),
charge_status: Set(record.charge_status.as_str().to_string()),
charge_evidence_json: Set(charge_evidence_json),
reconciliation_status: Set(record.reconciliation_status.as_str().to_string()),
reconciliation_attempts: Set(0),
reconciliation_last_error: Set(None),
reconciliation_last_attempt_at: Set(None),
authoritative_settled_at: Set(None),
authoritative_receipt_json: Set(None),
estimated_charge_micro_usd: Set(record.estimated_charge_micro_usd),
streamed: Set(record.streamed as i32),
latency_ms: Set(record.latency_ms as i64),
generation_time_ms: Set(record.generation_time_ms as i64),
error: Set(record.error),
created_at: Set(Utc::now().to_rfc3339()),
};
requests::Entity::insert(row)
.on_conflict(
OnConflict::column(requests::Column::RequestId)
.update_columns([
requests::Column::UserId,
requests::Column::ApiKeyId,
requests::Column::LaunchId,
requests::Column::AgentHarness,
requests::Column::ControllerInstanceId,
requests::Column::AcpSessionId,
requests::Column::NativeRootSessionId,
requests::Column::NativeAgentThreadId,
requests::Column::NativeParentAgentThreadId,
requests::Column::NativeTurnId,
requests::Column::RouteLeaseId,
requests::Column::SessionIdentityJson,
requests::Column::RouterId,
requests::Column::BindingDigest,
requests::Column::OriginalSelector,
requests::Column::ModelId,
requests::Column::ProviderId,
requests::Column::PromptTokens,
requests::Column::CompletionTokens,
requests::Column::ReasoningTokens,
requests::Column::CacheReadTokens,
requests::Column::CacheWriteTokens,
requests::Column::UncachedInputTokens,
requests::Column::OutputTokens,
requests::Column::UsageOrigin,
requests::Column::RawUsageJson,
requests::Column::ChargeStatus,
requests::Column::ChargeEvidenceJson,
requests::Column::EstimatedChargeMicroUsd,
requests::Column::Streamed,
requests::Column::LatencyMs,
requests::Column::GenerationTimeMs,
requests::Column::Error,
])
.to_owned(),
)
.exec(&self.db)
.await
.map_err(|e| BitrouterError::internal(format!("record_request: {e}")))?;
Ok(())
}
}
fn authoritative_usage(receipt: &SettlementReceipt, raw: serde_json::Value) -> Result<Usage> {
fn tokens(value: i64, bucket: &str) -> Result<u64> {
u64::try_from(value).map_err(|_| {
BitrouterError::bad_request(format!("authoritative receipt has negative {bucket}"))
})
}
let uncached = tokens(receipt.usage.uncached_input_tokens, "uncached_input_tokens")?;
let cache_read = tokens(receipt.usage.cache_read_tokens, "cache_read_tokens")?;
let cache_write = tokens(receipt.usage.cache_write_tokens, "cache_write_tokens")?;
let output = tokens(receipt.usage.output_tokens, "output_tokens")?;
let reasoning = tokens(receipt.usage.reasoning_tokens, "reasoning_tokens")?;
let prompt_tokens = uncached
.checked_add(cache_read)
.and_then(|sum| sum.checked_add(cache_write))
.ok_or_else(|| BitrouterError::bad_request("authoritative input token overflow"))?;
let completion_tokens = output
.checked_add(reasoning)
.ok_or_else(|| BitrouterError::bad_request("authoritative output token overflow"))?;
if prompt_tokens > i64::MAX as u64 || completion_tokens > i64::MAX as u64 {
return Err(BitrouterError::bad_request(
"authoritative token total exceeds local storage range",
));
}
Ok(Usage {
prompt_tokens,
completion_tokens,
reasoning_tokens: reasoning,
cache_read_tokens: cache_read,
cache_write_tokens: cache_write,
origin: UsageOrigin::AuthoritativeReceipt,
raw: Some(Box::new(raw)),
..Default::default()
})
}
fn truncate_error(error: &str) -> String {
const MAX_CHARS: usize = 512;
error.chars().take(MAX_CHARS).collect()
}
impl From<requests::Model> for MeteringUsageRecord {
fn from(row: requests::Model) -> Self {
let status = if row.error.is_some() {
"failed"
} else {
"completed"
};
let charge_status = ChargeStatus::from_persisted(&row.charge_status);
let charge_evidence = row
.charge_evidence_json
.as_deref()
.and_then(|json| serde_json::from_str(json).ok());
let final_charge_micro_usd = if charge_status == ChargeStatus::Computed {
charge_evidence
.as_ref()
.and_then(|evidence: &ChargeEvidence| evidence.charge_micro_usd)
.map(|charge| charge.max(0) as u64)
} else {
None
};
let mut usage_origin = match row.usage_origin.as_str() {
"provider_reported" => UsageOrigin::ProviderReported,
"authoritative_receipt" => UsageOrigin::AuthoritativeReceipt,
"estimated" => UsageOrigin::Estimated,
_ => UsageOrigin::Unknown,
};
let mut raw_usage = row
.raw_usage_json
.as_deref()
.and_then(|json| serde_json::from_str(json).ok());
let error_code = row.error.as_deref().map(|error| {
if error == "upstream_rate_limited" || error.contains("upstream rate limited") {
"upstream_rate_limited".to_string()
} else if error == "upstream_policy_violation"
|| error.contains("upstream policy violation")
|| error.contains("upstream content policy violation")
{
"upstream_policy_violation".to_string()
} else {
"request_failed".to_string()
}
});
let zero_usage_rejection_code = error_code
.as_deref()
.filter(|code| matches!(*code, "upstream_policy_violation" | "upstream_rate_limited"));
let legacy_zero_usage_rejection = usage_origin == UsageOrigin::Unknown
&& raw_usage.is_none()
&& row.prompt_tokens == 0
&& row.completion_tokens == 0
&& row.reasoning_tokens == 0
&& row.cache_read_tokens == 0
&& row.cache_write_tokens == 0;
if let (Some(error_code), true) = (zero_usage_rejection_code, legacy_zero_usage_rejection) {
usage_origin = UsageOrigin::ProviderReported;
raw_usage = Some(serde_json::json!({
"error": { "code": error_code },
"usage": null
}));
}
let authoritative_receipt = row
.authoritative_receipt_json
.as_deref()
.and_then(|json| serde_json::from_str(json).ok());
Self {
id: Some(row.request_id.clone()),
request_id: Some(row.request_id),
provider_id: row.provider_id,
model_id: row.model_id,
prompt_tokens: row.prompt_tokens.max(0) as u64,
completion_tokens: row.completion_tokens.max(0) as u64,
reasoning_tokens: row.reasoning_tokens.max(0) as u64,
uncached_input_tokens: row.uncached_input_tokens.max(0) as u64,
cache_read_tokens: row.cache_read_tokens.max(0) as u64,
cache_write_tokens: row.cache_write_tokens.max(0) as u64,
output_tokens: row.output_tokens.max(0) as u64,
usage_origin,
raw_usage,
final_charge_micro_usd,
charge_status,
charge_evidence,
reconciliation_status: ReconciliationStatus::from_persisted(&row.reconciliation_status),
reconciliation_attempts: row.reconciliation_attempts.max(0) as u32,
authoritative_receipt,
status: Some(status.to_string()),
error_code,
}
}
}
fn window_start(window: TimeWindow) -> DateTime<Utc> {
let now = Utc::now();
match window {
TimeWindow::LastMinute => now - Duration::minutes(1),
TimeWindow::LastHour => now - Duration::hours(1),
TimeWindow::Today => Utc
.with_ymd_and_hms(now.year(), now.month(), now.day(), 0, 0, 0)
.single()
.unwrap_or(now),
TimeWindow::ThisWeek => {
let days_from_monday = now.weekday().num_days_from_monday() as i64;
let midnight = Utc
.with_ymd_and_hms(now.year(), now.month(), now.day(), 0, 0, 0)
.single()
.unwrap_or(now);
midnight - Duration::days(days_from_monday)
}
TimeWindow::ThisMonth => Utc
.with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
.single()
.unwrap_or(now),
TimeWindow::Custom { start, .. } => start,
}
}
fn stored_tariff(row: &requests::Model) -> Result<Option<super::tariff::FrozenTariff>> {
row.charge_evidence_json
.as_deref()
.map(|json| {
serde_json::from_str::<ChargeEvidence>(json)
.map(|evidence| evidence.tariff_snapshot)
.map_err(|_| BitrouterError::internal("stored charge evidence is invalid"))
})
.transpose()
.map(Option::flatten)
}