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

ai_usagebar/moonshot/
vendor.rs
