usage_monitor_cli/provider/
kimi.rs

1use async_trait::async_trait;
2use chrono::{DateTime, Utc};
3
4use crate::error::SpendPanelError;
5use crate::model::{NamedRateWindow, RateWindow, UsageSnapshot};
6use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
7
8#[derive(Debug, serde::Deserialize)]
9struct KimiUsageResponse {
10    #[serde(default)]
11    usages: Vec<KimiUsage>,
12    /// Overall quota for the billing cycle ("uso total" in the Kimi UI).
13    #[serde(default, rename = "totalQuota")]
14    total_quota: Option<KimiUsageDetail>,
15}
16
17#[derive(Debug, serde::Deserialize)]
18struct KimiUsage {
19    #[serde(default)]
20    scope: String,
21    detail: KimiUsageDetail,
22    #[serde(default)]
23    limits: Option<Vec<KimiRateLimit>>,
24}
25
26#[derive(Debug, serde::Deserialize)]
27struct KimiRateLimit {
28    /// Rolling window descriptor (e.g. 300 TIME_UNIT_MINUTE = 5h session).
29    #[serde(default)]
30    window: Option<KimiWindow>,
31    detail: KimiUsageDetail,
32}
33
34#[derive(Debug, serde::Deserialize)]
35struct KimiWindow {
36    #[serde(default)]
37    duration: Option<u64>,
38    #[serde(default, rename = "timeUnit")]
39    time_unit: Option<String>,
40}
41
42impl KimiWindow {
43    /// Window length in minutes, converted from the declared time unit.
44    fn minutes(&self) -> Option<u32> {
45        let d = self.duration?;
46        let factor: u64 = match self.time_unit.as_deref() {
47            Some("TIME_UNIT_SECOND") => return u32::try_from(d.div_ceil(60)).ok(),
48            Some("TIME_UNIT_HOUR") => 60,
49            Some("TIME_UNIT_DAY") => 1440,
50            Some("TIME_UNIT_WEEK") => 10080,
51            _ => 1, // TIME_UNIT_MINUTE or unknown: assume minutes
52        };
53        u32::try_from(d.saturating_mul(factor)).ok()
54    }
55}
56
57#[derive(Debug, serde::Deserialize)]
58struct KimiUsageDetail {
59    #[serde(default)]
60    limit: String,
61    #[serde(default)]
62    used: Option<String>,
63    #[serde(default)]
64    remaining: Option<String>,
65    #[serde(default, rename = "resetTime")]
66    reset_time: Option<String>,
67}
68
69impl KimiUsageDetail {
70    /// (used, limit) request counts.
71    fn counts(&self) -> (u64, u64) {
72        let limit = self.limit.parse::<i64>().unwrap_or(0).max(0) as u64;
73        let used = match self.used.as_deref().and_then(|s| s.parse::<i64>().ok()) {
74            Some(u) => u.max(0) as u64,
75            None => {
76                let remaining = self
77                    .remaining
78                    .as_deref()
79                    .and_then(|s| s.parse::<i64>().ok())
80                    .unwrap_or(0);
81                limit.saturating_sub(remaining.max(0) as u64)
82            }
83        };
84        (used, limit)
85    }
86
87    fn resets_at(&self) -> Option<DateTime<Utc>> {
88        let raw = self.reset_time.as_deref()?;
89        if let Ok(secs) = raw.parse::<i64>() {
90            let secs = if secs > 1_000_000_000_000 {
91                secs / 1000
92            } else {
93                secs
94            };
95            return chrono::TimeZone::timestamp_opt(&Utc, secs, 0).single();
96        }
97        DateTime::parse_from_rfc3339(raw)
98            .ok()
99            .map(|d| d.with_timezone(&Utc))
100    }
101}
102
103/// Short human label for a window length in minutes (e.g. "5h", "7d").
104fn window_label(minutes: u32) -> String {
105    if minutes >= 1440 && minutes % 1440 == 0 {
106        format!("{}d", minutes / 1440)
107    } else if minutes >= 60 && minutes % 60 == 0 {
108        format!("{}h", minutes / 60)
109    } else {
110        format!("{}min", minutes)
111    }
112}
113
114/// Kimi coding usage provider (kimi.com, JWT auth token).
115pub struct KimiProvider {
116    metadata: ProviderMetadata,
117    base_url: Option<String>,
118}
119
120impl KimiProvider {
121    pub fn new() -> Self {
122        Self {
123            metadata: ProviderMetadata {
124                id: "kimi",
125                name: "Kimi",
126                description: "Kimi coding weekly/rate-limit usage monitor",
127                auth_methods: &["token", "api_key", "env"],
128                website: Some("https://www.kimi.com"),
129            },
130            base_url: None,
131        }
132    }
133
134    pub fn with_base_url(url: &str) -> Self {
135        let mut p = Self::new();
136        p.base_url = Some(url.to_string());
137        p
138    }
139
140    fn api_base(&self) -> &str {
141        self.base_url.as_deref().unwrap_or("https://www.kimi.com")
142    }
143
144    fn clean(raw: &str) -> String {
145        let mut v = raw.trim();
146        if v.len() >= 2
147            && ((v.starts_with('"') && v.ends_with('"'))
148                || (v.starts_with('\'') && v.ends_with('\'')))
149        {
150            v = &v[1..v.len() - 1];
151        }
152        v.trim().to_string()
153    }
154
155    fn resolve_token(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
156        for key in ["token", "api_key", "cookie"] {
157            if let Some(v) = ctx.config.get(key) {
158                let c = Self::clean(v);
159                if !c.is_empty() {
160                    return Ok(c);
161                }
162            }
163        }
164        for env in ["KIMI_AUTH_TOKEN", "KIMI_API_KEY"] {
165            if let Ok(v) = std::env::var(env) {
166                let c = Self::clean(&v);
167                if !c.is_empty() {
168                    return Ok(c);
169                }
170            }
171        }
172        Err(SpendPanelError::AuthFailed(
173            "kimi".into(),
174            "no auth token in token/api_key config, KIMI_AUTH_TOKEN, or KIMI_API_KEY".into(),
175        ))
176    }
177
178    fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
179        reqwest::Client::builder()
180            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
181            .build()
182            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
183    }
184
185    fn parse(body: &str) -> Result<UsageSnapshot, SpendPanelError> {
186        let resp: KimiUsageResponse = serde_json::from_str(body)
187            .map_err(|e| SpendPanelError::ParseError("kimi".into(), e.to_string()))?;
188        let coding = resp
189            .usages
190            .iter()
191            .find(|u| u.scope == "FEATURE_CODING")
192            .or_else(|| resp.usages.first())
193            .ok_or_else(|| {
194                SpendPanelError::ParseError("kimi".into(), "no usage scope in response".into())
195            })?;
196
197        let mut snapshot = UsageSnapshot::new("kimi");
198        let mut ordered: Vec<RateWindow> = Vec::new();
199
200        // Rolling rate limits (e.g. the 5h session) come first: they are the
201        // most urgent window. `limits[0]` is the session in practice.
202        if let Some(limits) = &coding.limits {
203            for (i, rate) in limits.iter().enumerate() {
204                let (u, l) = rate.detail.counts();
205                let minutes = rate
206                    .window
207                    .as_ref()
208                    .and_then(|w| w.minutes())
209                    .unwrap_or(300);
210                let label = if i == 0 {
211                    "Sessão".to_string()
212                } else {
213                    format!("Limite {}", window_label(minutes))
214                };
215                let mut window = RateWindow::new(u, l, label, minutes);
216                window.resets_at = rate.detail.resets_at();
217                ordered.push(window);
218            }
219        }
220
221        // The scope-level detail is the 7-day (weekly) quota.
222        let (wused, wlimit) = coding.detail.counts();
223        let mut weekly = RateWindow::new(wused, wlimit, "Semanal", 10080);
224        weekly.resets_at = coding.detail.resets_at();
225        ordered.push(weekly);
226
227        // Overall cycle quota ("uso total" in the Kimi UI).
228        if let Some(total) = &resp.total_quota {
229            let (tused, tlimit) = total.counts();
230            ordered.push(RateWindow::new(tused, tlimit, "Total", 0));
231        }
232
233        // Most urgent (shortest) window first; the cycle total (no window)
234        // always sorts last.
235        ordered.sort_by_key(|w| {
236            if w.window_minutes == 0 {
237                u32::MAX
238            } else {
239                w.window_minutes
240            }
241        });
242
243        for (idx, window) in ordered.into_iter().enumerate() {
244            match idx {
245                0 => snapshot.primary_rate_window = Some(window),
246                1 => snapshot.secondary_rate_window = Some(window),
247                2 => snapshot.tertiary_rate_window = Some(window),
248                n => snapshot.extra_rate_windows.push(NamedRateWindow {
249                    id: format!("window_{}", n),
250                    label: window.label.clone(),
251                    window,
252                }),
253            }
254        }
255        Ok(snapshot)
256    }
257}
258
259impl Default for KimiProvider {
260    fn default() -> Self {
261        Self::new()
262    }
263}
264
265#[async_trait]
266impl UsageProvider for KimiProvider {
267    fn metadata(&self) -> &ProviderMetadata {
268        &self.metadata
269    }
270
271    fn detect_credentials(&self) -> bool {
272        ["KIMI_AUTH_TOKEN", "KIMI_API_KEY"].iter().any(|e| {
273            std::env::var(e)
274                .map(|v| !v.trim().is_empty())
275                .unwrap_or(false)
276        })
277    }
278
279    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
280        let token = Self::resolve_token(ctx)?;
281        let client = Self::build_client(ctx)?;
282        let url = format!(
283            "{}/apiv2/kimi.gateway.billing.v1.BillingService/GetUsages",
284            self.api_base().trim_end_matches('/')
285        );
286        let resp = client
287            .post(url)
288            .header("Authorization", format!("Bearer {}", token))
289            .header("Cookie", format!("kimi-auth={}", token))
290            .header("Content-Type", "application/json")
291            .header("Accept", "*/*")
292            .header("connect-protocol-version", "1")
293            .body(r#"{"scope":["FEATURE_CODING"]}"#)
294            .send()
295            .await
296            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
297        let status = resp.status();
298        let body = resp
299            .text()
300            .await
301            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
302        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
303            return Err(SpendPanelError::AuthFailed(
304                "kimi".into(),
305                format!("invalid auth token (HTTP {})", status.as_u16()),
306            ));
307        }
308        if !status.is_success() {
309            return Err(SpendPanelError::ProviderError(
310                "kimi".into(),
311                format!("HTTP {}: {}", status, body),
312            ));
313        }
314        Self::parse(&body)
315    }
316}
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321    use pretty_assertions::assert_eq;
322    use wiremock::matchers::{method, path};
323    use wiremock::{Mock, MockServer, ResponseTemplate};
324
325    /// Mirrors the real GetUsages response shape (RFC3339 resetTime, window
326    /// descriptor on the session limit, top-level totalQuota).
327    const SAMPLE: &str = r#"{
328      "usages": [
329        {"scope": "FEATURE_CODING",
330         "detail": {"limit": "100", "used": "43", "remaining": "57",
331                    "resetTime": "2026-07-31T15:29:07.734013Z"},
332         "limits": [
333           {"window": {"duration": 300, "timeUnit": "TIME_UNIT_MINUTE"},
334            "detail": {"limit": "100", "used": "15", "remaining": "85",
335                       "resetTime": "2026-07-25T23:29:07.734013Z"}}
336         ]}
337      ],
338      "totalQuota": {"limit": "100", "used": "30", "remaining": "70"}
339    }"#;
340
341    #[test]
342    fn test_metadata() {
343        assert_eq!(KimiProvider::new().metadata().id, "kimi");
344    }
345
346    #[test]
347    fn test_parse_session_weekly_total() {
348        let snap = KimiProvider::parse(SAMPLE).unwrap();
349        let session = snap.primary_rate_window.unwrap();
350        assert_eq!(session.label, "Sessão");
351        assert_eq!(session.used, Some(15));
352        assert_eq!(session.limit, Some(100));
353        assert_eq!(session.window_minutes, 300);
354        assert!(session.resets_at.is_some());
355
356        let weekly = snap.secondary_rate_window.unwrap();
357        assert_eq!(weekly.label, "Semanal");
358        assert_eq!(weekly.used, Some(43));
359        assert_eq!(weekly.window_minutes, 10080);
360        assert!(weekly.resets_at.is_some());
361
362        let total = snap.tertiary_rate_window.unwrap();
363        assert_eq!(total.label, "Total");
364        assert_eq!(total.used, Some(30));
365        assert_eq!(total.limit, Some(100));
366        assert!(total.resets_at.is_none());
367    }
368
369    #[test]
370    fn test_reset_time_rfc3339_parses() {
371        let snap = KimiProvider::parse(SAMPLE).unwrap();
372        let weekly = snap.secondary_rate_window.unwrap();
373        let reset = weekly.resets_at.unwrap();
374        assert_eq!(reset.to_rfc3339(), "2026-07-31T15:29:07.734013+00:00");
375    }
376
377    #[test]
378    fn test_reset_time_epoch_still_parses() {
379        let body = r#"{"usages":[{"scope":"FEATURE_CODING",
380          "detail":{"limit":"10","used":"3","resetTime":"1788000000"}}]}"#;
381        let snap = KimiProvider::parse(body).unwrap();
382        assert!(snap.primary_rate_window.unwrap().resets_at.is_some());
383    }
384
385    #[test]
386    fn test_no_limits_weekly_becomes_primary() {
387        let body = r#"{"usages":[{"scope":"FEATURE_CODING",
388          "detail":{"limit":"1000","remaining":"600"}}]}"#;
389        let snap = KimiProvider::parse(body).unwrap();
390        let weekly = snap.primary_rate_window.unwrap();
391        // used = limit - remaining = 400
392        assert_eq!(weekly.used, Some(400));
393        assert_eq!(weekly.label, "Semanal");
394        assert!(snap.secondary_rate_window.is_none());
395    }
396
397    #[test]
398    fn test_no_total_quota_drops_tertiary() {
399        let body = r#"{"usages":[{"scope":"FEATURE_CODING",
400          "detail":{"limit":"100","used":"43"},
401          "limits":[{"window":{"duration":300,"timeUnit":"TIME_UNIT_MINUTE"},
402                     "detail":{"limit":"100","used":"15"}}]}]}"#;
403        let snap = KimiProvider::parse(body).unwrap();
404        assert_eq!(snap.primary_rate_window.unwrap().label, "Sessão");
405        assert_eq!(snap.secondary_rate_window.unwrap().label, "Semanal");
406        assert!(snap.tertiary_rate_window.is_none());
407    }
408
409    #[test]
410    fn test_window_unit_conversion() {
411        let w = |d: u64, u: &str| KimiWindow {
412            duration: Some(d),
413            time_unit: Some(u.into()),
414        };
415        assert_eq!(w(300, "TIME_UNIT_MINUTE").minutes(), Some(300));
416        assert_eq!(w(5, "TIME_UNIT_HOUR").minutes(), Some(300));
417        assert_eq!(w(1, "TIME_UNIT_DAY").minutes(), Some(1440));
418        assert_eq!(w(1, "TIME_UNIT_WEEK").minutes(), Some(10080));
419        assert_eq!(w(90, "TIME_UNIT_SECOND").minutes(), Some(2));
420    }
421
422    #[test]
423    fn test_falls_back_to_first_scope() {
424        let body = r#"{"usages":[{"scope":"OTHER",
425          "detail":{"limit":"10","used":"3"}}]}"#;
426        let snap = KimiProvider::parse(body).unwrap();
427        assert_eq!(snap.primary_rate_window.unwrap().used, Some(3));
428    }
429
430    #[test]
431    fn test_extra_limits_beyond_session() {
432        let body = r#"{"usages":[{"scope":"FEATURE_CODING",
433          "detail":{"limit":"100","used":"10"},
434          "limits":[
435            {"window":{"duration":300,"timeUnit":"TIME_UNIT_MINUTE"},
436             "detail":{"limit":"100","used":"5"}},
437            {"window":{"duration":1,"timeUnit":"TIME_UNIT_DAY"},
438             "detail":{"limit":"50","used":"20"}}
439          ]}]}"#;
440        let snap = KimiProvider::parse(body).unwrap();
441        // Sorted by urgency: 5h session, 1d limit, then weekly.
442        assert_eq!(snap.primary_rate_window.unwrap().label, "Sessão");
443        let daily = snap.secondary_rate_window.unwrap();
444        assert_eq!(daily.label, "Limite 1d");
445        assert_eq!(daily.used, Some(20));
446        assert_eq!(snap.tertiary_rate_window.unwrap().label, "Semanal");
447        assert!(snap.extra_rate_windows.is_empty());
448    }
449
450    #[tokio::test]
451    async fn test_fetch_success() {
452        let server = MockServer::start().await;
453        Mock::given(method("POST"))
454            .and(path(
455                "/apiv2/kimi.gateway.billing.v1.BillingService/GetUsages",
456            ))
457            .respond_with(ResponseTemplate::new(200).set_body_raw(SAMPLE, "application/json"))
458            .mount(&server)
459            .await;
460        let provider = KimiProvider::with_base_url(&server.uri());
461        let mut ctx = ProviderContext::new();
462        ctx.config.insert("token".into(), "jwt".into());
463        let snap = provider.fetch_usage(&ctx).await.unwrap();
464        assert_eq!(snap.primary_rate_window.unwrap().used, Some(15));
465        assert_eq!(snap.secondary_rate_window.unwrap().used, Some(43));
466        assert_eq!(snap.tertiary_rate_window.unwrap().used, Some(30));
467    }
468
469    #[tokio::test]
470    async fn test_fetch_401() {
471        let server = MockServer::start().await;
472        Mock::given(method("POST"))
473            .and(path(
474                "/apiv2/kimi.gateway.billing.v1.BillingService/GetUsages",
475            ))
476            .respond_with(ResponseTemplate::new(401))
477            .mount(&server)
478            .await;
479        let provider = KimiProvider::with_base_url(&server.uri());
480        let mut ctx = ProviderContext::new();
481        ctx.config.insert("token".into(), "bad".into());
482        assert!(matches!(
483            provider.fetch_usage(&ctx).await.unwrap_err(),
484            SpendPanelError::AuthFailed(_, _)
485        ));
486    }
487}
usage_monitor_cli/provider/kimi.rs

usage_monitor_cli/provider/
kimi.rs
