use crate::models::{
ClaudeQuotaSnapshot, CodexQuotaSnapshot, CopilotQuotaSnapshot, CursorQuotaSnapshot,
GrokQuotaSnapshot,
};
use crate::utils::{
get_claude_usage_cache_path, get_codex_usage_cache_path, get_copilot_usage_cache_path,
get_cursor_usage_cache_path, get_grok_usage_cache_path, write_json_atomic,
};
use anyhow::Result;
use serde::de::DeserializeOwned;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
trait CachedQuota: Serialize + DeserializeOwned {
const SCHEMA_VERSION: u32;
}
impl CachedQuota for ClaudeQuotaSnapshot {
const SCHEMA_VERSION: u32 = 1;
}
impl CachedQuota for CodexQuotaSnapshot {
const SCHEMA_VERSION: u32 = 1;
}
impl CachedQuota for CopilotQuotaSnapshot {
const SCHEMA_VERSION: u32 = 1;
}
impl CachedQuota for CursorQuotaSnapshot {
const SCHEMA_VERSION: u32 = 1;
}
impl CachedQuota for GrokQuotaSnapshot {
const SCHEMA_VERSION: u32 = 1;
}
#[derive(Serialize, Deserialize)]
struct VersionedCache<T> {
schema_version: u32,
snapshot: T,
}
fn load_cache<T: CachedQuota>(path: Result<PathBuf>) -> Option<T> {
let path = path.ok()?;
let body = std::fs::read_to_string(&path).ok()?;
let cached: VersionedCache<T> = match serde_json::from_str(&body) {
Ok(cached) => cached,
Err(error) => {
log::debug!("ignoring quota cache {}: {error}", path.display());
return None;
}
};
if cached.schema_version != T::SCHEMA_VERSION {
log::debug!(
"ignoring quota cache {}: written by schema v{}, current is v{}",
path.display(),
cached.schema_version,
T::SCHEMA_VERSION
);
return None;
}
Some(cached.snapshot)
}
fn save_cache<T: CachedQuota>(path: Result<PathBuf>, snapshot: &T) -> Result<()> {
write_json_atomic(
path?,
&VersionedCache {
schema_version: T::SCHEMA_VERSION,
snapshot,
},
)
}
pub fn load_claude_cache() -> Option<ClaudeQuotaSnapshot> {
load_cache(get_claude_usage_cache_path())
}
pub fn save_claude_cache(snap: &ClaudeQuotaSnapshot) -> Result<()> {
save_cache(get_claude_usage_cache_path(), snap)
}
pub fn load_codex_cache() -> Option<CodexQuotaSnapshot> {
load_cache(get_codex_usage_cache_path())
}
pub fn save_codex_cache(snap: &CodexQuotaSnapshot) -> Result<()> {
save_cache(get_codex_usage_cache_path(), snap)
}
pub fn load_copilot_cache() -> Option<CopilotQuotaSnapshot> {
load_cache(get_copilot_usage_cache_path())
}
pub fn save_copilot_cache(snap: &CopilotQuotaSnapshot) -> Result<()> {
save_cache(get_copilot_usage_cache_path(), snap)
}
pub fn load_cursor_cache() -> Option<CursorQuotaSnapshot> {
load_cache(get_cursor_usage_cache_path())
}
pub fn save_cursor_cache(snap: &CursorQuotaSnapshot) -> Result<()> {
save_cache(get_cursor_usage_cache_path(), snap)
}
pub fn load_grok_cache() -> Option<GrokQuotaSnapshot> {
load_cache(get_grok_usage_cache_path())
}
pub fn save_grok_cache(snap: &GrokQuotaSnapshot) -> Result<()> {
save_cache(get_grok_usage_cache_path(), snap)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::models::{QuotaSource, QuotaWindow};
use tempfile::TempDir;
fn cache_file() -> (TempDir, PathBuf) {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("provider_usage.json");
(dir, path)
}
fn round_trip<T: CachedQuota>(snapshot: &T) -> Option<T> {
let (_dir, path) = cache_file();
save_cache(Ok(path.clone()), snapshot).unwrap();
load_cache::<T>(Ok(path))
}
fn cursor_snapshot() -> CursorQuotaSnapshot {
CursorQuotaSnapshot {
source: QuotaSource::Api,
fetched_at: 1_700_000_000,
plan_type: Some("pro".into()),
total: Some(QuotaWindow {
used_percent: 12.5,
resets_at_unix: Some(1_700_600_000),
}),
auto: Some(QuotaWindow {
used_percent: 3.0,
resets_at_unix: Some(1_700_600_000),
}),
api: None,
on_demand_dollars: Some(1.25),
limit_reached: false,
needs_login: false,
}
}
#[test]
fn round_trip_preserves_cursor_snapshot() {
let loaded = round_trip(&cursor_snapshot()).expect("a just-written cache must load");
assert_eq!(loaded.source, QuotaSource::Api);
assert_eq!(loaded.fetched_at, 1_700_000_000);
assert_eq!(loaded.plan_type.as_deref(), Some("pro"));
let total = loaded.total.expect("total window");
assert_eq!(total.used_percent, 12.5);
assert_eq!(total.resets_at_unix, Some(1_700_600_000));
assert_eq!(loaded.auto.expect("auto window").used_percent, 3.0);
assert!(loaded.api.is_none());
assert_eq!(loaded.on_demand_dollars, Some(1.25));
assert!(!loaded.limit_reached);
}
#[test]
fn round_trip_preserves_other_providers() {
let claude = round_trip(&ClaudeQuotaSnapshot {
fetched_at: 11,
five_hour: Some(QuotaWindow {
used_percent: 42.0,
resets_at_unix: None,
}),
..Default::default()
})
.expect("claude cache");
assert_eq!(claude.fetched_at, 11);
assert_eq!(claude.five_hour.unwrap().used_percent, 42.0);
let codex = round_trip(&CodexQuotaSnapshot {
fetched_at: 22,
plan_type: Some("plus".into()),
reset_credits_available: Some(2),
..Default::default()
})
.expect("codex cache");
assert_eq!(codex.fetched_at, 22);
assert_eq!(codex.plan_type.as_deref(), Some("plus"));
assert_eq!(codex.reset_credits_available, Some(2));
let copilot = round_trip(&CopilotQuotaSnapshot {
fetched_at: 33,
premium_remaining: Some(120),
needs_login: true,
..Default::default()
})
.expect("copilot cache");
assert_eq!(copilot.fetched_at, 33);
assert_eq!(copilot.premium_remaining, Some(120));
assert!(copilot.needs_login);
let grok = round_trip(&GrokQuotaSnapshot {
fetched_at: 44,
period_label: Some("week".into()),
prepaid_balance_dollars: Some(7.5),
..Default::default()
})
.expect("grok cache");
assert_eq!(grok.fetched_at, 44);
assert_eq!(grok.period_label.as_deref(), Some("week"));
assert_eq!(grok.prepaid_balance_dollars, Some(7.5));
}
#[test]
fn saved_file_carries_the_current_schema_version() {
let (_dir, path) = cache_file();
save_cache(Ok(path.clone()), &cursor_snapshot()).unwrap();
let raw = std::fs::read_to_string(&path).unwrap();
let body: serde_json::Value = serde_json::from_str(&raw).unwrap();
assert_eq!(
body["schema_version"],
serde_json::json!(CursorQuotaSnapshot::SCHEMA_VERSION)
);
assert_eq!(body["snapshot"]["plan_type"], "pro");
assert!(serde_json::from_str::<CursorQuotaSnapshot>(&raw).is_err());
}
#[test]
fn unversioned_cache_is_rejected() {
let (_dir, path) = cache_file();
let bare = serde_json::to_string(&cursor_snapshot()).unwrap();
std::fs::write(&path, &bare).unwrap();
assert!(serde_json::from_str::<CursorQuotaSnapshot>(&bare).is_ok());
assert!(load_cache::<CursorQuotaSnapshot>(Ok(path)).is_none());
}
#[test]
fn mismatched_schema_version_is_rejected() {
let (_dir, path) = cache_file();
std::fs::write(
&path,
serde_json::to_string(&VersionedCache {
schema_version: CursorQuotaSnapshot::SCHEMA_VERSION + 1,
snapshot: cursor_snapshot(),
})
.unwrap(),
)
.unwrap();
assert!(load_cache::<CursorQuotaSnapshot>(Ok(path)).is_none());
}
#[test]
fn corrupt_or_absent_cache_is_none() {
let (_dir, path) = cache_file();
assert!(load_cache::<CursorQuotaSnapshot>(Ok(path.clone())).is_none());
std::fs::write(&path, "{not json").unwrap();
assert!(load_cache::<CursorQuotaSnapshot>(Ok(path)).is_none());
}
}