usage_monitor_cli/provider/
moonshot.rs1use async_trait::async_trait;
2
3use crate::error::SpendPanelError;
4use crate::model::{CreditsSnapshot, PlanInfo, UsageSnapshot};
5use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq)]
8enum MoonshotRegion {
9 International,
10 China,
11}
12
13impl MoonshotRegion {
14 fn from_raw(raw: &str) -> Self {
15 match raw.trim().to_ascii_lowercase().as_str() {
16 "china" | "cn" => Self::China,
17 _ => Self::International,
18 }
19 }
20
21 fn base_url(self) -> &'static str {
22 match self {
23 Self::International => "https://api.moonshot.ai",
24 Self::China => "https://api.moonshot.cn",
25 }
26 }
27
28 fn display_name(self) -> &'static str {
29 match self {
30 Self::International => "International (api.moonshot.ai)",
31 Self::China => "China (api.moonshot.cn)",
32 }
33 }
34}
35
36#[derive(Debug, serde::Deserialize)]
37struct BalanceResponse {
38 code: i64,
39 data: BalanceData,
40 scode: String,
41 status: bool,
42}
43
44#[derive(Debug, serde::Deserialize)]
45struct BalanceData {
46 available_balance: f64,
47 voucher_balance: f64,
48 cash_balance: f64,
49}
50
51#[derive(Debug, Clone, PartialEq)]
52struct MoonshotUsage {
53 available_balance: f64,
54 voucher_balance: f64,
55 cash_balance: f64,
56 region: MoonshotRegion,
57}
58
59pub struct MoonshotProvider {
60 metadata: ProviderMetadata,
61 base_url: Option<String>,
62}
63
64impl MoonshotProvider {
65 pub fn new() -> Self {
66 Self {
67 metadata: ProviderMetadata {
68 id: "moonshot",
69 name: "Moonshot / Kimi API",
70 description: "Moonshot / Kimi API balance monitor",
71 auth_methods: &["api_key", "env"],
72 website: Some("https://platform.moonshot.ai"),
73 },
74 base_url: None,
75 }
76 }
77
78 pub fn with_base_url(url: &str) -> Self {
79 let mut p = Self::new();
80 p.base_url = Some(url.to_string());
81 p
82 }
83
84 fn clean(raw: &str) -> String {
85 let mut value = raw.trim();
86 if value.len() >= 2
87 && ((value.starts_with('"') && value.ends_with('"'))
88 || (value.starts_with('\'') && value.ends_with('\'')))
89 {
90 value = &value[1..value.len() - 1];
91 }
92 value.trim().to_string()
93 }
94
95 fn resolve_api_key(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
96 for key in ["api_key", "token"] {
97 if let Some(value) = ctx.config.get(key) {
98 let cleaned = Self::clean(value);
99 if !cleaned.is_empty() {
100 return Ok(cleaned);
101 }
102 }
103 }
104 for env in ["MOONSHOT_API_KEY", "MOONSHOT_KEY"] {
105 if let Ok(value) = std::env::var(env) {
106 let cleaned = Self::clean(&value);
107 if !cleaned.is_empty() {
108 return Ok(cleaned);
109 }
110 }
111 }
112 Err(SpendPanelError::AuthFailed(
113 "moonshot".into(),
114 "no API key found in config, token, MOONSHOT_API_KEY, or MOONSHOT_KEY".into(),
115 ))
116 }
117
118 fn resolve_region(ctx: &ProviderContext) -> MoonshotRegion {
119 ctx.config
120 .get("region")
121 .map(|v| Self::clean(v))
122 .or_else(|| {
123 std::env::var("MOONSHOT_REGION")
124 .ok()
125 .map(|v| Self::clean(&v))
126 })
127 .map(|v| MoonshotRegion::from_raw(&v))
128 .unwrap_or(MoonshotRegion::International)
129 }
130
131 fn balance_url(&self, ctx: &ProviderContext, region: MoonshotRegion) -> String {
132 let base = ctx
133 .config
134 .get("api_url")
135 .or_else(|| ctx.config.get("base_url"))
136 .map(String::as_str)
137 .filter(|v| !v.is_empty())
138 .map(Self::clean)
139 .or_else(|| self.base_url.clone())
140 .unwrap_or_else(|| region.base_url().into());
141 let base = if base.starts_with("http://") || base.starts_with("https://") {
142 base
143 } else {
144 format!("https://{}", base)
145 };
146 format!("{}/v1/users/me/balance", base.trim_end_matches('/'))
147 }
148
149 fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
150 reqwest::Client::builder()
151 .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
152 .build()
153 .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
154 }
155
156 fn parse_response(
157 body: &str,
158 region: MoonshotRegion,
159 ) -> Result<MoonshotUsage, SpendPanelError> {
160 let response: BalanceResponse = serde_json::from_str(body)
161 .map_err(|e| SpendPanelError::ParseError("moonshot".into(), e.to_string()))?;
162 if response.code != 0 || !response.status {
163 return Err(SpendPanelError::ProviderError(
164 "moonshot".into(),
165 format!("code {}, scode {}", response.code, response.scode),
166 ));
167 }
168 Ok(MoonshotUsage {
169 available_balance: response.data.available_balance,
170 voucher_balance: response.data.voucher_balance,
171 cash_balance: response.data.cash_balance,
172 region,
173 })
174 }
175
176 async fn fetch_balance(
177 client: &reqwest::Client,
178 url: String,
179 api_key: &str,
180 region: MoonshotRegion,
181 ) -> Result<MoonshotUsage, SpendPanelError> {
182 let resp = client
183 .get(url)
184 .header("Authorization", format!("Bearer {}", api_key))
185 .header("Accept", "application/json")
186 .send()
187 .await
188 .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
189 let status = resp.status();
190 let body = resp
191 .text()
192 .await
193 .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
194 if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
195 return Err(SpendPanelError::AuthFailed(
196 "moonshot".into(),
197 format!("invalid API key (HTTP {})", status.as_u16()),
198 ));
199 }
200 if !status.is_success() {
201 return Err(SpendPanelError::ProviderError(
202 "moonshot".into(),
203 format!("HTTP {}", status),
204 ));
205 }
206 Self::parse_response(&body, region)
207 }
208
209 fn snapshot_from_usage(usage: MoonshotUsage) -> UsageSnapshot {
210 let mut snapshot = UsageSnapshot::new("moonshot");
211 let mut credits = CreditsSnapshot::new(usage.available_balance, "USD");
212 credits.bonus = Some(usage.voucher_balance);
213 credits.purchased = Some(usage.cash_balance);
214 snapshot.credits = Some(credits);
215 let mut features = vec![format!("region: {}", usage.region.display_name())];
216 if usage.cash_balance < 0.0 {
217 features.push(format!("cash deficit: ${:.2}", usage.cash_balance.abs()));
218 }
219 snapshot.plan = Some(PlanInfo {
220 name: format!("Balance: ${:.2}", usage.available_balance),
221 tier: None,
222 features,
223 price: None,
224 currency: Some("USD".into()),
225 billing_period: None,
226 });
227 snapshot
228 }
229}
230
231impl Default for MoonshotProvider {
232 fn default() -> Self {
233 Self::new()
234 }
235}
236
237#[async_trait]
238impl UsageProvider for MoonshotProvider {
239 fn metadata(&self) -> &ProviderMetadata {
240 &self.metadata
241 }
242
243 fn detect_credentials(&self) -> bool {
244 ["MOONSHOT_API_KEY", "MOONSHOT_KEY"]
245 .iter()
246 .any(|env| std::env::var(env).is_ok_and(|v| !Self::clean(&v).is_empty()))
247 }
248
249 async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
250 let api_key = Self::resolve_api_key(ctx)?;
251 let region = Self::resolve_region(ctx);
252 let url = self.balance_url(ctx, region);
253 let client = Self::build_client(ctx)?;
254 let usage = Self::fetch_balance(&client, url, &api_key, region).await?;
255 Ok(Self::snapshot_from_usage(usage))
256 }
257}
258
259#[cfg(test)]
260mod tests {
261 use super::*;
262 use pretty_assertions::assert_eq;
263 use wiremock::matchers::{header, method, path};
264 use wiremock::{Mock, MockServer, ResponseTemplate};
265
266 const SAMPLE: &str = r#"{
267 "code": 0,
268 "data": {"available_balance": 49.58, "voucher_balance": 50.00, "cash_balance": 12.34},
269 "scode": "0x0",
270 "status": true
271 }"#;
272
273 #[test]
274 fn test_provider_metadata() {
275 let meta = MoonshotProvider::new().metadata().clone();
276 assert_eq!(meta.id, "moonshot");
277 assert_eq!(meta.name, "Moonshot / Kimi API");
278 }
279
280 #[test]
281 fn test_regions() {
282 assert_eq!(
283 MoonshotRegion::International.base_url(),
284 "https://api.moonshot.ai"
285 );
286 assert_eq!(MoonshotRegion::China.base_url(), "https://api.moonshot.cn");
287 assert_eq!(MoonshotRegion::from_raw("cn"), MoonshotRegion::China);
288 }
289
290 #[test]
291 fn test_parse_documented_response() {
292 let usage =
293 MoonshotProvider::parse_response(SAMPLE, MoonshotRegion::International).unwrap();
294 assert_eq!(usage.available_balance, 49.58);
295 assert_eq!(usage.voucher_balance, 50.0);
296 assert_eq!(usage.cash_balance, 12.34);
297 let snapshot = MoonshotProvider::snapshot_from_usage(usage);
298 assert_eq!(snapshot.credits.as_ref().unwrap().balance, 49.58);
299 assert_eq!(snapshot.credits.as_ref().unwrap().bonus, Some(50.0));
300 assert_eq!(snapshot.credits.as_ref().unwrap().purchased, Some(12.34));
301 assert_eq!(snapshot.plan.unwrap().name, "Balance: $49.58");
302 }
303
304 #[test]
305 fn test_negative_cash_balance_is_deficit_feature() {
306 let json = r#"{"code":0,"data":{"available_balance":49.58,"voucher_balance":50.0,"cash_balance":-0.42},"scode":"0x0","status":true}"#;
307 let usage = MoonshotProvider::parse_response(json, MoonshotRegion::International).unwrap();
308 let snapshot = MoonshotProvider::snapshot_from_usage(usage);
309 assert!(
310 snapshot
311 .plan
312 .unwrap()
313 .features
314 .iter()
315 .any(|f| f.contains("deficit"))
316 );
317 }
318
319 #[test]
320 fn test_api_code_failure_returns_provider_error() {
321 let json = r#"{"code":401,"data":{"available_balance":0,"voucher_balance":0,"cash_balance":0},"scode":"unauthorized","status":false}"#;
322 let err =
323 MoonshotProvider::parse_response(json, MoonshotRegion::International).unwrap_err();
324 assert!(matches!(err, SpendPanelError::ProviderError(_, _)));
325 }
326
327 #[tokio::test]
328 async fn test_fetch_usage_sends_bearer_token() {
329 let server = MockServer::start().await;
330 Mock::given(method("GET"))
331 .and(path("/v1/users/me/balance"))
332 .and(header("authorization", "Bearer live-token"))
333 .and(header("accept", "application/json"))
334 .respond_with(ResponseTemplate::new(200).set_body_raw(SAMPLE, "application/json"))
335 .mount(&server)
336 .await;
337 let provider = MoonshotProvider::with_base_url(&server.uri());
338 let mut ctx = ProviderContext::with_api_key(" live-token ");
339 ctx.config.insert("region".into(), "china".into());
340 let snapshot = provider.fetch_usage(&ctx).await.unwrap();
341 assert_eq!(snapshot.credits.unwrap().balance, 49.58);
342 }
343
344 #[tokio::test]
345 async fn test_fetch_usage_401_is_auth_failed() {
346 let server = MockServer::start().await;
347 Mock::given(method("GET"))
348 .and(path("/v1/users/me/balance"))
349 .respond_with(ResponseTemplate::new(401))
350 .mount(&server)
351 .await;
352 let provider = MoonshotProvider::with_base_url(&server.uri());
353 let err = provider
354 .fetch_usage(&ProviderContext::with_api_key("bad"))
355 .await
356 .unwrap_err();
357 assert!(matches!(err, SpendPanelError::AuthFailed(_, _)));
358 }
359}