Files
rtk-ai__rtk/src/analytics/ccusage.rs
David Alecrim f68fa0087c fix(ccusage): add --yes flag and warn when falling back to npx
Fixes silent hang when ccusage is not installed globally. npx prompts
for confirmation before installing the package, blocking the subprocess
indefinitely. Adding --yes auto-confirms and the info message sets user
expectations.

Fixes rtk-ai/rtk#1226
2026-04-11 12:43:54 -03:00

329 lines
9.1 KiB
Rust

//! Parses Claude Code spending data for economics reporting.
//!
//! Provides isolated interface to ccusage (npm package) for fetching
//! Claude Code API usage metrics. Handles subprocess execution, JSON parsing,
//! and graceful degradation when ccusage is unavailable.
use crate::core::utils::{resolved_command, tool_exists};
use anyhow::{Context, Result};
use serde::Deserialize;
use std::process::Command;
// ── Public Types ──
/// Metrics from ccusage for a single period (day/week/month)
#[derive(Debug, Deserialize)]
pub struct CcusageMetrics {
#[serde(rename = "inputTokens")]
pub input_tokens: u64,
#[serde(rename = "outputTokens")]
pub output_tokens: u64,
#[serde(rename = "cacheCreationTokens", default)]
pub cache_creation_tokens: u64,
#[serde(rename = "cacheReadTokens", default)]
pub cache_read_tokens: u64,
#[serde(rename = "totalTokens")]
pub total_tokens: u64,
#[serde(rename = "totalCost")]
pub total_cost: f64,
}
/// Period data with key (date/month/week) and metrics
#[derive(Debug)]
pub struct CcusagePeriod {
pub key: String, // "2026-01-30" (daily), "2026-01" (monthly), "2026-01-20" (weekly ISO monday)
pub metrics: CcusageMetrics,
}
/// Time granularity for ccusage reports
#[derive(Debug, Clone, Copy)]
pub enum Granularity {
Daily,
Weekly,
Monthly,
}
// ── Internal Types for JSON Deserialization ──
#[derive(Debug, Deserialize)]
struct DailyResponse {
daily: Vec<DailyEntry>,
}
#[derive(Debug, Deserialize)]
struct DailyEntry {
date: String,
#[serde(flatten)]
metrics: CcusageMetrics,
}
#[derive(Debug, Deserialize)]
struct WeeklyResponse {
weekly: Vec<WeeklyEntry>,
}
#[derive(Debug, Deserialize)]
struct WeeklyEntry {
week: String, // ISO week start (Monday)
#[serde(flatten)]
metrics: CcusageMetrics,
}
#[derive(Debug, Deserialize)]
struct MonthlyResponse {
monthly: Vec<MonthlyEntry>,
}
#[derive(Debug, Deserialize)]
struct MonthlyEntry {
month: String,
#[serde(flatten)]
metrics: CcusageMetrics,
}
// ── Public API ──
/// Check if ccusage binary exists in PATH
fn binary_exists() -> bool {
tool_exists("ccusage")
}
/// Build the ccusage command, falling back to npx if binary not in PATH
fn build_command() -> Option<Command> {
if binary_exists() {
return Some(resolved_command("ccusage"));
}
// Fallback: try npx
eprintln!("[info] ccusage not installed globally, fetching via npx...");
let npx_check = resolved_command("npx")
.arg("--yes")
.arg("ccusage")
.arg("--help")
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.status();
if npx_check.map(|s| s.success()).unwrap_or(false) {
let mut cmd = resolved_command("npx");
cmd.arg("--yes");
cmd.arg("ccusage");
return Some(cmd);
}
None
}
/// Fetch usage data from ccusage for the last 90 days
///
/// Returns `Ok(None)` if ccusage is unavailable (graceful degradation)
/// Returns `Ok(Some(vec))` with parsed data on success
/// Returns `Err` only on unexpected failures (JSON parse, etc.)
pub fn fetch(granularity: Granularity) -> Result<Option<Vec<CcusagePeriod>>> {
let mut cmd = match build_command() {
Some(cmd) => cmd,
None => {
eprintln!("[warn] ccusage not found. Install: npm i -g ccusage (or use npx ccusage)");
return Ok(None);
}
};
let subcommand = match granularity {
Granularity::Daily => "daily",
Granularity::Weekly => "weekly",
Granularity::Monthly => "monthly",
};
let output = cmd
.arg(subcommand)
.arg("--json")
.arg("--since")
.arg("20250101") // 90 days back approx
.output();
let output = match output {
Err(e) => {
eprintln!("[warn] ccusage execution failed: {}", e);
return Ok(None);
}
Ok(o) => o,
};
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
eprintln!(
"[warn] ccusage exited with {}: {}",
output.status,
stderr.trim()
);
return Ok(None);
}
let stdout = String::from_utf8_lossy(&output.stdout);
let periods =
parse_json(&stdout, granularity).context("Failed to parse ccusage JSON output")?;
Ok(Some(periods))
}
// ── Internal Helpers ──
fn parse_json(json: &str, granularity: Granularity) -> Result<Vec<CcusagePeriod>> {
match granularity {
Granularity::Daily => {
let resp: DailyResponse =
serde_json::from_str(json).context("Invalid JSON structure for daily data")?;
Ok(resp
.daily
.into_iter()
.map(|e| CcusagePeriod {
key: e.date,
metrics: e.metrics,
})
.collect())
}
Granularity::Weekly => {
let resp: WeeklyResponse =
serde_json::from_str(json).context("Invalid JSON structure for weekly data")?;
Ok(resp
.weekly
.into_iter()
.map(|e| CcusagePeriod {
key: e.week,
metrics: e.metrics,
})
.collect())
}
Granularity::Monthly => {
let resp: MonthlyResponse =
serde_json::from_str(json).context("Invalid JSON structure for monthly data")?;
Ok(resp
.monthly
.into_iter()
.map(|e| CcusagePeriod {
key: e.month,
metrics: e.metrics,
})
.collect())
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_monthly_valid() {
let json = r#"{
"monthly": [
{
"month": "2026-01",
"inputTokens": 1000,
"outputTokens": 500,
"cacheCreationTokens": 100,
"cacheReadTokens": 200,
"totalTokens": 1800,
"totalCost": 12.34
}
]
}"#;
let result = parse_json(json, Granularity::Monthly);
assert!(result.is_ok());
let periods = result.unwrap();
assert_eq!(periods.len(), 1);
assert_eq!(periods[0].key, "2026-01");
assert_eq!(periods[0].metrics.input_tokens, 1000);
assert_eq!(periods[0].metrics.total_cost, 12.34);
}
#[test]
fn test_parse_daily_valid() {
let json = r#"{
"daily": [
{
"date": "2026-01-30",
"inputTokens": 100,
"outputTokens": 50,
"cacheCreationTokens": 0,
"cacheReadTokens": 0,
"totalTokens": 150,
"totalCost": 0.15
}
]
}"#;
let result = parse_json(json, Granularity::Daily);
assert!(result.is_ok());
let periods = result.unwrap();
assert_eq!(periods.len(), 1);
assert_eq!(periods[0].key, "2026-01-30");
}
#[test]
fn test_parse_weekly_valid() {
let json = r#"{
"weekly": [
{
"week": "2026-01-20",
"inputTokens": 500,
"outputTokens": 250,
"cacheCreationTokens": 50,
"cacheReadTokens": 100,
"totalTokens": 900,
"totalCost": 5.67
}
]
}"#;
let result = parse_json(json, Granularity::Weekly);
assert!(result.is_ok());
let periods = result.unwrap();
assert_eq!(periods.len(), 1);
assert_eq!(periods[0].key, "2026-01-20");
}
#[test]
fn test_parse_malformed_json() {
let json = r#"{ "monthly": [ { "broken": }"#;
let result = parse_json(json, Granularity::Monthly);
assert!(result.is_err());
}
#[test]
fn test_parse_missing_required_fields() {
let json = r#"{
"monthly": [
{
"month": "2026-01",
"inputTokens": 100
}
]
}"#;
let result = parse_json(json, Granularity::Monthly);
assert!(result.is_err()); // Missing required fields like totalTokens
}
#[test]
fn test_parse_default_cache_fields() {
let json = r#"{
"monthly": [
{
"month": "2026-01",
"inputTokens": 100,
"outputTokens": 50,
"totalTokens": 150,
"totalCost": 1.0
}
]
}"#;
let result = parse_json(json, Granularity::Monthly);
assert!(result.is_ok());
let periods = result.unwrap();
assert_eq!(periods[0].metrics.cache_creation_tokens, 0); // default
assert_eq!(periods[0].metrics.cache_read_tokens, 0);
}
}