//! Quota monitor — polls the local LS `GetUserStatus` to track //! prompt/flow credits and per-model rate limits without touching Google servers. use serde::Serialize; use std::sync::Arc; use tokio::sync::RwLock; use tracing::{debug, warn}; /// How often to poll the LS for fresh quota data (seconds). const POLL_INTERVAL_SECS: u64 = 60; // ─── Public types ──────────────────────────────────────────────────────────── #[derive(Debug, Clone, Serialize, Default)] pub struct QuotaSnapshot { /// When this snapshot was last refreshed (ISO-8601 UTC). pub last_updated: String, /// Overall plan info. pub plan: PlanInfo, /// Monthly credit balances. pub credits: CreditInfo, /// Per-model rate limits. pub models: Vec, } #[derive(Debug, Clone, Serialize, Default)] pub struct PlanInfo { pub plan_name: String, pub tier_id: String, pub tier_name: String, } #[derive(Debug, Clone, Serialize, Default)] pub struct CreditInfo { pub prompt_available: i64, pub prompt_total: i64, pub prompt_used_pct: f64, pub flow_available: i64, pub flow_total: i64, pub flow_used_pct: f64, pub flex_purchasable: i64, pub can_buy_more: bool, } #[derive(Debug, Clone, Serialize, Default)] pub struct ModelQuota { pub label: String, pub model_id: String, /// 0.0–1.0 remaining fraction (1.0 = full quota). pub remaining_fraction: f64, /// Percentage remaining (0–100). pub remaining_pct: f64, /// ISO-8601 UTC reset time. pub reset_time: String, /// Seconds until reset (negative = already reset). pub reset_in_secs: i64, /// Human-readable countdown. pub reset_in_human: String, } // ─── Quota Store ───────────────────────────────────────────────────────────── #[derive(Clone)] pub struct QuotaStore { inner: Arc>, } impl QuotaStore { pub fn new() -> Self { Self { inner: Arc::new(RwLock::new(QuotaSnapshot::default())), } } /// Get the latest cached snapshot. pub async fn snapshot(&self) -> QuotaSnapshot { self.inner.read().await.clone() } /// Start the background polling loop. Call once at startup. pub fn start_polling(self, backend: Arc) { tokio::spawn(async move { // Initial poll immediately. self.poll_once(&backend).await; let mut interval = tokio::time::interval( std::time::Duration::from_secs(POLL_INTERVAL_SECS), ); interval.tick().await; // consume the first immediate tick loop { interval.tick().await; self.poll_once(&backend).await; } }); } async fn poll_once(&self, backend: &crate::backend::Backend) { match backend .call_json("GetUserStatus", &serde_json::json!({})) .await { Ok((200, data)) => { let snapshot = parse_user_status(&data); debug!( "Quota poll: prompt {}/{} flow {}/{}", snapshot.credits.prompt_available, snapshot.credits.prompt_total, snapshot.credits.flow_available, snapshot.credits.flow_total, ); *self.inner.write().await = snapshot; } Ok((status, data)) => { warn!("GetUserStatus returned {status}: {data}"); } Err(e) => { warn!("GetUserStatus poll failed: {e}"); } } } } // ─── Parsing ───────────────────────────────────────────────────────────────── fn parse_user_status(data: &serde_json::Value) -> QuotaSnapshot { let now = chrono::Utc::now(); let us = &data["userStatus"]; let ps = &us["planStatus"]; let pi = &ps["planInfo"]; let ut = &us["userTier"]; let prompt_total = pi["monthlyPromptCredits"].as_i64().unwrap_or(0); let prompt_avail = ps["availablePromptCredits"].as_i64().unwrap_or(0); let flow_total = pi["monthlyFlowCredits"].as_i64().unwrap_or(0); let flow_avail = ps["availableFlowCredits"].as_i64().unwrap_or(0); let prompt_used_pct = if prompt_total > 0 { ((prompt_total - prompt_avail) as f64 / prompt_total as f64) * 100.0 } else { 0.0 }; let flow_used_pct = if flow_total > 0 { ((flow_total - flow_avail) as f64 / flow_total as f64) * 100.0 } else { 0.0 }; let models = us["cascadeModelConfigData"]["clientModelConfigs"] .as_array() .map(|arr| { arr.iter() .map(|m| { let label = m["label"].as_str().unwrap_or("").to_string(); let model_id = m["modelOrAlias"]["model"] .as_str() .unwrap_or("") .to_string(); let frac = m["quotaInfo"]["remainingFraction"] .as_f64() .unwrap_or(0.0); let reset_str = m["quotaInfo"]["resetTime"] .as_str() .unwrap_or("") .to_string(); let reset_in_secs = if !reset_str.is_empty() { chrono::DateTime::parse_from_rfc3339(&reset_str) .map(|dt| (dt.with_timezone(&chrono::Utc) - now).num_seconds()) .unwrap_or(0) } else { 0 }; let reset_in_human = if reset_in_secs > 0 { let h = reset_in_secs / 3600; let m = (reset_in_secs % 3600) / 60; format!("{h}h {m}m") } else { "available".to_string() }; ModelQuota { label, model_id, remaining_fraction: frac, remaining_pct: frac * 100.0, reset_time: reset_str, reset_in_secs, reset_in_human, } }) .collect() }) .unwrap_or_default(); QuotaSnapshot { last_updated: now.to_rfc3339(), plan: PlanInfo { plan_name: pi["planName"].as_str().unwrap_or("").to_string(), tier_id: ut["id"].as_str().unwrap_or("").to_string(), tier_name: ut["name"].as_str().unwrap_or("").to_string(), }, credits: CreditInfo { prompt_available: prompt_avail, prompt_total, prompt_used_pct, flow_available: flow_avail, flow_total, flow_used_pct, flex_purchasable: pi["monthlyFlexCreditPurchaseAmount"] .as_i64() .unwrap_or(0), can_buy_more: pi["canBuyMoreCredits"].as_bool().unwrap_or(false), }, models, } }