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::{color_span, escape, severity_color, severity_for};
11use crate::theme::Theme;
12use crate::tooltip::{Line as TooltipLine, push_window, 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, None);
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, None);
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, None);
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
181impl From<FetchOutcome> for VendorOutcome {
182    fn from(o: FetchOutcome) -> Self {
183        Self {
184            snapshot: crate::usage::VendorSnapshot::Zai(o.snapshot),
185            stale: o.stale,
186            last_error: o.last_error,
187            cache_age: o.cache_age,
188        }
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use crate::usage::{UsageWindow, ZaiSnapshot};
196
197    fn sample_snap() -> ZaiSnapshot {
198        let now = Utc::now();
199        ZaiSnapshot {
200            plan: "GLM Coding Pro".into(),
201            session: Some(UsageWindow {
202                utilization_pct: 42,
203                resets_at: Some(now + chrono::Duration::hours(2)),
204                window_duration: chrono::Duration::hours(5),
205            }),
206            weekly: Some(UsageWindow {
207                utilization_pct: 15,
208                resets_at: Some(now + chrono::Duration::days(3)),
209                window_duration: chrono::Duration::days(7),
210            }),
211            mcp: None,
212        }
213    }
214
215    fn outcome(s: ZaiSnapshot) -> VendorOutcome {
216        VendorOutcome {
217            snapshot: crate::usage::VendorSnapshot::Zai(s),
218            stale: false,
219            last_error: None,
220            cache_age: Some(std::time::Duration::from_secs(10)),
221        }
222    }
223
224    fn opts() -> RenderOpts {
225        RenderOpts {
226            format: None,
227            tooltip_format: None,
228            icon: None,
229            pace_tolerance: 5,
230            format_pace_color: false,
231            tooltip_pace_pts: false,
232        }
233    }
234
235    #[test]
236    fn default_format_renders_session_pct() {
237        let snap = sample_snap();
238        let oc = outcome(snap.clone());
239        let out = render(&oc, &snap, &Theme::default(), &opts(), Utc::now());
240        assert!(out.text.contains("42%"));
241    }
242
243    #[test]
244    fn tooltip_contains_all_windows_present() {
245        let snap = sample_snap();
246        let oc = outcome(snap.clone());
247        let out = render(&oc, &snap, &Theme::default(), &opts(), Utc::now());
248        assert!(out.tooltip.contains("Session"));
249        assert!(out.tooltip.contains("Weekly"));
250        assert!(!out.tooltip.contains("MCP"));
251    }
252
253    #[test]
254    fn empty_snapshot_renders_no_windows_message() {
255        let snap = ZaiSnapshot {
256            plan: "GLM Coding Unknown".into(),
257            session: None,
258            weekly: None,
259            mcp: None,
260        };
261        let oc = outcome(snap.clone());
262        let out = render(&oc, &snap, &Theme::default(), &opts(), Utc::now());
263        assert!(out.tooltip.contains("no usage windows reported"));
264    }
265
266    #[test]
267    fn severity_picks_worst_window() {
268        let mut snap = sample_snap();
269        snap.weekly.as_mut().unwrap().utilization_pct = 95;
270        assert_eq!(severity(&snap), PaceSeverity::Critical);
271    }
272
273    #[test]
274    fn custom_tooltip_uses_placeholders() {
275        let snap = sample_snap();
276        let oc = outcome(snap.clone());
277        let mut o = opts();
278        o.tooltip_format = Some("S:{zai_session_pct} W:{zai_weekly_pct}".into());
279        let out = render(&oc, &snap, &Theme::default(), &o, Utc::now());
280        assert_eq!(out.tooltip, "S:42 W:15");
281    }
282}
ai_usagebar/zai/vendor.rs

ai_usagebar/zai/
vendor.rs
