claude_code_stats/
lib.rs

1pub mod source;
2pub mod types;
3
4use std::fs;
5use std::path::PathBuf;
6
7use anyhow::Result;
8use chrono::{DateTime, Utc};
9
10use source::cli_probe::CliProbeSource;
11use source::oauth::OauthSource;
12use source::web::WebSource;
13use source::{UsageSource, fetch_with_failover};
14use types::{
15    ActivePayload, ApiResponse, CachedResponse, CostData, DataSource, ExtraUsage, WidgetPayload,
16    parse_reset_time, to_usage_window,
17};
18
19pub const CACHE_MAX_AGE_SECS: i64 = 4 * 60;
20pub const CACHE_FILE_NAME: &str = "claude-code-stats-cache.json";
21
22pub fn collect_widget_payload() -> Result<WidgetPayload> {
23    let now = Utc::now();
24    let (usage, cached_at, data_source, cost_data) = get_usage_with_cache(now)?;
25
26    let data = ActivePayload {
27        five_hour: usage
28            .five_hour
29            .map(|w| to_usage_window(&w, now, Some(300.0))),
30        seven_day: usage
31            .seven_day
32            .map(|w| to_usage_window(&w, now, Some(10080.0))),
33        seven_day_sonnet: usage
34            .seven_day_sonnet
35            .map(|w| to_usage_window(&w, now, Some(10080.0))),
36        seven_day_opus: usage
37            .seven_day_opus
38            .map(|w| to_usage_window(&w, now, Some(10080.0))),
39        extra_usage: usage.extra_usage.map(|e| ExtraUsage {
40            is_enabled: e.is_enabled,
41            monthly_limit: e.monthly_limit,
42            used_credits: e.used_credits,
43            utilization: e.utilization,
44        }),
45        updated_at: cached_at.to_rfc3339(),
46        source: data_source,
47        cost_data,
48    };
49
50    Ok(WidgetPayload::Active {
51        data: Box::new(data),
52    })
53}
54
55pub fn collect_widget_payload_json() -> String {
56    let payload = match collect_widget_payload() {
57        Ok(payload) => payload,
58        Err(err) => {
59            eprintln!("claude-code-stats error: {err:?}");
60            WidgetPayload::Error {
61                title: "Claude usage error".to_string(),
62                message: err.to_string(),
63            }
64        }
65    };
66
67    match serde_json::to_string(&payload) {
68        Ok(json) => json,
69        Err(err) => {
70            eprintln!("claude-code-stats serialization error: {err:?}");
71            "{\"status\":\"error\",\"title\":\"claude-code-stats\",\"message\":\"serialization failure\"}".to_string()
72        }
73    }
74}
75
76fn get_cache_path() -> PathBuf {
77    dirs::cache_dir()
78        .unwrap_or_else(|| PathBuf::from("/tmp"))
79        .join(CACHE_FILE_NAME)
80}
81
82fn get_usage_with_cache(
83    now: DateTime<Utc>,
84) -> Result<(ApiResponse, DateTime<Utc>, DataSource, Option<CostData>)> {
85    let cache_path = get_cache_path();
86
87    if let Ok(contents) = fs::read_to_string(&cache_path) {
88        if let Ok(cached) = serde_json::from_str::<CachedResponse>(&contents) {
89            if let Some(cached_at) = parse_reset_time(&cached.cached_at) {
90                let age = now.signed_duration_since(cached_at).num_seconds();
91                if age < CACHE_MAX_AGE_SECS {
92                    return Ok((cached.response, cached_at, cached.source, cached.cost_data));
93                }
94            }
95        }
96    }
97
98    let sources: Vec<&dyn UsageSource> = vec![&OauthSource, &WebSource, &CliProbeSource];
99    let (response, source_name) = fetch_with_failover(&sources)?;
100
101    let data_source = match source_name {
102        "oauth_api" => DataSource::OauthApi,
103        "web_api" => DataSource::WebApi,
104        "cli_probe" => DataSource::CliProbe,
105        _ => DataSource::OauthApi,
106    };
107
108    let cost_data = source::cost_scanner::scan_cost_data();
109
110    let cached = CachedResponse {
111        cached_at: now.to_rfc3339(),
112        response: response.clone(),
113        source: data_source,
114        cost_data: cost_data.clone(),
115    };
116
117    if let Ok(json) = serde_json::to_string(&cached) {
118        let _ = fs::write(&cache_path, json);
119    }
120
121    Ok((response, now, data_source, cost_data))
122}
claude_code_stats/lib.rs

claude_code_stats/
lib.rs
