use std::path::Path;
use anyhow::{Context, Result};
use rusqlite::{params, Connection, OpenFlags};
#[derive(Debug, Clone)]
pub(crate) struct TelemetryCostRecord {
pub(crate) session_id: String,
pub(crate) recorded_at_epoch_s: u64,
pub(crate) next_step_key: String,
pub(crate) next_step_title: Option<String>,
pub(crate) model: Option<String>,
pub(crate) session_cost_usd: f64,
pub(crate) input_tokens: u64,
pub(crate) output_tokens: u64,
pub(crate) cache_creation_input_tokens: u64,
pub(crate) cache_read_input_tokens: u64,
pub(crate) blended_total_tokens: Option<u64>,
}
pub(crate) fn insert(conn: &Connection, record: &TelemetryCostRecord) -> Result<()> {
conn.execute(
"INSERT OR REPLACE INTO telemetry_cost (
session_id,
recorded_at_epoch_s,
next_step_key,
next_step_title,
model,
session_cost_usd,
input_tokens,
output_tokens,
cache_creation_input_tokens,
cache_read_input_tokens,
blended_total_tokens
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
params![
record.session_id,
record.recorded_at_epoch_s,
record.next_step_key,
record.next_step_title,
record.model,
record.session_cost_usd,
i64::try_from(record.input_tokens)
.context("input_tokens overflow while writing telemetry cost")?,
i64::try_from(record.output_tokens)
.context("output_tokens overflow while writing telemetry cost")?,
i64::try_from(record.cache_creation_input_tokens)
.context("cache_creation_input_tokens overflow while writing telemetry cost")?,
i64::try_from(record.cache_read_input_tokens)
.context("cache_read_input_tokens overflow while writing telemetry cost")?,
record
.blended_total_tokens
.map(|value| {
i64::try_from(value)
.context("blended_total_tokens overflow while writing telemetry cost")
})
.transpose()?,
],
)?;
Ok(())
}
pub(crate) fn sum_for_focus(
path: &Path,
next_step_key: &str,
exclude_session_id: Option<&str>,
) -> Result<f64> {
if next_step_key.is_empty() || !path.is_file() {
return Ok(0.0);
}
let conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
.with_context(|| format!("open telemetry cost DB: {}", path.display()))?;
if let Some(session_id) = exclude_session_id {
let total = conn.query_row(
"SELECT COALESCE(SUM(session_cost_usd), 0.0)
FROM telemetry_cost
WHERE next_step_key = ?1 AND session_id != ?2",
params![next_step_key, session_id],
|row| row.get::<_, f64>(0),
)?;
return Ok(total);
}
let total = conn.query_row(
"SELECT COALESCE(SUM(session_cost_usd), 0.0)
FROM telemetry_cost
WHERE next_step_key = ?1",
params![next_step_key],
|row| row.get::<_, f64>(0),
)?;
Ok(total)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::schema;
#[test]
fn sum_for_focus_returns_zero_when_key_is_empty() {
let temp = tempfile::tempdir().unwrap();
let total = sum_for_focus(&temp.path().join("state.db"), "", None).unwrap();
assert_eq!(total, 0.0);
}
#[test]
fn sum_for_focus_returns_zero_when_db_is_missing() {
let temp = tempfile::tempdir().unwrap();
let total = sum_for_focus(
&temp.path().join("state.db"),
"github-issues:issue:98",
None,
)
.unwrap();
assert_eq!(total, 0.0);
}
#[test]
fn insert_and_sum_round_trip() {
let conn = Connection::open_in_memory().unwrap();
schema::initialize(&conn).unwrap();
insert(
&conn,
&TelemetryCostRecord {
session_id: "ses_1".to_owned(),
recorded_at_epoch_s: 1_000,
next_step_key: "handoff_title:seeded".to_owned(),
next_step_title: Some("Cost estimation in telemetry".to_owned()),
model: Some("gpt-5".to_owned()),
session_cost_usd: 1.25,
input_tokens: 100,
output_tokens: 50,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 25,
blended_total_tokens: None,
},
)
.unwrap();
let path = tempfile::NamedTempFile::new().unwrap();
let disk_conn = Connection::open(path.path()).unwrap();
schema::initialize(&disk_conn).unwrap();
insert(
&disk_conn,
&TelemetryCostRecord {
session_id: "ses_1".to_owned(),
recorded_at_epoch_s: 1_000,
next_step_key: "handoff_title:seeded".to_owned(),
next_step_title: Some("Cost estimation in telemetry".to_owned()),
model: Some("gpt-5".to_owned()),
session_cost_usd: 1.25,
input_tokens: 100,
output_tokens: 50,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 25,
blended_total_tokens: None,
},
)
.unwrap();
insert(
&disk_conn,
&TelemetryCostRecord {
session_id: "ses_2".to_owned(),
recorded_at_epoch_s: 2_000,
next_step_key: "handoff_title:seeded".to_owned(),
next_step_title: Some("Cost estimation in telemetry".to_owned()),
model: Some("gpt-5".to_owned()),
session_cost_usd: 2.0,
input_tokens: 200,
output_tokens: 100,
cache_creation_input_tokens: 0,
cache_read_input_tokens: 50,
blended_total_tokens: None,
},
)
.unwrap();
let total = sum_for_focus(path.path(), "handoff_title:seeded", Some("ses_2")).unwrap();
assert!((total - 1.25).abs() < f64::EPSILON);
}
}