235 lines
8.3 KiB
Rust
235 lines
8.3 KiB
Rust
//! Quota monitor — polls the local LS `GetUserStatus` to track
|
||
//! prompt/flow credits and per-model rate limits without touching Google servers.
|
||
//!
|
||
//! The LS's `GetUserStatus` response contains `cascadeModelConfigData` with
|
||
//! per-model `quotaInfo` (remaining fraction, reset time) and plan-level credit
|
||
//! balances (prompt, flow, flex). This data originates from Google's
|
||
//! `PredictionService/RetrieveUserQuota` / `v1internal:retrieveUserQuota` endpoint,
|
||
//! but asking the LS is simpler since it caches this data locally.
|
||
//!
|
||
//! The credit system uses the `google/internal/cloud/code/v1internal/credits`
|
||
//! proto package with `Credits_CreditType` enum. The `CASCADE_ENFORCE_QUOTA`
|
||
//! config key controls whether quotas are actually enforced.
|
||
|
||
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)) => {
|
||
// Profile picture fetch fails through iptables — harmless, suppress
|
||
let data_str = data.to_string();
|
||
if data_str.contains("profile picture") {
|
||
tracing::debug!("GetUserStatus: profile picture fetch failed (expected with iptables)");
|
||
} else {
|
||
warn!("GetUserStatus returned {status}: {data_str}");
|
||
}
|
||
}
|
||
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,
|
||
}
|
||
}
|