feat: initial commit — antigravity proxy with MITM, standalone LS, and snapshot tooling

src/quota.rs

@@ -0,0 +1,218 @@
+//! 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<ModelQuota>,
+}
+
+#[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<RwLock<QuotaSnapshot>>,
+}
+
+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<crate::backend::Backend>) {
+        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,
+    }
+}