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        if s.utilization_pct > max {
207            max = s.utilization_pct;
208        }
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 {
218        if let Some(extra) = snap.extra {
219            let p = extra.percent();
220            if p > max {
221                max = p;
222            }
223        }
224    }
225    crate::pango::severity_for(max)
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231    use crate::pacing::PaceSeverity;
232    use chrono::Duration;
233
234    fn w(pct: i32) -> UsageWindow {
235        UsageWindow {
236            utilization_pct: pct,
237            resets_at: None,
238            window_duration: Duration::hours(5),
239        }
240    }
241
242    fn snap(s: i32, w_: i32, sonnet: Option<i32>, extra: Option<(i64, i64)>) -> AnthropicSnapshot {
243        AnthropicSnapshot {
244            plan: "Max 5x".into(),
245            session: w(s),
246            weekly: w(w_),
247            sonnet: sonnet.map(w),
248            extra: extra.map(|(limit, spent)| ExtraUsage {
249                limit: Cents(limit),
250                spent: Cents(spent),
251            }),
252        }
253    }
254
255    #[test]
256    fn cents_format_positive() {
257        assert_eq!(Cents(0).fmt_dollars(), "$0.00");
258        assert_eq!(Cents(50).fmt_dollars(), "$0.50");
259        assert_eq!(Cents(250).fmt_dollars(), "$2.50");
260        assert_eq!(Cents(5000).fmt_dollars(), "$50.00");
261    }
262
263    #[test]
264    fn cents_format_negative_uses_leading_sign() {
265        // claudebar bug-fix: never "$-1.-50" — sign goes before the dollar sign.
266        assert_eq!(Cents(-150).fmt_dollars(), "-$1.50");
267        assert_eq!(Cents(-1).fmt_dollars(), "-$0.01");
268    }
269
270    #[test]
271    fn extra_percent_with_zero_limit_is_zero() {
272        assert_eq!(
273            ExtraUsage {
274                limit: Cents(0),
275                spent: Cents(100)
276            }
277            .percent(),
278            0
279        );
280    }
281
282    #[test]
283    fn extra_percent_truncates() {
284        // Bash integer division — 33/100 -> 33%, 50/100 -> 50%.
285        assert_eq!(
286            ExtraUsage {
287                limit: Cents(10000),
288                spent: Cents(3333)
289            }
290            .percent(),
291            33
292        );
293    }
294
295    #[test]
296    fn severity_picks_worst_of_three_windows() {
297        let s = snap(40, 60, Some(80), None);
298        assert_eq!(anthropic_severity(&s), PaceSeverity::High); // 80 → high
299    }
300
301    #[test]
302    fn severity_ignores_extra_when_no_cap_hit() {
303        // Extra at 95% but no rate-limit at 100% → extra is NOT promoted.
304        let s = snap(50, 60, None, Some((10000, 9500)));
305        assert_eq!(anthropic_severity(&s), PaceSeverity::Mid); // capped at 60
306    }
307
308    #[test]
309    fn severity_promotes_extra_when_session_at_100() {
310        let s = snap(100, 50, None, Some((10000, 9500)));
311        assert_eq!(anthropic_severity(&s), PaceSeverity::Critical); // 100 → critical
312    }
313
314    #[test]
315    fn severity_falls_through_to_extra_when_extra_higher_than_capped_window() {
316        // session = 100, weekly = 50, extra = 100% → max should be 100.
317        let s = snap(100, 50, None, Some((10000, 10000)));
318        assert_eq!(anthropic_severity(&s), PaceSeverity::Critical);
319    }
320}
ai_usagebar/usage.rs

ai_usagebar/
usage.rs
