ai_usagebar/kimi/
fetch.rs

1//! Fetch Kimi usage from `/coding/v1/usages`.
2
3use std::time::Duration;
4
5use chrono::{DateTime, Utc};
6
7use crate::cache::{Cache, acquire_lock};
8use crate::error::{AppError, Result};
9use crate::usage::KimiSnapshot;
10
11use super::types::UsagesResponse;
12
13pub const BASE_URL: &str = "https://api.kimi.com";
14const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
15const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
16/// Stable marker stored alongside code 0 for a successful HTTP response whose
17/// payload no longer matches Kimi's undocumented usage schema.
18pub const SCHEMA_DRIFT_MESSAGE: &str = "Kimi API schema drift";
19
20#[derive(Debug, Clone)]
21pub struct Endpoints {
22    pub usages: String,
23}
24
25impl Default for Endpoints {
26    fn default() -> Self {
27        Self {
28            usages: format!("{BASE_URL}/coding/v1/usages"),
29        }
30    }
31}
32
33#[derive(Debug, Clone)]
34pub struct FetchOutcome {
35    pub snapshot: KimiSnapshot,
36    pub stale: bool,
37    pub last_error: Option<(u16, String)>,
38    pub cache_age: Option<Duration>,
39}
40
41pub async fn fetch_snapshot(
42    client: &reqwest::Client,
43    api_key: &str,
44    cache: &Cache,
45    endpoints: &Endpoints,
46    cache_ttl: Duration,
47) -> Result<FetchOutcome> {
48    cache.ensure_dir()?;
49    let _lock = acquire_lock(&cache.lock_path(), LOCK_TIMEOUT)?;
50
51    if let Some(bytes) = cache.fresh_payload(cache_ttl)?
52        && let Ok(outcome) = reuse_cache(bytes, cache, false)
53    {
54        return Ok(outcome);
55    }
56    // Corrupt fresh cache: fall through to live fetch rather than return a
57    // fabricated zero snapshot.
58
59    match fetch_live(client, &endpoints.usages, api_key).await {
60        Ok(snap) => {
61            let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap_or_default();
62            cache.write_payload(&bytes)?;
63            Ok(FetchOutcome {
64                snapshot: snap,
65                stale: false,
66                last_error: None,
67                cache_age: Some(Duration::ZERO),
68            })
69        }
70        Err(e) if e.is_transient() => fallback_silent(cache, e),
71        Err(e) => {
72            cache.mark_stale();
73            if let Some((code, msg)) = error_to_pair(&e) {
74                cache.write_last_error(code, &msg);
75            }
76            fallback_with_error(cache, e)
77        }
78    }
79}
80
81fn fallback_silent(cache: &Cache, original: AppError) -> Result<FetchOutcome> {
82    let Some(bytes) = cache.maybe_payload()? else {
83        return Err(original);
84    };
85    match reuse_cache(bytes, cache, true) {
86        Ok(outcome) => Ok(outcome),
87        Err(_) => Err(original),
88    }
89}
90
91fn fallback_with_error(cache: &Cache, original: AppError) -> Result<FetchOutcome> {
92    let Some(bytes) = cache.maybe_payload()? else {
93        return Err(original);
94    };
95    match reuse_cache(bytes, cache, true) {
96        Ok(mut outcome) => {
97            outcome.last_error = error_to_pair(&original);
98            Ok(outcome)
99        }
100        Err(_) => Err(original),
101    }
102}
103
104fn error_to_pair(e: &AppError) -> Option<(u16, String)> {
105    match e {
106        AppError::Http { status, body } => Some((*status, body.clone())),
107        // A 2xx response with an unknown shape is not an HTTP 422 response.
108        AppError::Schema(_) => Some((0, SCHEMA_DRIFT_MESSAGE.into())),
109        e => Some((0, e.to_string())),
110    }
111}
112
113fn reuse_cache(bytes: Vec<u8>, cache: &Cache, stale: bool) -> Result<FetchOutcome> {
114    let snap = parse_cache(&bytes)?;
115    Ok(FetchOutcome {
116        snapshot: snap,
117        stale,
118        last_error: cache.read_last_error(),
119        cache_age: cache.payload_age(),
120    })
121}
122
123fn parse_cache(bytes: &[u8]) -> Result<KimiSnapshot> {
124    let v: serde_json::Value = serde_json::from_slice(bytes)?;
125    Ok(KimiSnapshot {
126        plan: v["plan"].as_str().map(|s| s.to_string()),
127        weekly_limit: parse_cache_u64(&v["weekly_limit"], "weekly_limit")?,
128        weekly_used: parse_cache_u64(&v["weekly_used"], "weekly_used")?,
129        weekly_remaining: parse_cache_u64(&v["weekly_remaining"], "weekly_remaining")?,
130        weekly_reset_at: parse_cache_datetime(&v["weekly_reset_at"])?,
131        window_limit: parse_cache_u64(&v["window_limit"], "window_limit")?,
132        window_used: parse_cache_u64(&v["window_used"], "window_used")?,
133        window_remaining: parse_cache_u64(&v["window_remaining"], "window_remaining")?,
134        window_reset_at: parse_cache_datetime(&v["window_reset_at"])?,
135    })
136}
137
138fn parse_cache_u64(v: &serde_json::Value, name: &str) -> Result<u64> {
139    v.as_u64()
140        .ok_or_else(|| AppError::Schema(format!("kimi cache: invalid {name}")))
141}
142
143fn parse_cache_datetime(v: &serde_json::Value) -> Result<Option<DateTime<Utc>>> {
144    match v {
145        serde_json::Value::Null => Ok(None),
146        serde_json::Value::String(s) => DateTime::parse_from_rfc3339(s)
147            .map(|dt| Some(dt.into()))
148            .map_err(|e| AppError::Schema(format!("kimi cache: invalid reset timestamp: {e}"))),
149        _ => Err(AppError::Schema(
150            "kimi cache: invalid reset timestamp".into(),
151        )),
152    }
153}
154
155fn snap_to_json(snap: &KimiSnapshot) -> serde_json::Value {
156    serde_json::json!({
157        "plan": snap.plan,
158        "weekly_limit": snap.weekly_limit,
159        "weekly_used": snap.weekly_used,
160        "weekly_remaining": snap.weekly_remaining,
161        "weekly_reset_at": snap.weekly_reset_at.map(|dt| dt.to_rfc3339()),
162        "window_limit": snap.window_limit,
163        "window_used": snap.window_used,
164        "window_remaining": snap.window_remaining,
165        "window_reset_at": snap.window_reset_at.map(|dt| dt.to_rfc3339()),
166    })
167}
168
169async fn fetch_live(client: &reqwest::Client, url: &str, api_key: &str) -> Result<KimiSnapshot> {
170    let resp = tokio::time::timeout(
171        HTTP_TIMEOUT,
172        client
173            .get(url)
174            .header("Authorization", format!("Bearer {api_key}"))
175            .header("Accept", "application/json")
176            .send(),
177    )
178    .await
179    .map_err(|_| AppError::Transport(format!("kimi timeout: {url}")))??;
180
181    let status = resp.status();
182
183    if !status.is_success() {
184        // Never surface upstream/proxy bodies: they can contain credentials or
185        // arbitrary markup. Keep the cached diagnostic useful but generic.
186        let body = if matches!(status.as_u16(), 401 | 403) {
187            "Kimi authentication failed".into()
188        } else {
189            format!("Kimi API returned HTTP {}", status.as_u16())
190        };
191        return Err(AppError::Http {
192            status: status.as_u16(),
193            body,
194        });
195    }
196
197    let bytes = resp.bytes().await?;
198    let r: UsagesResponse = serde_json::from_slice(&bytes)
199        .map_err(|e| AppError::Schema(format!("kimi usages response: {e}")))?;
200    r.into_snapshot()
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use tempfile::TempDir;
207
208    fn cache_fixture() -> (TempDir, Cache) {
209        let td = TempDir::new().unwrap();
210        let cache = Cache::at(td.path().join("kimi"));
211        cache.ensure_dir().unwrap();
212        (td, cache)
213    }
214
215    fn sample_json() -> &'static str {
216        r#"{
217            "user": { "membership": { "level": "LEVEL_INTERMEDIATE" } },
218            "usage": { "limit": "100", "used": "26", "remaining": "74", "resetTime": "2026-02-11T17:32:50.757941Z" },
219            "limits": [
220                {
221                    "window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" },
222                    "detail": { "limit": "100", "used": "15", "remaining": "85", "resetTime": "2026-02-07T12:32:50.757941Z" }
223                }
224            ]
225        }"#
226    }
227
228    fn sample_seed() -> serde_json::Value {
229        serde_json::json!({
230            "plan": "LEVEL_INTERMEDIATE",
231            "weekly_limit": 100,
232            "weekly_used": 30,
233            "weekly_remaining": 70,
234            "weekly_reset_at": "2026-02-11T17:32:50.757941Z",
235            "window_limit": 100,
236            "window_used": 20,
237            "window_remaining": 80,
238            "window_reset_at": "2026-02-07T12:32:50.757941Z"
239        })
240    }
241
242    #[tokio::test]
243    async fn live_200_returns_snapshot_and_sends_headers() {
244        let mut server = mockito::Server::new_async().await;
245        let m = server
246            .mock("GET", "/coding/v1/usages")
247            .with_status(200)
248            .with_body(sample_json())
249            .match_header("authorization", "Bearer sk-test")
250            .match_header("accept", "application/json")
251            .create_async()
252            .await;
253
254        let (_td, cache) = cache_fixture();
255        let client = reqwest::Client::new();
256        let endpoints = Endpoints {
257            usages: format!("{}/coding/v1/usages", server.url()),
258        };
259        let out = fetch_snapshot(
260            &client,
261            "sk-test",
262            &cache,
263            &endpoints,
264            Duration::from_secs(0),
265        )
266        .await
267        .unwrap();
268        m.assert_async().await;
269        assert_eq!(out.snapshot.plan, Some("LEVEL_INTERMEDIATE".into()));
270        assert_eq!(out.snapshot.weekly_limit, 100);
271        assert_eq!(out.snapshot.weekly_used, 26);
272        assert_eq!(out.snapshot.weekly_remaining, 74);
273        assert_eq!(out.snapshot.window_limit, 100);
274        assert_eq!(out.snapshot.window_used, 15);
275        assert!(!out.stale);
276    }
277
278    #[tokio::test]
279    async fn http_401_falls_back_to_cache() {
280        let mut server = mockito::Server::new_async().await;
281        server
282            .mock("GET", "/coding/v1/usages")
283            .with_status(401)
284            .with_body(r#"{"error": "invalid api key"}"#)
285            .create_async()
286            .await;
287
288        let (_td, cache) = cache_fixture();
289        cache
290            .write_payload(sample_seed().to_string().as_bytes())
291            .unwrap();
292
293        let client = reqwest::Client::new();
294        let endpoints = Endpoints {
295            usages: format!("{}/coding/v1/usages", server.url()),
296        };
297        let out = fetch_snapshot(
298            &client,
299            "bad-key",
300            &cache,
301            &endpoints,
302            Duration::from_secs(0),
303        )
304        .await
305        .unwrap();
306        assert!(out.stale);
307        assert_eq!(out.snapshot.weekly_used, 30);
308        assert_eq!(out.last_error.as_ref().map(|(c, _)| *c), Some(401));
309    }
310
311    #[tokio::test]
312    async fn http_500_falls_back_to_cache() {
313        let mut server = mockito::Server::new_async().await;
314        server
315            .mock("GET", "/coding/v1/usages")
316            .with_status(500)
317            .with_body(r#"{"error": "internal server error"}"#)
318            .create_async()
319            .await;
320
321        let (_td, cache) = cache_fixture();
322        cache
323            .write_payload(sample_seed().to_string().as_bytes())
324            .unwrap();
325
326        let client = reqwest::Client::new();
327        let endpoints = Endpoints {
328            usages: format!("{}/coding/v1/usages", server.url()),
329        };
330        let out = fetch_snapshot(
331            &client,
332            "sk-test",
333            &cache,
334            &endpoints,
335            Duration::from_secs(0),
336        )
337        .await
338        .unwrap();
339        assert!(out.stale);
340        assert_eq!(out.last_error.as_ref().map(|(c, _)| *c), Some(500));
341    }
342
343    #[tokio::test]
344    async fn http_401_without_cache_returns_http_error() {
345        let mut server = mockito::Server::new_async().await;
346        server
347            .mock("GET", "/coding/v1/usages")
348            .with_status(401)
349            .with_body(r#"{"error": "invalid api key"}"#)
350            .create_async()
351            .await;
352
353        let (_td, cache) = cache_fixture();
354        let client = reqwest::Client::new();
355        let endpoints = Endpoints {
356            usages: format!("{}/coding/v1/usages", server.url()),
357        };
358        let err = fetch_snapshot(
359            &client,
360            "bad-key",
361            &cache,
362            &endpoints,
363            Duration::from_secs(0),
364        )
365        .await
366        .unwrap_err();
367        match err {
368            AppError::Http { status, .. } => assert_eq!(status, 401),
369            other => panic!("expected Http 401, got {other:?}"),
370        }
371    }
372
373    #[tokio::test]
374    async fn malformed_numeric_200_returns_schema_error() {
375        let mut server = mockito::Server::new_async().await;
376        server
377            .mock("GET", "/coding/v1/usages")
378            .with_status(200)
379            .with_body(r#"{"usage": {"limit": "100", "used": "garbage"}}"#)
380            .create_async()
381            .await;
382
383        let (_td, cache) = cache_fixture();
384        let client = reqwest::Client::new();
385        let endpoints = Endpoints {
386            usages: format!("{}/coding/v1/usages", server.url()),
387        };
388        let err = fetch_snapshot(
389            &client,
390            "sk-test",
391            &cache,
392            &endpoints,
393            Duration::from_secs(0),
394        )
395        .await
396        .unwrap_err();
397        assert!(
398            err.to_string().contains("used") || err.to_string().contains("Schema"),
399            "expected schema error, got {err}"
400        );
401    }
402
403    #[tokio::test]
404    async fn malformed_numeric_200_with_seeded_cache_returns_stale_snapshot_and_preserves_cache() {
405        let mut server = mockito::Server::new_async().await;
406        server
407            .mock("GET", "/coding/v1/usages")
408            .with_status(200)
409            .with_body(r#"{"usage": {"limit": "100", "used": "garbage"}}"#)
410            .create_async()
411            .await;
412
413        let (_td, cache) = cache_fixture();
414        let seeded = sample_seed().to_string();
415        cache.write_payload(seeded.as_bytes()).unwrap();
416
417        let client = reqwest::Client::new();
418        let endpoints = Endpoints {
419            usages: format!("{}/coding/v1/usages", server.url()),
420        };
421        let out = fetch_snapshot(
422            &client,
423            "sk-test",
424            &cache,
425            &endpoints,
426            Duration::from_secs(0),
427        )
428        .await
429        .unwrap();
430
431        assert!(out.stale);
432        assert_eq!(out.snapshot.weekly_used, 30);
433        assert_eq!(out.snapshot.window_used, 20);
434        assert_eq!(out.last_error, Some((0, SCHEMA_DRIFT_MESSAGE.into())));
435
436        // The payload file must still contain the original seeded snapshot.
437        let payload = std::fs::read_to_string(cache.payload_path()).unwrap();
438        assert_eq!(payload, seeded);
439    }
440
441    #[tokio::test]
442    async fn error_object_200_returns_schema_error() {
443        let mut server = mockito::Server::new_async().await;
444        server
445            .mock("GET", "/coding/v1/usages")
446            .with_status(200)
447            .with_body(r#"{"error": "invalid token"}"#)
448            .create_async()
449            .await;
450
451        let (_td, cache) = cache_fixture();
452        let client = reqwest::Client::new();
453        let endpoints = Endpoints {
454            usages: format!("{}/coding/v1/usages", server.url()),
455        };
456        let err = fetch_snapshot(
457            &client,
458            "sk-test",
459            &cache,
460            &endpoints,
461            Duration::from_secs(0),
462        )
463        .await
464        .unwrap_err();
465        assert!(err.to_string().contains("usage block"), "got {err}");
466    }
467
468    #[tokio::test]
469    async fn corrupt_fresh_cache_ignored() {
470        let mut server = mockito::Server::new_async().await;
471        server
472            .mock("GET", "/coding/v1/usages")
473            .with_status(200)
474            .with_body(sample_json())
475            .create_async()
476            .await;
477
478        let (_td, cache) = cache_fixture();
479        cache.write_payload(b"not valid json".as_slice()).unwrap();
480
481        let client = reqwest::Client::new();
482        let endpoints = Endpoints {
483            usages: format!("{}/coding/v1/usages", server.url()),
484        };
485        let out = fetch_snapshot(
486            &client,
487            "sk-test",
488            &cache,
489            &endpoints,
490            Duration::from_secs(60),
491        )
492        .await
493        .unwrap();
494        assert_eq!(out.snapshot.weekly_used, 26);
495        assert!(!out.stale);
496    }
497
498    #[tokio::test]
499    async fn corrupt_stale_cache_returns_error() {
500        let mut server = mockito::Server::new_async().await;
501        server
502            .mock("GET", "/coding/v1/usages")
503            .with_status(401)
504            .with_body(r#"{"error": "invalid api key"}"#)
505            .create_async()
506            .await;
507
508        let (_td, cache) = cache_fixture();
509        cache.write_payload(b"not valid json".as_slice()).unwrap();
510
511        let client = reqwest::Client::new();
512        let endpoints = Endpoints {
513            usages: format!("{}/coding/v1/usages", server.url()),
514        };
515        let err = fetch_snapshot(
516            &client,
517            "bad-key",
518            &cache,
519            &endpoints,
520            Duration::from_secs(0),
521        )
522        .await
523        .unwrap_err();
524        assert!(
525            matches!(err, AppError::Http { status, .. } if status == 401),
526            "expected 401, got {err:?}"
527        );
528    }
529
530    #[tokio::test]
531    async fn transport_error_with_stale_cache_uses_cache() {
532        // Use a URL that will not resolve to trigger a transport error.
533        let (_td, cache) = cache_fixture();
534        cache
535            .write_payload(sample_seed().to_string().as_bytes())
536            .unwrap();
537
538        let client = reqwest::Client::new();
539        let endpoints = Endpoints {
540            usages: "http://localhost:1/coding/v1/usages".into(),
541        };
542        let out = fetch_snapshot(
543            &client,
544            "sk-test",
545            &cache,
546            &endpoints,
547            Duration::from_secs(0),
548        )
549        .await
550        .unwrap();
551        assert!(out.stale);
552        assert_eq!(out.snapshot.weekly_used, 30);
553    }
554
555    #[tokio::test]
556    async fn missing_counters_with_seeded_cache_preserves_snapshot() {
557        let mut server = mockito::Server::new_async().await;
558        server
559            .mock("GET", "/coding/v1/usages")
560            .with_status(200)
561            .with_body(r#"{"usage":{"limit":100}}"#)
562            .create_async()
563            .await;
564        let (_td, cache) = cache_fixture();
565        let seeded = sample_seed().to_string();
566        cache.write_payload(seeded.as_bytes()).unwrap();
567        let out = fetch_snapshot(
568            &reqwest::Client::new(),
569            "sk-test",
570            &cache,
571            &Endpoints {
572                usages: format!("{}/coding/v1/usages", server.url()),
573            },
574            Duration::ZERO,
575        )
576        .await
577        .unwrap();
578        assert!(out.stale);
579        assert_eq!(out.snapshot.weekly_used, 30);
580        assert_eq!(
581            std::fs::read_to_string(cache.payload_path()).unwrap(),
582            seeded
583        );
584    }
585
586    #[tokio::test]
587    async fn unrecognized_window_with_seeded_cache_preserves_snapshot() {
588        let mut server = mockito::Server::new_async().await;
589        server.mock("GET", "/coding/v1/usages").with_status(200)
590            .with_body(r#"{"usage":{"limit":100,"used":10},"limits":[{"window":{"duration":4,"timeUnit":"TIME_UNIT_HOUR"},"detail":{"limit":100,"used":10}}]}"#).create_async().await;
591        let (_td, cache) = cache_fixture();
592        let seeded = sample_seed().to_string();
593        cache.write_payload(seeded.as_bytes()).unwrap();
594        let out = fetch_snapshot(
595            &reqwest::Client::new(),
596            "sk-test",
597            &cache,
598            &Endpoints {
599                usages: format!("{}/coding/v1/usages", server.url()),
600            },
601            Duration::ZERO,
602        )
603        .await
604        .unwrap();
605        assert!(out.stale);
606        assert_eq!(out.snapshot.window_used, 20);
607        assert_eq!(
608            std::fs::read_to_string(cache.payload_path()).unwrap(),
609            seeded
610        );
611    }
612
613    #[tokio::test]
614    async fn http_error_body_is_redacted() {
615        let mut server = mockito::Server::new_async().await;
616        server
617            .mock("GET", "/coding/v1/usages")
618            .with_status(500)
619            .with_body("proxy secret: <token>")
620            .create_async()
621            .await;
622        let (_td, cache) = cache_fixture();
623        let err = fetch_snapshot(
624            &reqwest::Client::new(),
625            "sk-test",
626            &cache,
627            &Endpoints {
628                usages: format!("{}/coding/v1/usages", server.url()),
629            },
630            Duration::ZERO,
631        )
632        .await
633        .unwrap_err();
634        assert!(
635            matches!(err, AppError::Http { status: 500, ref body } if body == "Kimi API returned HTTP 500")
636        );
637    }
638}
ai_usagebar/kimi/fetch.rs

ai_usagebar/kimi/
fetch.rs
