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, 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", "grk".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    if snap.balance < 1.0 {
37        PaceSeverity::Critical
38    } else if snap.balance < 5.0 {
39        PaceSeverity::High
40    } else if snap.balance < 20.0 {
41        PaceSeverity::Mid
42    } else {
43        PaceSeverity::Low
44    }
45}
46
47pub fn render(
48    outcome: &VendorOutcome,
49    snap: &GrokSnapshot,
50    theme: &Theme,
51    opts: &RenderOpts,
52    now: DateTime<Utc>,
53) -> WaybarOutput {
54    let class = Class::from(severity(snap));
55    let format = opts
56        .format
57        .clone()
58        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
59    let values = build_placeholders(snap);
60
61    let mut text = substitute(&format, &values);
62    if outcome.stale {
63        text.push_str(" ⏸");
64    }
65
66    let wrapper_color = severity_color(severity(snap), theme).to_string();
67    let icon_prefix = match opts.icon.as_deref() {
68        Some(ic) if !ic.is_empty() => format!("{ic} "),
69        _ => String::new(),
70    };
71    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
72
73    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
74        substitute(fmt, &values)
75    } else {
76        render_tooltip(outcome, snap, theme, now)
77    };
78
79    WaybarOutput {
80        text: bar_text,
81        tooltip,
82        class,
83    }
84}
85
86fn render_tooltip(
87    outcome: &VendorOutcome,
88    snap: &GrokSnapshot,
89    theme: &Theme,
90    now: DateTime<Utc>,
91) -> String {
92    let blue = &theme.blue;
93    let dim = &theme.dim;
94    let fg = &theme.fg;
95    let color = severity_color(severity(snap), theme);
96
97    let mut lines: Vec<TooltipLine> = Vec::new();
98    lines.push(TooltipLine::Center(format!(
99        "<span font_weight='bold' foreground='{blue}'>Grok (xAI)</span>"
100    )));
101    lines.push(TooltipLine::Sep);
102    lines.push(TooltipLine::Body("".into()));
103
104    lines.push(TooltipLine::Body(format!(
105        " <span foreground='{fg}'>  󰢗  Prepaid balance</span>"
106    )));
107    lines.push(TooltipLine::Body(format!(
108        "   <span font_weight='bold' foreground='{color}'>{bal}</span>",
109        bal = escape(&usd(snap.balance))
110    )));
111
112    if let Some((code, msg)) = outcome.last_error.as_ref()
113        && *code != 0
114    {
115        let (icon, ecolor) = if *code >= 500 {
116            ("󰅚", theme.red.as_str())
117        } else {
118            ("󰀪", theme.orange.as_str())
119        };
120        lines.push(TooltipLine::Body("".into()));
121        lines.push(TooltipLine::Sep);
122        lines.push(TooltipLine::Body(format!(
123            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
124        )));
125        lines.push(TooltipLine::Body(format!(
126            "     <span foreground='{dim}'>{}</span>",
127            escape(msg)
128        )));
129    }
130
131    let updated = updated_at_hm(now, outcome.cache_age);
132    lines.push(TooltipLine::Body("".into()));
133    lines.push(TooltipLine::Sep);
134    lines.push(TooltipLine::Body(format!(
135        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
136    )));
137
138    render_bordered(&lines, theme)
139}
140
141impl From<FetchOutcome> for VendorOutcome {
142    fn from(o: FetchOutcome) -> Self {
143        Self {
144            snapshot: crate::usage::VendorSnapshot::Grok(o.snapshot),
145            stale: o.stale,
146            last_error: o.last_error,
147            cache_age: o.cache_age,
148        }
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use crate::usage::GrokSnapshot;
156
157    fn outcome(balance: f64) -> (GrokSnapshot, VendorOutcome) {
158        let snap = GrokSnapshot { balance };
159        let o = VendorOutcome {
160            snapshot: crate::usage::VendorSnapshot::Grok(snap.clone()),
161            stale: false,
162            last_error: None,
163            cache_age: Some(std::time::Duration::from_secs(10)),
164        };
165        (snap, o)
166    }
167
168    fn opts() -> RenderOpts {
169        RenderOpts {
170            format: None,
171            tooltip_format: None,
172            icon: None,
173            pace_tolerance: 5,
174            format_pace_color: false,
175            tooltip_pace_pts: false,
176        }
177    }
178
179    #[test]
180    fn renders_balance() {
181        let (snap, o) = outcome(25.0);
182        let out = render(&o, &snap, &Theme::default(), &opts(), Utc::now());
183        assert!(out.text.contains("$25.00"));
184        assert!(out.tooltip.contains("Prepaid balance"));
185    }
186
187    #[test]
188    fn low_balance_is_critical() {
189        let (snap, _) = outcome(0.5);
190        assert_eq!(severity(&snap), PaceSeverity::Critical);
191    }
192}
ai_usagebar/grok/vendor.rs

ai_usagebar/grok/
vendor.rs
