usage_monitor_cli/provider/
kimik2.rs

1use async_trait::async_trait;
2
3use crate::error::SpendPanelError;
4use crate::model::{CreditsSnapshot, UsageSnapshot};
5use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
6
7/// Kimi K2 credits provider (kimi-k2.ai, API-key auth).
8pub struct KimiK2Provider {
9    metadata: ProviderMetadata,
10    base_url: Option<String>,
11}
12
13impl KimiK2Provider {
14    pub fn new() -> Self {
15        Self {
16            metadata: ProviderMetadata {
17                id: "kimik2",
18                name: "Kimi K2",
19                description: "Kimi K2 credits monitor (kimi-k2.ai)",
20                auth_methods: &["api_key", "env"],
21                website: Some("https://kimi-k2.ai"),
22            },
23            base_url: None,
24        }
25    }
26
27    pub fn with_base_url(url: &str) -> Self {
28        let mut p = Self::new();
29        p.base_url = Some(url.to_string());
30        p
31    }
32
33    fn api_base(&self) -> &str {
34        self.base_url.as_deref().unwrap_or("https://kimi-k2.ai")
35    }
36
37    fn clean(raw: &str) -> String {
38        let mut v = raw.trim();
39        if v.len() >= 2
40            && ((v.starts_with('"') && v.ends_with('"'))
41                || (v.starts_with('\'') && v.ends_with('\'')))
42        {
43            v = &v[1..v.len() - 1];
44        }
45        v.trim().to_string()
46    }
47
48    fn resolve_key(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
49        for key in ["api_key", "token"] {
50            if let Some(v) = ctx.config.get(key) {
51                let c = Self::clean(v);
52                if !c.is_empty() {
53                    return Ok(c);
54                }
55            }
56        }
57        for env in ["KIMI_K2_API_KEY", "KIMIK2_API_KEY"] {
58            if let Ok(v) = std::env::var(env) {
59                let c = Self::clean(&v);
60                if !c.is_empty() {
61                    return Ok(c);
62                }
63            }
64        }
65        Err(SpendPanelError::AuthFailed(
66            "kimik2".into(),
67            "no API key in api_key/token config, KIMI_K2_API_KEY, or KIMIK2_API_KEY".into(),
68        ))
69    }
70
71    fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
72        reqwest::Client::builder()
73            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
74            .build()
75            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
76    }
77
78    /// Extracts a number from any of the candidate dotted paths, searching the
79    /// root plus the common `data`/`result`/`usage`/`credits` wrappers.
80    fn find_number(root: &serde_json::Value, paths: &[&[&str]]) -> Option<f64> {
81        let mut contexts: Vec<&serde_json::Value> = vec![root];
82        for key in ["data", "result", "usage", "credits"] {
83            if let Some(v) = root.get(key) {
84                contexts.push(v);
85                // Descend one more level so `data.credits` / `result.usage` etc. resolve.
86                for nested in ["usage", "credits"] {
87                    if let Some(n) = v.get(nested) {
88                        contexts.push(n);
89                    }
90                }
91            }
92        }
93        for path in paths {
94            for ctx in &contexts {
95                let mut cursor = *ctx;
96                let mut ok = true;
97                for key in *path {
98                    match cursor.get(key) {
99                        Some(next) => cursor = next,
100                        None => {
101                            ok = false;
102                            break;
103                        }
104                    }
105                }
106                let value = ok
107                    .then(|| {
108                        cursor
109                            .as_f64()
110                            .or_else(|| cursor.as_str().and_then(|s| s.parse().ok()))
111                    })
112                    .flatten();
113                if let Some(n) = value {
114                    return Some(n);
115                }
116            }
117        }
118        None
119    }
120
121    fn parse(body: &str) -> Result<UsageSnapshot, SpendPanelError> {
122        let json: serde_json::Value = serde_json::from_str(body)
123            .map_err(|e| SpendPanelError::ParseError("kimik2".into(), e.to_string()))?;
124
125        let consumed = Self::find_number(
126            &json,
127            &[
128                &["total_credits_consumed"],
129                &["totalCreditsConsumed"],
130                &["credits_consumed"],
131                &["consumedCredits"],
132                &["usedCredits"],
133                &["total"],
134            ],
135        )
136        .unwrap_or(0.0);
137        let remaining = Self::find_number(
138            &json,
139            &[
140                &["credits_remaining"],
141                &["creditsRemaining"],
142                &["remaining_credits"],
143                &["available_credits"],
144                &["credits_left"],
145            ],
146        )
147        .unwrap_or(0.0)
148        .max(0.0);
149
150        let mut snapshot = UsageSnapshot::new("kimik2");
151        let mut credits = CreditsSnapshot::new(remaining, "credits");
152        credits.used = Some(consumed);
153        if consumed > 0.0 || remaining > 0.0 {
154            credits.total = Some(consumed + remaining);
155        }
156        snapshot.credits = Some(credits);
157        Ok(snapshot)
158    }
159}
160
161impl Default for KimiK2Provider {
162    fn default() -> Self {
163        Self::new()
164    }
165}
166
167#[async_trait]
168impl UsageProvider for KimiK2Provider {
169    fn metadata(&self) -> &ProviderMetadata {
170        &self.metadata
171    }
172
173    fn detect_credentials(&self) -> bool {
174        ["KIMI_K2_API_KEY", "KIMIK2_API_KEY"].iter().any(|e| {
175            std::env::var(e)
176                .map(|v| !v.trim().is_empty())
177                .unwrap_or(false)
178        })
179    }
180
181    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
182        let key = Self::resolve_key(ctx)?;
183        let client = Self::build_client(ctx)?;
184        let url = format!("{}/api/user/credits", self.api_base().trim_end_matches('/'));
185        let resp = client
186            .get(url)
187            .header("Authorization", format!("Bearer {}", key))
188            .header("Accept", "application/json")
189            .send()
190            .await
191            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
192        let status = resp.status();
193        let body = resp
194            .text()
195            .await
196            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
197        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
198            return Err(SpendPanelError::AuthFailed(
199                "kimik2".into(),
200                format!("invalid API key (HTTP {})", status.as_u16()),
201            ));
202        }
203        if !status.is_success() {
204            return Err(SpendPanelError::ProviderError(
205                "kimik2".into(),
206                format!("HTTP {}: {}", status, body),
207            ));
208        }
209        Self::parse(&body)
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use pretty_assertions::assert_eq;
217    use wiremock::matchers::{header, method, path};
218    use wiremock::{Mock, MockServer, ResponseTemplate};
219
220    #[test]
221    fn test_metadata() {
222        assert_eq!(KimiK2Provider::new().metadata().id, "kimik2");
223    }
224
225    #[test]
226    fn test_resolve_key_missing() {
227        assert!(matches!(
228            KimiK2Provider::resolve_key(&ProviderContext::new()).unwrap_err(),
229            SpendPanelError::AuthFailed(_, _)
230        ));
231    }
232
233    #[test]
234    fn test_parse_flexible_paths() {
235        let snap = KimiK2Provider::parse(
236            r#"{"data":{"credits":{"total_credits_consumed":40,"credits_remaining":"60"}}}"#,
237        )
238        .unwrap();
239        let c = snap.credits.unwrap();
240        assert_eq!(c.balance, 60.0);
241        assert_eq!(c.used, Some(40.0));
242        assert_eq!(c.total, Some(100.0));
243    }
244
245    #[tokio::test]
246    async fn test_fetch_success() {
247        let server = MockServer::start().await;
248        Mock::given(method("GET"))
249            .and(path("/api/user/credits"))
250            .and(header("authorization", "Bearer k2"))
251            .respond_with(ResponseTemplate::new(200).set_body_raw(
252                r#"{"credits_remaining":75,"total_credits_consumed":25}"#,
253                "application/json",
254            ))
255            .mount(&server)
256            .await;
257        let provider = KimiK2Provider::with_base_url(&server.uri());
258        let mut ctx = ProviderContext::new();
259        ctx.config.insert("api_key".into(), "k2".into());
260        let snap = provider.fetch_usage(&ctx).await.unwrap();
261        assert_eq!(snap.credits.unwrap().balance, 75.0);
262    }
263
264    #[tokio::test]
265    async fn test_fetch_401() {
266        let server = MockServer::start().await;
267        Mock::given(method("GET"))
268            .and(path("/api/user/credits"))
269            .respond_with(ResponseTemplate::new(401))
270            .mount(&server)
271            .await;
272        let provider = KimiK2Provider::with_base_url(&server.uri());
273        let mut ctx = ProviderContext::new();
274        ctx.config.insert("api_key".into(), "bad".into());
275        assert!(matches!(
276            provider.fetch_usage(&ctx).await.unwrap_err(),
277            SpendPanelError::AuthFailed(_, _)
278        ));
279    }
280}
usage_monitor_cli/provider/kimik2.rs

usage_monitor_cli/provider/
kimik2.rs
