llmstat 0.1.0

Token usage distribution and cost across local LLM CLIs
//! Incremental cache for append-only JSONL logs (Claude transcripts, Codex
//! rollouts).
//!
//! A file is consumed up to a line boundary: an unterminated tail line that
//! doesn't parse as complete JSON is left for the next run, so a line can
//! never be half-read. Before resuming mid-file we memcmp a 64-byte probe
//! ending at the stored offset — if the prefix was rewritten the file is
//! reparsed from scratch.
//!
//! Cache file: ~/.cache/llmstat/<name>.bin (bincode, versioned).

use serde::{Deserialize, Serialize, de::DeserializeOwned};
use std::collections::HashMap;
use std::io::{BufRead, Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use xxhash_rust::xxh3::Xxh3;

use crate::report::{Call, Usage};

const VERSION: u32 = 2;
const PROBE: usize = 64;

/// 128-bit dedup key. Sources hash their identifying parts — a false
/// collision would drop a real call, so 64 bits is not enough.
pub fn key_of(parts: &[&[u8]]) -> u128 {
    let mut h = Xxh3::new();
    for p in parts {
        // length prefix keeps tuple boundaries unambiguous
        h.update(&(p.len() as u32).to_le_bytes());
        h.update(p);
    }
    h.digest128()
}

/// Per-file string table: session/model names repeat per call, so entries
/// store u32 indices into this instead of copies.
#[derive(Default)]
pub struct Dict {
    strings: Vec<String>,
    index: HashMap<String, u32>,
}

impl Dict {
    /// Rebuild a Dict from a cached `Entry.dict` so resumed parses keep
    /// existing indices stable.
    pub fn from_vec(strings: Vec<String>) -> Self {
        let index = strings
            .iter()
            .enumerate()
            .map(|(i, s)| (s.clone(), i as u32))
            .collect();
        Self { strings, index }
    }

    pub fn intern(&mut self, s: &str) -> u32 {
        if let Some(&i) = self.index.get(s) {
            return i;
        }
        let i = self.strings.len() as u32;
        self.strings.push(s.to_string());
        self.index.insert(s.to_string(), i);
        i
    }

    pub fn into_strings(self) -> Vec<String> {
        self.strings
    }
}

/// Serializable form of `Call` (`Call.source` is `&'static str`, set by the
/// owning source on rehydration; `session`/`model` index the entry's dict).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CachedCall {
    /// xxh3-128 dedup key (see `key_of`).
    pub key: u128,
    /// Index into `Entry::dict`.
    pub session: u32,
    /// Index into `Entry::dict`.
    pub model: u32,
    /// unix seconds
    pub ts: Option<i64>,
    pub usage: Usage,
    pub estimated: bool,
}

impl CachedCall {
    pub fn to_call(&self, source: &'static str, dict: &[Arc<str>]) -> Call {
        Call {
            source,
            session: dict[self.session as usize].clone(),
            model: dict[self.model as usize].clone(),
            ts: self.ts.and_then(|s| chrono::DateTime::from_timestamp(s, 0)),
            usage: self.usage,
            estimated: self.estimated,
        }
    }
}

/// Everything remembered about one scanned file. `S` is the parser state
/// needed to resume mid-file (e.g. Codex's current model/session id).
#[derive(Serialize, Deserialize)]
pub struct Entry<S> {
    /// Bytes consumed, always at a line boundary.
    pub offset: u64,
    /// File mtime (unix secs) at parse time.
    pub mtime: i64,
    /// Up to `PROBE` bytes ending at `offset` — append-safety probe.
    pub boundary: Vec<u8>,
    pub state: S,
    /// String table that `CachedCall.session`/`model` index into.
    pub dict: Vec<String>,
    /// Calls parsed from the file.
    pub entries: Vec<CachedCall>,
}

#[derive(Serialize)]
struct CacheFile<'a, S: Serialize> {
    version: u32,
    files: &'a HashMap<PathBuf, Entry<S>>,
}

/// On-disk shape for reading back.
#[derive(Deserialize)]
struct CacheFileRead<S> {
    version: u32,
    files: HashMap<PathBuf, Entry<S>>,
}

/// The cache is scoped to the directory set being scanned — a `--*-dir`
/// override must not read or clobber the default location's cache.
fn cache_path(name: &str, scope: &[&Path]) -> PathBuf {
    let mut h = Xxh3::new();
    for p in scope {
        let s = p.to_string_lossy();
        h.update(&(s.len() as u32).to_le_bytes());
        h.update(s.as_bytes());
    }
    std::env::home_dir()
        .unwrap_or_else(|| PathBuf::from("~"))
        .join(format!(".cache/llmstat/{name}-{:032x}.bin", h.digest128()))
}

pub fn load<S: DeserializeOwned>(name: &str, scope: &[&Path]) -> HashMap<PathBuf, Entry<S>> {
    let t0 = std::time::Instant::now();
    let Ok(bytes) = std::fs::read(cache_path(name, scope)) else {
        return HashMap::new();
    };
    let n = bytes.len();
    let out = match bincode::deserialize::<CacheFileRead<S>>(&bytes) {
        Ok(c) if c.version == VERSION => c.files,
        _ => HashMap::new(),
    };
    tracing::debug!(name, bytes = n, files = out.len(), elapsed = ?t0.elapsed(), "filecache load");
    out
}

pub fn save<S: Serialize>(name: &str, scope: &[&Path], files: &HashMap<PathBuf, Entry<S>>) {
    let file = cache_path(name, scope);
    if let Some(dir) = file.parent()
        && std::fs::create_dir_all(dir).is_err()
    {
        return;
    }
    let c = CacheFile {
        version: VERSION,
        files,
    };
    let Ok(bytes) = bincode::serialize(&c) else {
        return;
    };
    tracing::debug!(
        name,
        bytes = bytes.len(),
        files = files.len(),
        "filecache save"
    );
    let tmp = file.with_extension("tmp");
    if std::fs::write(&tmp, bytes).is_ok() {
        let _ = std::fs::rename(tmp, file);
    }
}

/// What to do with a file this run.
pub enum Plan<S> {
    /// offset == len and mtime unchanged — reuse cached entries, no reads.
    Reuse,
    /// File grew past the stored offset and the prefix probe verified —
    /// parse only the tail starting at this offset with this state.
    Resume(u64, S),
    /// New, shrunk, rewritten, or unverifiable — parse from byte 0.
    Full,
}

/// Bytes a `Plan` will read — drives the progress-bar threshold.
pub fn plan_bytes<S>(len: u64, p: &Plan<S>) -> u64 {
    match p {
        Plan::Reuse => 0,
        Plan::Resume(o, _) => len - o,
        Plan::Full => len,
    }
}

/// Show a parse progress bar only above this much fresh input — below it the
/// scan finishes before a human can read the bar anyway.
pub const BAR_MIN_BYTES: u64 = 32 << 20;

pub fn mtime(path: &Path) -> i64 {
    std::fs::metadata(path)
        .and_then(|m| m.modified())
        .ok()
        .and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0)
}

/// Classify `path` against its cached entry. `len` comes from the directory
/// walk's metadata. Opens the file only when a tail parse is plausible.
pub fn plan<S: Clone>(path: &Path, len: u64, cached: Option<&Entry<S>>) -> Plan<S> {
    let Some(e) = cached else { return Plan::Full };
    if len == e.offset && e.mtime == mtime(path) {
        return Plan::Reuse;
    }
    if len > e.offset && probe_ok(path, e) {
        return Plan::Resume(e.offset, e.state.clone());
    }
    Plan::Full
}

fn probe_ok<S>(path: &Path, e: &Entry<S>) -> bool {
    let n = e.boundary.len() as u64;
    let Ok(mut f) = std::fs::File::open(path) else {
        return false;
    };
    let mut buf = vec![0u8; n as usize];
    f.seek(SeekFrom::Start(e.offset - n)).is_ok()
        && f.read_exact(&mut buf).is_ok()
        && buf == e.boundary
}

/// Read the ≤`PROBE` bytes ending at `offset` (the probe stored for the next
/// run's append check).
pub fn boundary(path: &Path, offset: u64) -> Vec<u8> {
    let n = offset.min(PROBE as u64);
    let Ok(mut f) = std::fs::File::open(path) else {
        return Vec::new();
    };
    let mut buf = vec![0u8; n as usize];
    if f.seek(SeekFrom::Start(offset - n)).is_err() || f.read_exact(&mut buf).is_err() {
        return Vec::new();
    }
    buf
}

/// Iterate lines of `path` starting at byte `offset`, handing each to `f`.
/// Returns bytes consumed, always at a line boundary. A final line without a
/// trailing newline is consumed only if `is_complete` says it's already a
/// well-formed record (a partial append is left for the next run).
pub fn read_lines(
    path: &Path,
    offset: u64,
    is_complete: impl Fn(&str) -> bool,
    mut f: impl FnMut(&str),
) -> std::io::Result<u64> {
    let mut file = std::fs::File::open(path)?;
    file.seek(SeekFrom::Start(offset))?;
    let mut reader = std::io::BufReader::new(file);
    let mut consumed = offset;
    let mut buf = String::new();
    loop {
        buf.clear();
        let n = match reader.read_line(&mut buf) {
            Ok(0) | Err(_) => break,
            Ok(n) => n,
        };
        if buf.ends_with('\n') {
            consumed += n as u64;
            f(&buf);
        } else {
            if is_complete(&buf) {
                consumed += n as u64;
                f(&buf);
            }
            break;
        }
    }
    Ok(consumed)
}