ai_usagebar/zai/
vendor.rs

1//! Z.AI renderer — bar text + bordered Pango tooltip.
2
3use std::collections::HashMap;
4
5use chrono::{DateTime, Utc};
6
7use crate::countdown;
8use crate::format::{placeholders, substitute, updated_at_hm};
9use crate::pacing::PaceSeverity;
10use crate::pango::{self, color_span, escape, severity_color, severity_for};
11use crate::theme::Theme;
12use crate::tooltip::{Line as TooltipLine, render_bordered};
13use crate::usage::{UsageWindow, ZaiSnapshot};
14use crate::vendor::{RenderOpts, VendorOutcome};
15use crate::waybar::{Class, WaybarOutput};
16
17use super::fetch::FetchOutcome;
18
19pub const DEFAULT_FORMAT: &str = "{zai_session_pct}% · {zai_session_reset}";
20
21pub fn build_placeholders(snap: &ZaiSnapshot, now: DateTime<Utc>) -> HashMap<&'static str, String> {
22    let session_pct = snap
23        .session
24        .as_ref()
25        .map(|w| w.utilization_pct)
26        .unwrap_or(0);
27    let weekly_pct = snap.weekly.as_ref().map(|w| w.utilization_pct).unwrap_or(0);
28    let mcp_pct = snap.mcp.as_ref().map(|w| w.utilization_pct).unwrap_or(0);
29    placeholders(vec![
30        ("icon", "󰚩".to_string()),
31        ("vendor_short", "zai".to_string()),
32        // Cross-vendor aliases for scroll-cycle friendly formats.
33        ("session_pct", session_pct.to_string()),
34        (
35            "session_reset",
36            countdown::format(window_reset(&snap.session), now),
37        ),
38        ("weekly_pct", weekly_pct.to_string()),
39        (
40            "weekly_reset",
41            countdown::format(window_reset(&snap.weekly), now),
42        ),
43        ("plan", snap.plan.clone()),
44        ("zai_plan", snap.plan.clone()),
45        ("zai_session_pct", session_pct.to_string()),
46        (
47            "zai_session_reset",
48            countdown::format(window_reset(&snap.session), now),
49        ),
50        ("zai_weekly_pct", weekly_pct.to_string()),
51        (
52            "zai_weekly_reset",
53            countdown::format(window_reset(&snap.weekly), now),
54        ),
55        ("zai_mcp_pct", mcp_pct.to_string()),
56        (
57            "zai_mcp_reset",
58            countdown::format(window_reset(&snap.mcp), now),
59        ),
60    ])
61}
62
63fn window_reset(w: &Option<UsageWindow>) -> Option<DateTime<Utc>> {
64    w.as_ref().and_then(|w| w.resets_at)
65}
66
67pub fn severity(snap: &ZaiSnapshot) -> PaceSeverity {
68    let session = snap
69        .session
70        .as_ref()
71        .map(|w| w.utilization_pct)
72        .unwrap_or(0);
73    let weekly = snap.weekly.as_ref().map(|w| w.utilization_pct).unwrap_or(0);
74    let mcp = snap.mcp.as_ref().map(|w| w.utilization_pct).unwrap_or(0);
75    severity_for([session, weekly, mcp].into_iter().max().unwrap_or(0))
76}
77
78pub fn render(
79    outcome: &VendorOutcome,
80    snap: &ZaiSnapshot,
81    theme: &Theme,
82    opts: &RenderOpts,
83    now: DateTime<Utc>,
84) -> WaybarOutput {
85    let class = Class::from(severity(snap));
86    let format = opts
87        .format
88        .clone()
89        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
90    let values = build_placeholders(snap, now);
91
92    let mut text = substitute(&format, &values);
93    if outcome.stale {
94        text.push_str(" ⏸");
95    }
96    let wrapper_color = severity_color(severity(snap), theme).to_string();
97    let icon_prefix = match opts.icon.as_deref() {
98        Some(ic) if !ic.is_empty() => format!("{ic} "),
99        _ => String::new(),
100    };
101    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
102
103    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
104        substitute(fmt, &values)
105    } else {
106        render_tooltip(outcome, snap, theme, now)
107    };
108
109    WaybarOutput {
110        text: bar_text,
111        tooltip,
112        class,
113    }
114}
115
116fn render_tooltip(
117    outcome: &VendorOutcome,
118    snap: &ZaiSnapshot,
119    theme: &Theme,
120    now: DateTime<Utc>,
121) -> String {
122    let blue = &theme.blue;
123    let dim = &theme.dim;
124    let mut lines: Vec<TooltipLine> = Vec::new();
125    lines.push(TooltipLine::Center(format!(
126        "<span font_weight='bold' foreground='{blue}'>{plan}</span>",
127        plan = escape(&snap.plan)
128    )));
129    lines.push(TooltipLine::Sep);
130    lines.push(TooltipLine::Body("".into()));
131
132    if let Some(w) = snap.session.as_ref() {
133        push_window(&mut lines, "  󰔟  Session (5h)", w, theme, now);
134    }
135    if let Some(w) = snap.weekly.as_ref() {
136        if snap.session.is_some() {
137            lines.push(TooltipLine::Body("".into()));
138        }
139        push_window(&mut lines, "  󰃰  Weekly", w, theme, now);
140    }
141    if let Some(w) = snap.mcp.as_ref() {
142        lines.push(TooltipLine::Body("".into()));
143        lines.push(TooltipLine::Sep);
144        push_window(&mut lines, "  󰓹  MCP tools (monthly)", w, theme, now);
145    }
146    if snap.session.is_none() && snap.weekly.is_none() && snap.mcp.is_none() {
147        lines.push(TooltipLine::Body(format!(
148            " <span foreground='{dim}'>no usage windows reported</span>"
149        )));
150    }
151
152    if let Some((code, msg)) = outcome.last_error.as_ref()
153        && *code != 0
154    {
155        let (icon, ecolor) = if *code >= 500 {
156            ("󰅚", theme.red.as_str())
157        } else {
158            ("󰀪", theme.orange.as_str())
159        };
160        lines.push(TooltipLine::Body("".into()));
161        lines.push(TooltipLine::Sep);
162        lines.push(TooltipLine::Body(format!(
163            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
164        )));
165        lines.push(TooltipLine::Body(format!(
166            "     <span foreground='{dim}'>{}</span>",
167            escape(msg)
168        )));
169    }
170
171    let updated = updated_at_hm(now, outcome.cache_age);
172    lines.push(TooltipLine::Body("".into()));
173    lines.push(TooltipLine::Sep);
174    lines.push(TooltipLine::Body(format!(
175        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
176    )));
177
178    render_bordered(&lines, theme)
179}
180
181fn push_window(
182    lines: &mut Vec<TooltipLine>,
183    label: &str,
184    w: &UsageWindow,
185    theme: &Theme,
186    now: DateTime<Utc>,
187) {
188    let color = severity_color(severity_for(w.utilization_pct), theme);
189    let bar = pango::progress_bar(w.utilization_pct, color, theme, None);
190    let fg = &theme.fg;
191    let dim = &theme.dim;
192    lines.push(TooltipLine::Body(format!(
193        " <span foreground='{fg}'>{label}</span>"
194    )));
195    lines.push(TooltipLine::Body(format!(
196        "   {bar}  <span font_weight='bold' foreground='{color}'>{pct}%</span>",
197        pct = w.utilization_pct
198    )));
199    lines.push(TooltipLine::Body(format!(
200        " <span foreground='{dim}'>  ⏱  Resets in {cd}</span>",
201        cd = escape(&countdown::format(w.resets_at, now))
202    )));
203}
204
205impl From<FetchOutcome> for VendorOutcome {
206    fn from(o: FetchOutcome) -> Self {
207        Self {
208            snapshot: crate::usage::VendorSnapshot::Zai(o.snapshot),
209            stale: o.stale,
210            last_error: o.last_error,
211            cache_age: o.cache_age,
212        }
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use crate::usage::{UsageWindow, ZaiSnapshot};
220
221    fn sample_snap() -> ZaiSnapshot {
222        let now = Utc::now();
223        ZaiSnapshot {
224            plan: "GLM Coding Pro".into(),
225            session: Some(UsageWindow {
226                utilization_pct: 42,
227                resets_at: Some(now + chrono::Duration::hours(2)),
228                window_duration: chrono::Duration::hours(5),
229            }),
230            weekly: Some(UsageWindow {
231                utilization_pct: 15,
232                resets_at: Some(now + chrono::Duration::days(3)),
233                window_duration: chrono::Duration::days(7),
234            }),
235            mcp: None,
236        }
237    }
238
239    fn outcome(s: ZaiSnapshot) -> VendorOutcome {
240        VendorOutcome {
241            snapshot: crate::usage::VendorSnapshot::Zai(s),
242            stale: false,
243            last_error: None,
244            cache_age: Some(std::time::Duration::from_secs(10)),
245        }
246    }
247
248    fn opts() -> RenderOpts {
249        RenderOpts {
250            format: None,
251            tooltip_format: None,
252            icon: None,
253            pace_tolerance: 5,
254            format_pace_color: false,
255            tooltip_pace_pts: false,
256        }
257    }
258
259    #[test]
260    fn default_format_renders_session_pct() {
261        let snap = sample_snap();
262        let oc = outcome(snap.clone());
263        let out = render(&oc, &snap, &Theme::default(), &opts(), Utc::now());
264        assert!(out.text.contains("42%"));
265    }
266
267    #[test]
268    fn tooltip_contains_all_windows_present() {
269        let snap = sample_snap();
270        let oc = outcome(snap.clone());
271        let out = render(&oc, &snap, &Theme::default(), &opts(), Utc::now());
272        assert!(out.tooltip.contains("Session"));
273        assert!(out.tooltip.contains("Weekly"));
274        assert!(!out.tooltip.contains("MCP"));
275    }
276
277    #[test]
278    fn empty_snapshot_renders_no_windows_message() {
279        let snap = ZaiSnapshot {
280            plan: "GLM Coding Unknown".into(),
281            session: None,
282            weekly: None,
283            mcp: None,
284        };
285        let oc = outcome(snap.clone());
286        let out = render(&oc, &snap, &Theme::default(), &opts(), Utc::now());
287        assert!(out.tooltip.contains("no usage windows reported"));
288    }
289
290    #[test]
291    fn severity_picks_worst_window() {
292        let mut snap = sample_snap();
293        snap.weekly.as_mut().unwrap().utilization_pct = 95;
294        assert_eq!(severity(&snap), PaceSeverity::Critical);
295    }
296
297    #[test]
298    fn custom_tooltip_uses_placeholders() {
299        let snap = sample_snap();
300        let oc = outcome(snap.clone());
301        let mut o = opts();
302        o.tooltip_format = Some("S:{zai_session_pct} W:{zai_weekly_pct}".into());
303        let out = render(&oc, &snap, &Theme::default(), &o, Utc::now());
304        assert_eq!(out.tooltip, "S:42 W:15");
305    }
306}
ai_usagebar/zai/vendor.rs

ai_usagebar/zai/
vendor.rs
