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/// DeepSeek — credit balance from `/user/balance`.
81#[derive(Debug, Clone, PartialEq)]
82pub struct DeepseekSnapshot {
83    pub is_available: bool,
84    /// Current balance (prefer USD, fallback to CNY).
85    pub balance: f64,
86    /// Free-granted credits component.
87    pub granted: f64,
88    /// Topped-up (purchased) credits component.
89    pub topped_up: f64,
90    /// The currency of the above amounts ("USD", "CNY", etc.).
91    pub currency: String,
92}
93
94impl Eq for DeepseekSnapshot {}
95
96impl Default for DeepseekSnapshot {
97    fn default() -> Self {
98        Self {
99            is_available: false,
100            balance: 0.0,
101            granted: 0.0,
102            topped_up: 0.0,
103            currency: String::new(),
104        }
105    }
106}
107
108/// Discriminated union of vendor-specific snapshots. The widget and TUI match
109/// on this to pick a renderer.
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub enum VendorSnapshot {
112    Anthropic(AnthropicSnapshot),
113    Openai(OpenAiSnapshot),
114    Zai(ZaiSnapshot),
115    Openrouter(OpenRouterSnapshot),
116    Deepseek(DeepseekSnapshot),
117}
118
119/// OpenAI Codex OAuth — mirrors Anthropic's two-window + extras pattern.
120#[derive(Debug, Clone, PartialEq, Eq)]
121pub struct OpenAiSnapshot {
122    pub plan: String,
123    /// 5h window (Codex `rate_limit.primary_window`).
124    pub session: UsageWindow,
125    /// 7d window (Codex `rate_limit.secondary_window`).
126    pub weekly: UsageWindow,
127    /// Optional 7d code-review bucket.
128    pub code_review: Option<UsageWindow>,
129    /// Optional credit balance + approximate message-count ranges.
130    pub credits: Option<OpenAiCredits>,
131    /// Source of the snapshot — Codex OAuth vs admin-key fallback. Drives
132    /// the placeholder set and the "OpenAI does not expose this for Plus"
133    /// tooltip when the OAuth path isn't available.
134    pub source: OpenAiSource,
135}
136
137#[derive(Debug, Clone, Copy, PartialEq, Eq)]
138pub enum OpenAiSource {
139    CodexOauth,
140    AdminKeyMtd,
141    Unavailable,
142}
143
144#[derive(Debug, Clone, PartialEq, Eq)]
145pub struct OpenAiCredits {
146    /// Credit balance, formatted dollars ("$0.00", "$5.00", etc.) — kept as
147    /// a string because OpenAI returns it that way.
148    pub balance: String,
149    pub has_credits: bool,
150    pub unlimited: bool,
151    pub approx_local_messages: Option<(i64, i64)>,
152    pub approx_cloud_messages: Option<(i64, i64)>,
153}
154
155/// Z.AI / BigModel — list of buckets with discriminated types. We project the
156/// two we care about into named fields (5h tokens, weekly tokens, MCP).
157#[derive(Debug, Clone, PartialEq, Eq)]
158pub struct ZaiSnapshot {
159    pub plan: String,
160    pub session: Option<UsageWindow>,
161    pub weekly: Option<UsageWindow>,
162    pub mcp: Option<UsageWindow>,
163}
164
165/// OpenRouter — credit balance + lifetime/daily/weekly/monthly usage from
166/// `/api/v1/credits` and `/api/v1/key`.
167#[derive(Debug, Clone, PartialEq)]
168pub struct OpenRouterSnapshot {
169    pub label: String,
170    pub total_credits: f64,
171    pub total_usage: f64,
172    pub usage_daily: f64,
173    pub usage_weekly: f64,
174    pub usage_monthly: f64,
175    pub is_free_tier: bool,
176    pub limit: Option<f64>,
177    pub limit_remaining: Option<f64>,
178}
179
180impl Eq for OpenRouterSnapshot {}
181
182impl OpenRouterSnapshot {
183    pub fn balance(&self) -> f64 {
184        (self.total_credits - self.total_usage).max(0.0)
185    }
186    /// Percentage of total_credits consumed (0..=100). Returns 0 when
187    /// `total_credits` is 0 (free-tier-only accounts).
188    pub fn consumed_pct(&self) -> i32 {
189        if self.total_credits <= 0.0 {
190            return 0;
191        }
192        ((self.total_usage / self.total_credits) * 100.0)
193            .round()
194            .clamp(0.0, 100.0) as i32
195    }
196}
197
198/// Worst-of severity class for the Waybar bar text color. Mirrors
199/// claudebar:606-620 — "extra usage only matters when a rate limit hits 100%".
200pub fn anthropic_severity(snap: &AnthropicSnapshot) -> crate::pacing::PaceSeverity {
201    let mut max = snap.session.utilization_pct;
202    if snap.weekly.utilization_pct > max {
203        max = snap.weekly.utilization_pct;
204    }
205    if let Some(s) = &snap.sonnet
206        && s.utilization_pct > max
207    {
208        max = s.utilization_pct;
209    }
210    // Extra usage only promotes severity if a rate-limit window is at 100%.
211    let any_at_cap = snap.session.utilization_pct >= 100
212        || snap.weekly.utilization_pct >= 100
213        || snap
214            .sonnet
215            .as_ref()
216            .is_some_and(|s| s.utilization_pct >= 100);
217    if any_at_cap && let Some(extra) = snap.extra {
218        let p = extra.percent();
219        if p > max {
220            max = p;
221        }
222    }
223    crate::pango::severity_for(max)
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use crate::pacing::PaceSeverity;
230    use chrono::Duration;
231
232    fn w(pct: i32) -> UsageWindow {
233        UsageWindow {
234            utilization_pct: pct,
235            resets_at: None,
236            window_duration: Duration::hours(5),
237        }
238    }
239
240    fn snap(s: i32, w_: i32, sonnet: Option<i32>, extra: Option<(i64, i64)>) -> AnthropicSnapshot {
241        AnthropicSnapshot {
242            plan: "Max 5x".into(),
243            session: w(s),
244            weekly: w(w_),
245            sonnet: sonnet.map(w),
246            extra: extra.map(|(limit, spent)| ExtraUsage {
247                limit: Cents(limit),
248                spent: Cents(spent),
249            }),
250        }
251    }
252
253    #[test]
254    fn cents_format_positive() {
255        assert_eq!(Cents(0).fmt_dollars(), "$0.00");
256        assert_eq!(Cents(50).fmt_dollars(), "$0.50");
257        assert_eq!(Cents(250).fmt_dollars(), "$2.50");
258        assert_eq!(Cents(5000).fmt_dollars(), "$50.00");
259    }
260
261    #[test]
262    fn cents_format_negative_uses_leading_sign() {
263        // claudebar bug-fix: never "$-1.-50" — sign goes before the dollar sign.
264        assert_eq!(Cents(-150).fmt_dollars(), "-$1.50");
265        assert_eq!(Cents(-1).fmt_dollars(), "-$0.01");
266    }
267
268    #[test]
269    fn extra_percent_with_zero_limit_is_zero() {
270        assert_eq!(
271            ExtraUsage {
272                limit: Cents(0),
273                spent: Cents(100)
274            }
275            .percent(),
276            0
277        );
278    }
279
280    #[test]
281    fn extra_percent_truncates() {
282        // Bash integer division — 33/100 -> 33%, 50/100 -> 50%.
283        assert_eq!(
284            ExtraUsage {
285                limit: Cents(10000),
286                spent: Cents(3333)
287            }
288            .percent(),
289            33
290        );
291    }
292
293    #[test]
294    fn severity_picks_worst_of_three_windows() {
295        let s = snap(40, 60, Some(80), None);
296        assert_eq!(anthropic_severity(&s), PaceSeverity::High); // 80 → high
297    }
298
299    #[test]
300    fn severity_ignores_extra_when_no_cap_hit() {
301        // Extra at 95% but no rate-limit at 100% → extra is NOT promoted.
302        let s = snap(50, 60, None, Some((10000, 9500)));
303        assert_eq!(anthropic_severity(&s), PaceSeverity::Mid); // capped at 60
304    }
305
306    #[test]
307    fn severity_promotes_extra_when_session_at_100() {
308        let s = snap(100, 50, None, Some((10000, 9500)));
309        assert_eq!(anthropic_severity(&s), PaceSeverity::Critical); // 100 → critical
310    }
311
312    #[test]
313    fn severity_falls_through_to_extra_when_extra_higher_than_capped_window() {
314        // session = 100, weekly = 50, extra = 100% → max should be 100.
315        let s = snap(100, 50, None, Some((10000, 10000)));
316        assert_eq!(anthropic_severity(&s), PaceSeverity::Critical);
317    }
318}
ai_usagebar/usage.rs

ai_usagebar/
usage.rs
