ai_usagebar/grok/
vendor.rs

1//! xAI (Grok) renderer — bar text + bordered Pango tooltip. Prepaid credit
2//! balance in USD, from the Management API.
3
4use std::collections::HashMap;
5
6use chrono::{DateTime, Utc};
7
8use crate::format::{placeholders, substitute, updated_at_hm, usd};
9use crate::pacing::PaceSeverity;
10use crate::pango::{color_span, escape, severity_color};
11use crate::theme::Theme;
12use crate::tooltip::{Line as TooltipLine, render_bordered};
13use crate::usage::GrokSnapshot;
14use crate::vendor::{RenderOpts, VendorId, VendorOutcome};
15use crate::waybar::{Class, WaybarOutput};
16
17use super::fetch::FetchOutcome;
18
19pub const DEFAULT_FORMAT: &str = "{grok_balance}";
20
21pub fn build_placeholders(snap: &GrokSnapshot) -> HashMap<&'static str, String> {
22    placeholders(vec![
23        ("icon", "󰇷".to_string()),
24        ("vendor_short", VendorId::Grok.short_name().to_string()),
25        ("session_pct", "0".to_string()),
26        ("session_reset", "—".to_string()),
27        ("weekly_pct", "0".to_string()),
28        ("weekly_reset", "—".to_string()),
29        ("plan", "Grok".to_string()),
30        ("grok_balance", usd(snap.balance)),
31    ])
32}
33
34/// Prepaid credit: running low = warmer, empty/negative = critical.
35pub fn severity(snap: &GrokSnapshot) -> PaceSeverity {
36    crate::pango::balance_severity(snap.balance, "USD")
37}
38
39pub fn render(
40    outcome: &VendorOutcome,
41    snap: &GrokSnapshot,
42    theme: &Theme,
43    opts: &RenderOpts,
44    now: DateTime<Utc>,
45) -> WaybarOutput {
46    let class = Class::from(severity(snap));
47    let format = opts
48        .format
49        .clone()
50        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
51    let values = build_placeholders(snap);
52
53    let mut text = substitute(&format, &values);
54    if outcome.stale {
55        text.push_str(" ⏸");
56    }
57
58    let wrapper_color = severity_color(severity(snap), theme).to_string();
59    let icon_prefix = match opts.icon.as_deref() {
60        Some(ic) if !ic.is_empty() => format!("{ic} "),
61        _ => String::new(),
62    };
63    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
64
65    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
66        substitute(fmt, &values)
67    } else {
68        render_tooltip(outcome, snap, theme, now)
69    };
70
71    WaybarOutput {
72        text: bar_text,
73        tooltip,
74        class,
75    }
76}
77
78fn render_tooltip(
79    outcome: &VendorOutcome,
80    snap: &GrokSnapshot,
81    theme: &Theme,
82    now: DateTime<Utc>,
83) -> String {
84    let blue = &theme.blue;
85    let dim = &theme.dim;
86    let fg = &theme.fg;
87    let color = severity_color(severity(snap), theme);
88
89    let mut lines: Vec<TooltipLine> = Vec::new();
90    lines.push(TooltipLine::Center(format!(
91        "<span font_weight='bold' foreground='{blue}'>Grok (xAI)</span>"
92    )));
93    lines.push(TooltipLine::Sep);
94    lines.push(TooltipLine::Body("".into()));
95
96    lines.push(TooltipLine::Body(format!(
97        " <span foreground='{fg}'>  󰢗  Prepaid balance</span>"
98    )));
99    lines.push(TooltipLine::Body(format!(
100        "   <span font_weight='bold' foreground='{color}'>{bal}</span>",
101        bal = escape(&usd(snap.balance))
102    )));
103
104    if let Some((code, msg)) = outcome.last_error.as_ref()
105        && *code != 0
106    {
107        let (icon, ecolor) = if *code >= 500 {
108            ("󰅚", theme.red.as_str())
109        } else {
110            ("󰀪", theme.orange.as_str())
111        };
112        lines.push(TooltipLine::Body("".into()));
113        lines.push(TooltipLine::Sep);
114        lines.push(TooltipLine::Body(format!(
115            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
116        )));
117        lines.push(TooltipLine::Body(format!(
118            "     <span foreground='{dim}'>{}</span>",
119            escape(msg)
120        )));
121    }
122
123    let updated = updated_at_hm(now, outcome.cache_age);
124    lines.push(TooltipLine::Body("".into()));
125    lines.push(TooltipLine::Sep);
126    lines.push(TooltipLine::Body(format!(
127        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
128    )));
129
130    render_bordered(&lines, theme)
131}
132
133impl From<FetchOutcome> for VendorOutcome {
134    fn from(o: FetchOutcome) -> Self {
135        o.map(crate::usage::VendorSnapshot::Grok)
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use crate::usage::GrokSnapshot;
143
144    fn outcome(balance: f64) -> (GrokSnapshot, VendorOutcome) {
145        let snap = GrokSnapshot { balance };
146        let o = VendorOutcome {
147            snapshot: crate::usage::VendorSnapshot::Grok(snap.clone()),
148            stale: false,
149            last_error: None,
150            cache_age: Some(std::time::Duration::from_secs(10)),
151        };
152        (snap, o)
153    }
154
155    fn opts() -> RenderOpts {
156        RenderOpts {
157            format: None,
158            tooltip_format: None,
159            icon: None,
160            pace_tolerance: 5,
161            format_pace_color: false,
162            tooltip_pace_pts: false,
163        }
164    }
165
166    #[test]
167    fn renders_balance() {
168        let (snap, o) = outcome(25.0);
169        let out = render(&o, &snap, &Theme::default(), &opts(), Utc::now());
170        assert!(out.text.contains("$25.00"));
171        assert!(out.tooltip.contains("Prepaid balance"));
172    }
173
174    #[test]
175    fn low_balance_is_critical() {
176        let (snap, _) = outcome(0.5);
177        assert_eq!(severity(&snap), PaceSeverity::Critical);
178    }
179}
ai_usagebar/grok/vendor.rs

ai_usagebar/grok/
vendor.rs
