1use serde::Deserialize;
36
37use crate::usage::{UsageWindow, ZaiSnapshot};
38
39#[derive(Debug, Clone, Deserialize)]
40pub struct Envelope {
41 #[serde(default)]
42 pub code: i64,
43 #[serde(default)]
44 pub data: Option<MonitorData>,
45 #[serde(default)]
46 pub success: bool,
47 #[serde(default)]
48 pub msg: String,
49}
50
51impl Envelope {
52 pub fn check_ok(&self) -> crate::error::Result<()> {
60 if !self.success || (self.code != 0 && self.code != 200) {
61 let msg = if self.msg.is_empty() {
62 "no message".to_string()
63 } else {
64 self.msg.clone()
65 };
66 return Err(crate::error::AppError::Schema(format!(
67 "zai: API reported failure (code {}, success {}): {msg}",
68 self.code, self.success
69 )));
70 }
71 let Some(data) = &self.data else {
72 return Err(crate::error::AppError::Schema(
73 "zai: success response carried no `data`".into(),
74 ));
75 };
76 let (session, weekly) = classify_token_buckets(&data.limits).map_err(|why| {
80 crate::error::AppError::Schema(format!("zai: unrecognised limits layout: {why}"))
81 })?;
82 let mcp = classify_time_bucket(&data.limits).map_err(|why| {
83 crate::error::AppError::Schema(format!("zai: unrecognised limits layout: {why}"))
84 })?;
85 for (label, bucket) in [("session", session), ("weekly", weekly), ("MCP", mcp)] {
86 if bucket.is_some_and(|entry| entry.percentage.is_none()) {
87 return Err(crate::error::AppError::Schema(format!(
88 "zai: {label} limit carried no percentage"
89 )));
90 }
91 }
92 Ok(())
93 }
94}
95
96const UNIT_SESSION: i64 = 3;
101const UNIT_WEEKLY: i64 = 6;
102
103fn is_usage_bucket(kind: &str) -> bool {
107 kind == "TOKENS_LIMIT" || kind == "CREDIT_LIMIT"
108}
109
110type TokenBuckets<'a> = (Option<&'a LimitEntry>, Option<&'a LimitEntry>);
111
112fn classify_token_buckets(limits: &[LimitEntry]) -> Result<TokenBuckets<'_>, String> {
120 let tokens: Vec<&LimitEntry> = limits.iter().filter(|l| is_usage_bucket(&l.kind)).collect();
121
122 if tokens.is_empty() {
123 return Ok((None, None));
124 }
125 if tokens.iter().all(|l| l.unit.is_none()) {
126 return Err("usage buckets carry no unit discriminator".into());
127 }
128
129 let session = unique_by_unit(&tokens, UNIT_SESSION)?;
130 let weekly = unique_by_unit(&tokens, UNIT_WEEKLY)?;
131 if session.is_none() && weekly.is_none() {
132 let seen: Vec<String> = tokens
133 .iter()
134 .filter_map(|l| l.unit)
135 .map(|u| u.to_string())
136 .collect();
137 return Err(format!(
138 "no usage bucket carries a known unit code (saw {})",
139 seen.join(", ")
140 ));
141 }
142 Ok((session, weekly))
143}
144
145fn classify_time_bucket(limits: &[LimitEntry]) -> Result<Option<&LimitEntry>, String> {
146 let mut matching = limits.iter().filter(|l| l.kind == "TIME_LIMIT");
147 let first = matching.next();
148 if matching.next().is_some() {
149 return Err("two TIME_LIMIT buckets are present".into());
150 }
151 Ok(first)
152}
153
154fn unique_by_unit<'a>(
155 tokens: &[&'a LimitEntry],
156 code: i64,
157) -> Result<Option<&'a LimitEntry>, String> {
158 let mut matching = tokens.iter().filter(|l| l.unit == Some(code));
159 let first = matching.next().copied();
160 if matching.next().is_some() {
161 return Err(format!("two usage buckets carry unit {code}"));
162 }
163 Ok(first)
164}
165
166#[derive(Debug, Default, Clone, Deserialize)]
167#[serde(default)]
168pub struct MonitorData {
169 pub limits: Vec<LimitEntry>,
170 pub level: String,
171}
172
173#[derive(Debug, Default, Clone, Deserialize)]
174#[serde(default)]
175pub struct LimitEntry {
176 #[serde(rename = "type")]
177 pub kind: String,
178 #[serde(default, deserialize_with = "de_percent_opt")]
179 pub percentage: Option<f64>,
180 #[serde(rename = "nextResetTime", default, deserialize_with = "de_opt_ms")]
182 pub next_reset_time: Option<i64>,
183 pub unit: Option<i64>,
184 pub number: Option<i64>,
185}
186
187fn de_opt_ms<'de, D>(d: D) -> Result<Option<i64>, D::Error>
188where
189 D: serde::Deserializer<'de>,
190{
191 let v = serde_json::Value::deserialize(d)?;
192 match v {
193 serde_json::Value::Null => Ok(None),
194 serde_json::Value::Number(n) => {
195 let millis = if let Some(i) = n.as_i64() {
196 i
197 } else if let Some(f) = n.as_f64()
198 && f.is_finite()
199 && f.fract() == 0.0
200 && f.abs() <= (1_u64 << 53) as f64
201 {
202 f as i64
203 } else {
204 return Err(serde::de::Error::custom(
205 "nextResetTime must be an integer in range",
206 ));
207 };
208 match millis {
209 0 => Ok(None),
210 1.. => Ok(Some(millis)),
211 _ => Err(serde::de::Error::custom("nextResetTime cannot be negative")),
212 }
213 }
214 other => Err(serde::de::Error::custom(format!(
215 "nextResetTime must be an integer or null, got {other:?}"
216 ))),
217 }
218}
219
220fn de_percent_opt<'de, D>(d: D) -> Result<Option<f64>, D::Error>
221where
222 D: serde::Deserializer<'de>,
223{
224 let value = Option::<f64>::deserialize(d)?;
225 value
226 .map(|pct| {
227 if pct.is_finite() && (0.0..=101.0).contains(&pct) {
228 Ok(pct)
229 } else {
230 Err(serde::de::Error::custom(format!(
231 "percentage {pct} outside 0..=100"
232 )))
233 }
234 })
235 .transpose()
236}
237
238impl Envelope {
239 pub fn into_snapshot(self, config_plan_tier: Option<&str>) -> ZaiSnapshot {
242 let data = self.data.unwrap_or_default();
243 let (session, weekly) = classify_token_buckets(&data.limits).unwrap_or((None, None));
246 let session = session.and_then(|l| to_window(l, chrono::Duration::hours(5)));
247 let weekly = weekly.and_then(|l| to_window(l, chrono::Duration::days(7)));
248 let mcp = classify_time_bucket(&data.limits)
249 .ok()
250 .flatten()
251 .and_then(|l| to_window(l, chrono::Duration::days(30)));
252
253 let level = if !data.level.is_empty() {
255 data.level
256 } else {
257 config_plan_tier.unwrap_or("unknown").to_string()
258 };
259 let plan = format!("GLM Coding {}", capitalize(&level));
260
261 ZaiSnapshot {
262 plan,
263 session,
264 weekly,
265 mcp,
266 }
267 }
268}
269
270fn to_window(l: &LimitEntry, dur: chrono::Duration) -> Option<UsageWindow> {
271 let utilization_pct = l.percentage?.round().clamp(0.0, 100.0) as i32;
272 let resets_at = l
273 .next_reset_time
274 .and_then(chrono::DateTime::<chrono::Utc>::from_timestamp_millis);
275 Some(UsageWindow {
276 utilization_pct,
277 resets_at,
278 window_duration: dur,
279 })
280}
281
282fn capitalize(s: &str) -> String {
283 let mut chars = s.chars();
284 match chars.next() {
285 Some(c) => {
286 let mut out = String::with_capacity(s.len());
287 for u in c.to_uppercase() {
288 out.push(u);
289 }
290 out.push_str(chars.as_str());
291 out
292 }
293 None => String::new(),
294 }
295}
296
297#[cfg(test)]
298mod tests {
299 use super::*;
300
301 const REAL_BODY: &str = r#"{"code":200,"msg":"Operation successful","data":{
302 "limits":[
303 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":0},
304 {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":0,"nextResetTime":1779792169974},
305 {"type":"TIME_LIMIT","unit":5,"number":1,"usage":1000,"currentValue":0,"remaining":1000,"percentage":0,"nextResetTime":1779964969979,
306 "usageDetails":[{"modelCode":"search-prime","usage":0}]}
307 ],
308 "level":"pro"
309 },"success":true}"#;
310
311 #[test]
312 fn parses_real_response_shape() {
313 let env: Envelope = serde_json::from_str(REAL_BODY).unwrap();
314 let snap = env.into_snapshot(None);
315 assert_eq!(snap.plan, "GLM Coding Pro");
316 assert!(snap.session.is_some());
317 assert!(snap.weekly.is_some());
318 assert!(snap.mcp.is_some());
319 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 0);
320 assert!(snap.weekly.as_ref().unwrap().resets_at.is_some());
321 }
322
323 #[test]
324 fn missing_data_yields_neutral_snapshot() {
325 let env: Envelope = serde_json::from_str(r#"{"code":500,"success":false}"#).unwrap();
326 let snap = env.into_snapshot(Some("lite"));
327 assert_eq!(snap.plan, "GLM Coding Lite");
328 assert!(snap.session.is_none());
329 }
330
331 #[test]
332 fn percentage_with_float_rounds() {
333 let body = r#"{"data":{"limits":[
334 {"type":"TOKENS_LIMIT","unit":3,"percentage":42.7}
335 ],"level":"max"},"success":true}"#;
336 let env: Envelope = serde_json::from_str(body).unwrap();
337 let snap = env.into_snapshot(None);
338 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
339 }
340
341 #[test]
342 fn benign_percentage_overshoot_clamps_to_hundred() {
343 let body = r#"{"data":{"limits":[
344 {"type":"TOKENS_LIMIT","unit":3,"percentage":100.6}
345 ]},"success":true}"#;
346 let env: Envelope = serde_json::from_str(body).unwrap();
347 let snap = env.into_snapshot(None);
348 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 100);
349 }
350
351 #[test]
352 fn only_time_limit_means_no_session_or_weekly() {
353 let body = r#"{"data":{"limits":[
354 {"type":"TIME_LIMIT","percentage":12}
355 ]},"success":true}"#;
356 let env: Envelope = serde_json::from_str(body).unwrap();
357 let snap = env.into_snapshot(None);
358 assert!(snap.session.is_none());
359 assert!(snap.weekly.is_none());
360 assert!(snap.mcp.is_some());
361 }
362
363 #[test]
364 fn config_plan_tier_used_when_level_empty() {
365 let body = r#"{"data":{"limits":[],"level":""},"success":true}"#;
366 let env: Envelope = serde_json::from_str(body).unwrap();
367 let snap = env.into_snapshot(Some("max"));
368 assert_eq!(snap.plan, "GLM Coding Max");
369 }
370
371 #[test]
374 fn buckets_are_identified_by_unit_not_by_position() {
375 let body = r#"{"data":{"limits":[
376 {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":15,"nextResetTime":1779792169974},
377 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":42}
378 ],"level":"pro"},"success":true}"#;
379 let env: Envelope = serde_json::from_str(body).unwrap();
380 env.check_ok().unwrap();
381 let snap = env.into_snapshot(None);
382 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 42);
383 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 15);
384 assert!(snap.weekly.as_ref().unwrap().resets_at.is_some());
385 assert!(snap.session.as_ref().unwrap().resets_at.is_none());
386 }
387
388 #[test]
390 fn unknown_extra_bucket_is_dropped_not_shown_as_session() {
391 let body = r#"{"data":{"limits":[
392 {"type":"TOKENS_LIMIT","unit":4,"number":1,"percentage":99},
393 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":42},
394 {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":15}
395 ],"level":"pro"},"success":true}"#;
396 let env: Envelope = serde_json::from_str(body).unwrap();
397 env.check_ok().unwrap();
398 let snap = env.into_snapshot(None);
399 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 42);
400 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 15);
401 }
402
403 #[test]
404 fn duplicate_unit_is_an_error_not_a_coin_flip() {
405 let body = r#"{"data":{"limits":[
406 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":42},
407 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":7}
408 ],"level":"pro"},"success":true}"#;
409 let env: Envelope = serde_json::from_str(body).unwrap();
410 let err = env.check_ok().unwrap_err().to_string();
411 assert!(err.contains("unit 3"), "unhelpful error: {err}");
412 let snap = env.into_snapshot(None);
414 assert!(snap.session.is_none());
415 assert!(snap.weekly.is_none());
416 }
417
418 #[test]
419 fn all_unknown_units_is_an_error() {
420 let body = r#"{"data":{"limits":[
421 {"type":"TOKENS_LIMIT","unit":4,"number":1,"percentage":42},
422 {"type":"TOKENS_LIMIT","unit":7,"number":1,"percentage":15}
423 ],"level":"pro"},"success":true}"#;
424 let env: Envelope = serde_json::from_str(body).unwrap();
425 let err = env.check_ok().unwrap_err().to_string();
426 assert!(err.contains("4, 7"), "unhelpful error: {err}");
427 assert!(env.into_snapshot(None).session.is_none());
428 }
429
430 #[test]
433 fn unit_less_bucket_alongside_a_known_one_is_dropped() {
434 let body = r#"{"data":{"limits":[
435 {"type":"TOKENS_LIMIT","percentage":99},
436 {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":15}
437 ],"level":"pro"},"success":true}"#;
438 let env: Envelope = serde_json::from_str(body).unwrap();
439 env.check_ok().unwrap();
440 let snap = env.into_snapshot(None);
441 assert!(snap.session.is_none());
442 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 15);
443 }
444
445 #[test]
446 fn bodies_without_any_unit_are_rejected_not_guessed_by_position() {
447 let body = r#"{"data":{"limits":[
448 {"type":"TOKENS_LIMIT","percentage":10},
449 {"type":"TOKENS_LIMIT","percentage":20}
450 ],"level":"lite"},"success":true}"#;
451 let env: Envelope = serde_json::from_str(body).unwrap();
452 let err = env.check_ok().unwrap_err().to_string();
453 assert!(err.contains("no unit discriminator"), "{err}");
454 let snap = env.into_snapshot(None);
455 assert!(snap.session.is_none());
456 assert!(snap.weekly.is_none());
457 }
458
459 #[test]
460 fn a_named_bucket_without_percentage_is_rejected_not_zeroed() {
461 let body = r#"{"data":{"limits":[
462 {"type":"TOKENS_LIMIT","unit":3,"number":5}
463 ]},"success":true}"#;
464 let env: Envelope = serde_json::from_str(body).unwrap();
465 let err = env.check_ok().unwrap_err().to_string();
466 assert!(err.contains("session limit carried no percentage"), "{err}");
467 assert!(env.into_snapshot(None).session.is_none());
468 }
469
470 #[test]
471 fn duplicate_time_limit_is_rejected_not_selected_by_position() {
472 let body = r#"{"data":{"limits":[
473 {"type":"TIME_LIMIT","unit":5,"percentage":10},
474 {"type":"TIME_LIMIT","unit":5,"percentage":20}
475 ]},"success":true}"#;
476 let env: Envelope = serde_json::from_str(body).unwrap();
477 let err = env.check_ok().unwrap_err().to_string();
478 assert!(err.contains("two TIME_LIMIT"), "{err}");
479 assert!(env.into_snapshot(None).mcp.is_none());
480 }
481
482 #[test]
483 fn invalid_percentage_and_reset_values_are_schema_drift() {
484 for percentage in ["-1", "101.5", "150"] {
485 let body = format!(
486 r#"{{"data":{{"limits":[{{"type":"TOKENS_LIMIT","unit":3,"percentage":{percentage}}}]}},"success":true}}"#
487 );
488 assert!(serde_json::from_str::<Envelope>(&body).is_err(), "{body}");
489 }
490 for reset in ["-1", "1.5", "true", r#""later""#] {
491 let body = format!(
492 r#"{{"data":{{"limits":[{{"type":"TOKENS_LIMIT","unit":3,"percentage":0,"nextResetTime":{reset}}}]}},"success":true}}"#
493 );
494 assert!(serde_json::from_str::<Envelope>(&body).is_err(), "{body}");
495 }
496 }
497
498 #[test]
499 fn check_ok_accepts_the_real_response_shape() {
500 let env: Envelope = serde_json::from_str(REAL_BODY).unwrap();
501 env.check_ok().unwrap();
502 }
503
504 #[test]
505 fn reset_time_zero_or_null_becomes_none() {
506 let body = r#"{"data":{"limits":[
507 {"type":"TOKENS_LIMIT","unit":3,"percentage":0,"nextResetTime":null}
508 ]},"success":true}"#;
509 let env: Envelope = serde_json::from_str(body).unwrap();
510 let snap = env.into_snapshot(None);
511 assert!(snap.session.as_ref().unwrap().resets_at.is_none());
512 }
513
514 #[test]
517 fn credit_limit_buckets_fill_session_and_weekly() {
518 let body = r#"{"data":{"limits":[
519 {"type":"CREDIT_LIMIT","unit":3,"number":5,"percentage":42},
520 {"type":"CREDIT_LIMIT","unit":6,"number":1,"percentage":15,"nextResetTime":1779792169974}
521 ],"level":"pro"},"success":true}"#;
522 let env: Envelope = serde_json::from_str(body).unwrap();
523 env.check_ok().unwrap();
524 let snap = env.into_snapshot(None);
525 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 42);
526 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 15);
527 assert!(snap.weekly.as_ref().unwrap().resets_at.is_some());
528 }
529
530 #[test]
533 fn mixed_token_and_credit_buckets_are_both_kept() {
534 let body = r#"{"data":{"limits":[
535 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":42},
536 {"type":"CREDIT_LIMIT","unit":6,"number":1,"percentage":15}
537 ],"level":"pro"},"success":true}"#;
538 let env: Envelope = serde_json::from_str(body).unwrap();
539 env.check_ok().unwrap();
540 let snap = env.into_snapshot(None);
541 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 42);
542 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 15);
543 }
544
545 #[test]
548 fn token_and_credit_sharing_a_unit_is_still_a_duplicate() {
549 let body = r#"{"data":{"limits":[
550 {"type":"TOKENS_LIMIT","unit":3,"percentage":42},
551 {"type":"CREDIT_LIMIT","unit":3,"percentage":15}
552 ],"level":"pro"},"success":true}"#;
553 let env: Envelope = serde_json::from_str(body).unwrap();
554 let err = env.check_ok().unwrap_err().to_string();
555 assert!(err.contains("two usage buckets carry unit 3"), "{err}");
556 let snap = env.into_snapshot(None);
557 assert!(snap.session.is_none());
558 assert!(snap.weekly.is_none());
559 }
560}