Skip to main content

ai_usagebar/grokbot/
vendor.rs

1//! Grok Bot renderer — bar text + bordered Pango tooltip. One weekly meter,
2//! plus the no-included-allowance state and the on-demand footnote.
3
4use std::collections::HashMap;
5
6use chrono::{DateTime, Utc};
7
8use crate::countdown;
9use crate::format::{placeholders, substitute, updated_at_hm};
10use crate::pacing::PaceSeverity;
11use crate::pango::{color_span, escape, severity_color, severity_for};
12use crate::theme::Theme;
13use crate::tooltip::{Line as TooltipLine, WindowRow, push_window_with_row, render_bordered};
14use crate::usage::{GrokbotSnapshot, UsageWindow};
15use crate::vendor::{RenderOpts, VendorId, VendorOutcome};
16use crate::waybar::{Class, WaybarOutput};
17
18use super::fetch::FetchOutcome;
19
20pub const DEFAULT_FORMAT: &str = "{gbt_weekly_pct}%";
21
22pub fn build_placeholders(
23    snap: &GrokbotSnapshot,
24    now: DateTime<Utc>,
25) -> HashMap<&'static str, String> {
26    let reset = countdown::format(snap.reset_at, now);
27    // With no included allowance the weekly placeholders resolve to empty
28    // strings (the missing-placeholder convention) rather than a fabricated 0.
29    let allowance = |value: String| {
30        if snap.has_included_allowance {
31            value
32        } else {
33            String::new()
34        }
35    };
36    placeholders(vec![
37        ("icon", VendorId::Grokbot.bar_icon().to_string()),
38        ("vendor_short", VendorId::Grokbot.short_name().to_string()),
39        // Cross-vendor aliases.
40        ("plan", snap.plan.clone()),
41        ("weekly_pct", allowance(snap.weekly_pct.to_string())),
42        ("weekly_reset", allowance(reset.clone())),
43        // Grok Bot-specific placeholders.
44        ("gbt_plan", snap.plan.clone()),
45        ("gbt_weekly_pct", allowance(snap.weekly_pct.to_string())),
46        ("gbt_weekly_reset", allowance(reset)),
47        (
48            "gbt_on_demand",
49            if snap.on_demand_enabled { "on" } else { "off" }.to_string(),
50        ),
51    ])
52}
53
54/// No included allowance reads as Low: there is no pool to be exhausting.
55pub fn severity(snap: &GrokbotSnapshot) -> PaceSeverity {
56    if snap.has_included_allowance {
57        severity_for(snap.weekly_pct)
58    } else {
59        PaceSeverity::Low
60    }
61}
62
63pub fn render(
64    outcome: &VendorOutcome,
65    snap: &GrokbotSnapshot,
66    theme: &Theme,
67    opts: &RenderOpts,
68    now: DateTime<Utc>,
69) -> WaybarOutput {
70    let class = Class::from(severity(snap));
71    let format = opts
72        .format
73        .clone()
74        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
75    let values = build_placeholders(snap, now);
76    // User formats are Pango markup after Waybar renders them; the plan label
77    // is API-controlled, so it is escaped at this projection boundary, once.
78    let mut pango_values = values.clone();
79    for key in ["plan", "gbt_plan"] {
80        if let Some(value) = pango_values.get_mut(key) {
81            *value = escape(value);
82        }
83    }
84
85    let mut text = substitute(&format, &pango_values);
86    if outcome.stale {
87        text.push_str(" ⏸");
88    }
89
90    let wrapper_color = severity_color(severity(snap), theme).to_string();
91    let icon_prefix = match opts.icon.as_deref() {
92        Some(ic) if !ic.is_empty() => format!("{ic} "),
93        _ => String::new(),
94    };
95    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
96
97    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
98        substitute(fmt, &pango_values)
99    } else {
100        render_tooltip(outcome, snap, theme, now)
101    };
102
103    WaybarOutput {
104        text: bar_text,
105        tooltip,
106        class,
107    }
108}
109
110fn render_tooltip(
111    outcome: &VendorOutcome,
112    snap: &GrokbotSnapshot,
113    theme: &Theme,
114    now: DateTime<Utc>,
115) -> String {
116    let blue = &theme.blue;
117    let dim = &theme.dim;
118    let fg = &theme.fg;
119    let plan_color = severity_color(severity(snap), theme);
120
121    let mut lines: Vec<TooltipLine> = Vec::new();
122    lines.push(TooltipLine::Center(format!(
123        "<span font_weight='bold' foreground='{blue}'>Grok Bot</span>"
124    )));
125    lines.push(TooltipLine::Sep);
126    lines.push(TooltipLine::Body("".into()));
127
128    lines.push(TooltipLine::Body(format!(
129        " <span foreground='{fg}'>  󰣖  Plan</span>"
130    )));
131    lines.push(TooltipLine::Body(format!(
132        "   <span font_weight='bold' foreground='{plan_color}'>{}</span>",
133        escape(&snap.plan)
134    )));
135
136    lines.push(TooltipLine::Body("".into()));
137    if snap.has_included_allowance {
138        // The window's length is derived from the reported instants, never
139        // assumed — and only the reset carries a countdown.
140        push_window_with_row(
141            &mut lines,
142            "  󰅄  Weekly included usage",
143            &UsageWindow {
144                utilization_pct: snap.weekly_pct,
145                resets_at: snap.reset_at,
146                window_duration: snap.window.unwrap_or_else(chrono::Duration::zero),
147            },
148            theme,
149            now,
150            WindowRow::default(),
151        );
152        if let Some(note) = snap.on_demand_note() {
153            lines.push(TooltipLine::Body(format!(
154                "     <span foreground='{dim}'>{note}</span>"
155            )));
156        }
157    } else {
158        lines.push(TooltipLine::Body(format!(
159            " <span foreground='{dim}'>  No included allowance on this account</span>"
160        )));
161    }
162
163    if let Some((code, msg)) = outcome.last_error.as_ref() {
164        // Code zero has never meant HTTP; the rest are real status codes.
165        let (label, icon, ecolor) = match *code {
166            0 => ("Grok Bot error".to_string(), "󰅚", theme.red.as_str()),
167            code if code >= 500 => (format!("HTTP {code}"), "󰅚", theme.red.as_str()),
168            code => (format!("HTTP {code}"), "󰀪", theme.orange.as_str()),
169        };
170        lines.push(TooltipLine::Body("".into()));
171        lines.push(TooltipLine::Sep);
172        lines.push(TooltipLine::Body(format!(
173            " <span foreground='{ecolor}'>  {icon}  {label}</span>"
174        )));
175        if msg != &label {
176            lines.push(TooltipLine::Body(format!(
177                "     <span foreground='{dim}'>{}</span>",
178                escape(msg)
179            )));
180        }
181    }
182
183    let updated = updated_at_hm(now, outcome.cache_age);
184    lines.push(TooltipLine::Body("".into()));
185    lines.push(TooltipLine::Sep);
186    lines.push(TooltipLine::Body(format!(
187        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
188    )));
189
190    render_bordered(&lines, theme)
191}
192
193impl From<FetchOutcome> for VendorOutcome {
194    fn from(o: FetchOutcome) -> Self {
195        o.map(crate::usage::VendorSnapshot::Grokbot)
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use chrono::TimeZone;
203
204    fn now() -> DateTime<Utc> {
205        Utc.with_ymd_and_hms(2026, 9, 14, 12, 0, 0).unwrap()
206    }
207
208    fn sample_snap() -> GrokbotSnapshot {
209        GrokbotSnapshot {
210            plan: "Grok Bot Plan".into(),
211            has_included_allowance: true,
212            weekly_pct: 42,
213            has_available_usage: true,
214            on_demand_enabled: false,
215            period_start: Some(now() - chrono::Duration::days(3)),
216            reset_at: Some(now() + chrono::Duration::days(4)),
217            window: Some(chrono::Duration::days(7)),
218        }
219    }
220
221    fn sample_outcome(snap: GrokbotSnapshot) -> VendorOutcome {
222        VendorOutcome {
223            snapshot: crate::usage::VendorSnapshot::Grokbot(snap),
224            stale: false,
225            last_error: None,
226            cache_age: Some(std::time::Duration::from_secs(10)),
227        }
228    }
229
230    fn opts() -> RenderOpts {
231        RenderOpts {
232            format: None,
233            tooltip_format: None,
234            icon: None,
235            pace_tolerance: 5,
236            format_pace_color: false,
237            tooltip_pace_pts: false,
238        }
239    }
240
241    #[test]
242    fn default_render_shows_the_weekly_percent() {
243        let snap = sample_snap();
244        let outcome = sample_outcome(snap.clone());
245        let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
246        assert!(out.text.contains("42%"), "text: {}", out.text);
247        assert!(!out.text.contains("%%"), "text: {}", out.text);
248    }
249
250    #[test]
251    fn placeholder_set_contains_all_keys() {
252        let snap = sample_snap();
253        let values = build_placeholders(&snap, now());
254        for key in [
255            "gbt_plan",
256            "gbt_weekly_pct",
257            "gbt_weekly_reset",
258            "gbt_on_demand",
259            "plan",
260            "weekly_pct",
261            "weekly_reset",
262            "vendor_short",
263        ] {
264            assert!(values.contains_key(key), "missing placeholder {key}");
265        }
266        assert_eq!(values["gbt_weekly_pct"], "42");
267        assert_eq!(values["weekly_pct"], "42");
268        assert_eq!(values["gbt_on_demand"], "off");
269        assert!(!values["gbt_weekly_reset"].is_empty());
270    }
271
272    #[test]
273    fn no_allowance_renders_empty_weekly_placeholders_not_a_zero() {
274        let snap = GrokbotSnapshot {
275            has_included_allowance: false,
276            weekly_pct: 0,
277            ..sample_snap()
278        };
279        let values = build_placeholders(&snap, now());
280        for key in [
281            "gbt_weekly_pct",
282            "gbt_weekly_reset",
283            "weekly_pct",
284            "weekly_reset",
285        ] {
286            assert_eq!(values[key], "", "{key} must render empty");
287        }
288        // …and the default format must not leave literal braces behind. The
289        // dangling `%` is the kimi monthly-shape convention: substitution
290        // tolerates the shape rather than fabricating a figure.
291        let outcome = sample_outcome(snap.clone());
292        let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
293        assert!(!out.text.contains('{'), "{}", out.text);
294        assert!(!out.text.contains('}'), "{}", out.text);
295        assert!(out.text.contains('%'), "{}", out.text);
296        assert!(
297            !out.text.contains('0'),
298            "a fabricated 0 must not appear: {}",
299            out.text
300        );
301        assert_eq!(severity(&snap), PaceSeverity::Low);
302    }
303
304    #[test]
305    fn tooltip_draws_the_weekly_bar_and_a_reset_countdown() {
306        let snap = sample_snap();
307        let outcome = sample_outcome(snap.clone());
308        let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
309        assert!(out.tooltip.contains("Grok Bot"), "{}", out.tooltip);
310        assert!(out.tooltip.contains("Grok Bot Plan"), "{}", out.tooltip);
311        assert!(
312            out.tooltip.contains("Weekly included usage"),
313            "{}",
314            out.tooltip
315        );
316        assert!(out.tooltip.contains("42%"), "{}", out.tooltip);
317        assert!(out.tooltip.contains("Resets in"), "{}", out.tooltip);
318        // A countdown, not the raw timestamp.
319        assert!(!out.tooltip.contains("2026-09-18"), "{}", out.tooltip);
320    }
321
322    #[test]
323    fn the_no_allowance_state_is_a_text_line_not_a_meter() {
324        let snap = GrokbotSnapshot {
325            has_included_allowance: false,
326            weekly_pct: 0,
327            ..sample_snap()
328        };
329        let outcome = sample_outcome(snap.clone());
330        let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
331        assert!(
332            out.tooltip.contains("No included allowance"),
333            "{}",
334            out.tooltip
335        );
336        assert!(
337            !out.tooltip.contains("Weekly included usage"),
338            "{}",
339            out.tooltip
340        );
341    }
342
343    #[test]
344    fn the_on_demand_footnote_appears_only_at_an_exhausted_pool_with_on_demand() {
345        let mut snap = sample_snap();
346        snap.weekly_pct = 100;
347        snap.has_available_usage = true;
348        snap.on_demand_enabled = true;
349        let outcome = sample_outcome(snap.clone());
350        let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
351        assert!(out.tooltip.contains("on-demand"), "{}", out.tooltip);
352
353        snap.on_demand_enabled = false;
354        let outcome = sample_outcome(snap.clone());
355        let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
356        assert!(!out.tooltip.contains("on-demand may"), "{}", out.tooltip);
357    }
358
359    #[test]
360    fn plan_is_pango_escaped() {
361        let mut snap = sample_snap();
362        snap.plan = "A&B <beta>".into();
363        let outcome = sample_outcome(snap.clone());
364        let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
365        assert!(
366            out.tooltip.contains("A&amp;B &lt;beta&gt;"),
367            "tooltip: {}",
368            out.tooltip
369        );
370        let mut o = opts();
371        o.tooltip_format = Some("{gbt_plan}".into());
372        let out = render(&outcome, &snap, &Theme::default(), &o, now());
373        assert_eq!(out.tooltip, "A&amp;B &lt;beta&gt;");
374    }
375
376    #[test]
377    fn stale_appends_pause() {
378        let snap = sample_snap();
379        let mut outcome = sample_outcome(snap.clone());
380        outcome.stale = true;
381        let out = render(&outcome, &snap, &Theme::default(), &opts(), now());
382        assert!(out.text.contains("⏸"));
383    }
384
385    #[test]
386    fn fetch_outcome_conversion_preserves_metadata() {
387        let snap = sample_snap();
388        let fetch = FetchOutcome {
389            snapshot: snap,
390            stale: true,
391            last_error: Some((401, "bad".into())),
392            cache_age: Some(std::time::Duration::from_secs(42)),
393        };
394        let vendor: VendorOutcome = fetch.into();
395        assert!(matches!(
396            vendor.snapshot,
397            crate::usage::VendorSnapshot::Grokbot(_)
398        ));
399        assert!(vendor.stale);
400        assert_eq!(vendor.last_error, Some((401, "bad".into())));
401        assert_eq!(vendor.cache_age, Some(std::time::Duration::from_secs(42)));
402    }
403}