use std::fs; use std::path::Path; use chrono::Utc; use serde::{Deserialize, Serialize}; use super::TokenUsage; /// A single token usage record persisted to disk. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct TokenUsageRecord { pub story_id: String, pub agent_name: String, pub timestamp: String, #[serde(default)] pub model: Option, pub usage: TokenUsage, } /// Append a token usage record to the persistent JSONL file. /// /// Each line is a self-contained JSON object, making appends atomic and /// reads simple. The file lives at `.storkit/token_usage.jsonl`. pub fn append_record(project_root: &Path, record: &TokenUsageRecord) -> Result<(), String> { let path = token_usage_path(project_root); if let Some(parent) = path.parent() { fs::create_dir_all(parent) .map_err(|e| format!("Failed to create token_usage directory: {e}"))?; } let mut line = serde_json::to_string(record).map_err(|e| format!("Failed to serialize record: {e}"))?; line.push('\n'); use std::io::Write; let file = fs::OpenOptions::new() .create(true) .append(true) .open(&path) .map_err(|e| format!("Failed to open token_usage file: {e}"))?; let mut writer = std::io::BufWriter::new(file); writer .write_all(line.as_bytes()) .map_err(|e| format!("Failed to write token_usage record: {e}"))?; writer .flush() .map_err(|e| format!("Failed to flush token_usage file: {e}"))?; Ok(()) } /// Read all token usage records from the persistent file. pub fn read_all(project_root: &Path) -> Result, String> { let path = token_usage_path(project_root); if !path.exists() { return Ok(Vec::new()); } let content = fs::read_to_string(&path).map_err(|e| format!("Failed to read token_usage file: {e}"))?; let mut records = Vec::new(); for line in content.lines() { let trimmed = line.trim(); if trimmed.is_empty() { continue; } match serde_json::from_str::(trimmed) { Ok(record) => records.push(record), Err(e) => { crate::slog_warn!("[token_usage] Skipping malformed line: {e}"); } } } Ok(records) } /// Build a `TokenUsageRecord` from the parts available at completion time. pub fn build_record( story_id: &str, agent_name: &str, model: Option, usage: TokenUsage, ) -> TokenUsageRecord { TokenUsageRecord { story_id: story_id.to_string(), agent_name: agent_name.to_string(), timestamp: Utc::now().to_rfc3339(), model, usage, } } fn token_usage_path(project_root: &Path) -> std::path::PathBuf { project_root.join(".storkit").join("token_usage.jsonl") } #[cfg(test)] mod tests { use super::*; use tempfile::TempDir; fn sample_usage() -> TokenUsage { TokenUsage { input_tokens: 100, output_tokens: 200, cache_creation_input_tokens: 5000, cache_read_input_tokens: 10000, total_cost_usd: 1.57, } } #[test] fn append_and_read_roundtrip() { let dir = TempDir::new().unwrap(); let root = dir.path(); let record = build_record("42_story_foo", "coder-1", None, sample_usage()); append_record(root, &record).unwrap(); let records = read_all(root).unwrap(); assert_eq!(records.len(), 1); assert_eq!(records[0].story_id, "42_story_foo"); assert_eq!(records[0].agent_name, "coder-1"); assert_eq!(records[0].usage, sample_usage()); } #[test] fn multiple_appends_accumulate() { let dir = TempDir::new().unwrap(); let root = dir.path(); let r1 = build_record("s1", "coder-1", None, sample_usage()); let r2 = build_record("s2", "coder-2", None, sample_usage()); append_record(root, &r1).unwrap(); append_record(root, &r2).unwrap(); let records = read_all(root).unwrap(); assert_eq!(records.len(), 2); assert_eq!(records[0].story_id, "s1"); assert_eq!(records[1].story_id, "s2"); } #[test] fn read_empty_returns_empty() { let dir = TempDir::new().unwrap(); let records = read_all(dir.path()).unwrap(); assert!(records.is_empty()); } #[test] fn malformed_lines_are_skipped() { let dir = TempDir::new().unwrap(); let root = dir.path(); let path = root.join(".storkit").join("token_usage.jsonl"); fs::create_dir_all(path.parent().unwrap()).unwrap(); fs::write(&path, "not json\n{\"bad\":true}\n").unwrap(); let records = read_all(root).unwrap(); assert!(records.is_empty()); } #[test] fn token_usage_from_result_event() { let json = serde_json::json!({ "type": "result", "total_cost_usd": 1.57, "usage": { "input_tokens": 7, "output_tokens": 475, "cache_creation_input_tokens": 185020, "cache_read_input_tokens": 810585 } }); let usage = TokenUsage::from_result_event(&json).unwrap(); assert_eq!(usage.input_tokens, 7); assert_eq!(usage.output_tokens, 475); assert_eq!(usage.cache_creation_input_tokens, 185020); assert_eq!(usage.cache_read_input_tokens, 810585); assert!((usage.total_cost_usd - 1.57).abs() < f64::EPSILON); } #[test] fn token_usage_from_result_event_missing_usage() { let json = serde_json::json!({"type": "result"}); assert!(TokenUsage::from_result_event(&json).is_none()); } #[test] fn token_usage_from_result_event_partial_fields() { let json = serde_json::json!({ "type": "result", "total_cost_usd": 0.5, "usage": { "input_tokens": 10, "output_tokens": 20 } }); let usage = TokenUsage::from_result_event(&json).unwrap(); assert_eq!(usage.input_tokens, 10); assert_eq!(usage.output_tokens, 20); assert_eq!(usage.cache_creation_input_tokens, 0); assert_eq!(usage.cache_read_input_tokens, 0); } }