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 don't have a documented mapping, so we classify
27//! limits by **position + type**: the first TOKENS_LIMIT entry is treated as
28//! the session bucket, the second as weekly, and the TIME_LIMIT entry as the
29//! monthly MCP tool ceiling. When the shape drifts the smoke test will catch
30//! it and we can update this mapping.
31
32use serde::Deserialize;
33
34use crate::usage::{UsageWindow, ZaiSnapshot};
35
36#[derive(Debug, Clone, Deserialize)]
37pub struct Envelope {
38    #[serde(default)]
39    pub code: i64,
40    #[serde(default)]
41    pub data: Option<MonitorData>,
42    #[serde(default)]
43    pub success: bool,
44    #[serde(default)]
45    pub msg: String,
46}
47
48#[derive(Debug, Default, Clone, Deserialize)]
49#[serde(default)]
50pub struct MonitorData {
51    pub limits: Vec<LimitEntry>,
52    pub level: String,
53}
54
55#[derive(Debug, Default, Clone, Deserialize)]
56#[serde(default)]
57pub struct LimitEntry {
58    #[serde(rename = "type")]
59    pub kind: String,
60    pub percentage: f64,
61    /// Unix milliseconds — `null` / `0` / missing → None.
62    #[serde(rename = "nextResetTime", default, deserialize_with = "de_opt_ms")]
63    pub next_reset_time: Option<i64>,
64    pub unit: Option<i64>,
65    pub number: Option<i64>,
66}
67
68fn de_opt_ms<'de, D>(d: D) -> Result<Option<i64>, D::Error>
69where
70    D: serde::Deserializer<'de>,
71{
72    let v = serde_json::Value::deserialize(d)?;
73    Ok(match v {
74        serde_json::Value::Null => None,
75        serde_json::Value::Number(n) => n.as_i64().or_else(|| n.as_f64().map(|f| f as i64)),
76        _ => None,
77    })
78}
79
80impl Envelope {
81    /// Project the envelope into the canonical [`ZaiSnapshot`]. Returns a
82    /// snapshot with all windows `None` when `data` is missing.
83    pub fn into_snapshot(self, config_plan_tier: Option<&str>) -> ZaiSnapshot {
84        let data = self.data.unwrap_or_default();
85        let mut tokens_iter = data.limits.iter().filter(|l| l.kind == "TOKENS_LIMIT");
86        let session = tokens_iter
87            .next()
88            .map(|l| to_window(l, chrono::Duration::hours(5)));
89        let weekly = tokens_iter
90            .next()
91            .map(|l| to_window(l, chrono::Duration::days(7)));
92        let mcp = data
93            .limits
94            .iter()
95            .find(|l| l.kind == "TIME_LIMIT")
96            .map(|l| to_window(l, chrono::Duration::days(30)));
97
98        // Prefer the response's `level` field, then any config-provided tier.
99        let level = if !data.level.is_empty() {
100            data.level
101        } else {
102            config_plan_tier.unwrap_or("unknown").to_string()
103        };
104        let plan = format!("GLM Coding {}", capitalize(&level));
105
106        ZaiSnapshot {
107            plan,
108            session,
109            weekly,
110            mcp,
111        }
112    }
113}
114
115fn to_window(l: &LimitEntry, dur: chrono::Duration) -> UsageWindow {
116    let utilization_pct = l.percentage.round().clamp(0.0, 100.0) as i32;
117    let resets_at = l
118        .next_reset_time
119        .and_then(chrono::DateTime::<chrono::Utc>::from_timestamp_millis);
120    UsageWindow {
121        utilization_pct,
122        resets_at,
123        window_duration: dur,
124    }
125}
126
127fn capitalize(s: &str) -> String {
128    let mut chars = s.chars();
129    match chars.next() {
130        Some(c) => {
131            let mut out = String::with_capacity(s.len());
132            for u in c.to_uppercase() {
133                out.push(u);
134            }
135            out.push_str(chars.as_str());
136            out
137        }
138        None => String::new(),
139    }
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    const REAL_BODY: &str = r#"{"code":200,"msg":"Operation successful","data":{
147        "limits":[
148            {"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":0},
149            {"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":0,"nextResetTime":1779792169974},
150            {"type":"TIME_LIMIT","unit":5,"number":1,"usage":1000,"currentValue":0,"remaining":1000,"percentage":0,"nextResetTime":1779964969979,
151             "usageDetails":[{"modelCode":"search-prime","usage":0}]}
152        ],
153        "level":"pro"
154    },"success":true}"#;
155
156    #[test]
157    fn parses_real_response_shape() {
158        let env: Envelope = serde_json::from_str(REAL_BODY).unwrap();
159        let snap = env.into_snapshot(None);
160        assert_eq!(snap.plan, "GLM Coding Pro");
161        assert!(snap.session.is_some());
162        assert!(snap.weekly.is_some());
163        assert!(snap.mcp.is_some());
164        assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 0);
165        assert!(snap.weekly.as_ref().unwrap().resets_at.is_some());
166    }
167
168    #[test]
169    fn missing_data_yields_neutral_snapshot() {
170        let env: Envelope = serde_json::from_str(r#"{"code":500,"success":false}"#).unwrap();
171        let snap = env.into_snapshot(Some("lite"));
172        assert_eq!(snap.plan, "GLM Coding Lite");
173        assert!(snap.session.is_none());
174    }
175
176    #[test]
177    fn percentage_with_float_rounds() {
178        let body = r#"{"data":{"limits":[
179            {"type":"TOKENS_LIMIT","percentage":42.7}
180        ],"level":"max"},"success":true}"#;
181        let env: Envelope = serde_json::from_str(body).unwrap();
182        let snap = env.into_snapshot(None);
183        assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
184    }
185
186    #[test]
187    fn percentage_clamps_to_hundred() {
188        let body = r#"{"data":{"limits":[
189            {"type":"TOKENS_LIMIT","percentage":150}
190        ]},"success":true}"#;
191        let env: Envelope = serde_json::from_str(body).unwrap();
192        let snap = env.into_snapshot(None);
193        assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 100);
194    }
195
196    #[test]
197    fn only_time_limit_means_no_session_or_weekly() {
198        let body = r#"{"data":{"limits":[
199            {"type":"TIME_LIMIT","percentage":12}
200        ]},"success":true}"#;
201        let env: Envelope = serde_json::from_str(body).unwrap();
202        let snap = env.into_snapshot(None);
203        assert!(snap.session.is_none());
204        assert!(snap.weekly.is_none());
205        assert!(snap.mcp.is_some());
206    }
207
208    #[test]
209    fn config_plan_tier_used_when_level_empty() {
210        let body = r#"{"data":{"limits":[],"level":""},"success":true}"#;
211        let env: Envelope = serde_json::from_str(body).unwrap();
212        let snap = env.into_snapshot(Some("max"));
213        assert_eq!(snap.plan, "GLM Coding Max");
214    }
215
216    #[test]
217    fn reset_time_zero_or_null_becomes_none() {
218        let body = r#"{"data":{"limits":[
219            {"type":"TOKENS_LIMIT","percentage":0,"nextResetTime":null}
220        ]},"success":true}"#;
221        let env: Envelope = serde_json::from_str(body).unwrap();
222        let snap = env.into_snapshot(None);
223        assert!(snap.session.as_ref().unwrap().resets_at.is_none());
224    }
225}
ai_usagebar/zai/types.rs

ai_usagebar/zai/
types.rs
