tokr 0.1.0

Persistent token-usage ledger for AI coding agents. Captures on write, queries forever.
use anyhow::Result;
use rusqlite::{Transaction, params};
use serde::Deserialize;

use crate::pricing::Pricing;

#[derive(Debug, Deserialize)]
pub struct ClaudeLine {
    pub uuid: Option<String>,
    pub timestamp: Option<String>,
    #[serde(default, rename = "isSidechain")]
    pub is_sidechain: bool,
    pub message: Option<ClaudeMessage>,
    pub cwd: Option<String>,
    #[serde(rename = "sessionId")]
    pub session_id: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct ClaudeMessage {
    pub id: Option<String>,
    pub role: Option<String>,
    pub model: Option<String>,
    pub usage: Option<ClaudeUsage>,
}

#[derive(Debug, Deserialize)]
pub struct ClaudeUsage {
    #[serde(default)]
    pub input_tokens: u64,
    #[serde(default)]
    pub output_tokens: u64,
    #[serde(default)]
    pub cache_creation_input_tokens: u64,
    #[serde(default)]
    pub cache_read_input_tokens: u64,
    #[serde(default)]
    pub cache_creation: Option<ClaudeCacheCreation>,
}

#[derive(Debug, Deserialize, Default)]
pub struct ClaudeCacheCreation {
    #[serde(default)]
    pub ephemeral_5m_input_tokens: u64,
    #[serde(default)]
    pub ephemeral_1h_input_tokens: u64,
}

#[derive(Debug)]
pub struct UsageRow {
    pub message_id: String,
    pub source: String,
    pub uuid: String,
    pub session_id: String,
    pub project_path: String,
    pub cwd: Option<String>,
    pub transcript_path: String,
    pub timestamp: String,
    pub model: String,
    pub is_sidechain: bool,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_creation_5m: u64,
    pub cache_creation_1h: u64,
    pub cache_read_tokens: u64,
    pub reasoning_tokens: u64,
}

pub struct FileCtx<'a> {
    pub transcript_path: &'a str,
    pub fallback_session_id: &'a str,
    pub fallback_cwd: Option<&'a str>,
    pub project_path: &'a str,
}

pub fn parse_claude_line(line: &str, ctx: &FileCtx) -> Result<Option<UsageRow>> {
    let line = line.trim();
    if line.is_empty() {
        return Ok(None);
    }
    let row: ClaudeLine = serde_json::from_str(line)?;
    let Some(msg) = row.message else {
        return Ok(None);
    };
    if msg.role.as_deref() != Some("assistant") {
        return Ok(None);
    }
    let Some(usage) = msg.usage else {
        return Ok(None);
    };
    let Some(message_id) = msg.id else {
        return Ok(None);
    };
    let Some(uuid) = row.uuid else {
        return Ok(None);
    };

    let (cc_5m, cc_1h) = match usage.cache_creation {
        Some(c) => (c.ephemeral_5m_input_tokens, c.ephemeral_1h_input_tokens),
        None => (usage.cache_creation_input_tokens, 0),
    };

    Ok(Some(UsageRow {
        message_id,
        source: "claude_code".into(),
        uuid,
        session_id: row
            .session_id
            .unwrap_or_else(|| ctx.fallback_session_id.to_string()),
        project_path: ctx.project_path.to_string(),
        cwd: row
            .cwd
            .or_else(|| ctx.fallback_cwd.map(std::string::ToString::to_string)),
        transcript_path: ctx.transcript_path.to_string(),
        timestamp: row.timestamp.unwrap_or_default(),
        model: msg.model.unwrap_or_else(|| "unknown".into()),
        is_sidechain: row.is_sidechain,
        input_tokens: usage.input_tokens,
        output_tokens: usage.output_tokens,
        cache_creation_5m: cc_5m,
        cache_creation_1h: cc_1h,
        cache_read_tokens: usage.cache_read_input_tokens,
        reasoning_tokens: 0,
    }))
}

pub fn insert_row(tx: &Transaction, row: &UsageRow, pricing: &Pricing) -> Result<bool> {
    let cost = pricing.compute(
        &row.model,
        &row.timestamp,
        row.input_tokens,
        row.output_tokens,
        row.cache_creation_5m,
        row.cache_creation_1h,
        row.cache_read_tokens,
    );

    let n = tx.execute(
        "INSERT OR IGNORE INTO usage_events (
            message_id, source, uuid, session_id, project_path, cwd, transcript_path,
            timestamp, model, is_sidechain,
            input_tokens, output_tokens, cache_creation_5m, cache_creation_1h,
            cache_read_tokens, reasoning_tokens, cost_usd, pricing_version
         ) VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11,?12,?13,?14,?15,?16,?17,?18)",
        params![
            row.message_id,
            row.source,
            row.uuid,
            row.session_id,
            row.project_path,
            row.cwd,
            row.transcript_path,
            row.timestamp,
            row.model,
            i64::from(row.is_sidechain),
            row.input_tokens as i64,
            row.output_tokens as i64,
            row.cache_creation_5m as i64,
            row.cache_creation_1h as i64,
            row.cache_read_tokens as i64,
            row.reasoning_tokens as i64,
            cost.usd,
            cost.version,
        ],
    )?;
    Ok(n > 0)
}