1use std::time::SystemTime;
12
13use async_trait::async_trait;
14use reqwest::header::{ACCEPT, AUTHORIZATION, HeaderMap, HeaderValue};
15
16use forge_primitives::usage::oauth::OauthUsageError;
17use forge_primitives::usage::zai::{QuotaLimitData, QuotaLimitEntry, QuotaLimitResponse};
18use forge_primitives::usage::{UsageSnapshot, UsageSourceKind, UsageWindow};
19
20use crate::helpers::{
21 BaseUrlCredential, MissingBase, OAUTH_TIMEOUT, base_url_credential, redacted_body,
22 redacted_message, system_time_from_epoch,
23};
24use crate::{AccountEnv, BillingModel, ProbeError, Provider, ProviderBackend, ProviderHost};
25
26pub struct Zai;
28
29#[async_trait]
30impl ProviderBackend for Zai {
31 fn token(&self) -> Provider {
32 Provider::Zai
33 }
34
35 fn billing(&self) -> BillingModel {
36 BillingModel::Windows
37 }
38
39 fn source(&self) -> UsageSourceKind {
40 UsageSourceKind::ZaiMonitor
41 }
42
43 async fn probe(
44 &self,
45 account: &AccountEnv<'_>,
46 host: &dyn ProviderHost,
47 ) -> Result<UsageSnapshot, ProbeError> {
48 match choose_mapper(base_url_credential(account.env)) {
49 Mapper::Monitor(credential) => {
50 let client = host.http_client(OAUTH_TIMEOUT).map_err(OauthUsageError::Network)?;
51 let payload = monitor_probe(&client, &credential.base_url, &credential.bearer)
52 .await
53 .map_err(ProbeError::Fetch)?;
54 snapshot_from_zai_quota(payload)
55 }
56 Mapper::MissingBase(missing) => Err(ProbeError::Unmappable(missing.to_string())),
57 }
58 }
59}
60
61fn choose_mapper(credential: Result<BaseUrlCredential, MissingBase>) -> Mapper {
67 match credential {
68 Ok(credential) => Mapper::Monitor(credential),
69 Err(missing) => Mapper::MissingBase(missing),
70 }
71}
72
73#[derive(Debug, PartialEq, Eq)]
74enum Mapper {
75 Monitor(BaseUrlCredential),
78 MissingBase(MissingBase),
79}
80
81fn monitor_host(base_url: &str) -> &str {
86 let trimmed = base_url.trim().trim_end_matches('/');
87 let after = trimmed.split_once("://").map_or(trimmed, |(_, rest)| rest);
88 let host = after.split('/').next().unwrap_or_default();
89 if host.is_empty() {
90 return trimmed;
91 }
92 &trimmed[..trimmed.len() - after.len() + host.len()]
93}
94
95fn monitor_url(base_url: &str) -> String {
96 format!("{}/api/monitor/usage/quota/limit", monitor_host(base_url))
97}
98
99fn monitor_headers(key: &str) -> Result<HeaderMap, OauthUsageError> {
100 let mut headers = HeaderMap::new();
101 headers.insert(ACCEPT, HeaderValue::from_static("application/json"));
102 let auth = HeaderValue::from_str(key)
103 .map_err(|error| OauthUsageError::Network(format!("bad auth header: {error}")))?;
104 headers.insert(AUTHORIZATION, auth);
105 Ok(headers)
106}
107
108async fn monitor_probe(
113 client: &reqwest::Client,
114 base_url: &str,
115 key: &str,
116) -> Result<QuotaLimitData, OauthUsageError> {
117 let response = client
118 .get(monitor_url(base_url))
119 .headers(monitor_headers(key)?)
120 .send()
121 .await
122 .map_err(|error| OauthUsageError::Network(error.to_string()))?;
123
124 let status = response.status().as_u16();
125 let body = response
126 .bytes()
127 .await
128 .map_err(|error| OauthUsageError::Network(format!("body read: {error}")))?;
129
130 if status == 200 {
131 tracing::trace!(
132 target: "forge_gateway::zai",
133 event_name = "zai_monitor_response",
134 status,
135 outcome = "ok",
136 body_bytes = body.len(),
137 );
138 } else {
139 tracing::warn!(
140 target: "forge_gateway::zai",
141 event_name = "zai_monitor_response",
142 status,
143 outcome = "non_ok",
144 body_suffix = %redacted_body(&body, key),
145 );
146 }
147
148 match status {
149 200 => quota_from_body(&body, key),
150 401 | 403 => Err(OauthUsageError::Unauthorized(status)),
151 429 => Err(OauthUsageError::RateLimited { retry_after: None, snapshot: None }),
152 _ => Err(OauthUsageError::HttpStatus(status, redacted_body(&body, key))),
153 }
154}
155
156fn quota_from_body(body: &[u8], key: &str) -> Result<QuotaLimitData, OauthUsageError> {
167 if String::from_utf8_lossy(body).trim().is_empty() {
168 return Err(OauthUsageError::Decode(
169 "Z.ai monitor answered 200 with an empty body".to_owned(),
170 ));
171 }
172 let envelope: QuotaLimitResponse = serde_json::from_slice(body).map_err(|error| {
173 OauthUsageError::Decode(redacted_message(
174 &format!("Z.ai monitor body did not decode: {error}"),
175 key,
176 ))
177 })?;
178 let msg = redacted_message(envelope.msg.as_deref().unwrap_or("no msg"), key);
179 if !envelope.success {
180 if envelope.code == Some(401) {
181 return Err(OauthUsageError::Unauthorized(401));
182 }
183 return Err(OauthUsageError::Decode(format!(
184 "Z.ai monitor reported failure (code {}): {msg}",
185 envelope.code.unwrap_or(0),
186 )));
187 }
188 if envelope.code != Some(200) {
189 return Err(OauthUsageError::Decode(format!(
190 "Z.ai monitor reported code {}: {msg}",
191 envelope.code.unwrap_or(0),
192 )));
193 }
194 envelope
195 .data
196 .ok_or_else(|| OauthUsageError::Decode("Z.ai monitor 200 carried no data".to_owned()))
197}
198
199fn snapshot_from_zai_quota(payload: QuotaLimitData) -> Result<UsageSnapshot, ProbeError> {
214 let mut five_hour = None;
215 let mut seven_day = None;
216 for entry in payload.limits {
217 if entry.kind.as_deref() != Some("CREDIT_LIMIT") {
218 continue;
219 }
220 match entry.unit {
221 Some(3) => five_hour = window_from_entry(&entry),
222 Some(6) => seven_day = window_from_entry(&entry),
223 _ => {}
224 }
225 }
226 if five_hour.is_none() && seven_day.is_none() {
227 return Err(ProbeError::Unmappable(
228 "Z.ai quota response carried no CREDIT_LIMIT window entries.".to_owned(),
229 ));
230 }
231 Ok(UsageSnapshot {
232 source: UsageSourceKind::ZaiMonitor,
233 fetched_at: SystemTime::now(),
234 five_hour,
235 seven_day,
236 seven_day_opus: None,
237 seven_day_sonnet: None,
238 extra_usage: None,
239 spend: None,
240 balance: None,
241 })
242}
243
244fn window_from_entry(entry: &QuotaLimitEntry) -> Option<UsageWindow> {
245 let usage = entry.usage?;
246 let remaining = entry.remaining?;
247 let utilization = if usage > 0.0 {
248 ((usage - remaining).max(0.0) / usage * 100.0).clamp(0.0, 100.0)
249 } else {
250 0.0
251 };
252 Some(UsageWindow {
253 utilization,
254 resets_at: entry.next_reset_time.and_then(system_time_from_epoch),
255 reset_description: None,
256 status: None,
257 })
258}
259
260#[cfg(test)]
261mod tests {
262 use std::collections::HashMap;
263 use std::time::Duration;
264
265 use async_trait::async_trait;
266
267 use super::*;
268
269 fn env_with_base(base: &str) -> HashMap<String, String> {
270 let mut env = HashMap::new();
271 env.insert("ANTHROPIC_BASE_URL".to_owned(), base.to_owned());
272 env.insert("ANTHROPIC_AUTH_TOKEN".to_owned(), "zai-key".to_owned());
273 env
274 }
275
276 fn green_envelope_body() -> &'static [u8] {
277 br#"{
278 "code": 200,
279 "msg": "success",
280 "data": {
281 "limits": [
282 {"type":"CREDIT_LIMIT","unit":3,"number":5,"usage":28000,
283 "remaining":27104,"percentage":3.2,"currentValue":0,
284 "nextResetTime":1757025600000},
285 {"type":"CREDIT_LIMIT","unit":6,"number":1,"usage":140000,
286 "remaining":139000,"percentage":0.71,"currentValue":0,
287 "nextResetTime":1757000000000}
288 ],
289 "level": "max"
290 },
291 "success": true
292 }"#
293 }
294
295 #[test]
296 fn zai_backend_is_windowed() {
297 assert_eq!(Zai.token(), Provider::Zai);
298 assert_eq!(Zai.billing(), BillingModel::Windows);
299 }
300
301 #[test]
307 fn a_base_url_credential_earns_the_monitor_mapper_and_a_missing_base_the_error_arm() {
308 let credential =
309 base_url_credential(&env_with_base("https://api.z.ai/api/anthropic")).expect("cred");
310 assert_eq!(
311 choose_mapper(Ok(credential.clone())),
312 Mapper::Monitor(credential),
313 "the monitor arm runs with the credential that earned it",
314 );
315 assert!(
316 matches!(choose_mapper(Err(MissingBase)), Mapper::MissingBase(_)),
317 "a missing base surfaces Unmappable, not a probe against a default host",
318 );
319 }
320
321 #[tokio::test]
324 async fn a_missing_base_is_unmappable_without_probing() {
325 let backend = Zai;
326 let account = AccountEnv { env: &HashMap::new() };
327 let result = backend.probe(&account, &UnreachableHost).await;
328 assert!(matches!(result, Err(ProbeError::Unmappable(_))), "got {result:?}");
329 }
330
331 #[test]
332 fn monitor_url_derives_the_host_root() {
333 assert_eq!(
334 monitor_url("https://api.z.ai/api/anthropic"),
335 "https://api.z.ai/api/monitor/usage/quota/limit",
336 );
337 assert_eq!(
338 monitor_url("https://api.z.ai/api/anthropic/"),
339 "https://api.z.ai/api/monitor/usage/quota/limit",
340 "trailing slash trimmed so base and base/ behave identically",
341 );
342 }
343
344 #[test]
345 fn monitor_headers_send_the_key_raw_without_a_bearer_prefix() {
346 let headers = monitor_headers("zai-key").expect("headers");
347 let auth = headers
348 .get(AUTHORIZATION)
349 .expect("auth header set")
350 .to_str()
351 .expect("ascii header value");
352 assert_eq!(auth, "zai-key", "the key goes out raw; no Bearer prefix");
353 }
354
355 #[test]
358 fn quota_from_body_maps_the_fresh_account_envelope() {
359 let body = br#"{
360 "code": 200,
361 "msg": "success",
362 "data": {
363 "limits": [
364 {"type":"CREDIT_LIMIT","unit":3,"number":5,"usage":28000,
365 "remaining":28000,"percentage":0,"currentValue":0},
366 {"type":"CREDIT_LIMIT","unit":6,"number":1,"usage":140000,
367 "remaining":140000,"percentage":0,"currentValue":0,
368 "nextResetTime":1757000000000}
369 ],
370 "level": "max"
371 },
372 "success": true
373 }"#;
374 let data = quota_from_body(body, "test-key").expect("a green envelope parses");
375 assert_eq!(data.limits.len(), 2);
376 assert_eq!(data.level.as_deref(), Some("max"));
377 }
378
379 #[test]
380 fn quota_from_body_fails_on_wrong_key_inside_a_200() {
381 let err = quota_from_body(
382 br#"{"code":401,"msg":"token expired or incorrect","success":false}"#,
383 "test-key",
384 )
385 .expect_err("wrong key fails");
386 assert!(
387 matches!(err, OauthUsageError::Unauthorized(401)),
388 "a body-level 401 must classify as unauthorized, got {err:?}",
389 );
390 }
391
392 #[test]
393 fn quota_from_body_fails_on_wrong_path_inside_a_200() {
394 let err =
395 quota_from_body(br#"{"code":500,"msg":"404 NOT_FOUND","success":false}"#, "test-key")
396 .expect_err("a wrong path fails");
397 assert!(err.to_string().contains("404 NOT_FOUND"), "the msg reaches the log: {err}");
398 }
399
400 #[test]
403 fn quota_from_body_masks_the_key_in_its_msg() {
404 let err = quota_from_body(
405 br#"{"code":500,"msg":"refused Authorization glm-0123456789abcdef","success":false}"#,
406 "glm-0123456789abcdef",
407 )
408 .expect_err("the refusal fails");
409 assert!(!err.to_string().contains("0123456789abcdef"), "the key never prints: {err}");
410 assert!(err.to_string().contains("refused"), "and the host's own words stay: {err}");
411 }
412
413 #[test]
416 fn quota_from_body_masks_the_key_a_serde_message_quotes() {
417 let err = quota_from_body(
418 br#"{"code":"glm-0123456789abcdef","msg":"x","success":true}"#,
419 "glm-0123456789abcdef",
420 )
421 .expect_err("a code of the wrong type fails to decode");
422 assert!(!err.to_string().contains("0123456789abcdef"), "the value never prints: {err}");
423 assert!(err.to_string().contains("invalid type"), "and the reason stays: {err}");
424 }
425
426 #[test]
429 fn quota_from_body_fails_on_empty_body() {
430 let err = quota_from_body(b"", "test-key").expect_err("empty body fails");
431 assert!(err.to_string().contains("empty"), "got {err}");
432 }
433
434 #[test]
437 fn quota_from_body_requires_the_success_flag() {
438 let err = quota_from_body(
439 br#"{"code":200,"msg":"success","data":{"limits":[],"level":"max"}}"#,
440 "test-key",
441 )
442 .expect_err("no success flag fails");
443 assert!(matches!(err, OauthUsageError::Decode(_)), "got {err:?}");
444 }
445
446 #[tokio::test]
454 async fn a_200_envelope_probes_through_to_a_windowed_snapshot() {
455 let body = green_envelope_body();
456 let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
457 let addr = listener.local_addr().expect("addr");
458 std::thread::spawn(move || {
459 use std::io::{Read, Write as _};
460 let Ok((mut sock, _)) = listener.accept() else { return };
461 let mut request = Vec::new();
465 let mut byte = [0u8; 1];
466 while !request.ends_with(b"\r\n\r\n") {
467 match sock.read(&mut byte) {
468 Ok(1) => request.push(byte[0]),
469 _ => break,
470 }
471 }
472 let request_text = String::from_utf8_lossy(&request).into_owned();
473 assert!(
474 request_text.contains("GET /api/monitor/usage/quota/limit HTTP/1.1"),
475 "the request must hit the host-root monitor path, got: {}",
476 request_text.lines().next().unwrap_or_default()
477 );
478 assert!(
479 request_text.contains("authorization: zai-key"),
480 "the env key goes out raw, no Bearer prefix, got: {request_text}",
481 );
482 let header = format!(
483 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
484 body.len()
485 );
486 let _ = sock.write_all(header.as_bytes());
487 let _ = sock.write_all(body);
488 let _ = sock.shutdown(std::net::Shutdown::Both);
489 });
490 let env = env_with_base(&format!("http://{addr}/api/anthropic"));
491 let account = AccountEnv { env: &env };
492 let snapshot = Zai.probe(&account, &LocalHost).await.expect("snapshot");
493 assert_eq!(snapshot.source, UsageSourceKind::ZaiMonitor);
494 let five = snapshot.five_hour.expect("5h window");
495 assert!(
496 (five.utilization - 3.2).abs() < 1e-9,
497 "896 of 28000 credits is 3.2%, got {}",
498 five.utilization,
499 );
500 }
501
502 struct LocalHost;
503
504 #[async_trait]
505 impl ProviderHost for LocalHost {
506 fn http_client(&self, timeout: Duration) -> Result<reqwest::Client, String> {
507 reqwest::Client::builder().timeout(timeout).build().map_err(|e| e.to_string())
508 }
509
510 fn streaming_http_client(
511 &self,
512 _connect_timeout: Duration,
513 _idle_timeout: Duration,
514 ) -> Result<reqwest::Client, String> {
515 unreachable!("no probe streams a response")
516 }
517
518 async fn user_agent(&self) -> Result<String, String> {
519 unreachable!("the zai monitor probe sends no User-Agent header")
520 }
521 }
522
523 struct UnreachableHost;
524
525 #[async_trait]
526 impl ProviderHost for UnreachableHost {
527 fn http_client(&self, _timeout: Duration) -> Result<reqwest::Client, String> {
528 unreachable!("the probe must not build a client for a missing base url")
529 }
530
531 fn streaming_http_client(
532 &self,
533 _connect_timeout: Duration,
534 _idle_timeout: Duration,
535 ) -> Result<reqwest::Client, String> {
536 unreachable!("the probe must not build a client for a missing base url")
537 }
538
539 async fn user_agent(&self) -> Result<String, String> {
540 unreachable!("the probe must not resolve a UA for a missing base url")
541 }
542 }
543
544 #[test]
549 fn zai_quota_maps_credit_windows_to_percentages() {
550 let payload: QuotaLimitData = serde_json::from_str(
551 r#"{
552 "limits": [
553 {"type":"CREDIT_LIMIT","unit":3,"number":5,"usage":28000,
554 "remaining":27104,"percentage":3.2,"currentValue":0,
555 "nextResetTime":1757025600000},
556 {"type":"CREDIT_LIMIT","unit":6,"number":1,"usage":140000,
557 "remaining":139000,"percentage":0.71,"currentValue":0,
558 "nextResetTime":1757000000000}
559 ],
560 "level": "max"
561 }"#,
562 )
563 .expect("decode");
564 let snapshot = snapshot_from_zai_quota(payload).expect("maps");
565 assert_eq!(snapshot.source, UsageSourceKind::ZaiMonitor);
566 let five = snapshot.five_hour.expect("5h window");
567 assert!(
568 (five.utilization - 3.2).abs() < 1e-9,
569 "896 of 28000 credits is 3.2%, got {}",
570 five.utilization,
571 );
572 assert_eq!(
573 five.resets_at,
574 Some(std::time::UNIX_EPOCH + std::time::Duration::from_secs(1_757_025_600)),
576 "epoch-ms nextResetTime becomes the reset moment",
577 );
578 let weekly = snapshot.seven_day.expect("weekly window");
579 assert!(
580 (weekly.utilization - 1000.0 / 1400.0).abs() < 1e-9,
581 "1000 of 140000 credits, got {}",
582 weekly.utilization,
583 );
584 assert!(snapshot.spend.is_none(), "a subscription carries no per-key spend");
585 }
586
587 #[test]
591 fn zai_quota_maps_a_fresh_account_with_no_five_hour_reset() {
592 let payload: QuotaLimitData = serde_json::from_str(
593 r#"{
594 "limits": [
595 {"type":"CREDIT_LIMIT","unit":3,"number":5,"usage":28000,
596 "remaining":28000},
597 {"type":"CREDIT_LIMIT","unit":6,"number":1,"usage":140000,
598 "remaining":140000,"nextResetTime":1757000000000}
599 ],
600 "level": "max"
601 }"#,
602 )
603 .expect("decode");
604 let snapshot = snapshot_from_zai_quota(payload).expect("maps");
605 let five = snapshot.five_hour.expect("5h window");
606 assert!(five.utilization.abs() < f64::EPSILON, "fresh account has consumed nothing");
607 assert_eq!(five.resets_at, None, "no nextResetTime means no reset moment yet");
608 }
609
610 #[test]
615 fn zai_quota_skips_an_entry_without_remaining() {
616 let payload: QuotaLimitData = serde_json::from_str(
617 r#"{
618 "limits": [
619 {"type":"CREDIT_LIMIT","unit":3,"number":5,"usage":28000},
620 {"type":"CREDIT_LIMIT","unit":6,"number":1,"usage":140000,
621 "remaining":139000,"nextResetTime":1757000000000}
622 ],
623 "level": "max"
624 }"#,
625 )
626 .expect("decode");
627 let snapshot = snapshot_from_zai_quota(payload).expect("the weekly entry maps");
628 assert!(
629 snapshot.five_hour.is_none(),
630 "an entry with no remaining must not map to 100% utilization",
631 );
632 assert!(snapshot.seven_day.is_some(), "its present sibling still maps");
633 }
634
635 #[test]
639 fn zai_quota_with_no_mappable_entries_is_an_error_not_zero() {
640 let empty: QuotaLimitData =
641 serde_json::from_str(r#"{"limits":[],"level":"max"}"#).expect("decode");
642 assert!(snapshot_from_zai_quota(empty).is_err(), "no windows must not read as a zero bill");
643 let foreign_kind: QuotaLimitData = serde_json::from_str(
644 r#"{"limits":[{"type":"TOKENS_LIMIT","unit":3,"number":5,"usage":1,"remaining":1}]}"#,
645 )
646 .expect("decode");
647 assert!(
648 snapshot_from_zai_quota(foreign_kind).is_err(),
649 "a non-CREDIT_LIMIT entry must not become a window",
650 );
651 }
652}