ai_usagebar/
usage.rs

1//! Canonical in-memory representation of "how much have I used my plan".
2//!
3//! Each vendor's snapshot lives in its own variant — this is deliberate.
4//! Anthropic exposes three windows + extra credits; OpenAI Codex exposes two
5//! windows + credit balance + message-count ranges; OpenRouter is a single
6//! credit-balance number with daily/weekly/monthly totals; Z.AI is a list of
7//! token + MCP buckets. Forcing them into a shared shape would either drop
8//! information or paper over genuine differences.
9//!
10//! Renderers (widget tooltip, TUI tab) consume a `VendorSnapshot` directly,
11//! not a flattened shape — so each vendor controls its own presentation while
12//! sharing the pacing math, color thresholds, and Pango primitives.
13
14use chrono::{DateTime, Utc};
15
16/// A single usage window — generic enough that every vendor with a notion of
17/// "% used vs. when does it reset" can express itself with it.
18///
19/// `utilization_pct` is `0..=100` (integer percent, matching claudebar's units).
20/// `resets_at` is `None` when the vendor doesn't report a reset time.
21#[derive(Debug, Clone, PartialEq, Eq)]
22pub struct UsageWindow {
23    pub utilization_pct: i32,
24    pub resets_at: Option<DateTime<Utc>>,
25    /// Window length (used for pacing math).
26    pub window_duration: chrono::Duration,
27}
28
29/// Money expressed in cents to dodge float roundoff.
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub struct Cents(pub i64);
32
33impl Cents {
34    /// Format as `[-]$D.CC`. Negative values render `-$D.CC` (not `$-D.CC`),
35    /// matching claudebar's `_fmt_dollars` (claudebar:532-537).
36    pub fn fmt_dollars(self) -> String {
37        let (sign, abs) = if self.0 < 0 {
38            ("-", -self.0)
39        } else {
40            ("", self.0)
41        };
42        format!("{sign}${}.{:02}", abs / 100, abs % 100)
43    }
44}
45
46/// Anthropic-specific snapshot — three rolling windows plus optional
47/// pay-as-you-go credit balance.
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct AnthropicSnapshot {
50    /// "Claude Pro", "Claude Max 5x", "Claude Max 20x", etc.
51    pub plan: String,
52    pub session: UsageWindow,
53    pub weekly: UsageWindow,
54    /// Some vendors of Claude (Pro, some Max tiers) don't have a separate
55    /// Sonnet bucket — in which case this is None.
56    pub sonnet: Option<UsageWindow>,
57    /// `None` when `extra_usage.is_enabled` is false or the block is absent.
58    pub extra: Option<ExtraUsage>,
59}
60
61/// "Extra usage" pay-as-you-go block (claudebar's `extra_usage`).
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub struct ExtraUsage {
64    pub limit: Cents,
65    pub spent: Cents,
66}
67
68impl ExtraUsage {
69    /// Integer percentage of the monthly limit consumed (0..=100, saturating
70    /// at 0 when limit is non-positive — matches claudebar:540-542).
71    pub fn percent(self) -> i32 {
72        if self.limit.0 <= 0 {
73            0
74        } else {
75            ((self.spent.0 * 100) / self.limit.0) as i32
76        }
77    }
78}
79
80/// Discriminated union of vendor-specific snapshots. The widget and TUI match
81/// on this to pick a renderer.
82#[derive(Debug, Clone, PartialEq, Eq)]
83pub enum VendorSnapshot {
84    Anthropic(AnthropicSnapshot),
85    Openai(OpenAiSnapshot),
86    Zai(ZaiSnapshot),
87    Openrouter(OpenRouterSnapshot),
88}
89
90/// OpenAI Codex OAuth — mirrors Anthropic's two-window + extras pattern.
91#[derive(Debug, Clone, PartialEq, Eq)]
92pub struct OpenAiSnapshot {
93    pub plan: String,
94    /// 5h window (Codex `rate_limit.primary_window`).
95    pub session: UsageWindow,
96    /// 7d window (Codex `rate_limit.secondary_window`).
97    pub weekly: UsageWindow,
98    /// Optional 7d code-review bucket.
99    pub code_review: Option<UsageWindow>,
100    /// Optional credit balance + approximate message-count ranges.
101    pub credits: Option<OpenAiCredits>,
102    /// Source of the snapshot — Codex OAuth vs admin-key fallback. Drives
103    /// the placeholder set and the "OpenAI does not expose this for Plus"
104    /// tooltip when the OAuth path isn't available.
105    pub source: OpenAiSource,
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
109pub enum OpenAiSource {
110    CodexOauth,
111    AdminKeyMtd,
112    Unavailable,
113}
114
115#[derive(Debug, Clone, PartialEq, Eq)]
116pub struct OpenAiCredits {
117    /// Credit balance, formatted dollars ("$0.00", "$5.00", etc.) — kept as
118    /// a string because OpenAI returns it that way.
119    pub balance: String,
120    pub has_credits: bool,
121    pub unlimited: bool,
122    pub approx_local_messages: Option<(i64, i64)>,
123    pub approx_cloud_messages: Option<(i64, i64)>,
124}
125
126/// Z.AI / BigModel — list of buckets with discriminated types. We project the
127/// two we care about into named fields (5h tokens, weekly tokens, MCP).
128#[derive(Debug, Clone, PartialEq, Eq)]
129pub struct ZaiSnapshot {
130    pub plan: String,
131    pub session: Option<UsageWindow>,
132    pub weekly: Option<UsageWindow>,
133    pub mcp: Option<UsageWindow>,
134}
135
136/// OpenRouter — credit balance + lifetime/daily/weekly/monthly usage from
137/// `/api/v1/credits` and `/api/v1/key`.
138#[derive(Debug, Clone, PartialEq)]
139pub struct OpenRouterSnapshot {
140    pub label: String,
141    pub total_credits: f64,
142    pub total_usage: f64,
143    pub usage_daily: f64,
144    pub usage_weekly: f64,
145    pub usage_monthly: f64,
146    pub is_free_tier: bool,
147    pub limit: Option<f64>,
148    pub limit_remaining: Option<f64>,
149}
150
151impl Eq for OpenRouterSnapshot {}
152
153impl OpenRouterSnapshot {
154    pub fn balance(&self) -> f64 {
155        (self.total_credits - self.total_usage).max(0.0)
156    }
157    /// Percentage of total_credits consumed (0..=100). Returns 0 when
158    /// `total_credits` is 0 (free-tier-only accounts).
159    pub fn consumed_pct(&self) -> i32 {
160        if self.total_credits <= 0.0 {
161            return 0;
162        }
163        ((self.total_usage / self.total_credits) * 100.0)
164            .round()
165            .clamp(0.0, 100.0) as i32
166    }
167}
168
169/// Worst-of severity class for the Waybar bar text color. Mirrors
170/// claudebar:606-620 — "extra usage only matters when a rate limit hits 100%".
171pub fn anthropic_severity(snap: &AnthropicSnapshot) -> crate::pacing::PaceSeverity {
172    let mut max = snap.session.utilization_pct;
173    if snap.weekly.utilization_pct > max {
174        max = snap.weekly.utilization_pct;
175    }
176    if let Some(s) = &snap.sonnet {
177        if s.utilization_pct > max {
178            max = s.utilization_pct;
179        }
180    }
181    // Extra usage only promotes severity if a rate-limit window is at 100%.
182    let any_at_cap = snap.session.utilization_pct >= 100
183        || snap.weekly.utilization_pct >= 100
184        || snap
185            .sonnet
186            .as_ref()
187            .is_some_and(|s| s.utilization_pct >= 100);
188    if any_at_cap {
189        if let Some(extra) = snap.extra {
190            let p = extra.percent();
191            if p > max {
192                max = p;
193            }
194        }
195    }
196    crate::pango::severity_for(max)
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::pacing::PaceSeverity;
203    use chrono::Duration;
204
205    fn w(pct: i32) -> UsageWindow {
206        UsageWindow {
207            utilization_pct: pct,
208            resets_at: None,
209            window_duration: Duration::hours(5),
210        }
211    }
212
213    fn snap(s: i32, w_: i32, sonnet: Option<i32>, extra: Option<(i64, i64)>) -> AnthropicSnapshot {
214        AnthropicSnapshot {
215            plan: "Max 5x".into(),
216            session: w(s),
217            weekly: w(w_),
218            sonnet: sonnet.map(w),
219            extra: extra.map(|(limit, spent)| ExtraUsage {
220                limit: Cents(limit),
221                spent: Cents(spent),
222            }),
223        }
224    }
225
226    #[test]
227    fn cents_format_positive() {
228        assert_eq!(Cents(0).fmt_dollars(), "$0.00");
229        assert_eq!(Cents(50).fmt_dollars(), "$0.50");
230        assert_eq!(Cents(250).fmt_dollars(), "$2.50");
231        assert_eq!(Cents(5000).fmt_dollars(), "$50.00");
232    }
233
234    #[test]
235    fn cents_format_negative_uses_leading_sign() {
236        // claudebar bug-fix: never "$-1.-50" — sign goes before the dollar sign.
237        assert_eq!(Cents(-150).fmt_dollars(), "-$1.50");
238        assert_eq!(Cents(-1).fmt_dollars(), "-$0.01");
239    }
240
241    #[test]
242    fn extra_percent_with_zero_limit_is_zero() {
243        assert_eq!(
244            ExtraUsage {
245                limit: Cents(0),
246                spent: Cents(100)
247            }
248            .percent(),
249            0
250        );
251    }
252
253    #[test]
254    fn extra_percent_truncates() {
255        // Bash integer division — 33/100 -> 33%, 50/100 -> 50%.
256        assert_eq!(
257            ExtraUsage {
258                limit: Cents(10000),
259                spent: Cents(3333)
260            }
261            .percent(),
262            33
263        );
264    }
265
266    #[test]
267    fn severity_picks_worst_of_three_windows() {
268        let s = snap(40, 60, Some(80), None);
269        assert_eq!(anthropic_severity(&s), PaceSeverity::High); // 80 → high
270    }
271
272    #[test]
273    fn severity_ignores_extra_when_no_cap_hit() {
274        // Extra at 95% but no rate-limit at 100% → extra is NOT promoted.
275        let s = snap(50, 60, None, Some((10000, 9500)));
276        assert_eq!(anthropic_severity(&s), PaceSeverity::Mid); // capped at 60
277    }
278
279    #[test]
280    fn severity_promotes_extra_when_session_at_100() {
281        let s = snap(100, 50, None, Some((10000, 9500)));
282        assert_eq!(anthropic_severity(&s), PaceSeverity::Critical); // 100 → critical
283    }
284
285    #[test]
286    fn severity_falls_through_to_extra_when_extra_higher_than_capped_window() {
287        // session = 100, weekly = 50, extra = 100% → max should be 100.
288        let s = snap(100, 50, None, Some((10000, 10000)));
289        assert_eq!(anthropic_severity(&s), PaceSeverity::Critical);
290    }
291}
ai_usagebar/usage.rs

ai_usagebar/
usage.rs
