tokr 0.1.0

Persistent token-usage ledger for AI coding agents. Captures on write, queries forever.
use anyhow::{Context, Result};
use serde::Deserialize;
use std::io::{Read, Seek, SeekFrom};
use std::path::Path;

use crate::{codex, db::Db, ingest, paths, pricing::Pricing};

#[derive(Deserialize)]
struct HookInput {
    transcript_path: Option<String>,
    session_id: Option<String>,
    cwd: Option<String>,
}

enum Source {
    ClaudeCode,
    Codex,
}

fn detect_source(transcript_path: &str) -> Source {
    let normalized = transcript_path.replace('\\', "/");
    if normalized.split('/').any(|part| part == ".codex") {
        return Source::Codex;
    }
    let name_is_rollout = transcript_path
        .trim_end_matches(['/', '\\'])
        .rsplit(['/', '\\'])
        .find(|part| !part.is_empty())
        .is_some_and(|n| n.starts_with("rollout-"));
    if name_is_rollout {
        Source::Codex
    } else {
        Source::ClaudeCode
    }
}

pub fn run(final_pass: bool) -> Result<()> {
    let mut buf = String::new();
    std::io::stdin()
        .read_to_string(&mut buf)
        .context("reading hook stdin")?;

    let input: HookInput = match serde_json::from_str(&buf) {
        Ok(v) => v,
        Err(e) => {
            if std::env::var_os("TOKR_DEBUG").is_some() {
                eprintln!("tokr hook: ignoring malformed stdin: {e}");
            }
            return Ok(());
        }
    };
    let Some(transcript_path) = input.transcript_path else {
        if std::env::var_os("TOKR_DEBUG").is_some() {
            eprintln!("tokr hook: stdin payload has no transcript_path");
        }
        return Ok(());
    };
    if !Path::new(&transcript_path).exists() {
        if std::env::var_os("TOKR_DEBUG").is_some() {
            eprintln!("tokr hook: transcript_path does not exist: {transcript_path}");
        }
        return Ok(());
    }

    let mut db = Db::open()?;
    let pricing = Pricing::load()?;

    match detect_source(&transcript_path) {
        Source::Codex => run_codex(&transcript_path, &mut db, &pricing),
        Source::ClaudeCode => run_claude(
            &transcript_path,
            &input.session_id.unwrap_or_default(),
            input.cwd.as_deref(),
            final_pass,
            &mut db,
            &pricing,
        ),
    }
}

fn run_claude(
    transcript_path: &str,
    session_id: &str,
    cwd: Option<&str>,
    final_pass: bool,
    db: &mut Db,
    pricing: &Pricing,
) -> Result<()> {
    let project_path = paths::project_from_transcript(transcript_path);

    let cursor = if final_pass {
        crate::db::Cursor::default()
    } else {
        db.cursor_for(transcript_path)?
    };

    let mut f = std::fs::File::open(transcript_path)
        .with_context(|| format!("opening {transcript_path}"))?;
    let len = f.metadata()?.len();
    if cursor.byte_offset > len {
        f.seek(SeekFrom::Start(0))?;
    } else {
        f.seek(SeekFrom::Start(cursor.byte_offset))?;
    }

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

    let consume_to = match bytes.iter().rposition(|b| *b == b'\n') {
        Some(i) => i + 1,
        None => 0,
    };
    let safe = &bytes[..consume_to];

    let ctx = ingest::FileCtx {
        transcript_path,
        fallback_session_id: session_id,
        fallback_cwd: cwd,
        project_path: &project_path,
    };

    let tx = db.conn.transaction()?;
    let mut last_uuid: Option<String> = cursor.last_uuid.clone();
    for line in safe.split(|b| *b == b'\n') {
        if line.is_empty() {
            continue;
        }
        let Ok(line_str) = std::str::from_utf8(line) else {
            continue;
        };
        if let Ok(Some(row)) = ingest::parse_claude_line(line_str, &ctx) {
            last_uuid = Some(row.uuid.clone());
            let _ = ingest::insert_row(&tx, &row, pricing)?;
        }
    }
    tx.commit()?;

    let new_offset = if final_pass {
        cursor.byte_offset.max(consume_to as u64)
    } else {
        cursor.byte_offset + consume_to as u64
    };
    db.set_cursor(transcript_path, new_offset, last_uuid.as_deref())?;

    Ok(())
}

fn run_codex(transcript_path: &str, db: &mut Db, pricing: &Pricing) -> Result<()> {
    let path = Path::new(transcript_path);
    let Ok(rows) = codex::parse_rollout(path, None) else {
        return Ok(());
    };
    let tx = db.conn.transaction()?;
    for row in rows {
        ingest::insert_row(&tx, &row, pricing)?;
    }
    tx.commit()?;

    if let Ok(meta) = std::fs::metadata(path) {
        db.set_cursor(transcript_path, meta.len(), None)?;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{Source, detect_source};

    #[test]
    fn detects_codex_from_unix_path() {
        assert!(matches!(
            detect_source("/home/alice/.codex/sessions/2026/04/rollout-demo.jsonl"),
            Source::Codex
        ));
    }

    #[test]
    fn detects_codex_from_windows_path() {
        assert!(matches!(
            detect_source(r"C:\Users\Alice\.codex\sessions\2026\04\rollout-demo.jsonl"),
            Source::Codex
        ));
    }

    #[test]
    fn detects_claude_from_transcript_path() {
        assert!(matches!(
            detect_source("/home/alice/.claude/projects/-home-alice-repo/session.jsonl"),
            Source::ClaudeCode
        ));
    }
}