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
103type TokenBuckets<'a> = (Option<&'a LimitEntry>, Option<&'a LimitEntry>);
104
105fn classify_token_buckets(limits: &[LimitEntry]) -> Result<TokenBuckets<'_>, String> {
112 let tokens: Vec<&LimitEntry> = limits.iter().filter(|l| l.kind == "TOKENS_LIMIT").collect();
113
114 if tokens.is_empty() {
115 return Ok((None, None));
116 }
117 if tokens.iter().all(|l| l.unit.is_none()) {
118 return Err("TOKENS_LIMIT buckets carry no unit discriminator".into());
119 }
120
121 let session = unique_by_unit(&tokens, UNIT_SESSION)?;
122 let weekly = unique_by_unit(&tokens, UNIT_WEEKLY)?;
123 if session.is_none() && weekly.is_none() {
124 let seen: Vec<String> = tokens
125 .iter()
126 .filter_map(|l| l.unit)
127 .map(|u| u.to_string())
128 .collect();
129 return Err(format!(
130 "no TOKENS_LIMIT bucket carries a known unit code (saw {})",
131 seen.join(", ")
132 ));
133 }
134 Ok((session, weekly))
135}
136
137fn classify_time_bucket(limits: &[LimitEntry]) -> Result<Option<&LimitEntry>, String> {
138 let mut matching = limits.iter().filter(|l| l.kind == "TIME_LIMIT");
139 let first = matching.next();
140 if matching.next().is_some() {
141 return Err("two TIME_LIMIT buckets are present".into());
142 }
143 Ok(first)
144}
145
146fn unique_by_unit<'a>(
147 tokens: &[&'a LimitEntry],
148 code: i64,
149) -> Result<Option<&'a LimitEntry>, String> {
150 let mut matching = tokens.iter().filter(|l| l.unit == Some(code));
151 let first = matching.next().copied();
152 if matching.next().is_some() {
153 return Err(format!("two TOKENS_LIMIT buckets carry unit {code}"));
154 }
155 Ok(first)
156}
157
158#[derive(Debug, Default, Clone, Deserialize)]
159#[serde(default)]
160pub struct MonitorData {
161 pub limits: Vec<LimitEntry>,
162 pub level: String,
163}
164
165#[derive(Debug, Default, Clone, Deserialize)]
166#[serde(default)]
167pub struct LimitEntry {
168 #[serde(rename = "type")]
169 pub kind: String,
170 #[serde(default, deserialize_with = "de_percent_opt")]
171 pub percentage: Option<f64>,
172 #[serde(rename = "nextResetTime", default, deserialize_with = "de_opt_ms")]
174 pub next_reset_time: Option<i64>,
175 pub unit: Option<i64>,
176 pub number: Option<i64>,
177}
178
179fn de_opt_ms<'de, D>(d: D) -> Result<Option<i64>, D::Error>
180where
181 D: serde::Deserializer<'de>,
182{
183 let v = serde_json::Value::deserialize(d)?;
184 match v {
185 serde_json::Value::Null => Ok(None),
186 serde_json::Value::Number(n) => {
187 let millis = if let Some(i) = n.as_i64() {
188 i
189 } else if let Some(f) = n.as_f64()
190 && f.is_finite()
191 && f.fract() == 0.0
192 && f.abs() <= (1_u64 << 53) as f64
193 {
194 f as i64
195 } else {
196 return Err(serde::de::Error::custom(
197 "nextResetTime must be an integer in range",
198 ));
199 };
200 match millis {
201 0 => Ok(None),
202 1.. => Ok(Some(millis)),
203 _ => Err(serde::de::Error::custom("nextResetTime cannot be negative")),
204 }
205 }
206 other => Err(serde::de::Error::custom(format!(
207 "nextResetTime must be an integer or null, got {other:?}"
208 ))),
209 }
210}
211
212fn de_percent_opt<'de, D>(d: D) -> Result<Option<f64>, D::Error>
213where
214 D: serde::Deserializer<'de>,
215{
216 let value = Option::<f64>::deserialize(d)?;
217 value
218 .map(|pct| {
219 if pct.is_finite() && (0.0..=101.0).contains(&pct) {
220 Ok(pct)
221 } else {
222 Err(serde::de::Error::custom(format!(
223 "percentage {pct} outside 0..=100"
224 )))
225 }
226 })
227 .transpose()
228}
229
230impl Envelope {
231 pub fn into_snapshot(self, config_plan_tier: Option<&str>) -> ZaiSnapshot {
234 let data = self.data.unwrap_or_default();
235 let (session, weekly) = classify_token_buckets(&data.limits).unwrap_or((None, None));
238 let session = session.and_then(|l| to_window(l, chrono::Duration::hours(5)));
239 let weekly = weekly.and_then(|l| to_window(l, chrono::Duration::days(7)));
240 let mcp = classify_time_bucket(&data.limits)
241 .ok()
242 .flatten()
243 .and_then(|l| to_window(l, chrono::Duration::days(30)));
244
245 let level = if !data.level.is_empty() {
247 data.level
248 } else {
249 config_plan_tier.unwrap_or("unknown").to_string()
250 };
251 let plan = format!("GLM Coding {}", capitalize(&level));
252
253 ZaiSnapshot {
254 plan,
255 session,
256 weekly,
257 mcp,
258 }
259 }
260}
261
262fn to_window(l: &LimitEntry, dur: chrono::Duration) -> Option<UsageWindow> {
263 let utilization_pct = l.percentage?.round().clamp(0.0, 100.0) as i32;
264 let resets_at = l
265 .next_reset_time
266 .and_then(chrono::DateTime::<chrono::Utc>::from_timestamp_millis);
267 Some(UsageWindow {
268 utilization_pct,
269 resets_at,
270 window_duration: dur,
271 })
272}
273
274fn capitalize(s: &str) -> String {
275 let mut chars = s.chars();
276 match chars.next() {
277 Some(c) => {
278 let mut out = String::with_capacity(s.len());
279 for u in c.to_uppercase() {
280 out.push(u);
281 }
282 out.push_str(chars.as_str());
283 out
284 }
285 None => String::new(),
286 }
287}
288
289#[cfg(test)]
290mod tests {
291 use super::*;
292
293 const REAL_BODY: &str = r#"{"code":200,"msg":"Operation successful","data":{
294 "limits":[
295 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":0},
296 {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":0,"nextResetTime":1779792169974},
297 {"type":"TIME_LIMIT","unit":5,"number":1,"usage":1000,"currentValue":0,"remaining":1000,"percentage":0,"nextResetTime":1779964969979,
298 "usageDetails":[{"modelCode":"search-prime","usage":0}]}
299 ],
300 "level":"pro"
301 },"success":true}"#;
302
303 #[test]
304 fn parses_real_response_shape() {
305 let env: Envelope = serde_json::from_str(REAL_BODY).unwrap();
306 let snap = env.into_snapshot(None);
307 assert_eq!(snap.plan, "GLM Coding Pro");
308 assert!(snap.session.is_some());
309 assert!(snap.weekly.is_some());
310 assert!(snap.mcp.is_some());
311 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 0);
312 assert!(snap.weekly.as_ref().unwrap().resets_at.is_some());
313 }
314
315 #[test]
316 fn missing_data_yields_neutral_snapshot() {
317 let env: Envelope = serde_json::from_str(r#"{"code":500,"success":false}"#).unwrap();
318 let snap = env.into_snapshot(Some("lite"));
319 assert_eq!(snap.plan, "GLM Coding Lite");
320 assert!(snap.session.is_none());
321 }
322
323 #[test]
324 fn percentage_with_float_rounds() {
325 let body = r#"{"data":{"limits":[
326 {"type":"TOKENS_LIMIT","unit":3,"percentage":42.7}
327 ],"level":"max"},"success":true}"#;
328 let env: Envelope = serde_json::from_str(body).unwrap();
329 let snap = env.into_snapshot(None);
330 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
331 }
332
333 #[test]
334 fn benign_percentage_overshoot_clamps_to_hundred() {
335 let body = r#"{"data":{"limits":[
336 {"type":"TOKENS_LIMIT","unit":3,"percentage":100.6}
337 ]},"success":true}"#;
338 let env: Envelope = serde_json::from_str(body).unwrap();
339 let snap = env.into_snapshot(None);
340 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 100);
341 }
342
343 #[test]
344 fn only_time_limit_means_no_session_or_weekly() {
345 let body = r#"{"data":{"limits":[
346 {"type":"TIME_LIMIT","percentage":12}
347 ]},"success":true}"#;
348 let env: Envelope = serde_json::from_str(body).unwrap();
349 let snap = env.into_snapshot(None);
350 assert!(snap.session.is_none());
351 assert!(snap.weekly.is_none());
352 assert!(snap.mcp.is_some());
353 }
354
355 #[test]
356 fn config_plan_tier_used_when_level_empty() {
357 let body = r#"{"data":{"limits":[],"level":""},"success":true}"#;
358 let env: Envelope = serde_json::from_str(body).unwrap();
359 let snap = env.into_snapshot(Some("max"));
360 assert_eq!(snap.plan, "GLM Coding Max");
361 }
362
363 #[test]
366 fn buckets_are_identified_by_unit_not_by_position() {
367 let body = r#"{"data":{"limits":[
368 {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":15,"nextResetTime":1779792169974},
369 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":42}
370 ],"level":"pro"},"success":true}"#;
371 let env: Envelope = serde_json::from_str(body).unwrap();
372 env.check_ok().unwrap();
373 let snap = env.into_snapshot(None);
374 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 42);
375 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 15);
376 assert!(snap.weekly.as_ref().unwrap().resets_at.is_some());
377 assert!(snap.session.as_ref().unwrap().resets_at.is_none());
378 }
379
380 #[test]
382 fn unknown_extra_bucket_is_dropped_not_shown_as_session() {
383 let body = r#"{"data":{"limits":[
384 {"type":"TOKENS_LIMIT","unit":4,"number":1,"percentage":99},
385 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":42},
386 {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":15}
387 ],"level":"pro"},"success":true}"#;
388 let env: Envelope = serde_json::from_str(body).unwrap();
389 env.check_ok().unwrap();
390 let snap = env.into_snapshot(None);
391 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 42);
392 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 15);
393 }
394
395 #[test]
396 fn duplicate_unit_is_an_error_not_a_coin_flip() {
397 let body = r#"{"data":{"limits":[
398 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":42},
399 {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":7}
400 ],"level":"pro"},"success":true}"#;
401 let env: Envelope = serde_json::from_str(body).unwrap();
402 let err = env.check_ok().unwrap_err().to_string();
403 assert!(err.contains("unit 3"), "unhelpful error: {err}");
404 let snap = env.into_snapshot(None);
406 assert!(snap.session.is_none());
407 assert!(snap.weekly.is_none());
408 }
409
410 #[test]
411 fn all_unknown_units_is_an_error() {
412 let body = r#"{"data":{"limits":[
413 {"type":"TOKENS_LIMIT","unit":4,"number":1,"percentage":42},
414 {"type":"TOKENS_LIMIT","unit":7,"number":1,"percentage":15}
415 ],"level":"pro"},"success":true}"#;
416 let env: Envelope = serde_json::from_str(body).unwrap();
417 let err = env.check_ok().unwrap_err().to_string();
418 assert!(err.contains("4, 7"), "unhelpful error: {err}");
419 assert!(env.into_snapshot(None).session.is_none());
420 }
421
422 #[test]
425 fn unit_less_bucket_alongside_a_known_one_is_dropped() {
426 let body = r#"{"data":{"limits":[
427 {"type":"TOKENS_LIMIT","percentage":99},
428 {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":15}
429 ],"level":"pro"},"success":true}"#;
430 let env: Envelope = serde_json::from_str(body).unwrap();
431 env.check_ok().unwrap();
432 let snap = env.into_snapshot(None);
433 assert!(snap.session.is_none());
434 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 15);
435 }
436
437 #[test]
438 fn bodies_without_any_unit_are_rejected_not_guessed_by_position() {
439 let body = r#"{"data":{"limits":[
440 {"type":"TOKENS_LIMIT","percentage":10},
441 {"type":"TOKENS_LIMIT","percentage":20}
442 ],"level":"lite"},"success":true}"#;
443 let env: Envelope = serde_json::from_str(body).unwrap();
444 let err = env.check_ok().unwrap_err().to_string();
445 assert!(err.contains("no unit discriminator"), "{err}");
446 let snap = env.into_snapshot(None);
447 assert!(snap.session.is_none());
448 assert!(snap.weekly.is_none());
449 }
450
451 #[test]
452 fn a_named_bucket_without_percentage_is_rejected_not_zeroed() {
453 let body = r#"{"data":{"limits":[
454 {"type":"TOKENS_LIMIT","unit":3,"number":5}
455 ]},"success":true}"#;
456 let env: Envelope = serde_json::from_str(body).unwrap();
457 let err = env.check_ok().unwrap_err().to_string();
458 assert!(err.contains("session limit carried no percentage"), "{err}");
459 assert!(env.into_snapshot(None).session.is_none());
460 }
461
462 #[test]
463 fn duplicate_time_limit_is_rejected_not_selected_by_position() {
464 let body = r#"{"data":{"limits":[
465 {"type":"TIME_LIMIT","unit":5,"percentage":10},
466 {"type":"TIME_LIMIT","unit":5,"percentage":20}
467 ]},"success":true}"#;
468 let env: Envelope = serde_json::from_str(body).unwrap();
469 let err = env.check_ok().unwrap_err().to_string();
470 assert!(err.contains("two TIME_LIMIT"), "{err}");
471 assert!(env.into_snapshot(None).mcp.is_none());
472 }
473
474 #[test]
475 fn invalid_percentage_and_reset_values_are_schema_drift() {
476 for percentage in ["-1", "101.5", "150"] {
477 let body = format!(
478 r#"{{"data":{{"limits":[{{"type":"TOKENS_LIMIT","unit":3,"percentage":{percentage}}}]}},"success":true}}"#
479 );
480 assert!(serde_json::from_str::<Envelope>(&body).is_err(), "{body}");
481 }
482 for reset in ["-1", "1.5", "true", r#""later""#] {
483 let body = format!(
484 r#"{{"data":{{"limits":[{{"type":"TOKENS_LIMIT","unit":3,"percentage":0,"nextResetTime":{reset}}}]}},"success":true}}"#
485 );
486 assert!(serde_json::from_str::<Envelope>(&body).is_err(), "{body}");
487 }
488 }
489
490 #[test]
491 fn check_ok_accepts_the_real_response_shape() {
492 let env: Envelope = serde_json::from_str(REAL_BODY).unwrap();
493 env.check_ok().unwrap();
494 }
495
496 #[test]
497 fn reset_time_zero_or_null_becomes_none() {
498 let body = r#"{"data":{"limits":[
499 {"type":"TOKENS_LIMIT","unit":3,"percentage":0,"nextResetTime":null}
500 ]},"success":true}"#;
501 let env: Envelope = serde_json::from_str(body).unwrap();
502 let snap = env.into_snapshot(None);
503 assert!(snap.session.as_ref().unwrap().resets_at.is_none());
504 }
505}