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    crate::pango::balance_severity(snap.available, &snap.currency)
46}
47
48pub fn render(
49    outcome: &VendorOutcome,
50    snap: &MoonshotSnapshot,
51    theme: &Theme,
52    opts: &RenderOpts,
53    now: DateTime<Utc>,
54) -> WaybarOutput {
55    let class = Class::from(severity(snap));
56    let format = opts
57        .format
58        .clone()
59        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
60    let values = build_placeholders(snap);
61
62    let mut text = substitute(&format, &values);
63    if outcome.stale {
64        text.push_str(" ⏸");
65    }
66
67    let wrapper_color = severity_color(severity(snap), theme).to_string();
68    let icon_prefix = match opts.icon.as_deref() {
69        Some(ic) if !ic.is_empty() => format!("{ic} "),
70        _ => String::new(),
71    };
72    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
73
74    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
75        substitute(fmt, &values)
76    } else {
77        render_tooltip(outcome, snap, theme, now)
78    };
79
80    WaybarOutput {
81        text: bar_text,
82        tooltip,
83        class,
84    }
85}
86
87fn render_tooltip(
88    outcome: &VendorOutcome,
89    snap: &MoonshotSnapshot,
90    theme: &Theme,
91    now: DateTime<Utc>,
92) -> String {
93    let blue = &theme.blue;
94    let dim = &theme.dim;
95    let fg = &theme.fg;
96    let color = severity_color(severity(snap), theme);
97
98    let mut lines: Vec<TooltipLine> = Vec::new();
99    lines.push(TooltipLine::Center(format!(
100        "<span font_weight='bold' foreground='{blue}'>Kimi (Moonshot)</span>"
101    )));
102    lines.push(TooltipLine::Sep);
103    lines.push(TooltipLine::Body("".into()));
104
105    lines.push(TooltipLine::Body(format!(
106        " <span foreground='{fg}'>  󰢗  Balance</span>"
107    )));
108    lines.push(TooltipLine::Body(format!(
109        "   <span font_weight='bold' foreground='{color}'>{bal}</span>",
110        bal = escape(&money(snap.available, &snap.currency))
111    )));
112    lines.push(TooltipLine::Body(format!(
113        " <span foreground='{dim}'>     cash {cash} · voucher {voucher}</span>",
114        cash = escape(&money(snap.cash, &snap.currency)),
115        voucher = escape(&money(snap.voucher, &snap.currency))
116    )));
117
118    if snap.available <= 0.0 {
119        lines.push(TooltipLine::Body("".into()));
120        lines.push(TooltipLine::Body(format!(
121            " <span foreground='{}'>  󰀪  out of credit — inference blocked</span>",
122            theme.orange.as_str()
123        )));
124    }
125
126    if let Some((code, msg)) = outcome.last_error.as_ref()
127        && *code != 0
128    {
129        let (icon, ecolor) = if *code >= 500 {
130            ("󰅚", theme.red.as_str())
131        } else {
132            ("󰀪", theme.orange.as_str())
133        };
134        lines.push(TooltipLine::Body("".into()));
135        lines.push(TooltipLine::Sep);
136        lines.push(TooltipLine::Body(format!(
137            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
138        )));
139        lines.push(TooltipLine::Body(format!(
140            "     <span foreground='{dim}'>{}</span>",
141            escape(msg)
142        )));
143    }
144
145    let updated = updated_at_hm(now, outcome.cache_age);
146    lines.push(TooltipLine::Body("".into()));
147    lines.push(TooltipLine::Sep);
148    lines.push(TooltipLine::Body(format!(
149        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
150    )));
151
152    render_bordered(&lines, theme)
153}
154
155impl From<FetchOutcome> for VendorOutcome {
156    fn from(o: FetchOutcome) -> Self {
157        Self {
158            snapshot: crate::usage::VendorSnapshot::Moonshot(o.snapshot),
159            stale: o.stale,
160            last_error: o.last_error,
161            cache_age: o.cache_age,
162        }
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use crate::usage::MoonshotSnapshot;
170
171    fn sample_snap() -> MoonshotSnapshot {
172        MoonshotSnapshot {
173            available: 49.58,
174            voucher: 46.58,
175            cash: 3.0,
176            currency: "USD".into(),
177        }
178    }
179
180    fn sample_outcome(snap: MoonshotSnapshot) -> VendorOutcome {
181        VendorOutcome {
182            snapshot: crate::usage::VendorSnapshot::Moonshot(snap),
183            stale: false,
184            last_error: None,
185            cache_age: Some(std::time::Duration::from_secs(10)),
186        }
187    }
188
189    fn opts() -> RenderOpts {
190        RenderOpts {
191            format: None,
192            tooltip_format: None,
193            icon: None,
194            pace_tolerance: 5,
195            format_pace_color: false,
196            tooltip_pace_pts: false,
197        }
198    }
199
200    #[test]
201    fn default_render_shows_balance() {
202        let snap = sample_snap();
203        let out = render(
204            &sample_outcome(snap.clone()),
205            &snap,
206            &Theme::default(),
207            &opts(),
208            Utc::now(),
209        );
210        assert!(out.text.contains("$49.58"));
211    }
212
213    /// `cash_balance` is documented as negative (debt). It used to render as
214    /// "$-5.71" here while OpenRouter rendered the same debt as "-$5.71";
215    /// both now go through `format::money`.
216    #[test]
217    fn a_cash_debt_carries_its_sign_ahead_of_the_symbol() {
218        let mut snap = sample_snap();
219        snap.cash = -5.71;
220        let out = render(
221            &sample_outcome(snap.clone()),
222            &snap,
223            &Theme::default(),
224            &opts(),
225            Utc::now(),
226        );
227        assert!(out.tooltip.contains("cash -$5.71"), "{}", out.tooltip);
228        assert!(!out.tooltip.contains("$-5.71"), "{}", out.tooltip);
229
230        snap.currency = "CNY".into();
231        let out = render(
232            &sample_outcome(snap.clone()),
233            &snap,
234            &Theme::default(),
235            &opts(),
236            Utc::now(),
237        );
238        assert!(out.tooltip.contains("cash -¥5.71"), "{}", out.tooltip);
239    }
240
241    #[test]
242    fn tooltip_includes_cash_and_voucher() {
243        let snap = sample_snap();
244        let out = render(
245            &sample_outcome(snap.clone()),
246            &snap,
247            &Theme::default(),
248            &opts(),
249            Utc::now(),
250        );
251        assert!(out.tooltip.contains("cash $3.00"));
252        assert!(out.tooltip.contains("voucher $46.58"));
253    }
254
255    #[test]
256    fn cny_uses_yuan_symbol() {
257        let mut snap = sample_snap();
258        snap.currency = "CNY".into();
259        let out = render(
260            &sample_outcome(snap.clone()),
261            &snap,
262            &Theme::default(),
263            &opts(),
264            Utc::now(),
265        );
266        assert!(out.text.contains("¥49.58"));
267    }
268
269    #[test]
270    fn zero_balance_is_critical() {
271        let mut snap = sample_snap();
272        snap.available = 0.0;
273        assert_eq!(severity(&snap), PaceSeverity::Critical);
274    }
275}