Skip to main content

ai_usagebar/moonshot/
vendor.rs

1//! Moonshot / Kimi renderer — bar text + bordered Pango tooltip. Balance-only,
2//! currency-aware (USD on `.ai`, CNY on `.cn`).
3
4use std::collections::HashMap;
5
6use chrono::{DateTime, Utc};
7
8use crate::format::{money, 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::MoonshotSnapshot;
14use crate::vendor::{RenderOpts, VendorId, VendorOutcome};
15use crate::waybar::{Class, WaybarOutput};
16
17use super::fetch::FetchOutcome;
18
19pub const DEFAULT_FORMAT: &str = "{km_balance}";
20
21pub fn build_placeholders(snap: &MoonshotSnapshot) -> HashMap<&'static str, String> {
22    placeholders(vec![
23        ("icon", "󰚩".to_string()),
24        ("vendor_short", VendorId::Moonshot.short_name().to_string()),
25        // Cross-vendor aliases — Kimi has no rate-limit windows here.
26        ("session_pct", "0".to_string()),
27        ("session_reset", "—".to_string()),
28        ("weekly_pct", "0".to_string()),
29        ("weekly_reset", "—".to_string()),
30        ("plan", "Kimi".to_string()),
31        ("km_balance", money(snap.available, &snap.currency)),
32        ("km_voucher", money(snap.voucher, &snap.currency)),
33        ("km_cash", money(snap.cash, &snap.currency)),
34        ("currency", snap.currency.clone()),
35    ])
36}
37
38/// `available_balance <= 0` blocks the inference API, so that's critical.
39/// Otherwise scale the low/high/mid thresholds by currency (CNY ≈ 7× USD),
40/// mirroring DeepSeek.
41pub fn severity(snap: &MoonshotSnapshot) -> PaceSeverity {
42    if snap.available <= 0.0 {
43        return PaceSeverity::Critical;
44    }
45    let (t_critical, t_high, t_mid) = match snap.currency.as_str() {
46        "CNY" => (7.0_f64, 35.0, 140.0),
47        _ => (1.0_f64, 5.0, 20.0),
48    };
49    if snap.available < t_critical {
50        PaceSeverity::Critical
51    } else if snap.available < t_high {
52        PaceSeverity::High
53    } else if snap.available < t_mid {
54        PaceSeverity::Mid
55    } else {
56        PaceSeverity::Low
57    }
58}
59
60pub fn render(
61    outcome: &VendorOutcome,
62    snap: &MoonshotSnapshot,
63    theme: &Theme,
64    opts: &RenderOpts,
65    now: DateTime<Utc>,
66) -> WaybarOutput {
67    let class = Class::from(severity(snap));
68    let format = opts
69        .format
70        .clone()
71        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
72    let values = build_placeholders(snap);
73
74    let mut text = substitute(&format, &values);
75    if outcome.stale {
76        text.push_str(" ⏸");
77    }
78
79    let wrapper_color = severity_color(severity(snap), theme).to_string();
80    let icon_prefix = match opts.icon.as_deref() {
81        Some(ic) if !ic.is_empty() => format!("{ic} "),
82        _ => String::new(),
83    };
84    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
85
86    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
87        substitute(fmt, &values)
88    } else {
89        render_tooltip(outcome, snap, theme, now)
90    };
91
92    WaybarOutput {
93        text: bar_text,
94        tooltip,
95        class,
96    }
97}
98
99fn render_tooltip(
100    outcome: &VendorOutcome,
101    snap: &MoonshotSnapshot,
102    theme: &Theme,
103    now: DateTime<Utc>,
104) -> String {
105    let blue = &theme.blue;
106    let dim = &theme.dim;
107    let fg = &theme.fg;
108    let color = severity_color(severity(snap), theme);
109
110    let mut lines: Vec<TooltipLine> = Vec::new();
111    lines.push(TooltipLine::Center(format!(
112        "<span font_weight='bold' foreground='{blue}'>Kimi (Moonshot)</span>"
113    )));
114    lines.push(TooltipLine::Sep);
115    lines.push(TooltipLine::Body("".into()));
116
117    lines.push(TooltipLine::Body(format!(
118        " <span foreground='{fg}'>  󰢗  Balance</span>"
119    )));
120    lines.push(TooltipLine::Body(format!(
121        "   <span font_weight='bold' foreground='{color}'>{bal}</span>",
122        bal = escape(&money(snap.available, &snap.currency))
123    )));
124    lines.push(TooltipLine::Body(format!(
125        " <span foreground='{dim}'>     cash {cash} · voucher {voucher}</span>",
126        cash = escape(&money(snap.cash, &snap.currency)),
127        voucher = escape(&money(snap.voucher, &snap.currency))
128    )));
129
130    if snap.available <= 0.0 {
131        lines.push(TooltipLine::Body("".into()));
132        lines.push(TooltipLine::Body(format!(
133            " <span foreground='{}'>  󰀪  out of credit — inference blocked</span>",
134            theme.orange.as_str()
135        )));
136    }
137
138    if let Some((code, msg)) = outcome.last_error.as_ref()
139        && *code != 0
140    {
141        let (icon, ecolor) = if *code >= 500 {
142            ("󰅚", theme.red.as_str())
143        } else {
144            ("󰀪", theme.orange.as_str())
145        };
146        lines.push(TooltipLine::Body("".into()));
147        lines.push(TooltipLine::Sep);
148        lines.push(TooltipLine::Body(format!(
149            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
150        )));
151        lines.push(TooltipLine::Body(format!(
152            "     <span foreground='{dim}'>{}</span>",
153            escape(msg)
154        )));
155    }
156
157    let updated = updated_at_hm(now, outcome.cache_age);
158    lines.push(TooltipLine::Body("".into()));
159    lines.push(TooltipLine::Sep);
160    lines.push(TooltipLine::Body(format!(
161        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
162    )));
163
164    render_bordered(&lines, theme)
165}
166
167impl From<FetchOutcome> for VendorOutcome {
168    fn from(o: FetchOutcome) -> Self {
169        Self {
170            snapshot: crate::usage::VendorSnapshot::Moonshot(o.snapshot),
171            stale: o.stale,
172            last_error: o.last_error,
173            cache_age: o.cache_age,
174        }
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181    use crate::usage::MoonshotSnapshot;
182
183    fn sample_snap() -> MoonshotSnapshot {
184        MoonshotSnapshot {
185            available: 49.58,
186            voucher: 46.58,
187            cash: 3.0,
188            currency: "USD".into(),
189        }
190    }
191
192    fn sample_outcome(snap: MoonshotSnapshot) -> VendorOutcome {
193        VendorOutcome {
194            snapshot: crate::usage::VendorSnapshot::Moonshot(snap),
195            stale: false,
196            last_error: None,
197            cache_age: Some(std::time::Duration::from_secs(10)),
198        }
199    }
200
201    fn opts() -> RenderOpts {
202        RenderOpts {
203            format: None,
204            tooltip_format: None,
205            icon: None,
206            pace_tolerance: 5,
207            format_pace_color: false,
208            tooltip_pace_pts: false,
209        }
210    }
211
212    #[test]
213    fn default_render_shows_balance() {
214        let snap = sample_snap();
215        let out = render(
216            &sample_outcome(snap.clone()),
217            &snap,
218            &Theme::default(),
219            &opts(),
220            Utc::now(),
221        );
222        assert!(out.text.contains("$49.58"));
223    }
224
225    /// `cash_balance` is documented as negative (debt). It used to render as
226    /// "$-5.71" here while OpenRouter rendered the same debt as "-$5.71";
227    /// both now go through `format::money`.
228    #[test]
229    fn a_cash_debt_carries_its_sign_ahead_of_the_symbol() {
230        let mut snap = sample_snap();
231        snap.cash = -5.71;
232        let out = render(
233            &sample_outcome(snap.clone()),
234            &snap,
235            &Theme::default(),
236            &opts(),
237            Utc::now(),
238        );
239        assert!(out.tooltip.contains("cash -$5.71"), "{}", out.tooltip);
240        assert!(!out.tooltip.contains("$-5.71"), "{}", out.tooltip);
241
242        snap.currency = "CNY".into();
243        let out = render(
244            &sample_outcome(snap.clone()),
245            &snap,
246            &Theme::default(),
247            &opts(),
248            Utc::now(),
249        );
250        assert!(out.tooltip.contains("cash -¥5.71"), "{}", out.tooltip);
251    }
252
253    #[test]
254    fn tooltip_includes_cash_and_voucher() {
255        let snap = sample_snap();
256        let out = render(
257            &sample_outcome(snap.clone()),
258            &snap,
259            &Theme::default(),
260            &opts(),
261            Utc::now(),
262        );
263        assert!(out.tooltip.contains("cash $3.00"));
264        assert!(out.tooltip.contains("voucher $46.58"));
265    }
266
267    #[test]
268    fn cny_uses_yuan_symbol() {
269        let mut snap = sample_snap();
270        snap.currency = "CNY".into();
271        let out = render(
272            &sample_outcome(snap.clone()),
273            &snap,
274            &Theme::default(),
275            &opts(),
276            Utc::now(),
277        );
278        assert!(out.text.contains("¥49.58"));
279    }
280
281    #[test]
282    fn zero_balance_is_critical() {
283        let mut snap = sample_snap();
284        snap.available = 0.0;
285        assert_eq!(severity(&snap), PaceSeverity::Critical);
286    }
287}