From 2efb29a2626ebeb9aad0ad40c3369a608359794c Mon Sep 17 00:00:00 2001 From: Yanuar Date: Mon, 31 Aug 2026 15:29:16 +0700 Subject: [PATCH 1/6] feat(cli): add stats command --- src/main.rs | 88 ++++++++ src/stats.rs | 605 +++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 693 insertions(+) create mode 100644 src/stats.rs diff --git a/src/main.rs b/src/main.rs index c48eefb..b7aa4f7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -23,6 +23,7 @@ mod remote_mcp; mod session; mod skill; mod sound; +mod stats; mod streaming; mod terminal_title; mod theme; @@ -781,6 +782,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)] @@ -1025,6 +1045,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 { @@ -1325,6 +1371,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 = @@ -1363,6 +1450,7 @@ mod tests { let help = root_help().unwrap(); assert!(help.contains("Usage: crabcode")); assert!(help.contains("completion Generate shell completion script")); + assert!(help.contains("stats Show token usage and cost statistics")); assert!( help.contains("serve Host the current workspace for browser and CLI clients") ); diff --git a/src/stats.rs b/src/stats.rs new file mode 100644 index 0000000..114fc08 --- /dev/null +++ b/src/stats.rs @@ -0,0 +1,605 @@ +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 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, + total_cost: f64, +} + +#[derive(Clone, Debug)] +struct MessageRow { + session_id: i64, + timestamp: i64, + parts: String, + model: Option, + provider: Option, + output_tokens: u64, +} + +#[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.output = usage.output.saturating_add(message.output_tokens); + *session_tokens.entry(message.session_id).or_default() = session_tokens + .get(&message.session_id) + .copied() + .unwrap_or_default() + .saturating_add(message.output_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.output = stats.usage.output.saturating_add(message.output_tokens); + } + } + + usage.cost = report_sessions + .iter() + .filter_map(|id| selected_sessions.get(id)) + .map(|session| session.total_cost) + .sum(); + + 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, s.total_cost + 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)?, + total_cost: row.get(2)?, + }) + })?; + 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)?; + Ok(MessageRow { + session_id: row.get(0)?, + timestamp: row.get(1)?, + parts: row.get(2)?, + model: row.get(3)?, + provider: row.get(4)?, + output_tokens: output_tokens.max(0) as u64, + }) + })?; + 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", report.sessions.to_string()), + ("Messages", report.messages.to_string()), + ("Days", report.days.to_string()), + ], + ) +} + +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", &stats.messages.to_string())); + 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 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"}]"#, + 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.output, 2_000); + assert_eq!(report.average_tokens_per_session, 1_000); + assert_eq!(report.median_tokens_per_session, 1_000); + 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.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 + .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"); + } +} From 0ccb647b3c3d97f944cb2b98fcb7e66963c2d278 Mon Sep 17 00:00:00 2001 From: Yanuar Date: Mon, 31 Aug 2026 16:27:28 +0700 Subject: [PATCH 2/6] fix(stats): persist model token usage --- src/acp/service.rs | 12 ++++ src/agent/subagent.rs | 10 +++ src/aisdk/README.md | 1 + src/aisdk/chunk.rs | 18 +++++ src/aisdk/providers/anthropic.rs | 26 ++++++++ src/aisdk/providers/compatible.rs | 27 +++++--- src/aisdk/providers/openai.rs | 15 +++-- src/aisdk/response.rs | 7 ++ src/app.rs | 29 +++++++++ src/llm/client.rs | 7 ++ src/llm/mod.rs | 1 + src/main.rs | 1 + src/session/types.rs | 13 ++++ src/stats.rs | 87 ++++++++++++++++++++----- src/ui/components/chat.rs | 105 +++++++++++++++++++++++++++++- 15 files changed, 327 insertions(+), 32 deletions(-) diff --git a/src/acp/service.rs b/src/acp/service.rs index 1277e7e..841a910 100644 --- a/src/acp/service.rs +++ b/src/acp/service.rs @@ -747,6 +747,18 @@ impl AcpService { base_context_tokens.saturating_add(token_count), )?; } + crate::llm::ChunkMessage::Usage(usage) => { + assistant + .parts + .push(crate::session::types::MessagePart::usage( + usage.input, + usage.output, + usage.cache_read, + usage.cache_write, + 0.0, + )); + assistant.output_tokens = Some(usage.output as usize); + } crate::llm::ChunkMessage::Cancelled => cancelled = true, crate::llm::ChunkMessage::Failed(error) => failed = Some(error), crate::llm::ChunkMessage::PermissionRequest(prompt) => { diff --git a/src/agent/subagent.rs b/src/agent/subagent.rs index 00c46d2..75c9bc4 100644 --- a/src/agent/subagent.rs +++ b/src/agent/subagent.rs @@ -225,6 +225,11 @@ pub async fn run_subagent( message ); } + ChunkType::Usage(usage) => { + if let Some(sender) = sender.as_ref() { + let _ = sender.send(crate::llm::ChunkMessage::Usage(usage)); + } + } ChunkType::Retry(status) => { if let Some(sender) = sender.as_ref() { let _ = sender.send(crate::llm::ChunkMessage::Retry(status)); @@ -318,6 +323,11 @@ pub async fn run_subagent( message ); } + ChunkType::Usage(usage) => { + if let Some(sender) = sender.as_ref() { + let _ = sender.send(crate::llm::ChunkMessage::Usage(usage)); + } + } ChunkType::Retry(status) => { if let Some(sender) = sender.as_ref() { let _ = sender.send(crate::llm::ChunkMessage::Retry(status)); diff --git a/src/aisdk/README.md b/src/aisdk/README.md index 8ba7f3c..6538337 100644 --- a/src/aisdk/README.md +++ b/src/aisdk/README.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/chunk.rs b/src/aisdk/chunk.rs index c062ccd..0e9c1a3 100644 --- a/src/aisdk/chunk.rs +++ b/src/aisdk/chunk.rs @@ -20,6 +20,7 @@ pub enum ChunkType { end_turn: Option, reasoning_items: Vec, doom_loop_triggers: Vec, + usage: Option, }, Retry(crate::retry::RetryStatus), StreamRollback { @@ -28,6 +29,7 @@ pub enum ChunkType { }, Warning(String), Metadata(String), + Usage(TokenUsage), End { reason: Option, }, @@ -37,6 +39,21 @@ pub enum ChunkType { NotSupported(String), } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct TokenUsage { + /// Non-cached input tokens. + pub input: u64, + pub output: u64, + pub cache_read: u64, + pub cache_write: u64, +} + +impl TokenUsage { + pub fn is_empty(self) -> bool { + self.input == 0 && self.output == 0 && self.cache_read == 0 && self.cache_write == 0 + } +} + #[derive(Debug, Clone, PartialEq, Eq, Default)] pub struct ReasoningReplayItem { pub id: Option, @@ -58,6 +75,7 @@ impl ChunkType { end_turn, reasoning_items: Vec::new(), doom_loop_triggers: Vec::new(), + usage: None, } } } diff --git a/src/aisdk/providers/anthropic.rs b/src/aisdk/providers/anthropic.rs index c34c585..a04d55c 100644 --- a/src/aisdk/providers/anthropic.rs +++ b/src/aisdk/providers/anthropic.rs @@ -20,6 +20,28 @@ pub struct Anthropic { reasoning_effort: Option, } +fn anthropic_usage(usage: &serde_json::Value) -> Option { + let usage = crate::chunk::TokenUsage { + input: usage + .get("input_tokens") + .and_then(|value| value.as_u64()) + .unwrap_or(0), + output: usage + .get("output_tokens") + .and_then(|value| value.as_u64()) + .unwrap_or(0), + cache_read: usage + .get("cache_read_input_tokens") + .and_then(|value| value.as_u64()) + .unwrap_or(0), + cache_write: usage + .get("cache_creation_input_tokens") + .and_then(|value| value.as_u64()) + .unwrap_or(0), + }; + (!usage.is_empty()).then_some(usage) +} + impl Anthropic { pub fn builder() -> AnthropicBuilder { AnthropicBuilder::default() @@ -237,6 +259,7 @@ fn anthropic_stream_chunk( // Partial usage early in the stream (cache fields may already appear). if let Some(usage) = value.get("message").and_then(|m| m.get("usage")) { log_anthropic_usage(usage); + return anthropic_usage(usage).map(ChunkType::Usage).map(Ok); } None } @@ -256,6 +279,9 @@ fn anthropic_stream_chunk( // Final usage wins for cache_read / cache_creation. if let Some(usage) = value.get("usage") { log_anthropic_usage(usage); + if let Some(usage) = anthropic_usage(usage) { + return Some(Ok(ChunkType::Usage(usage))); + } } anthropic_message_delta(value).map(Ok) } diff --git a/src/aisdk/providers/compatible.rs b/src/aisdk/providers/compatible.rs index bbd0aa3..552603c 100644 --- a/src/aisdk/providers/compatible.rs +++ b/src/aisdk/providers/compatible.rs @@ -436,7 +436,7 @@ fn debug_log(msg: &str) { /// Log OpenAI-compatible / AI Gateway usage via the host logger. /// Looks for `prompt_tokens_details.cached_tokens` and Anthropic-style fields /// that some gateways forward. -fn log_openai_compatible_usage(usage: &serde_json::Value) { +fn openai_compatible_usage(usage: &serde_json::Value) -> Option { let prompt = usage.get("prompt_tokens").and_then(|v| v.as_u64()); let completion = usage.get("completion_tokens").and_then(|v| v.as_u64()); let cached = usage @@ -459,7 +459,7 @@ fn log_openai_compatible_usage(usage: &serde_json::Value) { && cache_read == 0 && cache_creation == 0 { - return; + return None; } // Prefer OpenAI-style cached_tokens; fall back to Anthropic-style cache_read. @@ -489,6 +489,13 @@ fn log_openai_compatible_usage(usage: &serde_json::Value) { cache_creation, hit_pct )); + + Some(crate::chunk::TokenUsage { + input: prompt_v.saturating_sub(effective_cached), + output: completion.unwrap_or(0), + cache_read: effective_cached, + cache_write: cache_creation, + }) } fn process_sse_data(data: &str) -> Vec> { @@ -525,26 +532,30 @@ fn process_sse_data(data: &str) -> Vec> { // Final usage often arrives on a choices-empty (or choices-missing) chunk. // Log cache-related fields so gateway Anthropic hits are verifiable. - if let Some(usage) = value.get("usage") { - log_openai_compatible_usage(usage); - } + let usage = value.get("usage").and_then(openai_compatible_usage); let Some(choices) = value["choices"].as_array() else { debug_log(&format!( "[SSE] No choices array. JSON keys: {:?}", value.as_object().map(|o| o.keys().collect::>()) )); - return vec![]; + return usage + .map(|usage| vec![Ok(ChunkType::Usage(usage))]) + .unwrap_or_default(); }; if choices.is_empty() { debug_log("[SSE] choices array is empty"); - return vec![]; + return usage + .map(|usage| vec![Ok(ChunkType::Usage(usage))]) + .unwrap_or_default(); } let choice = &choices[0]; let finish_reason = choice["finish_reason"].as_str().unwrap_or(""); - let mut chunks = Vec::new(); + let mut chunks = usage + .map(|usage| vec![Ok(ChunkType::Usage(usage))]) + .unwrap_or_default(); // Log the full choice structure for debugging debug_log(&format!( diff --git a/src/aisdk/providers/openai.rs b/src/aisdk/providers/openai.rs index 45e60fe..f5eeed4 100644 --- a/src/aisdk/providers/openai.rs +++ b/src/aisdk/providers/openai.rs @@ -817,7 +817,7 @@ impl OpenAI { progress.record_chunk(chunk); } if tx.send(chunk).is_err() { - return; + return None; } if is_completed { let mut state = websocket_state.lock().await; @@ -1545,14 +1545,12 @@ fn response_sse_data_to_chunk(data: &str) -> Option> { return Some(Ok(responses_error_chunk(&value, event_type))); } let resp = &value["response"]; - if let Some(usage) = resp.get("usage") { - log_openai_responses_usage(usage); - } log_openai_responses_completed(resp); Some(Ok(ChunkType::ResponseCompleted { end_turn: resp.get("end_turn").and_then(|value| value.as_bool()), reasoning_items: reasoning_items_from_response_output(resp), doom_loop_triggers: doom_loop_triggers_from(resp), + usage: resp.get("usage").and_then(openai_responses_usage), })) } // Grok Build / cli-chat-proxy: `response.doom_loop_check` with @@ -1589,7 +1587,7 @@ fn response_sse_data_to_chunk(data: &str) -> Option> { /// Log Responses API usage for prompt-cache visibility. /// Looks for `input_tokens_details.cached_tokens` (OpenAI/xAI shape). -fn log_openai_responses_usage(usage: &serde_json::Value) { +fn openai_responses_usage(usage: &serde_json::Value) -> Option { let input = usage .get("input_tokens") .or_else(|| usage.get("prompt_tokens")) @@ -1623,6 +1621,13 @@ fn log_openai_responses_usage(usage: &serde_json::Value) { cached, hit_pct )); + + Some(crate::chunk::TokenUsage { + input: input_v.saturating_sub(cached), + output: output.unwrap_or(0), + cache_read: cached, + cache_write: 0, + }) } /// Attribute a `response.completed` payload: status, incomplete reason, and diff --git a/src/aisdk/response.rs b/src/aisdk/response.rs index cef392f..9145157 100644 --- a/src/aisdk/response.rs +++ b/src/aisdk/response.rs @@ -260,9 +260,13 @@ pub async fn stream_with_tools( end_turn, reasoning_items, doom_loop_triggers, + usage, }) => { saw_terminal_event = true; response_end_turn = end_turn; + if let Some(usage) = usage { + let _ = tx_loop.send(ChunkType::Usage(usage)); + } for item in reasoning_items { merge_reasoning_replay_item(&mut reasoning_replay_items, item); } @@ -342,6 +346,9 @@ pub async fn stream_with_tools( } let _ = tx_loop.send(ChunkType::Metadata(msg)); } + Ok(ChunkType::Usage(usage)) => { + let _ = tx_loop.send(ChunkType::Usage(usage)); + } Ok(ChunkType::Warning(msg)) => { let _ = tx_loop.send(ChunkType::Warning(msg)); } diff --git a/src/app.rs b/src/app.rs index b213216..f80edb1 100644 --- a/src/app.rs +++ b/src/app.rs @@ -9728,6 +9728,35 @@ impl App { push_toast(Toast::new(msg, ToastLevel::Warning, None)); true } + crate::llm::ChunkMessage::Usage(usage) => { + let cost = self + .discovery + .as_ref() + .and_then(|discovery| { + discovery.get_model_pricing(&self.provider_name.to_lowercase(), &self.model) + }) + .map(|pricing| { + let per_million = 1_000_000.0; + usage.input as f64 / per_million * pricing.input + + usage.output as f64 / per_million * pricing.output + + usage.cache_read as f64 / per_million + * pricing.cache_read.unwrap_or(pricing.input) + + usage.cache_write as f64 / per_million + * pricing.cache_write.unwrap_or(pricing.input) + }) + .unwrap_or(0.0); + if let Some(chat) = self.chat_for_session_mut(session_id) { + chat.record_usage( + usage.input, + usage.output, + usage.cache_read, + usage.cache_write, + cost, + ); + } + self.mark_streaming_snapshot_pending(session_id); + true + } crate::llm::ChunkMessage::End => { self.finish_streaming_session(session_id); false diff --git a/src/llm/client.rs b/src/llm/client.rs index ff30a07..0273806 100644 --- a/src/llm/client.rs +++ b/src/llm/client.rs @@ -970,6 +970,7 @@ pub async fn summarize_for_compaction( | ChunkType::RetryableFailure(_) | ChunkType::Warning(_) | ChunkType::Metadata(_) + | ChunkType::Usage(_) | ChunkType::Start | ChunkType::Incomplete(_) => {} ChunkType::StreamRollback { text, .. } => { @@ -1029,6 +1030,7 @@ pub async fn generate_session_title( | ChunkType::RetryableFailure(_) | ChunkType::Warning(_) | ChunkType::Metadata(_) + | ChunkType::Usage(_) | ChunkType::Start | ChunkType::Incomplete(_) => {} ChunkType::StreamRollback { text, .. } => { @@ -1978,6 +1980,11 @@ async fn relay_stream_to_sender( stats.record_metadata(&message); crate::emit_log!("[RELAY] Metadata {}", message); } + ChunkType::Usage(usage) => { + let elapsed_ms = start_time.elapsed().as_millis(); + stats.record_chunk("Usage", elapsed_ms); + let _ = sender.send(crate::llm::ChunkMessage::Usage(usage)); + } ChunkType::Retry(status) => { let elapsed_ms = start_time.elapsed().as_millis(); stats.record_chunk("Retry", elapsed_ms); diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 780d79d..8b264dc 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -17,6 +17,7 @@ pub enum ChunkMessage { reasoning: String, }, Warning(String), + Usage(crate::aisdk::chunk::TokenUsage), ToolCalls(Vec), ToolResult(ToolCallResult), SubagentStarted { diff --git a/src/main.rs b/src/main.rs index b7aa4f7..9ee638b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -551,6 +551,7 @@ async fn run_print_mode( } crate::llm::ChunkMessage::ToolCalls(_) | crate::llm::ChunkMessage::ToolResult(_) + | crate::llm::ChunkMessage::Usage(_) | crate::llm::ChunkMessage::Metrics { .. } | crate::llm::ChunkMessage::Cancelled | crate::llm::ChunkMessage::Reasoning(_) diff --git a/src/session/types.rs b/src/session/types.rs index 2ace65a..7719283 100644 --- a/src/session/types.rs +++ b/src/session/types.rs @@ -88,6 +88,19 @@ impl MessagePart { } } + pub fn usage(input: u64, output: u64, cache_read: u64, cache_write: u64, cost: f64) -> Self { + Self { + part_type: "usage".to_string(), + data: serde_json::json!({ + "input": input, + "output": output, + "cache_read": cache_read, + "cache_write": cache_write, + "cost": cost, + }), + } + } + pub fn text_value(&self) -> Option<&str> { self.data.get("text").and_then(|value| value.as_str()) } diff --git a/src/stats.rs b/src/stats.rs index 114fc08..f7b5030 100644 --- a/src/stats.rs +++ b/src/stats.rs @@ -17,6 +17,47 @@ pub struct StatsOptions { 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; @@ -46,7 +87,6 @@ impl UsageTotals { struct SessionRow { id: i64, workspace_path: Option, - total_cost: f64, } #[derive(Clone, Debug)] @@ -56,7 +96,7 @@ struct MessageRow { parts: String, model: Option, provider: Option, - output_tokens: u64, + usage: UsageTotals, } #[derive(Clone, Debug, Default, PartialEq)] @@ -131,12 +171,16 @@ fn collect(conn: &Connection, options: &StatsOptions, now: i64) -> Result Result = session_tokens.into_values().collect(); per_session.sort_unstable(); let average_tokens_per_session = if per_session.is_empty() { @@ -222,7 +270,7 @@ fn collect(conn: &Connection, options: &StatsOptions, now: i64) -> Result Result> { let mut statement = conn.prepare( - "SELECT s.id, w.root_path, s.total_cost + "SELECT s.id, w.root_path FROM sessions s LEFT JOIN workspaces w ON w.id = s.workspace_id", )?; @@ -230,7 +278,6 @@ fn load_sessions(conn: &Connection) -> Result> { Ok(SessionRow { id: row.get(0)?, workspace_path: row.get(1)?, - total_cost: row.get(2)?, }) })?; rows.collect::>>() @@ -245,13 +292,14 @@ fn load_messages(conn: &Connection) -> Result> { )?; 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)?, - output_tokens: output_tokens.max(0) as u64, + usage, }) })?; rows.collect::>>() @@ -513,7 +561,7 @@ mod tests { ('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"}]"#, + 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"}]"# ], ) @@ -535,13 +583,18 @@ mod tests { 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.average_tokens_per_session, 1_000); - assert_eq!(report.median_tokens_per_session, 1_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); } diff --git a/src/ui/components/chat.rs b/src/ui/components/chat.rs index 4978403..4f473a5 100644 --- a/src/ui/components/chat.rs +++ b/src/ui/components/chat.rs @@ -45,6 +45,7 @@ fn assistant_tool_part_info( if result_ids.contains(id) { return None; } + let mut info = parsed_tool_message_from_object(part.data.as_object()?, false); if part.data.get("status").is_none() { info.status = "running".to_string(); @@ -1688,6 +1689,85 @@ fn plan_update_display( } impl Chat { + pub fn record_usage( + &mut self, + input: u64, + output: u64, + cache_read: u64, + cache_write: u64, + cost: f64, + ) { + if self.streaming_assistant_idx().is_none() { + self.messages.push(Message::incomplete("")); + } + let Some(message) = self + .messages + .iter_mut() + .rfind(|message| message.role == MessageRole::Assistant && !message.is_complete) + else { + return; + }; + + if let Some(part) = message + .parts + .iter_mut() + .find(|part| part.part_type == "usage") + { + let current_input = part + .data + .get("input") + .and_then(JsonValue::as_u64) + .unwrap_or(0); + let current_output = part + .data + .get("output") + .and_then(JsonValue::as_u64) + .unwrap_or(0); + let current_cache_read = part + .data + .get("cache_read") + .and_then(JsonValue::as_u64) + .unwrap_or(0); + let current_cache_write = part + .data + .get("cache_write") + .and_then(JsonValue::as_u64) + .unwrap_or(0); + let current_cost = part + .data + .get("cost") + .and_then(JsonValue::as_f64) + .unwrap_or(0.0); + part.data = serde_json::json!({ + "input": current_input.saturating_add(input), + "output": current_output.saturating_add(output), + "cache_read": current_cache_read.saturating_add(cache_read), + "cache_write": current_cache_write.saturating_add(cache_write), + "cost": current_cost + cost, + }); + } else { + message + .parts + .push(crate::session::types::MessagePart::usage( + input, + output, + cache_read, + cache_write, + cost, + )); + } + + if output > 0 { + message.output_tokens = Some( + message + .output_tokens + .unwrap_or(0) + .saturating_add(output as usize), + ); + message.token_count = message.output_tokens; + } + } + pub fn new() -> Self { Self { messages: Vec::new(), @@ -2593,8 +2673,8 @@ impl Chat { .rposition(|m| m.role == MessageRole::Assistant) { if let Some(msg) = self.messages.get_mut(idx) { - msg.output_tokens = Some(token_count); - msg.token_count = Some(token_count); + msg.output_tokens = Some(msg.output_tokens.unwrap_or(token_count)); + msg.token_count = msg.output_tokens; msg.duration_ms = Some(decode_duration_ms); msg.tokens_per_sec = final_tps; msg.finish_reasoning_timer(finalized_at); @@ -8078,6 +8158,27 @@ mod tests { assert_eq!(chat.messages[2].content, " assistant"); } + #[test] + fn record_usage_accumulates_provider_steps_on_streaming_assistant() { + let mut chat = Chat::new(); + + chat.record_usage(1_000, 100, 500, 50, 0.01); + chat.record_usage(2_000, 200, 750, 25, 0.02); + + let message = chat.messages.last().unwrap(); + let usage = message + .parts + .iter() + .find(|part| part.part_type == "usage") + .unwrap(); + assert_eq!(usage.data["input"], 3_000); + assert_eq!(usage.data["output"], 300); + assert_eq!(usage.data["cache_read"], 1_250); + assert_eq!(usage.data["cache_write"], 75); + assert!((usage.data["cost"].as_f64().unwrap() - 0.03).abs() < f64::EPSILON); + assert_eq!(message.output_tokens, Some(300)); + } + #[test] fn click_hit_test_maps_visible_row_to_message_index() { let mut chat = Chat::with_messages(vec![Message::user("hello"), Message::assistant("hi")]); From 07be9105618dfd0764ce3030b5ea61210cd92e00 Mon Sep 17 00:00:00 2001 From: Yanuar Date: Mon, 31 Aug 2026 16:57:44 +0700 Subject: [PATCH 3/6] fix(stats): request compatible provider usage --- src/aisdk/providers/compatible.rs | 28 +++++++++++++++++++++++----- 1 file changed, 23 insertions(+), 5 deletions(-) diff --git a/src/aisdk/providers/compatible.rs b/src/aisdk/providers/compatible.rs index 552603c..245e705 100644 --- a/src/aisdk/providers/compatible.rs +++ b/src/aisdk/providers/compatible.rs @@ -151,11 +151,7 @@ impl Provider for OpenAICompatible { } } - let mut body = serde_json::json!({ - "model": self.model_name, - "messages": chat_messages, - "stream": true, - }); + let mut body = openai_compatible_request_body(&self.model_name, chat_messages); if !tool_params.is_empty() { body["tools"] = serde_json::Value::Array(tool_params); @@ -237,6 +233,20 @@ impl Provider for OpenAICompatible { } } +fn openai_compatible_request_body( + model_name: &str, + messages: Vec, +) -> serde_json::Value { + serde_json::json!({ + "model": model_name, + "messages": messages, + "stream": true, + "stream_options": { + "include_usage": true + } + }) +} + fn openai_compatible_user_content(user: &crate::message::UserMessage) -> serde_json::Value { if user.images.is_empty() { return serde_json::json!(user.content); @@ -662,6 +672,14 @@ mod tests { assert!(provider.api_key.is_empty()); } + #[test] + fn request_asks_streaming_gateways_for_token_usage() { + let body = openai_compatible_request_body("test-model", Vec::new()); + + assert_eq!(body["stream"], true); + assert_eq!(body["stream_options"]["include_usage"], true); + } + #[test] fn emits_tool_call_delta_without_finish_reason() { let data = r#"{"choices":[{"index":0,"delta":{"tool_calls":[{"id":"tool-1","index":0,"type":"function","function":{"name":"question","arguments":"{\"questions\":[{\"header\":\"Hobbies\",\"options\":[]}]}"}}]}}]}"#; From b1bce54bc4a555d5c0c338159ff83fbc25e7fd33 Mon Sep 17 00:00:00 2001 From: Yanuar Date: Wed, 2 Sep 2026 08:38:44 +0700 Subject: [PATCH 4/6] fix(aisdk): preserve Anthropic stop reasons --- src/aisdk/providers/anthropic.rs | 122 +++++++++++++++++++++++-------- src/aisdk/providers/openai.rs | 4 +- src/aisdk/response.rs | 1 + 3 files changed, 93 insertions(+), 34 deletions(-) diff --git a/src/aisdk/providers/anthropic.rs b/src/aisdk/providers/anthropic.rs index a04d55c..1694f5c 100644 --- a/src/aisdk/providers/anthropic.rs +++ b/src/aisdk/providers/anthropic.rs @@ -221,28 +221,28 @@ impl Provider for Anthropic { let stream = response .bytes_stream() .eventsource() - .filter_map(|ev| match ev { + .flat_map(|ev| match ev { Ok(event) => { let event_type = event.event.as_str(); let data = &event.data; if data.is_empty() { - return futures::future::ready(None); + return futures::stream::iter(Vec::new()); } match serde_json::from_str::(data) { Ok(value) => { - futures::future::ready(anthropic_stream_chunk(event_type, &value)) + futures::stream::iter(anthropic_stream_chunks(event_type, &value)) } - Err(e) => futures::future::ready(Some(Ok(ChunkType::Failed(format!( + Err(e) => futures::stream::iter(vec![Ok(ChunkType::Failed(format!( "Invalid SSE data: {}", e - ))))), + )))]), } } - Err(e) => futures::future::ready(Some(Ok(ChunkType::RetryableFailure( + Err(e) => futures::stream::iter(vec![Ok(ChunkType::RetryableFailure( RetryError::from_message(format!("SSE error: {}", e)), - )))), + ))]), }) .boxed(); @@ -254,45 +254,61 @@ fn anthropic_stream_chunk( event_type: &str, value: &serde_json::Value, ) -> Option> { + anthropic_stream_chunks(event_type, value) + .into_iter() + .next() +} + +fn anthropic_stream_chunks(event_type: &str, value: &serde_json::Value) -> Vec> { match event_type { "message_start" => { - // Partial usage early in the stream (cache fields may already appear). + // This is partial usage. Log it for cache observability, but only + // persist the final message_delta usage to avoid double-counting. if let Some(usage) = value.get("message").and_then(|m| m.get("usage")) { log_anthropic_usage(usage); - return anthropic_usage(usage).map(ChunkType::Usage).map(Ok); } - None + Vec::new() } "content_block_start" => { if let Some(payload) = anthropic_hosted_search_start(value) { - Some(Ok(ChunkType::ProviderToolCall(payload))) + vec![Ok(ChunkType::ProviderToolCall(payload))] } else if let Some(payload) = anthropic_hosted_search_result(value) { - Some(Ok(ChunkType::ProviderToolCall(payload))) + vec![Ok(ChunkType::ProviderToolCall(payload))] } else { anthropic_tool_call_start(value) .map(ChunkType::ToolCall) .map(Ok) + .into_iter() + .collect() } } - "content_block_delta" => anthropic_content_block_delta(value).map(Ok), + "content_block_delta" => anthropic_content_block_delta(value) + .map(Ok) + .into_iter() + .collect(), "message_delta" => { - // Final usage wins for cache_read / cache_creation. + let mut chunks = Vec::with_capacity(2); + // Final usage wins for cache_read / cache_creation. Emit it before + // the terminal chunk so persistence sees usage even for failures. if let Some(usage) = value.get("usage") { log_anthropic_usage(usage); if let Some(usage) = anthropic_usage(usage) { - return Some(Ok(ChunkType::Usage(usage))); + chunks.push(Ok(ChunkType::Usage(usage))); } } - anthropic_message_delta(value).map(Ok) + if let Some(terminal) = anthropic_message_delta(value) { + chunks.push(Ok(terminal)); + } + chunks } - "message_stop" => Some(Ok(ChunkType::End { reason: None })), + "message_stop" => vec![Ok(ChunkType::End { reason: None })], "error" => { let error_msg = value["error"]["message"] .as_str() .unwrap_or("Unknown error"); - Some(Ok(ChunkType::Failed(error_msg.to_string()))) + vec![Ok(ChunkType::Failed(error_msg.to_string()))] } - _ => None, + _ => Vec::new(), } } @@ -836,40 +852,82 @@ mod tests { } #[test] - fn max_tokens_stop_reason_emits_incomplete_chunk() { + 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 chunk = anthropic_stream_chunk("message_delta", &value) - .expect("event should produce a chunk") - .expect("chunk should parse"); + let chunks = anthropic_stream_chunks("message_delta", &value); - assert!(matches!(chunk, ChunkType::Incomplete(_))); + 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_stop_reason_emits_terminal_reason() { + 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 chunk = anthropic_stream_chunk("message_delta", &value) - .expect("event should produce a chunk") - .expect("chunk should parse"); + let chunks = anthropic_stream_chunks("message_delta", &value); assert!(matches!( - chunk, - ChunkType::End { - reason: Some(FinishReason::EndTurn) - } + 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 message_start_usage_is_not_emitted_or_double_counted() { + let value = serde_json::json!({ + "type": "message_start", + "message": { + "usage": { + "input_tokens": 12, + "output_tokens": 0, + "cache_read_input_tokens": 5, + "cache_creation_input_tokens": 2, + }, + }, + }); + + assert!(anthropic_stream_chunks("message_start", &value).is_empty()); + } + #[test] fn groups_adjacent_tool_calls_and_results() { let messages = vec![ diff --git a/src/aisdk/providers/openai.rs b/src/aisdk/providers/openai.rs index f5eeed4..e09fb88 100644 --- a/src/aisdk/providers/openai.rs +++ b/src/aisdk/providers/openai.rs @@ -817,7 +817,7 @@ impl OpenAI { progress.record_chunk(chunk); } if tx.send(chunk).is_err() { - return None; + return; } if is_completed { let mut state = websocket_state.lock().await; @@ -1604,7 +1604,7 @@ fn openai_responses_usage(usage: &serde_json::Value) -> Option vec![ From 24175b71a3d685156765465b9261e8008219b6d1 Mon Sep 17 00:00:00 2001 From: Blankeos Date: Thu, 3 Sep 2026 02:50:39 +0800 Subject: [PATCH 5/6] fix(stats): comma-group overview counts like OpenCode Format sessions, messages, days, and per-model message counts with thousands separators. Keep token totals compact (K/M/B) and leave tool-usage counts ungrouped. --- src/stats.rs | 68 ++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 64 insertions(+), 4 deletions(-) diff --git a/src/stats.rs b/src/stats.rs index f7b5030..cbc6e0b 100644 --- a/src/stats.rs +++ b/src/stats.rs @@ -375,9 +375,9 @@ fn render_overview(report: &StatsReport) -> String { render_table( "OVERVIEW", &[ - ("Sessions", report.sessions.to_string()), - ("Messages", report.messages.to_string()), - ("Days", report.days.to_string()), + ("Sessions", comma_number(report.sessions as u64)), + ("Messages", comma_number(report.messages as u64)), + ("Days", comma_number(report.days as u64)), ], ) } @@ -413,7 +413,7 @@ 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", &stats.messages.to_string())); + lines.push(metric_row(" Messages", &comma_number(stats.messages))); lines.push(metric_row( " Input Tokens", &compact_number(stats.usage.input), @@ -514,6 +514,18 @@ fn truncate(value: &str, width: usize) -> String { 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"), @@ -641,6 +653,7 @@ mod tests { 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()) @@ -655,4 +668,51 @@ mod tests { 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%)")); + } } From b11eef1bbe1beb2f7c4422069624d481b9662bde Mon Sep 17 00:00:00 2001 From: Blankeos Date: Thu, 3 Sep 2026 03:07:46 +0800 Subject: [PATCH 6/6] fix(acp): price persisted usage from the models.dev catalog ACP was storing billed tokens with cost 0.0, so stats/history from ACP sessions looked token-rich and cost-poor. Reuse the same host catalog estimate as the TUI instead of leaving cost blank. --- src/acp/service.rs | 45 +++++++++++++++++++++++++++++++++++++++++- src/app.rs | 9 ++++----- src/model/discovery.rs | 14 +++++++++++++ 3 files changed, 62 insertions(+), 6 deletions(-) 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/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/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)?;