marver 0.0.28

//! What an agent has spent, read from Claude Code's transcript.

use std::io::{Read, Seek, SeekFrom};
use std::path::Path;

use serde_json::Value;

/// What one read of a transcript found.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct Usage {
    /// The model named by the most recent assistant message.
    pub model: Option<String>,
    /// How full the context window was on that message: its own input, plus
    /// the cache it read, plus the cache it wrote.
    pub context_tokens: Option<u64>,
    /// Output tokens across every assistant message in this read.
    pub output_tokens: u64,
    /// Where reading stopped, to be handed back to the next call.
    pub offset: u64,
    /// Where this read started, so a write can be made conditional on nobody
    /// else having read the same bytes first.
    pub start: u64,
}

/// Read a transcript from `offset`, returning what is new.
pub fn read_from(path: &Path, offset: u64) -> std::io::Result<Usage> {
    let mut file = std::fs::File::open(path)?;
    let len = file.metadata()?.len();

    // Replaced rather than appended to — compaction, a new session writing to
    // the same path — so anything remembered about it is about a different
    // file.
    let start = if len < offset { 0 } else { offset };
    file.seek(SeekFrom::Start(start))?;

    let mut bytes = Vec::new();
    file.read_to_end(&mut bytes)?;

    let mut usage = Usage {
        offset: start,
        start: offset,
        ..Usage::default()
    };
    let mut consumed = 0usize;
    // Split in the file's own bytes, not in a decoded string. Measuring the
    // lossy decode counted three bytes for every invalid one, so the offset
    // stored was past the offset read; once it passed the end, the rewind above
    // took it back to zero and `record_usage` added the whole transcript again
    // — and again on every hook after that, since the overshoot repeats.
    for chunk in bytes.split_inclusive(|&b| b == b'\n') {
        if !chunk.ends_with(b"\n") {
            // Incomplete: leave it for the next read.
            break;
        }
        consumed += chunk.len();
        // Lossy: one invalid byte should cost the read nothing, and the numbers
        // wanted here are ASCII.
        let line = String::from_utf8_lossy(chunk);
        let Ok(value) = serde_json::from_str::<Value>(line.trim_end()) else {
            continue;
        };
        absorb(&mut usage, &value);
    }
    usage.offset = start + consumed as u64;
    Ok(usage)
}

/// Fold one transcript line into a running total.
fn absorb(usage: &mut Usage, line: &Value) {
    if line.get("type").and_then(Value::as_str) != Some("assistant") {
        return;
    }
    let Some(message) = line.get("message") else {
        return;
    };
    if let Some(model) = message.get("model").and_then(Value::as_str) {
        usage.model = Some(model.to_string());
    }
    let Some(u) = message.get("usage") else {
        return;
    };
    let count = |key: &str| u.get(key).and_then(Value::as_u64).unwrap_or(0);

    usage.output_tokens += count("output_tokens");
    // The level as of this message, replacing rather than adding to what an
    // earlier one reported.
    usage.context_tokens = Some(
        count("input_tokens")
            + count("cache_read_input_tokens")
            + count("cache_creation_input_tokens"),
    );
}

/// A model id shortened for a header, e.g. `claude-opus-5` → `opus-5`.
pub fn short_model(model: &str) -> String {
    let stem = model.strip_prefix("claude-").unwrap_or(model);
    // `claude-haiku-4-5-20251001` → `haiku-4-5`
    match stem.rsplit_once('-') {
        Some((head, tail)) if tail.len() == 8 && tail.chars().all(|c| c.is_ascii_digit()) => {
            head.to_string()
        }
        _ => stem.to_string(),
    }
}

/// `303650` → `304k`, for a header column that cannot afford six digits.
pub fn compact(tokens: u64) -> String {
    match tokens {
        0..=9_999 => tokens.to_string(),
        10_000..=999_999 => format!("{}k", tokens.div_ceil(1_000)),
        _ => format!("{:.1}M", tokens as f64 / 1_000_000.0),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write;
    use tempfile::TempDir;

    /// One assistant line, in the shape a real transcript uses.
    fn assistant(
        model: &str,
        input: u64,
        cache_read: u64,
        cache_write: u64,
        output: u64,
    ) -> String {
        format!(
            r#"{{"type":"assistant","message":{{"model":"{model}","usage":{{"input_tokens":{input},"cache_read_input_tokens":{cache_read},"cache_creation_input_tokens":{cache_write},"output_tokens":{output}}}}}}}"#
        )
    }

    fn write(dir: &TempDir, name: &str, lines: &[String]) -> std::path::PathBuf {
        let path = dir.path().join(name);
        let mut file = std::fs::File::create(&path).unwrap();
        for line in lines {
            writeln!(file, "{line}").unwrap();
        }
        path
    }

    #[test]
    fn context_is_the_last_messages_level_and_output_is_a_total() {
        // The two numbers behave differently and confusing them is the easy
        // mistake: a context window that summed across turns would read as
        // hundreds of thousands of tokens over the model's limit.
        let dir = TempDir::new().unwrap();
        let path = write(
            &dir,
            "t.jsonl",
            &[
                assistant("claude-opus-5", 2, 100, 50, 255),
                assistant("claude-opus-5", 1, 300, 20, 1051),
            ],
        );

        let usage = read_from(&path, 0).unwrap();

        assert_eq!(usage.context_tokens, Some(1 + 300 + 20));
        assert_eq!(usage.output_tokens, 255 + 1051);
        assert_eq!(usage.model.as_deref(), Some("claude-opus-5"));
    }

    #[test]
    fn a_second_read_starts_where_the_first_stopped() {
        let dir = TempDir::new().unwrap();
        let path = write(
            &dir,
            "t.jsonl",
            &[assistant("claude-opus-5", 2, 100, 0, 255)],
        );
        let first = read_from(&path, 0).unwrap();
        assert_eq!(first.output_tokens, 255);

        let mut file = std::fs::OpenOptions::new()
            .append(true)
            .open(&path)
            .unwrap();
        writeln!(file, "{}", assistant("claude-opus-5", 1, 300, 0, 40)).unwrap();

        let second = read_from(&path, first.offset).unwrap();

        assert_eq!(
            second.output_tokens, 40,
            "only what is new, or totals would double with every hook"
        );
        assert_eq!(second.context_tokens, Some(301));
    }

    #[test]
    fn a_half_written_line_is_left_for_next_time() {
        // Hooks arrive while Claude Code is writing.
        let dir = TempDir::new().unwrap();
        let path = dir.path().join("t.jsonl");
        let mut file = std::fs::File::create(&path).unwrap();
        writeln!(file, "{}", assistant("claude-opus-5", 2, 100, 0, 255)).unwrap();
        write!(file, r#"{{"type":"assistant","message":{{"mod"#).unwrap();
        file.flush().unwrap();

        let first = read_from(&path, 0).unwrap();
        assert_eq!(first.output_tokens, 255);

        // Now the rest of it arrives.
        let mut file = std::fs::OpenOptions::new()
            .append(true)
            .open(&path)
            .unwrap();
        writeln!(
            file,
            r#"el":"claude-opus-5","usage":{{"output_tokens":7}}}}}}"#
        )
        .unwrap();

        let second = read_from(&path, first.offset).unwrap();
        assert_eq!(second.output_tokens, 7, "the completed line is not lost");
    }

    #[test]
    fn a_replaced_transcript_is_read_from_the_beginning() {
        // Shorter than the remembered offset means a different file at the
        // same path.
        let dir = TempDir::new().unwrap();
        let path = write(
            &dir,
            "t.jsonl",
            &[
                assistant("claude-opus-5", 2, 100, 0, 255),
                assistant("claude-opus-5", 2, 100, 0, 255),
            ],
        );
        let long = read_from(&path, 0).unwrap();

        let path = write(&dir, "t.jsonl", &[assistant("claude-opus-5", 1, 5, 0, 9)]);
        let after = read_from(&path, long.offset).unwrap();

        assert_eq!(after.output_tokens, 9);
        assert_eq!(after.context_tokens, Some(6));
    }

    #[test]
    fn lines_that_are_not_assistant_turns_are_ignored() {
        let dir = TempDir::new().unwrap();
        let path = write(
            &dir,
            "t.jsonl",
            &[
                r#"{"type":"user","message":{"content":"do it"}}"#.to_string(),
                r#"{"type":"system","subtype":"init"}"#.to_string(),
                assistant("claude-opus-5", 2, 100, 0, 255),
                r#"{"type":"file-history-snapshot","snapshot":{}}"#.to_string(),
            ],
        );

        let usage = read_from(&path, 0).unwrap();

        assert_eq!(usage.output_tokens, 255);
    }

    #[test]
    fn nonsense_is_skipped_rather_than_failing_the_read() {
        // The format carries no compatibility promise.
        let dir = TempDir::new().unwrap();
        let path = write(
            &dir,
            "t.jsonl",
            &[
                "not json at all".to_string(),
                r#"{"type":"assistant"}"#.to_string(),
                r#"{"type":"assistant","message":{"model":"claude-opus-5"}}"#.to_string(),
                assistant("claude-opus-5", 2, 100, 0, 255),
                "{".to_string(),
            ],
        );

        let usage = read_from(&path, 0).unwrap();

        assert_eq!(usage.output_tokens, 255);
        assert_eq!(usage.model.as_deref(), Some("claude-opus-5"));
    }

    #[test]
    fn a_transcript_with_no_assistant_turns_reports_nothing_rather_than_zero() {
        // "Nothing yet" and "zero tokens" are different claims, and only the
        // first is true of a task whose agent has not answered.
        let dir = TempDir::new().unwrap();
        let path = write(
            &dir,
            "t.jsonl",
            &[r#"{"type":"user","message":{"content":"do it"}}"#.to_string()],
        );

        let usage = read_from(&path, 0).unwrap();

        assert_eq!(usage.context_tokens, None);
        assert_eq!(usage.output_tokens, 0);
        assert_eq!(usage.model, None);
    }

    #[test]
    fn a_missing_transcript_is_an_error_the_caller_can_ignore() {
        let dir = TempDir::new().unwrap();
        assert!(read_from(&dir.path().join("nope.jsonl"), 0).is_err());
    }

    #[test]
    fn model_names_shorten_without_being_guessed_at() {
        assert_eq!(short_model("claude-opus-5"), "opus-5");
        assert_eq!(short_model("claude-haiku-4-5-20251001"), "haiku-4-5");
        assert_eq!(short_model("claude-sonnet-5"), "sonnet-5");
        // Unknown shapes are left alone rather than mangled.
        assert_eq!(short_model("something-else"), "something-else");
        assert_eq!(short_model("gpt-4"), "gpt-4");
    }

    #[test]
    fn token_counts_shorten_for_a_header() {
        assert_eq!(compact(0), "0");
        assert_eq!(compact(9_999), "9999");
        assert_eq!(compact(10_000), "10k");
        assert_eq!(compact(303_650), "304k");
        assert_eq!(compact(1_500_000), "1.5M");
    }
}