Skip to main content

ai_usagebar/commandcode/
vendor.rs

1use std::collections::HashMap;
2use std::time::Duration;
3
4use chrono::{DateTime, Utc};
5
6use crate::countdown;
7use crate::format::{placeholders, substitute, updated_at_hm, usd};
8use crate::pacing::PaceSeverity;
9use crate::pango::{color_span, escape, severity_color, severity_for};
10use crate::theme::Theme;
11use crate::tooltip::{Line as TooltipLine, WindowRow, push_window_with_detail, render_bordered};
12use crate::usage::UsageWindow;
13use crate::vendor::{RenderOpts, VendorId, VendorOutcome};
14use crate::waybar::{Class, WaybarOutput};
15
16use super::fetch::FetchOutcome;
17use super::types::{Snapshot, SpendWindow};
18
19pub const DEFAULT_FORMAT: &str = "{cc_session_pct}% · {cc_session_reset}";
20const DEFAULT_PLAN: &str = "Command Code";
21const UNAVAILABLE: &str = "—";
22
23impl From<FetchOutcome> for VendorOutcome {
24    fn from(outcome: FetchOutcome) -> Self {
25        outcome.map(crate::usage::VendorSnapshot::CommandCode)
26    }
27}
28
29pub fn build_placeholders(snap: &Snapshot, now: DateTime<Utc>) -> HashMap<&'static str, String> {
30    let plan = sanitize(snap.plan.as_deref().unwrap_or(DEFAULT_PLAN));
31    let session = window_values(snap.five_hour.as_ref(), now);
32    let weekly = window_values(snap.weekly.as_ref(), now);
33    // The monthly allowance rendered as a window of its own: dollars drawn
34    // from the plan pool, refilling at the billing period end.
35    let monthly = snap.monthly_window();
36    let monthly_values = window_values(monthly.as_ref(), now);
37    let remaining = snap
38        .credits
39        .as_ref()
40        .map(|credits| usd(credits.remaining()))
41        .unwrap_or_else(|| UNAVAILABLE.to_string());
42    let pool = snap
43        .credit_pool
44        .map(usd)
45        .unwrap_or_else(|| UNAVAILABLE.to_string());
46    let spent = snap
47        .credits_spent()
48        .map(usd)
49        .unwrap_or_else(|| UNAVAILABLE.to_string());
50    // When the monthly ledger refills (billing period end). Absent until the
51    // subscription supplies it.
52    let credits_reset = snap
53        .period_end
54        .map(|at| countdown::format(Some(at), now))
55        .unwrap_or_else(|| UNAVAILABLE.to_string());
56
57    placeholders([
58        (
59            "vendor_short",
60            VendorId::CommandCode.short_name().to_string(),
61        ),
62        ("plan", plan.clone()),
63        ("cc_plan", plan),
64        // Generic names so a shared format string works across vendors.
65        ("session_pct", session.percent.clone()),
66        ("session_reset", session.reset.clone()),
67        ("weekly_pct", weekly.percent.clone()),
68        ("weekly_reset", weekly.reset.clone()),
69        ("cc_session_pct", session.percent),
70        ("cc_session_reset", session.reset),
71        ("cc_session_used", session.used),
72        ("cc_session_cap", session.cap),
73        ("cc_weekly_pct", weekly.percent),
74        ("cc_weekly_reset", weekly.reset),
75        ("cc_weekly_used", weekly.used),
76        ("cc_weekly_cap", weekly.cap),
77        ("cc_monthly_pct", monthly_values.percent.clone()),
78        ("cc_monthly_reset", monthly_values.reset.clone()),
79        ("cc_monthly_used", monthly_values.used.clone()),
80        ("cc_monthly_cap", monthly_values.cap.clone()),
81        ("cc_credits", remaining),
82        ("cc_credits_pool", pool),
83        ("cc_credits_spent", spent),
84        ("cc_credits_reset", credits_reset),
85    ])
86}
87
88#[derive(Debug)]
89struct WindowValues {
90    percent: String,
91    reset: String,
92    used: String,
93    cap: String,
94}
95
96fn window_values(window: Option<&SpendWindow>, now: DateTime<Utc>) -> WindowValues {
97    let Some(window) = window else {
98        return WindowValues {
99            percent: UNAVAILABLE.to_string(),
100            reset: UNAVAILABLE.to_string(),
101            used: UNAVAILABLE.to_string(),
102            cap: UNAVAILABLE.to_string(),
103        };
104    };
105    WindowValues {
106        percent: window.pct().to_string(),
107        reset: countdown::format(window.resets_at, now),
108        used: usd(window.used),
109        cap: usd(window.cap),
110    }
111}
112
113fn sanitize(value: &str) -> String {
114    crate::display::sanitize_untrusted_field(value)
115}
116
117pub fn severity(snap: &Snapshot) -> PaceSeverity {
118    severity_for(snap.worst_pct())
119}
120
121pub fn render(
122    outcome: &VendorOutcome,
123    snap: &Snapshot,
124    theme: &Theme,
125    opts: &RenderOpts,
126    now: DateTime<Utc>,
127) -> WaybarOutput {
128    render_with_meta(
129        snap,
130        outcome.stale,
131        outcome.last_error.as_ref(),
132        outcome.cache_age,
133        theme,
134        opts,
135        now,
136    )
137}
138
139fn render_with_meta(
140    snap: &Snapshot,
141    stale: bool,
142    last_error: Option<&(u16, String)>,
143    cache_age: Option<Duration>,
144    theme: &Theme,
145    opts: &RenderOpts,
146    now: DateTime<Utc>,
147) -> WaybarOutput {
148    let sev = severity(snap);
149    let format = opts.format.as_deref().unwrap_or(DEFAULT_FORMAT);
150    let values = escaped_placeholders(snap, now);
151    let mut text = substitute(format, &values);
152    if stale {
153        text.push_str(" ⏸");
154    }
155    let icon_prefix = match opts.icon.as_deref() {
156        Some(icon) if !icon.is_empty() => format!("{} ", escape(icon)),
157        _ => String::new(),
158    };
159    let bar_text = color_span(severity_color(sev, theme), &format!("{icon_prefix}{text}"));
160    let tooltip = opts
161        .tooltip_format
162        .as_deref()
163        .map(|format| substitute(format, &values))
164        .unwrap_or_else(|| render_tooltip(snap, stale, last_error, cache_age, theme, now));
165
166    WaybarOutput {
167        text: bar_text,
168        tooltip,
169        class: Class::from(sev),
170    }
171}
172
173fn escaped_placeholders(snap: &Snapshot, now: DateTime<Utc>) -> HashMap<&'static str, String> {
174    let mut values = build_placeholders(snap, now);
175    for key in ["plan", "cc_plan"] {
176        if let Some(value) = values.get_mut(key) {
177            *value = escape(value);
178        }
179    }
180    values
181}
182
183fn render_tooltip(
184    snap: &Snapshot,
185    stale: bool,
186    last_error: Option<&(u16, String)>,
187    cache_age: Option<Duration>,
188    theme: &Theme,
189    now: DateTime<Utc>,
190) -> String {
191    let plan = snap.plan.as_deref().unwrap_or(DEFAULT_PLAN);
192    let mut lines = vec![TooltipLine::Center(format!(
193        "<span font_weight='bold' foreground='{}'>{}</span>",
194        theme.blue,
195        escape(&format!("Command Code {}", sanitize(plan)))
196    ))];
197    lines.push(TooltipLine::Sep);
198    lines.push(TooltipLine::Body(String::new()));
199
200    let mut present = false;
201    for (label, window) in [
202        ("  󰔟  Session (5h)", snap.five_hour.as_ref()),
203        ("  󰃰  Weekly", snap.weekly.as_ref()),
204        ("  󰃰  Monthly", snap.monthly_window().as_ref()),
205    ] {
206        let Some(window) = window else {
207            continue;
208        };
209        if present {
210            lines.push(TooltipLine::Body(String::new()));
211        }
212        present = true;
213        let usage = UsageWindow {
214            utilization_pct: window.pct(),
215            resets_at: window.resets_at,
216            // The shared renderer does not pace this row, but a concrete
217            // value keeps the conversion honest for future callers.
218            window_duration: chrono::Duration::zero(),
219        };
220        let detail = format!("{} of {}", usd(window.used), usd(window.cap));
221        push_window_with_detail(
222            &mut lines,
223            label,
224            &usage,
225            theme,
226            now,
227            WindowRow::default(),
228            Some(&detail),
229        );
230    }
231    if !present {
232        lines.push(TooltipLine::Body(format!(
233            " <span foreground='{}'>no usage windows reported</span>",
234            theme.dim
235        )));
236    }
237
238    if let Some(credits) = snap.credits.as_ref() {
239        lines.push(TooltipLine::Body(String::new()));
240        lines.push(TooltipLine::Sep);
241        lines.push(TooltipLine::Body(format!(
242            " <span foreground='{}'>  󰄑  Credits</span>",
243            theme.fg
244        )));
245        lines.push(TooltipLine::Body(format!(
246            " <span foreground='{}'>     balance: {}</span>",
247            theme.dim,
248            escape(&usd(credits.remaining()))
249        )));
250        if credits.monthly != 0.0 {
251            lines.push(TooltipLine::Body(format!(
252                " <span foreground='{}'>     monthly: {}</span>",
253                theme.dim,
254                escape(&usd(credits.monthly))
255            )));
256        }
257        if credits.purchased != 0.0 {
258            lines.push(TooltipLine::Body(format!(
259                " <span foreground='{}'>     purchased: {}</span>",
260                theme.dim,
261                escape(&usd(credits.purchased))
262            )));
263        }
264        if credits.free != 0.0 {
265            lines.push(TooltipLine::Body(format!(
266                " <span foreground='{}'>     free: {}</span>",
267                theme.dim,
268                escape(&usd(credits.free))
269            )));
270        }
271        if let Some(at) = snap.period_end {
272            lines.push(TooltipLine::Body(format!(
273                " <span foreground='{}'>     resets in {}</span>",
274                theme.dim,
275                escape(&countdown::format(Some(at), now))
276            )));
277        }
278    }
279
280    if stale {
281        lines.push(TooltipLine::Body(String::new()));
282        lines.push(TooltipLine::Body(format!(
283            " <span foreground='{}'>  ⏸  Showing cached data</span>",
284            theme.orange
285        )));
286    }
287    if let Some((code, message)) = last_error
288        && *code != 0
289    {
290        lines.push(TooltipLine::Body(String::new()));
291        lines.push(TooltipLine::Sep);
292        lines.push(TooltipLine::Body(format!(
293            " <span foreground='{}'>  HTTP {code}: {}</span>",
294            theme.orange,
295            escape(message)
296        )));
297    }
298
299    lines.push(TooltipLine::Body(String::new()));
300    lines.push(TooltipLine::Sep);
301    lines.push(TooltipLine::Body(format!(
302        " <span foreground='{}'>  Updated {}</span>",
303        theme.dim,
304        updated_at_hm(now, cache_age)
305    )));
306    render_bordered(&lines, theme)
307}
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312    use crate::commandcode::types::{Credits, Snapshot, SpendWindow};
313
314    fn at(value: &str) -> DateTime<Utc> {
315        value.parse().expect("RFC3339 timestamp")
316    }
317
318    fn sample() -> Snapshot {
319        Snapshot {
320            plan: Some("GOAT".into()),
321            five_hour: Some(SpendWindow {
322                used: 1.23,
323                cap: 14.0,
324                resets_at: Some(at("2026-08-27T04:40:19Z")),
325            }),
326            weekly: Some(SpendWindow {
327                used: 5.24,
328                cap: 35.0,
329                resets_at: Some(at("2026-09-02T18:36:12Z")),
330            }),
331            credits: Some(Credits {
332                monthly: 49.28,
333                purchased: 0.0,
334                free: 0.0,
335            }),
336            credit_pool: Some(70.0),
337            period_end: Some(at("2026-09-17T14:28:52Z")),
338        }
339    }
340
341    #[test]
342    fn exposes_exact_and_generic_placeholders() {
343        let values = build_placeholders(&sample(), at("2026-08-27T02:30:00Z"));
344
345        assert_eq!(values["vendor_short"], "cmc");
346        assert_eq!(values["plan"], "GOAT");
347        assert_eq!(values["cc_session_pct"], "9");
348        assert_eq!(values["cc_weekly_pct"], "15");
349        // The monthly allowance is a derived window: $20.72 of the $70 pool.
350        assert_eq!(values["cc_monthly_pct"], "30");
351        assert_eq!(values["cc_monthly_used"], "$20.72");
352        assert_eq!(values["cc_monthly_cap"], "$70.00");
353        assert_eq!(values["cc_monthly_reset"], "21d 11h");
354        // Generic aliases keep a shared format string working across vendors.
355        assert_eq!(values["session_pct"], "9");
356        assert_eq!(values["weekly_pct"], "15");
357    }
358
359    #[test]
360    fn spend_figures_use_the_shared_money_formatter() {
361        let values = build_placeholders(&sample(), at("2026-08-27T02:30:00Z"));
362
363        assert_eq!(values["cc_session_used"], "$1.23");
364        assert_eq!(values["cc_session_cap"], "$14.00");
365        assert_eq!(values["cc_credits"], "$49.28");
366        assert_eq!(values["cc_credits_pool"], "$70.00");
367        assert_eq!(values["cc_credits_spent"], "$20.72");
368        // 2026-09-17 minus 2026-08-27 → 21 days and change.
369        assert_eq!(values["cc_credits_reset"], "21d 11h");
370    }
371
372    #[test]
373    fn monthly_window_needs_ledger_and_a_recognised_plan() {
374        // No ledger: nothing to derive the spend from.
375        let no_ledger = Snapshot {
376            credits: None,
377            credit_pool: Some(70.0),
378            period_end: Some(at("2026-09-17T14:28:52Z")),
379            ..sample()
380        };
381        assert!(no_ledger.monthly_window().is_none());
382        assert_eq!(no_ledger.worst_pct(), 15);
383
384        // No pool: no denominator.
385        let no_pool = Snapshot {
386            credit_pool: None,
387            ..sample()
388        };
389        assert!(no_pool.monthly_window().is_none());
390    }
391
392    #[test]
393    fn default_format_leads_with_the_session_window() {
394        assert_eq!(DEFAULT_FORMAT, "{cc_session_pct}% · {cc_session_reset}");
395    }
396
397    #[test]
398    fn absent_windows_and_ledger_are_unavailable_not_zero() {
399        let values = build_placeholders(&Snapshot::default(), at("2026-08-27T02:30:00Z"));
400
401        for key in [
402            "session_pct",
403            "weekly_pct",
404            "cc_session_pct",
405            "cc_session_reset",
406            "cc_weekly_pct",
407            "cc_session_used",
408            "cc_credits",
409            "cc_credits_pool",
410            "cc_credits_spent",
411        ] {
412            assert_eq!(values[key], UNAVAILABLE, "{key} should be unavailable");
413            assert_ne!(values[key], "0");
414        }
415        // With no plan, the vendor name stands in.
416        assert_eq!(values["plan"], DEFAULT_PLAN);
417    }
418
419    #[test]
420    fn severity_follows_the_window_closest_to_its_cap() {
421        let mut snapshot = sample();
422        assert_eq!(severity(&snapshot), severity_for(15));
423
424        snapshot.weekly = Some(SpendWindow {
425            used: 34.0,
426            cap: 35.0,
427            resets_at: None,
428        });
429        assert_eq!(severity(&snapshot), severity_for(97));
430    }
431
432    #[test]
433    fn plan_is_sanitized_before_it_reaches_the_bar() {
434        let snapshot = Snapshot {
435            plan: Some("GO\u{1b}[31mAT\u{7}".into()),
436            ..sample()
437        };
438
439        let values = build_placeholders(&snapshot, at("2026-08-27T02:30:00Z"));
440
441        assert!(!values["plan"].contains('\u{1b}'));
442        assert!(!values["plan"].contains('\u{7}'));
443        assert!(!values["cc_plan"].contains('\u{1b}'));
444    }
445
446    #[test]
447    fn tooltip_shows_both_windows_and_the_credit_ledger() {
448        let theme = Theme::default();
449        let tooltip = render_tooltip(
450            &sample(),
451            false,
452            None,
453            None,
454            &theme,
455            at("2026-08-27T02:30:00Z"),
456        );
457
458        assert!(tooltip.contains("Command Code GOAT"), "{tooltip}");
459        assert!(tooltip.contains("Session (5h)"), "{tooltip}");
460        assert!(tooltip.contains("$1.23 of $14.00"), "{tooltip}");
461        assert!(tooltip.contains("Resets in 2h 10m"), "{tooltip}");
462        assert!(tooltip.contains("Weekly"), "{tooltip}");
463        // The monthly allowance renders as a third window row.
464        assert!(tooltip.contains("Monthly"), "{tooltip}");
465        assert!(tooltip.contains("$20.72 of $70.00"), "{tooltip}");
466        assert!(tooltip.contains("$49.28"), "{tooltip}");
467        assert!(tooltip.contains("resets in 21d 11h"), "{tooltip}");
468    }
469
470    #[test]
471    fn tooltip_says_so_when_the_vendor_reports_no_windows() {
472        let theme = Theme::default();
473        let tooltip = render_tooltip(
474            &Snapshot::default(),
475            false,
476            None,
477            None,
478            &theme,
479            at("2026-08-27T02:30:00Z"),
480        );
481
482        assert!(tooltip.contains("no usage windows reported"), "{tooltip}");
483    }
484
485    #[test]
486    fn stale_and_http_errors_surface_in_the_tooltip() {
487        let theme = Theme::default();
488        let tooltip = render_tooltip(
489            &sample(),
490            true,
491            Some(&(503, "service unavailable".to_string())),
492            None,
493            &theme,
494            at("2026-08-27T02:30:00Z"),
495        );
496
497        assert!(tooltip.contains("Showing cached data"), "{tooltip}");
498        assert!(tooltip.contains("HTTP 503"), "{tooltip}");
499    }
500
501    #[test]
502    fn an_unknown_plan_still_renders_without_an_allowance_line() {
503        let snapshot = Snapshot {
504            plan: Some("individual-future".into()),
505            credit_pool: None,
506            ..sample()
507        };
508        let theme = Theme::default();
509
510        let tooltip = render_tooltip(
511            &snapshot,
512            false,
513            None,
514            None,
515            &theme,
516            at("2026-08-27T02:30:00Z"),
517        );
518
519        assert!(tooltip.contains("$49.28"), "{tooltip}");
520        // No recognised plan → no monthly window, no allowance line at all.
521        assert!(!tooltip.contains("Monthly"), "{tooltip}");
522        assert!(!tooltip.contains("$20.72 of $70.00"), "{tooltip}");
523    }
524}