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