Skip to main content

ai_usagebar/kilo/
vendor.rs

1//! Kilo Code renderer — bar text + bordered Pango tooltip. Balance-only, like
2//! DeepSeek (a remaining-credit number with no rate-limit windows).
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::KiloSnapshot;
14use crate::vendor::{RenderOpts, VendorOutcome};
15use crate::waybar::{Class, WaybarOutput};
16
17use super::fetch::FetchOutcome;
18
19pub const DEFAULT_FORMAT: &str = "{kilo_balance}";
20
21pub fn build_placeholders(snap: &KiloSnapshot) -> HashMap<&'static str, String> {
22    placeholders(vec![
23        ("icon", "󰭟".to_string()),
24        ("vendor_short", "klo".to_string()),
25        // Cross-vendor aliases — Kilo has no rate-limit windows.
26        ("session_pct", "0".to_string()),
27        ("session_reset", "—".to_string()),
28        ("weekly_pct", "0".to_string()),
29        ("weekly_reset", "—".to_string()),
30        ("plan", snap.label.clone()),
31        ("kilo_balance", format_money(snap.balance)),
32    ])
33}
34
35fn format_money(v: f64) -> String {
36    if v < 0.0 {
37        format!("-${:.2}", -v)
38    } else {
39        format!("${v:.2}")
40    }
41}
42
43/// Kilo has no purchased-total on this endpoint, so severity keys on the
44/// absolute remaining USD balance (mirrors DeepSeek's USD thresholds): running
45/// low = warmer color, empty = critical (the `402` boundary).
46pub fn severity(snap: &KiloSnapshot) -> PaceSeverity {
47    if snap.balance < 1.0 {
48        PaceSeverity::Critical
49    } else if snap.balance < 5.0 {
50        PaceSeverity::High
51    } else if snap.balance < 20.0 {
52        PaceSeverity::Mid
53    } else {
54        PaceSeverity::Low
55    }
56}
57
58pub fn render(
59    outcome: &VendorOutcome,
60    snap: &KiloSnapshot,
61    theme: &Theme,
62    opts: &RenderOpts,
63    now: DateTime<Utc>,
64) -> WaybarOutput {
65    let class = Class::from(severity(snap));
66    let format = opts
67        .format
68        .clone()
69        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
70    let values = build_placeholders(snap);
71
72    let mut text = substitute(&format, &values);
73    if outcome.stale {
74        text.push_str(" ⏸");
75    }
76
77    let wrapper_color = severity_color(severity(snap), theme).to_string();
78    let icon_prefix = match opts.icon.as_deref() {
79        Some(ic) if !ic.is_empty() => format!("{ic} "),
80        _ => String::new(),
81    };
82    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
83
84    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
85        substitute(fmt, &values)
86    } else {
87        render_tooltip(outcome, snap, theme, now)
88    };
89
90    WaybarOutput {
91        text: bar_text,
92        tooltip,
93        class,
94    }
95}
96
97fn render_tooltip(
98    outcome: &VendorOutcome,
99    snap: &KiloSnapshot,
100    theme: &Theme,
101    now: DateTime<Utc>,
102) -> String {
103    let blue = &theme.blue;
104    let dim = &theme.dim;
105    let fg = &theme.fg;
106    let color = severity_color(severity(snap), theme);
107
108    let mut lines: Vec<TooltipLine> = Vec::new();
109    lines.push(TooltipLine::Center(format!(
110        "<span font_weight='bold' foreground='{blue}'>{label}</span>",
111        label = escape(&snap.label)
112    )));
113    lines.push(TooltipLine::Sep);
114    lines.push(TooltipLine::Body("".into()));
115
116    lines.push(TooltipLine::Body(format!(
117        " <span foreground='{fg}'>  󰢗  Balance</span>"
118    )));
119    lines.push(TooltipLine::Body(format!(
120        "   <span font_weight='bold' foreground='{color}'>{bal}</span>",
121        bal = escape(&format_money(snap.balance))
122    )));
123
124    if let Some((code, msg)) = outcome.last_error.as_ref()
125        && *code != 0
126    {
127        let (icon, ecolor) = if *code >= 500 {
128            ("󰅚", theme.red.as_str())
129        } else {
130            ("󰀪", theme.orange.as_str())
131        };
132        lines.push(TooltipLine::Body("".into()));
133        lines.push(TooltipLine::Sep);
134        lines.push(TooltipLine::Body(format!(
135            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
136        )));
137        lines.push(TooltipLine::Body(format!(
138            "     <span foreground='{dim}'>{}</span>",
139            escape(msg)
140        )));
141    }
142
143    let updated = updated_at_hm(now, outcome.cache_age);
144    lines.push(TooltipLine::Body("".into()));
145    lines.push(TooltipLine::Sep);
146    lines.push(TooltipLine::Body(format!(
147        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
148    )));
149
150    render_bordered(&lines, theme)
151}
152
153impl From<FetchOutcome> for VendorOutcome {
154    fn from(o: FetchOutcome) -> Self {
155        Self {
156            snapshot: crate::usage::VendorSnapshot::Kilo(o.snapshot),
157            stale: o.stale,
158            last_error: o.last_error,
159            cache_age: o.cache_age,
160        }
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use crate::usage::KiloSnapshot;
168
169    fn sample_snap() -> KiloSnapshot {
170        KiloSnapshot {
171            label: "Kilo".into(),
172            balance: 8.42,
173        }
174    }
175
176    fn sample_outcome(snap: KiloSnapshot) -> VendorOutcome {
177        VendorOutcome {
178            snapshot: crate::usage::VendorSnapshot::Kilo(snap),
179            stale: false,
180            last_error: None,
181            cache_age: Some(std::time::Duration::from_secs(10)),
182        }
183    }
184
185    fn opts() -> RenderOpts {
186        RenderOpts {
187            format: None,
188            tooltip_format: None,
189            icon: None,
190            pace_tolerance: 5,
191            format_pace_color: false,
192            tooltip_pace_pts: false,
193        }
194    }
195
196    #[test]
197    fn default_render_shows_balance() {
198        let snap = sample_snap();
199        let outcome = sample_outcome(snap.clone());
200        let out = render(&outcome, &snap, &Theme::default(), &opts(), Utc::now());
201        assert!(out.text.contains("$8.42"));
202    }
203
204    #[test]
205    fn tooltip_includes_balance() {
206        let snap = sample_snap();
207        let outcome = sample_outcome(snap.clone());
208        let out = render(&outcome, &snap, &Theme::default(), &opts(), Utc::now());
209        assert!(out.tooltip.contains("Balance"));
210        assert!(out.tooltip.contains("$8.42"));
211    }
212
213    #[test]
214    fn severity_scales_with_balance() {
215        assert_eq!(
216            severity(&KiloSnapshot {
217                label: "".into(),
218                balance: 0.5
219            }),
220            PaceSeverity::Critical
221        );
222        assert_eq!(
223            severity(&KiloSnapshot {
224                label: "".into(),
225                balance: 3.0
226            }),
227            PaceSeverity::High
228        );
229        assert_eq!(
230            severity(&KiloSnapshot {
231                label: "".into(),
232                balance: 12.0
233            }),
234            PaceSeverity::Mid
235        );
236        assert_eq!(
237            severity(&KiloSnapshot {
238                label: "".into(),
239                balance: 50.0
240            }),
241            PaceSeverity::Low
242        );
243    }
244
245    #[test]
246    fn stale_appends_pause() {
247        let snap = sample_snap();
248        let mut outcome = sample_outcome(snap.clone());
249        outcome.stale = true;
250        let out = render(&outcome, &snap, &Theme::default(), &opts(), Utc::now());
251        assert!(out.text.contains("⏸"));
252    }
253
254    #[test]
255    fn custom_tooltip_uses_placeholder() {
256        let snap = sample_snap();
257        let outcome = sample_outcome(snap.clone());
258        let mut o = opts();
259        o.tooltip_format = Some("bal: {kilo_balance}".into());
260        let out = render(&outcome, &snap, &Theme::default(), &o, Utc::now());
261        assert_eq!(out.tooltip, "bal: $8.42");
262    }
263}