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 TOKENS_LIMIT 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 TOKENS_LIMIT 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
103type TokenBuckets<'a> = (Option<&'a LimitEntry>, Option<&'a LimitEntry>);
104
105/// Match the TOKENS_LIMIT entries to the (session, weekly) windows by `unit`.
106///
107/// Buckets carrying an unknown `unit` are dropped rather than shown under a
108/// label we can't justify, so Z.AI adding a third window is inert here. A
109/// non-empty set with no discriminator is drift, not a backwards-compatibility
110/// case: caches retain the raw `unit` field, so position would still be a guess.
111fn 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    /// Unix milliseconds — `null` / `0` / missing → None.
173    #[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    /// Project the envelope into the canonical [`ZaiSnapshot`]. Returns a
232    /// snapshot with all windows `None` when `data` is missing.
233    pub fn into_snapshot(self, config_plan_tier: Option<&str>) -> ZaiSnapshot {
234        let data = self.data.unwrap_or_default();
235        // On the fetch path `check_ok` has already turned an unnameable layout
236        // into an error; direct callers get empty windows for the same reason.
237        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        // Prefer the response's `level` field, then any config-provided tier.
246        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    /// The regression: session/weekly used to be "first TOKENS_LIMIT, second
364    /// TOKENS_LIMIT", so a reordered array swapped the two windows.
365    #[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    /// A third bucket must not be promoted to "session" just by leading the array.
381    #[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        // And the projection refuses to pick one rather than showing either.
405        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    /// A bucket whose `unit` went missing can't be named, so it is dropped —
423    /// never quietly slotted into whichever window is still free.
424    #[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}