path: root/crates/dinggerd/src/cost.rs

//! What pi agents cost: today, this week, and this month, in UTC (a day
//! starts at 00:00 UTC, a week on Monday, a month on its 1st).
//!
//! pi records each reply's cost in its session file (`message.usage.cost.
//! total`, at `message.timestamp`). A thread of its own reads every session
//! file changed this period, then only what was appended, and counts each
//! reply once: a fork (`/fork`, `--fork`) copies its parent's entries into a
//! new file unchanged, so a reply is known by its entry id and time, and
//! counts for the session it first appeared in (files are read oldest
//! first). The totals include agents since closed, or started over.
//!
//! Published as global tmux options (`@dingger_cost_today`, `_week`,
//! `_month`, dollars with two decimals); each pi pane's own today is `@dingger_cost`, written
//! by the watcher from [`session_today`].

use std::collections::{BTreeMap, HashMap, HashSet};
use std::fs::File;
use std::io::{Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::sync::{Mutex, MutexGuard, OnceLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use serde::Deserialize;
use tracing::{info, warn};

const DAY_MS: i64 = 86_400_000;

/// A UTC day: days since 1970-01-01.
type Day = i64;

/// The days a period covers, from its first day through today.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Periods {
    pub today: Day,
    pub week: Day,
    pub month: Day,
}

impl Periods {
    pub fn at(now_ms: i64) -> Self {
        let today = now_ms.div_euclid(DAY_MS);
        // 1970-01-01 was a Thursday: Monday is 3 days before it.
        let week = today - (today + 3).rem_euclid(7);
        let (year, month, _) = civil_from_days(today);
        Periods { today, week, month: days_from_civil(year, month, 1) }
    }

    /// The first day any period needs.
    fn start(&self) -> Day {
        self.week.min(self.month)
    }
}

/// Howard Hinnant's `civil_from_days`: (year, month, day).
fn civil_from_days(days: Day) -> (i64, u32, u32) {
    let z = days + 719_468;
    let era = z.div_euclid(146_097);
    let doe = z.rem_euclid(146_097);
    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
    let mp = (5 * doy + 2) / 153;
    let day = (doy - (153 * mp + 2) / 5 + 1) as u32;
    let month = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
    let year = yoe + era * 400 + i64::from(month <= 2);
    (year, month, day)
}

fn days_from_civil(year: i64, month: u32, day: u32) -> Day {
    let year = if month <= 2 { year - 1 } else { year };
    let era = year.div_euclid(400);
    let yoe = year.rem_euclid(400);
    let month = i64::from(month);
    let doy = (153 * (if month > 2 { month - 3 } else { month + 9 }) + 2) / 5 + i64::from(day) - 1;
    let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
    era * 146_097 + doe - 719_468
}

/// Dollars spent in each period.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct Totals {
    pub today: f64,
    pub week: f64,
    pub month: f64,
}

impl Totals {
    fn add(&mut self, day: Day, cost: f64, periods: &Periods) {
        if day == periods.today {
            self.today += cost;
        }
        if day >= periods.week && day <= periods.today {
            self.week += cost;
        }
        if day >= periods.month && day <= periods.today {
            self.month += cost;
        }
    }
}

/// One session file, read up to `offset` (the end of its last full line).
#[derive(Debug, Default)]
struct FileCost {
    offset: u64,
    /// The header's time, which orders files: parents before their forks.
    started: String,
    by_day: BTreeMap<Day, f64>,
}

/// Every reply counted, by session file.
#[derive(Debug, Default)]
pub struct Ledger {
    files: HashMap<PathBuf, FileCost>,
    /// Each reply counted, by its entry id and time: the file it counts for.
    counted: HashMap<(String, i64), PathBuf>,
}

#[derive(Deserialize)]
struct Entry {
    #[serde(default)]
    id: Option<String>,
    #[serde(default)]
    message: Option<Message>,
}

#[derive(Deserialize)]
struct Message {
    #[serde(default)]
    role: Option<String>,
    #[serde(default)]
    timestamp: Option<i64>,
    #[serde(default)]
    usage: Option<Usage>,
}

#[derive(Deserialize)]
struct Usage {
    #[serde(default)]
    cost: Option<CostField>,
}

#[derive(Deserialize)]
struct CostField {
    #[serde(default)]
    total: Option<f64>,
}

#[derive(Deserialize)]
struct Header {
    #[serde(default)]
    timestamp: Option<String>,
}

/// A reply's cost, its time, and its id ("" for old entries without one).
fn reply(line: &[u8]) -> Option<(String, i64, f64)> {
    // Most lines are not replies; a cost is rare enough to look for first.
    if !line.windows(6).any(|w| w == b"\"cost\"") {
        return None;
    }
    let entry: Entry = serde_json::from_slice(line).ok()?;
    let message = entry.message?;
    if message.role.as_deref() != Some("assistant") {
        return None;
    }
    let cost = message.usage?.cost?.total?;
    let at = message.timestamp?;
    (cost > 0.0 && cost.is_finite()).then(|| (entry.id.unwrap_or_default(), at, cost))
}

fn header_time(path: &Path) -> String {
    let mut first = String::new();
    if let Ok(mut file) = File::open(path) {
        let mut buf = vec![0u8; 4096];
        if let Ok(n) = file.read(&mut buf) {
            let line = buf[..n].split(|b| *b == b'\n').next().unwrap_or_default();
            first = serde_json::from_slice::<Header>(line).ok().and_then(|h| h.timestamp).unwrap_or_default();
        }
    }
    first
}

impl Ledger {
    /// Read what changed in `files` (each with its modification time):
    /// files changed before the periods began are left out.
    pub fn scan(&mut self, files: Vec<(PathBuf, SystemTime)>, periods: &Periods) {
        let start_ms = periods.start() * DAY_MS;
        let mut files: Vec<(String, PathBuf)> = files
            .into_iter()
            .filter(|(_, modified)| millis(*modified) >= start_ms)
            .map(|(path, _)| {
                let started = match self.files.get(&path) {
                    Some(known) => known.started.clone(),
                    None => header_time(&path),
                };
                (started, path)
            })
            .collect();
        // Parents before their forks: a fork's copies count for the parent.
        files.sort();
        for (started, path) in files {
            if let Err(err) = self.read(&path, started, periods) {
                warn!(path = %path.display(), %err, "cost: cannot read a session file");
            }
        }
        // Days before every period are no longer needed.
        for file in self.files.values_mut() {
            file.by_day.retain(|day, _| *day >= periods.start());
        }
        self.counted.retain(|(_, at), _| at.div_euclid(DAY_MS) >= periods.start());
    }

    fn read(&mut self, path: &Path, started: String, periods: &Periods) -> std::io::Result<()> {
        let mut file = File::open(path)?;
        let len = file.metadata()?.len();
        let known = self.files.get(path).map(|f| f.offset).unwrap_or(0);
        if len < known {
            // Rewritten: count it again from the start.
            self.forget(path);
        }
        let offset = self.files.get(path).map(|f| f.offset).unwrap_or(0);
        if len == offset {
            self.files.entry(path.to_path_buf()).or_insert_with(|| FileCost { started, ..Default::default() });
            return Ok(());
        }
        file.seek(SeekFrom::Start(offset))?;
        let mut bytes = Vec::with_capacity((len - offset) as usize);
        file.take(len - offset).read_to_end(&mut bytes)?;
        // Only full lines: the rest is still being written.
        let Some(end) = bytes.iter().rposition(|b| *b == b'\n') else {
            return Ok(());
        };
        let mut counted = Vec::new();
        for line in bytes[..end].split(|b| *b == b'\n') {
            let Some((id, at, cost)) = reply(line) else { continue };
            let day = at.div_euclid(DAY_MS);
            if day < periods.start() || self.counted.contains_key(&(id.clone(), at)) {
                continue;
            }
            self.counted.insert((id, at), path.to_path_buf());
            counted.push((day, cost));
        }
        let entry = self.files.entry(path.to_path_buf()).or_insert_with(|| FileCost { started, ..Default::default() });
        entry.offset = offset + end as u64 + 1;
        for (day, cost) in counted {
            *entry.by_day.entry(day).or_default() += cost;
        }
        Ok(())
    }

    fn forget(&mut self, path: &Path) {
        self.files.remove(path);
        self.counted.retain(|_, file| file != path);
    }

    pub fn totals(&self, periods: &Periods) -> Totals {
        let mut totals = Totals::default();
        for file in self.files.values() {
            for (day, cost) in &file.by_day {
                totals.add(*day, *cost, periods);
            }
        }
        totals
    }

    /// What a session spent today (its own replies, not a parent's).
    pub fn today(&self, path: &Path, periods: &Periods) -> f64 {
        self.files.get(path).and_then(|f| f.by_day.get(&periods.today)).copied().unwrap_or(0.0)
    }
}

fn millis(time: SystemTime) -> i64 {
    time.duration_since(UNIX_EPOCH).map(|d| d.as_millis() as i64).unwrap_or(0)
}

/// pi's session files: `<agent dir>/sessions/<directory>/<session>.jsonl`.
fn session_files(root: &Path) -> Vec<(PathBuf, SystemTime)> {
    let mut found = Vec::new();
    let Ok(dirs) = std::fs::read_dir(root) else { return found };
    for dir in dirs.flatten() {
        let Ok(entries) = std::fs::read_dir(dir.path()) else { continue };
        for entry in entries.flatten() {
            let path = entry.path();
            if path.extension().is_some_and(|ext| ext == "jsonl") {
                if let Ok(modified) = entry.metadata().and_then(|m| m.modified()) {
                    found.push((path, modified));
                }
            }
        }
    }
    found
}

/// Session files of running agents, wherever they are (reported by the
/// watcher): those outside pi's own directory are read too.
fn agent_sessions() -> MutexGuard<'static, HashSet<PathBuf>> {
    static SESSIONS: OnceLock<Mutex<HashSet<PathBuf>>> = OnceLock::new();
    SESSIONS.get_or_init(Default::default).lock().unwrap_or_else(|poisoned| poisoned.into_inner())
}

pub fn note_session(path: &str) {
    if !path.is_empty() {
        agent_sessions().insert(PathBuf::from(path));
    }
}

/// Today's spend of each session file, as of the last scan.
fn today_by_session() -> MutexGuard<'static, HashMap<PathBuf, f64>> {
    static TODAY: OnceLock<Mutex<HashMap<PathBuf, f64>>> = OnceLock::new();
    TODAY.get_or_init(Default::default).lock().unwrap_or_else(|poisoned| poisoned.into_inner())
}

/// What the session spent today, as of the last scan.
pub fn session_today(path: &str) -> f64 {
    today_by_session().get(Path::new(path)).copied().unwrap_or(0.0)
}

/// Dollars as a tmux option: "" for nothing.
pub fn option_value(dollars: f64) -> String {
    if dollars < 0.005 {
        String::new()
    } else {
        format!("{dollars:.2}")
    }
}

const SCAN_EVERY: Duration = Duration::from_secs(15);

/// Scan now and every [`SCAN_EVERY`] (tests: `DINGGER_COST_SCAN_SECS`),
/// publishing the totals when they change.
pub fn spawn(root: PathBuf) {
    let every = std::env::var("DINGGER_COST_SCAN_SECS")
        .ok()
        .and_then(|secs| secs.parse().ok())
        .map(Duration::from_secs_f64)
        .unwrap_or(SCAN_EVERY);
    std::thread::spawn(move || {
        let mut ledger = Ledger::default();
        let mut published: Option<[String; 3]> = None;
        loop {
            let started = std::time::Instant::now();
            let periods = Periods::at(millis(SystemTime::now()));
            let mut files = session_files(&root);
            let extra: Vec<PathBuf> = agent_sessions().iter().filter(|p| !p.starts_with(&root)).cloned().collect();
            for path in extra {
                if let Ok(modified) = std::fs::metadata(&path).and_then(|m| m.modified()) {
                    files.push((path, modified));
                }
            }
            ledger.scan(files, &periods);
            let totals = ledger.totals(&periods);
            let today: HashMap<PathBuf, f64> = ledger
                .files
                .keys()
                .map(|path| (path.clone(), ledger.today(path, &periods)))
                .filter(|(_, today)| *today > 0.0)
                .collect();
            let changed = {
                let mut known = today_by_session();
                let changed = *known != today;
                *known = today;
                changed
            };
            // The watcher writes each agent's amount on its pane.
            if changed {
                crate::hooks::wake();
            }
            // Always a value: tmux refuses to set an empty one.
            let values = [totals.today, totals.week, totals.month].map(|dollars| format!("{dollars:.2}"));
            if published.as_ref() != Some(&values) {
                if published.is_none() {
                    info!(files = ledger.files.len(), took_ms = started.elapsed().as_millis() as u64, "cost: first scan");
                }
                for (name, value) in ["@dingger_cost_today", "@dingger_cost_week", "@dingger_cost_month"].iter().zip(&values) {
                    if let Err(err) = crate::tmux::set_global(name, value) {
                        warn!(%err, "failed to write {name}");
                    }
                }
                published = Some(values);
            }
            std::thread::sleep(every);
        }
    });
}

#[cfg(test)]
mod tests {
    use super::*;

    fn ms(date: &str) -> i64 {
        // "2026-11-01T09:30" in UTC
        let (d, t) = date.split_once('T').unwrap();
        let p: Vec<i64> = d.split('-').map(|x| x.parse().unwrap()).collect();
        let (h, m) = t.split_once(':').unwrap();
        days_from_civil(p[0], p[1] as u32, p[2] as u32) * DAY_MS
            + h.parse::<i64>().unwrap() * 3_600_000
            + m.parse::<i64>().unwrap() * 60_000
    }

    #[test]
    fn periods_start_at_utc_midnight_monday_and_the_first() {
        let p = Periods::at(ms("2026-11-01T09:30")); // a Sunday
        assert_eq!(civil_from_days(p.today), (2026, 11, 1));
        assert_eq!(civil_from_days(p.week), (2026, 10, 26), "the Monday before");
        assert_eq!(civil_from_days(p.month), (2026, 11, 1));
        let p = Periods::at(ms("2026-11-02T00:00")); // Monday, at midnight UTC
        assert_eq!(civil_from_days(p.week), (2026, 11, 2));
        assert_eq!(civil_from_days(Periods::at(ms("2026-11-01T23:59")).today), (2026, 11, 1));
        let p = Periods::at(ms("2024-03-01T12:00"));
        assert_eq!(civil_from_days(p.month), (2024, 3, 1));
        assert_eq!(civil_from_days(p.today - 1), (2024, 2, 29), "a leap day");
        for day in [-1000, 0, 19_000, 20_759] {
            let (y, m, d) = civil_from_days(day);
            assert_eq!(days_from_civil(y, m, d), day);
        }
    }

    fn line(id: &str, at: i64, cost: f64) -> String {
        format!(
            "{{\"type\":\"message\",\"id\":\"{id}\",\"message\":{{\"role\":\"assistant\",\"timestamp\":{at},\"usage\":{{\"cost\":{{\"total\":{cost}}}}}}}}}\n"
        )
    }

    fn session(dir: &Path, name: &str, started: &str, lines: &[String]) -> PathBuf {
        let path = dir.join(name);
        let mut text = format!("{{\"type\":\"session\",\"timestamp\":\"{started}\"}}\n");
        text.extend(lines.iter().cloned());
        std::fs::write(&path, text).unwrap();
        path
    }

    fn files(paths: &[&PathBuf]) -> Vec<(PathBuf, SystemTime)> {
        paths.iter().map(|p| ((*p).clone(), SystemTime::now())).collect()
    }

    #[test]
    fn each_reply_counts_once_for_the_session_it_first_appeared_in() {
        let dir = std::env::temp_dir().join(format!("dingger-cost-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let now = ms("2026-11-01T12:00");
        let periods = Periods::at(now);
        let old = line("a1", ms("2026-10-27T10:00"), 2.0); // this week, last month
        let today = line("a2", ms("2026-11-01T08:00"), 1.5);
        let before = line("a0", ms("2026-10-01T08:00"), 100.0); // before both periods
        let user = "{\"type\":\"message\",\"id\":\"u1\",\"message\":{\"role\":\"user\",\"timestamp\":1,\"usage\":{\"cost\":{\"total\":9}}}}\n".to_string();
        // The fork copies the parent's entries, then has its own reply.
        let fork_own = line("f1", ms("2026-11-01T09:00"), 0.25);
        let fork = session(&dir, "b-fork.jsonl", "2026-11-01T08:30:00Z", &[before.clone(), old.clone(), today.clone(), fork_own]);
        let parent = session(&dir, "a-parent.jsonl", "2026-10-20T00:00:00Z", &[before, old, user, today]);
        let mut ledger = Ledger::default();
        // Listed fork first: the order read is by start, not by listing.
        ledger.scan(files(&[&fork, &parent]), &periods);
        assert_eq!(ledger.totals(&periods), Totals { today: 1.75, week: 3.75, month: 1.75 });
        assert_eq!(ledger.today(&parent, &periods), 1.5, "the parent's own reply");
        assert_eq!(ledger.today(&fork, &periods), 0.25, "the fork: only its own reply");

        // Appended: read from where it stopped; a line still being written waits.
        let more = line("a3", ms("2026-11-01T11:00"), 0.5);
        let mut f = std::fs::OpenOptions::new().append(true).open(&parent).unwrap();
        use std::io::Write;
        f.write_all(more.as_bytes()).unwrap();
        f.write_all(&line("a4", ms("2026-11-01T11:30"), 4.0).as_bytes()[..20]).unwrap();
        ledger.scan(files(&[&fork, &parent]), &periods);
        assert_eq!(ledger.today(&parent, &periods), 2.0);
        f.write_all(&line("a4", ms("2026-11-01T11:30"), 4.0).as_bytes()[20..]).unwrap();
        ledger.scan(files(&[&fork, &parent]), &periods);
        assert_eq!(ledger.today(&parent, &periods), 6.0, "the line, once finished");

        // The next day: today starts over, the week and month go on.
        let tomorrow = Periods::at(ms("2026-11-02T00:01"));
        let totals = ledger.totals(&tomorrow);
        assert_eq!((totals.today, totals.month), (0.0, 6.25));
        assert_eq!(totals.week, 0.0, "a new week on Monday");

        // A file changed before the periods began is not read.
        let stale = session(&dir, "c.jsonl", "2026-09-01T00:00:00Z", &[line("s1", ms("2026-11-01T01:00"), 7.0)]);
        ledger.scan(vec![(stale, UNIX_EPOCH + Duration::from_secs(1))], &periods);
        assert_eq!(ledger.totals(&periods).today, 6.25);

        // Rewritten shorter: counted again from the start, still once.
        session(&dir, "a-parent.jsonl", "2026-10-20T00:00:00Z", &[line("a2", ms("2026-11-01T08:00"), 1.5)]);
        ledger.scan(files(&[&fork, &parent]), &periods);
        assert_eq!(ledger.today(&parent, &periods), 1.5);
        assert_eq!(ledger.totals(&periods).today, 1.75);
        std::fs::remove_dir_all(&dir).unwrap();
    }

    #[test]
    fn dollars_as_an_option() {
        assert_eq!(option_value(0.0), "");
        assert_eq!(option_value(0.004), "");
        assert_eq!(option_value(19.849), "19.85");
        assert_eq!(option_value(1967.14), "1967.14");
    }
}