diff --git a/src/acp/service.rs b/src/acp/service.rs index 551161a..174c587 100644 --- a/src/acp/service.rs +++ b/src/acp/service.rs @@ -755,7 +755,7 @@ impl AcpService { usage.output, usage.cache_read, usage.cache_write, - 0.0, + estimate_session_usage_cost(&session, &usage), )); if usage.output > 0 { assistant.output_tokens = Some( @@ -992,6 +992,27 @@ fn model_reasoning_capability( .filter(|capability| !capability.values().is_empty()) } +fn estimate_session_usage_cost( + session: &AcpSession, + usage: &crate::aisdk::chunk::TokenUsage, +) -> f64 { + crate::model::discovery::Discovery::new_with_custom(Some( + session.config.merged_config.custom_providers.clone(), + )) + .ok() + .map(|discovery| { + discovery.estimate_usage_cost( + &session.provider, + &session.model, + usage.input, + usage.output, + usage.cache_read, + usage.cache_write, + ) + }) + .unwrap_or(0.0) +} + fn model_context_window(config: &LoadedConfig, provider: &str, model: &str) -> Option { let discovery = crate::model::discovery::Discovery::new_with_custom(Some( config.merged_config.custom_providers.clone(), @@ -1884,4 +1905,26 @@ mod tests { ] ); } + + #[test] + fn unknown_model_usage_cost_is_zero() { + let mut session = session_with_config(LoadedConfig { + merged_config: crate::config::configuration::MergedConfig::default(), + raw_merged: serde_json::Value::Null, + diagnostics: Default::default(), + inventory: Default::default(), + project_root: PathBuf::from("/tmp"), + cwd: PathBuf::from("/tmp"), + xdg_config_home: PathBuf::from("/tmp"), + }); + session.provider = "no-such-provider".to_string(); + session.model = "no-such-model".to_string(); + let usage = crate::aisdk::chunk::TokenUsage { + input: 1_000, + output: 1_000, + cache_read: 0, + cache_write: 0, + }; + assert_eq!(estimate_session_usage_cost(&session, &usage), 0.0); + } } diff --git a/src/aisdk/README.md b/src/aisdk/README.md index 8ba7f3c..4d3eb5c 100644 --- a/src/aisdk/README.md +++ b/src/aisdk/README.md @@ -14,7 +14,7 @@ In-tree AI SDK used by the host binary (`mod aisdk` in `src/main.rs`). 2. Do not reintroduce a parallel `aisdk/` workspace package until we intentionally extract and publish. 3. Host app imports via `crate::aisdk::...` (binary module path), not an external crate dep. 4. **No hard coupling to the host product.** Code here must not depend on product-specific concepts (TUI, sessions, tools registry, prefs DB, agent loop, product config, host `crate::` app modules, etc.). Keep this layer a generic multi-provider AI SDK — same spirit as Vercel AI SDK for Rust: providers, messages, streams, tools, retries. Product-specific behavior belongs outside this tree (call sites in the app). -5. **Capabilities ≠ host policy.** Expose provider-executed features as ordinary tools (e.g. `openai::tools::web_search()`, `xai::tools::x_search()` via `ToolTransport`). Do **not** encode host preferences like `websearch.native` / `.hosted_web_search(true)` inside aisdk. The host maps policy → `HostedSearchSelection` / tools to pass. +5. **Capabilities ≠ host policy.** Expose provider-executed features as ordinary tools (e.g. `openai::tools::web_search()`, `xai::tools::x_search()` via `ToolTransport`). Do **not** encode host preferences like `websearch.native` / `.hosted_web_search(true)` inside aisdk. The host maps policy → `HostedSearchSelection` / tools to pass. Providers emit billed `ChunkType::Usage(TokenUsage)` (and optional `ResponseCompleted.usage`); pricing stays in the host. 6. **Logging:** use `crate::log::log(...)` (host-injected via `aisdk::log::set_logger`). Never call host `emit_log!` from this tree. 7. **Debug SSE dumps:** feature-gated behind `aisdk-sse-debug` (optional path via `AISDK_SSE_DEBUG_LOG`). 8. **Boundary check:** run `scripts/check-aisdk-boundary.sh` (also documented in root `AGENTS.md`). @@ -33,6 +33,7 @@ In-tree AI SDK used by the host binary (`mod aisdk` in `src/main.rs`). Done for packaging/host hooks: - Neutral logging (`log` module + host `set_logger`) +- Provider-neutral token usage events (`ChunkType::Usage`) - No `crate::aisdk::...` inside the tree (`mod.rs` / re-exports use `super::`) - Absolute `crate::{chunk,error,...}` paths are crate-root-shaped (host re-exports them today) - Product-leaky debug path renamed/feature-gated diff --git a/src/aisdk/providers/anthropic.rs b/src/aisdk/providers/anthropic.rs index c5a12bd..24c0600 100644 --- a/src/aisdk/providers/anthropic.rs +++ b/src/aisdk/providers/anthropic.rs @@ -989,6 +989,68 @@ mod tests { )); } + #[test] + fn max_tokens_delta_emits_final_usage_then_incomplete() { + let value = serde_json::json!({ + "type": "message_delta", + "delta": { + "stop_reason": "max_tokens", + }, + "usage": { + "input_tokens": 12, + "output_tokens": 34, + "cache_read_input_tokens": 5, + "cache_creation_input_tokens": 2, + }, + }); + let mut pending = PartialAnthropicUsage::default(); + let chunks = anthropic_stream_chunks("message_delta", &value, &mut pending); + + assert!(matches!( + chunks.as_slice(), + [ + Ok(ChunkType::Usage(crate::chunk::TokenUsage { + input: 12, + output: 34, + cache_read: 5, + cache_write: 2, + })), + Ok(ChunkType::Incomplete(_)), + ] + )); + } + + #[test] + fn end_turn_delta_emits_final_usage_then_terminal_reason() { + let value = serde_json::json!({ + "type": "message_delta", + "delta": { + "stop_reason": "end_turn", + }, + "usage": { + "input_tokens": 7, + "output_tokens": 11, + }, + }); + let mut pending = PartialAnthropicUsage::default(); + let chunks = anthropic_stream_chunks("message_delta", &value, &mut pending); + + assert!(matches!( + chunks.as_slice(), + [ + Ok(ChunkType::Usage(crate::chunk::TokenUsage { + input: 7, + output: 11, + cache_read: 0, + cache_write: 0, + })), + Ok(ChunkType::End { + reason: Some(FinishReason::EndTurn) + }), + ] + )); + } + #[test] fn groups_adjacent_tool_calls_and_results() { let messages = vec![ diff --git a/src/app.rs b/src/app.rs index e13d2ba..5e6d868 100644 --- a/src/app.rs +++ b/src/app.rs @@ -9743,11 +9743,10 @@ impl App { let cost = self .discovery .as_ref() - .and_then(|discovery| { - discovery.get_model_pricing(&self.provider_name.to_lowercase(), &self.model) - }) - .map(|pricing| { - pricing.estimate_tokens( + .map(|discovery| { + discovery.estimate_usage_cost( + &self.provider_name, + &self.model, usage.input, usage.output, usage.cache_read, diff --git a/src/main.rs b/src/main.rs index 96ab384..9e50483 100644 --- a/src/main.rs +++ b/src/main.rs @@ -24,6 +24,7 @@ mod remote_mcp; mod session; mod skill; mod sound; +mod stats; mod streaming; mod terminal_title; mod theme; @@ -790,6 +791,25 @@ enum Command { target: Option, }, + /// Show token usage and cost statistics + Stats { + /// Show stats for the last N days (default: all time) + #[arg(long)] + days: Option, + + /// Number of tools to show (default: all) + #[arg(long)] + tools: Option, + + /// Show model statistics; optionally limit to the top N + #[arg(long, num_args = 0..=1, default_missing_value = "all")] + models: Option, + + /// Filter by project (default: all projects, empty string: current project) + #[arg(long, value_name = "PROJECT", num_args = 0..=1, default_missing_value = "")] + project: Option, + }, + /// Manage survive-quit background jobs (list / logs / stop) Jobs { #[command(subcommand)] @@ -1007,6 +1027,32 @@ async fn main() -> Result<()> { Some(Command::Upgrade { target }) => { return crate::upgrade::upgrade(target.as_deref()); } + Some(Command::Stats { + days, + tools, + models, + project, + }) => { + let models = models + .as_deref() + .map(|value| { + if value == "all" { + Ok(None) + } else { + value + .parse::() + .map(Some) + .context("--models must be a non-negative integer") + } + }) + .transpose()?; + return crate::stats::run(crate::stats::StatsOptions { + days: *days, + tools: *tools, + models, + project: project.clone(), + }); + } Some(Command::Jobs { command }) => { let cwd = std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")); match command { @@ -1307,6 +1353,47 @@ mod tests { } } + #[test] + fn parses_stats_command_and_compatible_options() { + let args = Args::try_parse_from([ + "crabcode", + "stats", + "--days", + "7", + "--tools", + "5", + "--models", + "3", + "--project", + "", + ]) + .unwrap(); + + match args.command { + Some(Command::Stats { + days, + tools, + models, + project, + }) => { + assert_eq!(days, Some(7)); + assert_eq!(tools, Some(5)); + assert_eq!(models.as_deref(), Some("3")); + assert_eq!(project.as_deref(), Some("")); + } + other => panic!("expected stats command, got {other:?}"), + } + + let args = Args::try_parse_from(["crabcode", "stats", "--models"]).unwrap(); + assert!(matches!( + args.command, + Some(Command::Stats { + models: Some(ref models), + .. + }) if models == "all" + )); + } + #[test] fn generates_bash_completion() { let script = String::from_utf8(crate::completion::generate_script(Shell::Bash)).unwrap(); @@ -1352,6 +1439,8 @@ mod tests { assert!(help.contains("Usage: crabcode")); assert!(help.contains("completion")); assert!(help.contains("Generate or install shell completions")); + assert!(help.contains("stats")); + assert!(help.contains("Show token usage and cost statistics")); assert!( help.contains("serve Host the current workspace for browser and CLI clients") ); diff --git a/src/model/discovery.rs b/src/model/discovery.rs index 78788ef..05d3814 100644 --- a/src/model/discovery.rs +++ b/src/model/discovery.rs @@ -796,6 +796,20 @@ impl Discovery { model.cost.clone() } + pub fn estimate_usage_cost( + &self, + provider_id: &str, + model_id: &str, + input: u64, + output: u64, + cache_read: u64, + cache_write: u64, + ) -> f64 { + self.get_model_pricing(&provider_id.to_lowercase(), model_id) + .map(|pricing| pricing.estimate_tokens(input, output, cache_read, cache_write)) + .unwrap_or(0.0) + } + pub fn get_model_limit(&self, provider_id: &str, model_id: &str) -> Option { let entry = self.load_cache_entry().ok()??; let provider = entry.data.get(provider_id)?; diff --git a/src/stats.rs b/src/stats.rs new file mode 100644 index 0000000..cbc6e0b --- /dev/null +++ b/src/stats.rs @@ -0,0 +1,718 @@ +use anyhow::{Context, Result}; +use chrono::{Local, TimeZone}; +use rusqlite::Connection; +use serde_json::Value; +use std::collections::{HashMap, HashSet}; +use std::path::Path; +use std::time::{SystemTime, UNIX_EPOCH}; + +const BOX_WIDTH: usize = 56; +const TOOL_BAR_WIDTH: usize = 20; + +#[derive(Clone, Debug, Default)] +pub struct StatsOptions { + pub days: Option, + pub tools: Option, + pub models: Option>, + pub project: Option, +} + +fn usage_from_parts(parts: &str, fallback_output: u64) -> UsageTotals { + let Ok(parts) = serde_json::from_str::>(parts) else { + return UsageTotals { + output: fallback_output, + ..UsageTotals::default() + }; + }; + let usage_parts: Vec<&Value> = parts + .iter() + .filter(|part| part.get("type").and_then(Value::as_str) == Some("usage")) + .collect(); + if usage_parts.is_empty() { + return UsageTotals { + output: fallback_output, + ..UsageTotals::default() + }; + } + + usage_parts + .into_iter() + .fold(UsageTotals::default(), |mut totals, usage| { + totals.input = totals + .input + .saturating_add(usage.get("input").and_then(Value::as_u64).unwrap_or(0)); + totals.output = totals + .output + .saturating_add(usage.get("output").and_then(Value::as_u64).unwrap_or(0)); + totals.cache_read = totals + .cache_read + .saturating_add(usage.get("cache_read").and_then(Value::as_u64).unwrap_or(0)); + totals.cache_write = totals.cache_write.saturating_add( + usage + .get("cache_write") + .and_then(Value::as_u64) + .unwrap_or(0), + ); + totals.cost += usage.get("cost").and_then(Value::as_f64).unwrap_or(0.0); + totals + }) +} + +fn model_title_row() -> String { + let text = "MODEL USAGE"; + let left = (BOX_WIDTH.saturating_sub(text.len())) / 2; + let right = BOX_WIDTH.saturating_sub(text.len() + left); + format!("│{}{}{}│", " ".repeat(left), text, " ".repeat(right)) +} + +#[derive(Clone, Debug, Default, PartialEq)] +struct UsageTotals { + input: u64, + output: u64, + cache_read: u64, + cache_write: u64, + cost: f64, +} + +impl UsageTotals { + fn tokens(&self) -> u64 { + self.input + .saturating_add(self.output) + .saturating_add(self.cache_read) + .saturating_add(self.cache_write) + } +} + +#[derive(Clone, Debug)] +struct SessionRow { + id: i64, + workspace_path: Option, +} + +#[derive(Clone, Debug)] +struct MessageRow { + session_id: i64, + timestamp: i64, + parts: String, + model: Option, + provider: Option, + usage: UsageTotals, +} + +#[derive(Clone, Debug, Default, PartialEq)] +struct ModelStats { + messages: u64, + usage: UsageTotals, +} + +#[derive(Clone, Debug, Default, PartialEq)] +struct StatsReport { + sessions: usize, + messages: usize, + days: usize, + usage: UsageTotals, + average_tokens_per_session: u64, + median_tokens_per_session: u64, + tool_total: u64, + tools: Vec<(String, u64)>, + models: Vec<(String, ModelStats)>, +} + +pub fn run(options: StatsOptions) -> Result<()> { + let conn = crate::persistence::db::get_db_conn()?; + let conn = conn.lock().unwrap(); + let report = collect(&conn, &options, now_timestamp())?; + print!("{}", render(&report, &options)); + Ok(()) +} + +fn now_timestamp() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() as i64 +} + +fn collect(conn: &Connection, options: &StatsOptions, now: i64) -> Result { + let sessions = load_sessions(conn)?; + let messages = load_messages(conn)?; + let project_filter = resolve_project_filter(options.project.as_deref())?; + let cutoff = options + .days + .map(|days| now.saturating_sub((days.saturating_mul(86_400)) as i64)); + + let selected_sessions: HashMap = sessions + .iter() + .filter(|session| project_matches(session, project_filter.as_deref())) + .map(|session| (session.id, session)) + .collect(); + + let filtered_messages: Vec<&MessageRow> = messages + .iter() + .filter(|message| selected_sessions.contains_key(&message.session_id)) + .filter(|message| cutoff.is_none_or(|cutoff| message.timestamp >= cutoff)) + .collect(); + + let active_session_ids: HashSet = filtered_messages + .iter() + .map(|message| message.session_id) + .collect(); + let report_sessions: HashSet = if options.days.is_some() { + active_session_ids + } else { + selected_sessions.keys().copied().collect() + }; + + let mut usage = UsageTotals::default(); + let mut session_tokens: HashMap = + report_sessions.iter().copied().map(|id| (id, 0)).collect(); + let mut tool_counts: HashMap = HashMap::new(); + let mut model_counts: HashMap = HashMap::new(); + let mut active_days = HashSet::new(); + + for message in &filtered_messages { + usage.input = usage.input.saturating_add(message.usage.input); + usage.output = usage.output.saturating_add(message.usage.output); + usage.cache_read = usage.cache_read.saturating_add(message.usage.cache_read); + usage.cache_write = usage.cache_write.saturating_add(message.usage.cache_write); + usage.cost += message.usage.cost; + *session_tokens.entry(message.session_id).or_default() = session_tokens + .get(&message.session_id) + .copied() + .unwrap_or_default() + .saturating_add(message.usage.tokens()); + + if let Some(day) = Local + .timestamp_opt(message.timestamp, 0) + .single() + .map(|timestamp| timestamp.date_naive()) + { + active_days.insert(day); + } + + for tool in tool_names(&message.parts) { + *tool_counts.entry(tool).or_default() += 1; + } + + if let Some(model) = message.model.as_deref().filter(|model| !model.is_empty()) { + let name = match message + .provider + .as_deref() + .filter(|provider| !provider.is_empty()) + { + Some(provider) if !model.starts_with(&format!("{provider}/")) => { + format!("{provider}/{model}") + } + _ => model.to_string(), + }; + let stats = model_counts.entry(name).or_default(); + stats.messages += 1; + stats.usage.input = stats.usage.input.saturating_add(message.usage.input); + stats.usage.output = stats.usage.output.saturating_add(message.usage.output); + stats.usage.cache_read = stats + .usage + .cache_read + .saturating_add(message.usage.cache_read); + stats.usage.cache_write = stats + .usage + .cache_write + .saturating_add(message.usage.cache_write); + stats.usage.cost += message.usage.cost; + } + } + + let mut per_session: Vec = session_tokens.into_values().collect(); + per_session.sort_unstable(); + let average_tokens_per_session = if per_session.is_empty() { + 0 + } else { + usage.tokens() / per_session.len() as u64 + }; + let median_tokens_per_session = median(&per_session); + + let mut tools: Vec<_> = tool_counts.into_iter().collect(); + tools.sort_by(|(name_a, count_a), (name_b, count_b)| { + count_b.cmp(count_a).then_with(|| name_a.cmp(name_b)) + }); + let tool_total = tools.iter().map(|(_, count)| count).sum(); + if let Some(limit) = options.tools { + tools.truncate(limit); + } + + let mut models: Vec<_> = model_counts.into_iter().collect(); + models.sort_by(|(name_a, stats_a), (name_b, stats_b)| { + stats_b + .usage + .tokens() + .cmp(&stats_a.usage.tokens()) + .then_with(|| stats_b.messages.cmp(&stats_a.messages)) + .then_with(|| name_a.cmp(name_b)) + }); + if let Some(Some(limit)) = options.models { + models.truncate(limit); + } + + Ok(StatsReport { + sessions: report_sessions.len(), + messages: filtered_messages.len(), + days: options + .days + .map(|days| days as usize) + .unwrap_or(active_days.len()), + usage, + average_tokens_per_session, + median_tokens_per_session, + tool_total, + tools, + models, + }) +} + +fn load_sessions(conn: &Connection) -> Result> { + let mut statement = conn.prepare( + "SELECT s.id, w.root_path + FROM sessions s + LEFT JOIN workspaces w ON w.id = s.workspace_id", + )?; + let rows = statement.query_map([], |row| { + Ok(SessionRow { + id: row.get(0)?, + workspace_path: row.get(1)?, + }) + })?; + rows.collect::>>() + .context("failed to load sessions for stats") +} + +fn load_messages(conn: &Connection) -> Result> { + let mut statement = conn.prepare( + "SELECT session_id, timestamp, parts, model, provider, + COALESCE(output_tokens, tokens_used, 0) + FROM messages", + )?; + let rows = statement.query_map([], |row| { + let output_tokens: i64 = row.get(5)?; + let usage = usage_from_parts(&row.get::<_, String>(2)?, output_tokens.max(0) as u64); + Ok(MessageRow { + session_id: row.get(0)?, + timestamp: row.get(1)?, + parts: row.get(2)?, + model: row.get(3)?, + provider: row.get(4)?, + usage, + }) + })?; + rows.collect::>>() + .context("failed to load messages for stats") +} + +fn resolve_project_filter(project: Option<&str>) -> Result> { + match project { + None => Ok(None), + Some("") => Ok(Some( + std::env::current_dir()? + .canonicalize() + .unwrap_or(std::env::current_dir()?) + .to_string_lossy() + .into_owned(), + )), + Some(project) => Ok(Some( + Path::new(project) + .canonicalize() + .unwrap_or_else(|_| Path::new(project).to_path_buf()) + .to_string_lossy() + .into_owned(), + )), + } +} + +fn project_matches(session: &SessionRow, project: Option<&str>) -> bool { + let Some(project) = project else { + return true; + }; + session.workspace_path.as_deref().is_some_and(|workspace| { + workspace == project + || Path::new(workspace) + .file_name() + .is_some_and(|name| name.to_string_lossy() == project) + }) +} + +fn tool_names(parts: &str) -> Vec { + serde_json::from_str::>(parts) + .unwrap_or_default() + .into_iter() + .filter(|part| part.get("type").and_then(Value::as_str) == Some("tool_call")) + .filter_map(|part| { + part.get("name") + .and_then(Value::as_str) + .filter(|name| !name.is_empty()) + .map(str::to_string) + }) + .collect() +} + +fn median(sorted: &[u64]) -> u64 { + match sorted.len() { + 0 => 0, + len if len % 2 == 1 => sorted[len / 2], + len => sorted[len / 2 - 1].saturating_add(sorted[len / 2]) / 2, + } +} + +fn render(report: &StatsReport, options: &StatsOptions) -> String { + let mut sections = vec![render_overview(report), render_cost_and_tokens(report)]; + if options.models.is_some() && !report.models.is_empty() { + sections.push(render_models(&report.models)); + } + if !report.tools.is_empty() { + sections.push(render_tools(&report.tools, report.tool_total)); + } + format!("{}\n", sections.join("\n\n")) +} + +fn render_overview(report: &StatsReport) -> String { + render_table( + "OVERVIEW", + &[ + ("Sessions", comma_number(report.sessions as u64)), + ("Messages", comma_number(report.messages as u64)), + ("Days", comma_number(report.days as u64)), + ], + ) +} + +fn render_cost_and_tokens(report: &StatsReport) -> String { + let average_cost = if report.days == 0 { + 0.0 + } else { + report.usage.cost / report.days as f64 + }; + render_table( + "COST & TOKENS", + &[ + ("Total Cost", format!("${:.2}", report.usage.cost)), + ("Avg Cost/Day", format!("${average_cost:.2}")), + ( + "Avg Tokens/Session", + compact_number(report.average_tokens_per_session), + ), + ( + "Median Tokens/Session", + compact_number(report.median_tokens_per_session), + ), + ("Input", compact_number(report.usage.input)), + ("Output", compact_number(report.usage.output)), + ("Cache Read", compact_number(report.usage.cache_read)), + ("Cache Write", compact_number(report.usage.cache_write)), + ], + ) +} + +fn render_models(models: &[(String, ModelStats)]) -> String { + let mut lines = vec![top_border(), model_title_row(), middle_border()]; + for (index, (name, stats)) in models.iter().enumerate() { + lines.push(text_row(&format!(" {name}"))); + lines.push(metric_row(" Messages", &comma_number(stats.messages))); + lines.push(metric_row( + " Input Tokens", + &compact_number(stats.usage.input), + )); + lines.push(metric_row( + " Output Tokens", + &compact_number(stats.usage.output), + )); + lines.push(metric_row( + " Cache Read", + &compact_number(stats.usage.cache_read), + )); + lines.push(metric_row( + " Cache Write", + &compact_number(stats.usage.cache_write), + )); + lines.push(metric_row(" Cost", &format!("${:.4}", stats.usage.cost))); + if index + 1 < models.len() { + lines.push(middle_border()); + } + } + lines.push(bottom_border()); + lines.join("\n") +} + +fn render_tools(tools: &[(String, u64)], total: u64) -> String { + let max = tools.first().map(|(_, count)| *count).unwrap_or(0); + let mut lines = vec![top_border(), centered_row("TOOL USAGE"), middle_border()]; + for (name, count) in tools { + let percentage = if total == 0 { + 0.0 + } else { + *count as f64 * 100.0 / total as f64 + }; + let bar_len = if max == 0 { + 0 + } else { + ((*count as f64 / max as f64) * TOOL_BAR_WIDTH as f64) + .round() + .max(1.0) as usize + }; + let name = truncate(name, 18); + let body = format!( + " {name:<18} {:<20} {count} ({percentage:>4.1}%)", + "█".repeat(bar_len) + ); + lines.push(text_row(&body)); + } + lines.push(bottom_border()); + lines.join("\n") +} + +fn render_table(title: &str, rows: &[(&str, String)]) -> String { + let mut lines = vec![top_border(), centered_row(title), middle_border()]; + lines.extend(rows.iter().map(|(label, value)| metric_row(label, value))); + lines.push(bottom_border()); + lines.join("\n") +} + +fn top_border() -> String { + format!("┌{}┐", "─".repeat(BOX_WIDTH)) +} + +fn middle_border() -> String { + format!("├{}┤", "─".repeat(BOX_WIDTH)) +} + +fn bottom_border() -> String { + format!("└{}┘", "─".repeat(BOX_WIDTH)) +} + +fn centered_row(text: &str) -> String { + let left = (BOX_WIDTH.saturating_sub(text.chars().count()) / 2).saturating_sub(1); + let right = BOX_WIDTH.saturating_sub(text.chars().count() + left); + format!("│{}{}{}│", " ".repeat(left), text, " ".repeat(right)) +} + +fn metric_row(label: &str, value: &str) -> String { + let usable = BOX_WIDTH - 1; + let label = truncate(label, usable.saturating_sub(value.chars().count())); + let spaces = usable.saturating_sub(label.chars().count() + value.chars().count()); + format!("│{label}{}{value} │", " ".repeat(spaces)) +} + +fn text_row(text: &str) -> String { + let text = truncate(text, BOX_WIDTH); + let padding = BOX_WIDTH.saturating_sub(text.chars().count()); + format!("│{text}{}│", " ".repeat(padding)) +} + +fn truncate(value: &str, width: usize) -> String { + if value.chars().count() <= width { + return value.to_string(); + } + if width <= 2 { + return value.chars().take(width).collect(); + } + format!("{}..", value.chars().take(width - 2).collect::()) +} + +fn comma_number(value: u64) -> String { + let digits = value.to_string(); + let mut formatted = String::with_capacity(digits.len() + digits.len() / 3); + for (index, digit) in digits.chars().enumerate() { + if index > 0 && (digits.len() - index).is_multiple_of(3) { + formatted.push(','); + } + formatted.push(digit); + } + formatted +} + +fn compact_number(value: u64) -> String { + const UNITS: [(u64, &str); 4] = [ + (1_000_000_000, "B"), + (1_000_000, "M"), + (1_000, "K"), + (1, ""), + ]; + for (divisor, suffix) in UNITS { + if value >= divisor { + return if divisor == 1 { + value.to_string() + } else { + format!("{:.1}{suffix}", value as f64 / divisor as f64) + }; + } + } + "0".to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::persistence::migrations::run_migrations; + use rusqlite::params; + + fn test_db() -> Connection { + let mut conn = Connection::open_in_memory().unwrap(); + run_migrations(&mut conn).unwrap(); + conn.execute( + "INSERT INTO workspaces (root_path, display_name) VALUES ('/tmp/one', 'one')", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO sessions (session_identifier, name, workspace_id, total_cost) + VALUES ('ses_1', 'One', 1, 1.25), ('ses_2', 'Two', 1, 0.75)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO messages + (id, session_id, role, parts, timestamp, tokens_used, output_tokens, model, provider) + VALUES + ('m1', 1, 'assistant', ?1, 1000, 1200, 1200, 'gpt-test', 'openai'), + ('m2', 1, 'user', '[]', 1001, 0, 0, NULL, NULL), + ('m3', 2, 'assistant', ?2, 90000, 800, 800, 'openai/gpt-test', 'openai')", + params![ + r#"[{"type":"tool_call","name":"read"},{"type":"tool_call","name":"bash"},{"type":"usage","input":4000,"output":1200,"cache_read":3000,"cache_write":500,"cost":0.125}]"#, + r#"[{"type":"tool_call","name":"read"}]"# + ], + ) + .unwrap(); + conn + } + + #[test] + fn collects_totals_tools_models_and_median() { + let report = collect( + &test_db(), + &StatsOptions { + models: Some(None), + ..StatsOptions::default() + }, + 100_000, + ) + .unwrap(); + + assert_eq!(report.sessions, 2); + assert_eq!(report.messages, 3); + assert_eq!(report.usage.input, 4_000); + assert_eq!(report.usage.output, 2_000); + assert_eq!(report.usage.cache_read, 3_000); + assert_eq!(report.usage.cache_write, 500); + assert_eq!(report.usage.cost, 0.125); + assert_eq!(report.average_tokens_per_session, 4_750); + assert_eq!(report.median_tokens_per_session, 4_750); + assert_eq!(report.tool_total, 3); + assert_eq!(report.tools, vec![("read".into(), 2), ("bash".into(), 1)]); + assert_eq!(report.models[0].0, "openai/gpt-test"); + assert_eq!(report.models[0].1.messages, 2); + assert_eq!(report.models[0].1.usage.input, 4_000); + assert_eq!(report.models[0].1.usage.output, 2_000); + } + + #[test] + fn days_filter_counts_only_active_sessions_and_uses_requested_days() { + let report = collect( + &test_db(), + &StatsOptions { + days: Some(1), + ..StatsOptions::default() + }, + 100_000, + ) + .unwrap(); + + assert_eq!(report.sessions, 1); + assert_eq!(report.messages, 1); + assert_eq!(report.days, 1); + assert_eq!(report.usage.output, 800); + } + + #[test] + fn renders_opencode_style_sections() { + let output = render( + &StatsReport { + sessions: 2, + messages: 3, + days: 2, + usage: UsageTotals { + output: 2_000, + cost: 2.0, + ..UsageTotals::default() + }, + average_tokens_per_session: 1_000, + median_tokens_per_session: 1_000, + tool_total: 3, + tools: vec![("read".into(), 2), ("bash".into(), 1)], + ..StatsReport::default() + }, + &StatsOptions::default(), + ); + + assert!(output.contains("│ OVERVIEW │")); + assert!(output.contains("│Sessions 2 │")); + assert!(output.contains("│Output 2.0K │")); + assert!(output.contains("│ TOOL USAGE │")); + assert!(output.contains(" read ████████████████████ 2 (66.7%)")); + assert!(output + .lines() + .filter(|line| !line.is_empty()) + .all(|line| line.chars().count() == 58)); + } + + #[test] + fn compact_numbers_match_stats_display() { + assert_eq!(compact_number(0), "0"); + assert_eq!(compact_number(999), "999"); + assert_eq!(compact_number(1_000), "1.0K"); + assert_eq!(compact_number(10_600_000), "10.6M"); + assert_eq!(compact_number(1_310_600_000), "1.3B"); + } + + #[test] + fn comma_numbers_group_thousands() { + assert_eq!(comma_number(0), "0"); + assert_eq!(comma_number(121), "121"); + assert_eq!(comma_number(999), "999"); + assert_eq!(comma_number(1_000), "1,000"); + assert_eq!(comma_number(4_499), "4,499"); + assert_eq!(comma_number(30_446), "30,446"); + assert_eq!(comma_number(1_234_567), "1,234,567"); + } + + #[test] + fn overview_and_model_counts_use_commas_not_compact() { + let output = render( + &StatsReport { + sessions: 4_499, + messages: 30_446, + days: 121, + models: vec![( + "openai/gpt-test".into(), + ModelStats { + messages: 12_345, + usage: UsageTotals { + input: 4_000, + ..UsageTotals::default() + }, + }, + )], + tools: vec![("read".into(), 1234)], + tool_total: 1234, + ..StatsReport::default() + }, + &StatsOptions { + models: Some(None), + ..StatsOptions::default() + }, + ); + + assert!(output.contains("│Sessions 4,499 │")); + assert!(output.contains("│Messages 30,446 │")); + assert!(output.contains("│Days 121 │")); + assert!(output.contains("│ Messages 12,345 │")); + assert!(output.contains("│ Input Tokens 4.0K │")); + assert!(output.contains(" read ████████████████████ 1234 (100.0%)")); + assert!(!output.contains("1,234 (100.0%)")); + } +}