Skip to main content

ai_usagebar/zai/
types.rs

1//! Wire types for the undocumented Z.AI / BigModel monitor endpoint
2//! `https://api.z.ai/api/monitor/usage/quota/limit`.
3//!
4//! Real response shape (captured 2026-05-23):
5//!
6//! ```json
7//! {
8//!   "code": 200,
9//!   "msg": "Operation successful",
10//!   "data": {
11//!     "limits": [
12//!       {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":0},
13//!       {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":0,
14//!        "nextResetTime":1779792169974},
15//!       {"type":"TIME_LIMIT","unit":5,"number":1,"usage":1000,
16//!        "currentValue":0,"remaining":1000,"percentage":0,
17//!        "nextResetTime":1779964969979,
18//!        "usageDetails":[{"modelCode":"search-prime","usage":0},...]}
19//!     ],
20//!     "level":"pro"
21//!   },
22//!   "success": true
23//! }
24//! ```
25//!
26//! The `unit`/`number` codes have no documented mapping, but `unit` is the one
27//! field that tells the two usage buckets apart independently of where
28//! they sit in the array: the 5h window carries `unit:3`, the 7d one `unit:6`.
29//! So the session/weekly split keys off `unit`, not off position — Z.AI is free
30//! to reorder `limits` or insert a bucket, and a positional split would silently
31//! swap the two windows or promote a stranger to "session". A layout we cannot
32//! name is an error via [`Envelope::check_ok`], never a guess. The TIME_LIMIT
33//! entry is the monthly MCP tool ceiling.
34
35use 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    /// Z.AI signals failure *inside* a 200 response: `success: false` with a
53    /// non-200 `code` and the reason in `msg`, and `data: null`. Without this
54    /// check such a body deserializes cleanly, overwrites a good cache, clears
55    /// the previous error, and renders as an unknown plan with empty windows —
56    /// indistinguishable from a real account with no usage.
57    ///
58    /// `code` is accepted when absent (0) or 200; anything else is a failure.
59    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        // A body whose token buckets we cannot name is drift, not usage: let it
77        // through and the widget would render one window's figure under the
78        // other's label, and cache it as if it were vouched for.
79        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
96/// `unit` codes of the two usage buckets in the captured response. The
97/// enum behind them is undocumented — `number` (5 and 1) is consistent with
98/// "5 hours" / "1 week", but we don't lean on that, so the window durations
99/// stay hardcoded and only the *identity* of each bucket comes from `unit`.
100const UNIT_SESSION: i64 = 3;
101const UNIT_WEEKLY: i64 = 6;
102
103/// The API renamed the bucket type from `TOKENS_LIMIT` to `CREDIT_LIMIT`;
104/// both spellings denote the same session/weekly windows, so either is
105/// accepted and they must never be counted as strangers of each other.
106fn 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
112/// Match the TOKENS_LIMIT/CREDIT_LIMIT entries to the (session, weekly)
113/// windows by `unit`.
114///
115/// Buckets carrying an unknown `unit` are dropped rather than shown under a
116/// label we can't justify, so Z.AI adding a third window is inert here. A
117/// non-empty set with no discriminator is drift, not a backwards-compatibility
118/// case: caches retain the raw `unit` field, so position would still be a guess.
119fn 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    /// Unix milliseconds — `null` / `0` / missing → None.
181    #[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    /// Project the envelope into the canonical [`ZaiSnapshot`]. Returns a
240    /// snapshot with all windows `None` when `data` is missing.
241    pub fn into_snapshot(self, config_plan_tier: Option<&str>) -> ZaiSnapshot {
242        let data = self.data.unwrap_or_default();
243        // On the fetch path `check_ok` has already turned an unnameable layout
244        // into an error; direct callers get empty windows for the same reason.
245        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        // Prefer the response's `level` field, then any config-provided tier.
254        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    /// The regression: session/weekly used to be "first TOKENS_LIMIT, second
372    /// TOKENS_LIMIT", so a reordered array swapped the two windows.
373    #[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    /// A third bucket must not be promoted to "session" just by leading the array.
389    #[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        // And the projection refuses to pick one rather than showing either.
413        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    /// A bucket whose `unit` went missing can't be named, so it is dropped —
431    /// never quietly slotted into whichever window is still free.
432    #[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    /// The API renamed TOKENS_LIMIT to CREDIT_LIMIT; both spellings must map
515    /// to the same session/weekly windows, keyed by `unit` as before.
516    #[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    /// A rename mid-rollout means one window may still arrive under the old
531    /// spelling while the other already uses the new one; both must be read.
532    #[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    /// Two spellings naming the *same* window are duplicates, not aliases to
546    /// be reconciled by kind — same rule as two TOKENS_LIMIT entries.
547    #[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}