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