1use std::time::Duration;
10
11use crate::cache::{Cache, acquire_lock_async};
12use crate::error::{AppError, Result};
13use crate::usage::{ModelStudioSnapshot, UsageWindow};
14use crate::vendor::{MAX_BODY_BYTES, read_body_capped};
15
16use super::creds::Credentials;
17use super::types::{
18 ConsoleRegion, ConsoleSite, FIVE_HOUR_WINDOW, MONTHLY_WINDOW, WEEKLY_WINDOW, form_body,
19 gateway_for, parse_response, usage_path,
20};
21
22const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
23const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
24
25#[derive(Debug, Clone)]
29pub struct Endpoints {
30 pub base: String,
31}
32
33impl Endpoints {
34 pub fn for_gateway(region: ConsoleRegion, site: ConsoleSite) -> Self {
35 Self {
36 base: format!("https://{}", gateway_for(region, site).host),
37 }
38 }
39
40 fn usage_url(&self, action: &str) -> String {
41 format!("{}{}", self.base, usage_path(action))
42 }
43}
44
45impl Default for Endpoints {
46 fn default() -> Self {
47 Self::for_gateway(ConsoleRegion::CnBeijing, ConsoleSite::Domestic)
48 }
49}
50
51pub type FetchOutcome = crate::outcome::Outcome<ModelStudioSnapshot>;
54
55pub async fn fetch_snapshot(
58 client: &reqwest::Client,
59 cfg: &crate::config::ModelStudioConfig,
60 cache: &Cache,
61 cache_ttl: Duration,
62) -> Result<FetchOutcome> {
63 let creds = super::resolve_credentials(cfg)?;
64 let endpoints = Endpoints::for_gateway(creds.region, creds.site);
65 fetch_snapshot_with(client, &creds, cache, &endpoints, cache_ttl).await
66}
67
68pub async fn fetch_snapshot_with(
70 client: &reqwest::Client,
71 creds: &Credentials,
72 cache: &Cache,
73 endpoints: &Endpoints,
74 cache_ttl: Duration,
75) -> Result<FetchOutcome> {
76 cache.ensure_dir()?;
77 let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
78
79 if let Some(bytes) = cache.fresh_payload(cache_ttl)?
82 && let Ok(outcome) = reuse_cache(&bytes, cache, false, &creds.fingerprint)
83 {
84 return Ok(outcome);
85 }
86
87 match usage_call(client, creds, endpoints).await {
88 Ok(snap) => {
89 let bytes = serde_json::to_vec(&snap_to_json(&snap, &creds.fingerprint))?;
90 cache.write_payload(&bytes)?;
91 Ok(crate::outcome::Outcome::fresh(snap))
92 }
93 Err(e) if e.is_transient() => fallback_silent(cache, &creds.fingerprint, e),
94 Err(e) => {
95 cache.mark_stale();
96 if let Some((code, msg)) = error_to_pair(&e) {
97 cache.write_last_error(code, &msg);
98 }
99 fallback_with_error(cache, &creds.fingerprint, e)
100 }
101 }
102}
103
104fn fallback_silent(cache: &Cache, fingerprint: &str, original: AppError) -> Result<FetchOutcome> {
105 crate::outcome::fallback(cache, None, original, |bytes| {
106 parse_cache_at(bytes, fingerprint)
107 })
108}
109
110fn fallback_with_error(
111 cache: &Cache,
112 fingerprint: &str,
113 original: AppError,
114) -> Result<FetchOutcome> {
115 let last_error = error_to_pair(&original);
116 crate::outcome::fallback(cache, last_error, original, |bytes| {
117 parse_cache_at(bytes, fingerprint)
118 })
119}
120
121fn error_to_pair(e: &AppError) -> Option<(u16, String)> {
125 match e {
126 AppError::Http { status, .. } if matches!(status, 401 | 403) => {
127 Some((*status, "Model Studio authentication failed".into()))
128 }
129 AppError::Http { status, body } => Some((*status, body.clone())),
130 AppError::Credentials(msg) => Some((0, msg.clone())),
131 e => Some((0, e.to_string())),
132 }
133}
134
135fn reuse_cache(
136 bytes: &[u8],
137 cache: &Cache,
138 stale: bool,
139 fingerprint: &str,
140) -> Result<FetchOutcome> {
141 let snap = parse_cache_at(bytes, fingerprint)?;
142 Ok(crate::outcome::Outcome::cached(snap, cache, stale))
143}
144
145async fn usage_call(
146 client: &reqwest::Client,
147 creds: &Credentials,
148 endpoints: &Endpoints,
149) -> Result<ModelStudioSnapshot> {
150 let action = gateway_for(creds.region, creds.site).action;
151 let resp = tokio::time::timeout(
152 HTTP_TIMEOUT,
153 client
154 .post(endpoints.usage_url(action))
155 .header("Authorization", format!("Bearer {}", creds.access_token))
156 .header("Content-Type", "application/x-www-form-urlencoded")
157 .body(form_body(creds.region))
158 .send(),
159 )
160 .await
161 .map_err(|_| AppError::Transport("modelstudio timeout: console gateway".into()))?
162 .map_err(|e| AppError::Transport(format!("modelstudio transport: {e}")))?;
163
164 let status = resp.status();
165 let bytes = read_body_capped(resp, MAX_BODY_BYTES).await?;
166 if !status.is_success() {
167 let body = if matches!(status.as_u16(), 401 | 403) {
170 "Model Studio authentication failed".into()
171 } else {
172 format!("Model Studio gateway returned HTTP {}", status.as_u16())
173 };
174 return Err(AppError::Http {
175 status: status.as_u16(),
176 body,
177 });
178 }
179 parse_response(&bytes)?.to_snapshot()
180}
181
182fn snap_to_json(snap: &ModelStudioSnapshot, fingerprint: &str) -> serde_json::Value {
183 let window = |w: &Option<UsageWindow>| match w {
184 Some(w) => serde_json::json!({
185 "pct": w.utilization_pct,
186 "reset_ms": w.resets_at.map(|t| t.timestamp_millis()),
187 }),
188 None => serde_json::Value::Null,
189 };
190 serde_json::json!({
191 "account": fingerprint,
192 "session": window(&snap.session),
193 "weekly": window(&snap.weekly),
194 "monthly": window(&snap.monthly),
195 })
196}
197
198fn parse_cache_at(bytes: &[u8], fingerprint: &str) -> Result<ModelStudioSnapshot> {
199 let v: serde_json::Value = serde_json::from_slice(bytes)?;
200 if v.get("account").and_then(serde_json::Value::as_str) != Some(fingerprint) {
201 return Err(AppError::Schema(
202 "modelstudio cache belongs to a different login; refetching".into(),
203 ));
204 }
205 let invalid = |field: &str| AppError::Schema(format!("modelstudio cache: invalid {field}"));
206 let window = |field: &str, duration| -> Result<Option<UsageWindow>> {
207 match v.get(field) {
208 None | Some(serde_json::Value::Null) => Ok(None),
209 Some(w) => {
210 let pct = w["pct"]
211 .as_i64()
212 .filter(|pct| (0..=100).contains(pct))
213 .ok_or_else(|| invalid(field))? as i32;
214 let resets_at = match w.get("reset_ms") {
215 None | Some(serde_json::Value::Null) => None,
216 Some(value) => {
217 let ms = value
218 .as_i64()
219 .filter(|ms| *ms >= 0)
220 .ok_or_else(|| invalid(field))?;
221 Some(
222 chrono::DateTime::from_timestamp_millis(ms)
223 .ok_or_else(|| invalid(field))?,
224 )
225 }
226 };
227 Ok(Some(UsageWindow {
228 utilization_pct: pct,
229 resets_at,
230 window_duration: duration,
231 }))
232 }
233 }
234 };
235 Ok(ModelStudioSnapshot {
236 session: window("session", FIVE_HOUR_WINDOW)?,
237 weekly: window("weekly", WEEKLY_WINDOW)?,
238 monthly: window("monthly", MONTHLY_WINDOW)?,
239 })
240}
241
242#[cfg(test)]
243mod tests {
244 use super::*;
245 use tempfile::TempDir;
246
247 fn cache_fixture() -> (TempDir, Cache) {
248 let td = TempDir::new().unwrap();
249 let cache = Cache::at(td.path().join("modelstudio"));
250 cache.ensure_dir().unwrap();
251 (td, cache)
252 }
253
254 fn creds(region: ConsoleRegion, site: ConsoleSite) -> Credentials {
255 Credentials {
256 access_token: "tok-test".into(),
257 site,
258 region,
259 fingerprint: crate::cache::fingerprint_of("tok-test"),
260 }
261 }
262
263 fn gateway_mock(server: &mut mockito::ServerGuard, creds: &Credentials) -> mockito::Mock {
269 let action = gateway_for(creds.region, creds.site).action;
270 server
271 .mock("POST", usage_path(action).as_str())
272 .match_header("authorization", "Bearer tok-test")
273 .match_header("content-type", "application/x-www-form-urlencoded")
274 .match_body(mockito::Matcher::Exact(form_body(creds.region)))
275 }
276
277 fn usage_body() -> String {
278 serde_json::json!({
283 "code": "200",
284 "data": { "DataV2": { "ret": ["SUCCESS::接口调用成功"], "data": {
285 "msg": "Success.",
286 "code": "SUCCESS",
287 "data": {
288 "per5HourPercentage": 0.4217,
289 "per5HourResetTime": 1789200000000_i64,
290 "per1WeekPercentage": 0.7356,
291 "per1WeekResetTime": 1789600000000_i64,
292 "per1MonthPercentage": 0.44,
293 "per1MonthResetTime": 1794067200000_i64,
294 },
295 "requestId": "rid-t",
296 "success": true,
297 }}},
298 "httpStatusCode": "200",
299 "successResponse": true,
300 })
301 .to_string()
302 }
303
304 #[tokio::test]
308 async fn all_four_gateway_cells_build_the_url_and_body() {
309 for (region, site) in [
310 (ConsoleRegion::CnBeijing, ConsoleSite::Domestic),
311 (ConsoleRegion::CnBeijing, ConsoleSite::International),
312 (ConsoleRegion::ApSoutheast1, ConsoleSite::Domestic),
313 (ConsoleRegion::ApSoutheast1, ConsoleSite::International),
314 ] {
315 let mut server = mockito::Server::new_async().await;
316 let c = creds(region, site);
317 let m = gateway_mock(&mut server, &c)
318 .with_status(200)
319 .with_body(usage_body())
320 .create_async()
321 .await;
322
323 let (_td, cache) = cache_fixture();
324 let out = fetch_snapshot_with(
325 &reqwest::Client::new(),
326 &c,
327 &cache,
328 &Endpoints { base: server.url() },
329 Duration::ZERO,
330 )
331 .await
332 .unwrap();
333
334 m.assert_async().await;
335 assert_eq!(
336 out.snapshot.session.as_ref().unwrap().utilization_pct,
337 42,
338 "{region:?}/{site:?}"
339 );
340 assert_eq!(
341 out.snapshot.weekly.as_ref().unwrap().utilization_pct,
342 74,
343 "{region:?}/{site:?}"
344 );
345 assert!(!out.stale);
346 }
347 }
348
349 #[tokio::test]
350 async fn live_fetch_round_trips_the_epoch_ms_resets() {
351 let mut server = mockito::Server::new_async().await;
352 let c = creds(ConsoleRegion::CnBeijing, ConsoleSite::Domestic);
353 gateway_mock(&mut server, &c)
354 .with_status(200)
355 .with_body(usage_body())
356 .create_async()
357 .await;
358
359 let (_td, cache) = cache_fixture();
360 let out = fetch_snapshot_with(
361 &reqwest::Client::new(),
362 &c,
363 &cache,
364 &Endpoints { base: server.url() },
365 Duration::ZERO,
366 )
367 .await
368 .unwrap();
369 assert_eq!(
370 out.snapshot
371 .session
372 .as_ref()
373 .unwrap()
374 .resets_at
375 .unwrap()
376 .timestamp_millis(),
377 1_789_200_000_000
378 );
379 assert_eq!(
380 out.snapshot
381 .weekly
382 .as_ref()
383 .unwrap()
384 .resets_at
385 .unwrap()
386 .timestamp_millis(),
387 1_789_600_000_000
388 );
389 assert_eq!(
390 out.snapshot.monthly.as_ref().unwrap().utilization_pct,
391 44,
392 "the monthly window rides beside the 5h/weekly ones"
393 );
394 assert_eq!(
395 out.snapshot
396 .monthly
397 .as_ref()
398 .unwrap()
399 .resets_at
400 .unwrap()
401 .timestamp_millis(),
402 1_794_067_200_000
403 );
404
405 let stored = std::fs::read(cache.payload_path()).unwrap();
407 assert!(!String::from_utf8_lossy(&stored).contains("tok-test"));
408 assert_eq!(
409 parse_cache_at(&stored, &c.fingerprint).unwrap(),
410 out.snapshot
411 );
412 }
413
414 #[tokio::test]
417 async fn not_logined_is_a_credentials_error_naming_the_fix() {
418 let mut server = mockito::Server::new_async().await;
419 let c = creds(ConsoleRegion::CnBeijing, ConsoleSite::Domestic);
420 gateway_mock(&mut server, &c)
421 .with_status(200)
422 .with_body(r#"{"success":false,"errorCode":"NotLogined"}"#)
423 .create_async()
424 .await;
425
426 let (_td, cache) = cache_fixture();
427 let err = fetch_snapshot_with(
428 &reqwest::Client::new(),
429 &c,
430 &cache,
431 &Endpoints { base: server.url() },
432 Duration::ZERO,
433 )
434 .await
435 .unwrap_err();
436 assert!(matches!(err, AppError::Credentials(_)), "{err:?}");
437 assert!(err.to_string().contains("bl auth login --console"), "{err}");
438 }
439
440 #[tokio::test]
443 async fn not_logined_with_a_warm_cache_serves_it_stale_with_the_fix() {
444 let mut server = mockito::Server::new_async().await;
445 let c = creds(ConsoleRegion::CnBeijing, ConsoleSite::Domestic);
446 gateway_mock(&mut server, &c)
447 .with_status(200)
448 .with_body(r#"{"success":false,"errorCode":"NotLogined","token":"tok-test"}"#)
449 .create_async()
450 .await;
451
452 let (_td, cache) = cache_fixture();
453 let snap = ModelStudioSnapshot {
454 session: Some(UsageWindow {
455 utilization_pct: 42,
456 resets_at: None,
457 window_duration: FIVE_HOUR_WINDOW,
458 }),
459 weekly: None,
460 monthly: None,
461 };
462 cache
463 .write_payload(&serde_json::to_vec(&snap_to_json(&snap, &c.fingerprint)).unwrap())
464 .unwrap();
465
466 let out = fetch_snapshot_with(
467 &reqwest::Client::new(),
468 &c,
469 &cache,
470 &Endpoints { base: server.url() },
471 Duration::ZERO,
472 )
473 .await
474 .unwrap();
475 assert!(out.stale);
476 assert_eq!(out.snapshot, snap);
477 let (code, body) = out.last_error.unwrap();
478 assert_eq!(code, 0);
479 assert!(!body.contains("tok-test"), "{body}");
480 assert!(body.contains("bl auth login --console"), "{body}");
481 }
482
483 #[tokio::test]
486 async fn other_gateway_failures_are_schema_errors() {
487 let mut server = mockito::Server::new_async().await;
488 let c = creds(ConsoleRegion::CnBeijing, ConsoleSite::Domestic);
489 gateway_mock(&mut server, &c)
490 .with_status(200)
491 .with_body(r#"{"success":false,"errorCode":"SystemExeptionError"}"#)
492 .create_async()
493 .await;
494
495 let (_td, cache) = cache_fixture();
496 let err = fetch_snapshot_with(
497 &reqwest::Client::new(),
498 &c,
499 &cache,
500 &Endpoints { base: server.url() },
501 Duration::ZERO,
502 )
503 .await
504 .unwrap_err();
505 assert!(matches!(err, AppError::Schema(_)), "{err:?}");
506 assert!(err.to_string().contains("SystemExeptionError"), "{err}");
507 }
508
509 #[tokio::test]
511 async fn an_http_401_is_redacted_and_falls_back_to_the_cache() {
512 let mut server = mockito::Server::new_async().await;
513 let c = creds(ConsoleRegion::CnBeijing, ConsoleSite::Domestic);
514 server
515 .mock("POST", usage_path("BroadScopeAspnGateway").as_str())
516 .with_status(401)
517 .with_body(r#"{"message":"Bearer tok-test echoed"}"#)
518 .create_async()
519 .await;
520
521 let (_td, cache) = cache_fixture();
522 let snap = ModelStudioSnapshot {
523 session: None,
524 weekly: Some(UsageWindow {
525 utilization_pct: 74,
526 resets_at: None,
527 window_duration: WEEKLY_WINDOW,
528 }),
529 monthly: None,
530 };
531 cache
532 .write_payload(&serde_json::to_vec(&snap_to_json(&snap, &c.fingerprint)).unwrap())
533 .unwrap();
534
535 let out = fetch_snapshot_with(
536 &reqwest::Client::new(),
537 &c,
538 &cache,
539 &Endpoints { base: server.url() },
540 Duration::ZERO,
541 )
542 .await
543 .unwrap();
544 assert!(out.stale);
545 assert_eq!(out.snapshot.weekly.as_ref().unwrap().utilization_pct, 74);
546 let (code, body) = out.last_error.unwrap();
547 assert_eq!(code, 401);
548 assert_eq!(body, "Model Studio authentication failed");
549 }
550
551 #[tokio::test]
552 async fn a_fresh_cache_is_served_without_a_network_call() {
553 let (_td, cache) = cache_fixture();
554 let c = creds(ConsoleRegion::CnBeijing, ConsoleSite::Domestic);
555 let snap = ModelStudioSnapshot {
556 session: None,
557 weekly: Some(UsageWindow {
558 utilization_pct: 74,
559 resets_at: chrono::DateTime::from_timestamp_millis(1_789_600_000_000),
560 window_duration: WEEKLY_WINDOW,
561 }),
562 monthly: None,
563 };
564 cache
565 .write_payload(&serde_json::to_vec(&snap_to_json(&snap, &c.fingerprint)).unwrap())
566 .unwrap();
567
568 let out = fetch_snapshot_with(
570 &reqwest::Client::new(),
571 &c,
572 &cache,
573 &Endpoints {
574 base: "http://127.0.0.1:1".into(),
575 },
576 Duration::from_secs(60),
577 )
578 .await
579 .unwrap();
580
581 assert!(!out.stale);
582 assert_eq!(out.snapshot, snap);
583 }
584
585 #[tokio::test]
587 async fn a_cache_from_a_previous_login_is_not_reused() {
588 let mut server = mockito::Server::new_async().await;
589 let c = creds(ConsoleRegion::CnBeijing, ConsoleSite::Domestic);
590 let m = gateway_mock(&mut server, &c)
591 .with_status(200)
592 .with_body(usage_body())
593 .create_async()
594 .await;
595
596 let (_td, cache) = cache_fixture();
597 let stale_login = ModelStudioSnapshot {
598 session: Some(UsageWindow {
599 utilization_pct: 99,
600 resets_at: None,
601 window_duration: FIVE_HOUR_WINDOW,
602 }),
603 weekly: None,
604 monthly: None,
605 };
606 cache
607 .write_payload(
608 &serde_json::to_vec(&snap_to_json(&stale_login, "someone-elses-fingerprint"))
609 .unwrap(),
610 )
611 .unwrap();
612
613 let out = fetch_snapshot_with(
615 &reqwest::Client::new(),
616 &c,
617 &cache,
618 &Endpoints { base: server.url() },
619 Duration::from_secs(3600),
620 )
621 .await
622 .unwrap();
623 m.assert_async().await;
624 assert!(!out.stale);
625 assert_eq!(
626 out.snapshot.session.as_ref().unwrap().utilization_pct,
627 42,
628 "refetched, not 99"
629 );
630 }
631
632 #[tokio::test]
633 async fn a_transport_error_with_a_stale_cache_uses_the_cache() {
634 let (_td, cache) = cache_fixture();
635 let c = creds(ConsoleRegion::CnBeijing, ConsoleSite::Domestic);
636 let snap = ModelStudioSnapshot {
637 session: Some(UsageWindow {
638 utilization_pct: 42,
639 resets_at: None,
640 window_duration: FIVE_HOUR_WINDOW,
641 }),
642 weekly: None,
643 monthly: None,
644 };
645 cache
646 .write_payload(&serde_json::to_vec(&snap_to_json(&snap, &c.fingerprint)).unwrap())
647 .unwrap();
648
649 let out = fetch_snapshot_with(
650 &reqwest::Client::new(),
651 &c,
652 &cache,
653 &Endpoints {
654 base: "http://127.0.0.1:1".into(),
655 },
656 Duration::ZERO,
657 )
658 .await
659 .unwrap();
660
661 assert!(out.stale);
662 assert_eq!(out.snapshot, snap);
663 }
664
665 #[test]
666 fn cache_round_trips_both_windows_and_validates() {
667 let snap = ModelStudioSnapshot {
668 session: Some(UsageWindow {
669 utilization_pct: 42,
670 resets_at: chrono::DateTime::from_timestamp_millis(1_789_200_000_000),
671 window_duration: FIVE_HOUR_WINDOW,
672 }),
673 weekly: Some(UsageWindow {
674 utilization_pct: 74,
675 resets_at: None,
676 window_duration: WEEKLY_WINDOW,
677 }),
678 monthly: Some(UsageWindow {
679 utilization_pct: 100,
680 resets_at: chrono::DateTime::from_timestamp_millis(1_794_067_200_000),
681 window_duration: MONTHLY_WINDOW,
682 }),
683 };
684 let bytes = serde_json::to_vec(&snap_to_json(&snap, "fp")).unwrap();
685 assert_eq!(parse_cache_at(&bytes, "fp").unwrap(), snap);
686
687 let absent = ModelStudioSnapshot {
689 session: None,
690 weekly: None,
691 monthly: None,
692 };
693 let bytes = serde_json::to_vec(&snap_to_json(&absent, "fp")).unwrap();
694 assert_eq!(parse_cache_at(&bytes, "fp").unwrap(), absent);
695
696 for bad in [
699 serde_json::json!({"account":"fp","session":{"pct":150},"weekly":null}),
700 serde_json::json!({"account":"fp","session":{"pct":"42"},"weekly":null}),
701 serde_json::json!({"account":"fp","session":{"pct":42,"reset_ms":-5},"weekly":null}),
702 serde_json::json!({"account":"other","session":null,"weekly":null}),
703 serde_json::json!({"session":null,"weekly":null}),
704 ] {
705 assert!(
706 parse_cache_at(bad.to_string().as_bytes(), "fp").is_err(),
707 "{bad} must not parse"
708 );
709 }
710 }
711
712 #[test]
714 fn endpoints_follow_the_gateway_matrix() {
715 assert_eq!(
716 Endpoints::default().base,
717 "https://bailian-cs.console.aliyun.com"
718 );
719 assert_eq!(
720 Endpoints::for_gateway(ConsoleRegion::ApSoutheast1, ConsoleSite::International).base,
721 "https://bailian-singapore-cs.alibabacloud.com"
722 );
723 assert_eq!(
724 Endpoints::default().usage_url("BroadScopeAspnGateway"),
725 "https://bailian-cs.console.aliyun.com/cli/api.json?\
726 action=BroadScopeAspnGateway&product=sfm_bailian\
727 &api=zeldaHttp.apikeyMgr.%2Ftokenplan%2Fpersonal%2Fapi%2Fv2%2Fusage"
728 );
729 }
730}