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