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

ai_usagebar/grok/
vendor.rs
