use std::fs::{self, File, OpenOptions};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};
use fs2::FileExt;
use crate::error::{AppError, Result};
pub const DEFAULT_TTL: Duration = Duration::from_secs(60);
pub const MAX_STALE: Duration = Duration::from_secs(7 * 24 * 3600);
#[derive(Debug, Clone)]
pub struct Cache {
dir: PathBuf,
}
impl Cache {
pub fn for_vendor(vendor: &str) -> Result<Self> {
let base = xdg_cache_dir()?.join("ai-usagebar").join(vendor);
Ok(Self { dir: base })
}
pub fn for_vendor_account(vendor: &str, label: &str) -> Result<Self> {
let base = xdg_cache_dir()?
.join("ai-usagebar")
.join(vendor)
.join(label);
Ok(Self { dir: base })
}
pub fn at(path: PathBuf) -> Self {
Self { dir: path }
}
pub fn ensure_dir(&self) -> Result<()> {
fs::create_dir_all(&self.dir).map_err(|e| AppError::io_at(&self.dir, e))
}
pub fn dir(&self) -> &Path {
&self.dir
}
pub fn payload_path(&self) -> PathBuf {
self.dir.join("usage.json")
}
pub fn stale_path(&self) -> PathBuf {
self.dir.join(".stale")
}
pub fn last_error_path(&self) -> PathBuf {
self.dir.join(".last_error")
}
pub fn lock_path(&self) -> PathBuf {
self.dir.join(".fetch.lock")
}
pub fn payload_age(&self) -> Option<Duration> {
let meta = fs::metadata(self.payload_path()).ok()?;
let mtime = meta.modified().ok()?;
SystemTime::now().duration_since(mtime).ok()
}
pub fn fresh_payload(&self, ttl: Duration) -> Result<Option<Vec<u8>>> {
let Some(age) = self.payload_age() else {
return Ok(None);
};
if age < ttl {
self.read_payload().map(Some)
} else {
Ok(None)
}
}
pub fn maybe_payload(&self) -> Result<Option<Vec<u8>>> {
if !self.payload_path().exists() {
return Ok(None);
}
self.read_payload().map(Some)
}
pub fn fallback_payload(&self, max_stale: Duration) -> Result<Option<Vec<u8>>> {
let Some(age) = self.payload_age() else {
return Ok(None);
};
if age > max_stale {
return Ok(None);
}
self.read_payload().map(Some)
}
fn read_payload(&self) -> Result<Vec<u8>> {
let p = self.payload_path();
let mut f = File::open(&p).map_err(|e| AppError::io_at(&p, e))?;
let mut buf = Vec::new();
f.read_to_end(&mut buf)
.map_err(|e| AppError::io_at(&p, e))?;
Ok(buf)
}
pub fn write_payload(&self, bytes: &[u8]) -> Result<()> {
self.ensure_dir()?;
let mut tmp = tempfile::Builder::new()
.prefix(".usage.")
.tempfile_in(&self.dir)
.map_err(|e| AppError::io_at(&self.dir, e))?;
tmp.write_all(bytes)
.map_err(|e| AppError::io_at(tmp.path(), e))?;
tmp.as_file_mut()
.sync_all()
.map_err(|e| AppError::io_at(tmp.path(), e))?;
tmp.persist(self.payload_path())
.map_err(|e| AppError::io_at(self.payload_path(), e.error))?;
let _ = fs::remove_file(self.stale_path());
let _ = fs::remove_file(self.last_error_path());
Ok(())
}
pub fn mark_stale(&self) {
let _ = self.ensure_dir();
let _ = File::create(self.stale_path());
}
pub fn is_stale(&self) -> bool {
self.stale_path().exists()
}
pub fn write_last_error(&self, code: u16, msg: &str) {
let _ = self.ensure_dir();
let path = self.last_error_path();
let body = format!("{code}\n{msg}");
let _ = atomic_write(&path, body.as_bytes());
}
pub fn clear_last_error(&self) {
let _ = fs::remove_file(self.last_error_path());
}
pub fn read_last_error(&self) -> Option<(u16, String)> {
let raw = fs::read_to_string(self.last_error_path()).ok()?;
let (code, msg) = raw.split_once('\n').unwrap_or((raw.as_str(), ""));
Some((code.parse::<u16>().ok()?, msg.to_string()))
}
}
pub async fn acquire_lock_async(path: &Path, timeout: Duration) -> Result<LockGuard> {
let path = path.to_path_buf();
tokio::task::spawn_blocking(move || acquire_lock(&path, timeout))
.await
.map_err(|e| AppError::Other(format!("cache lock task failed: {e}")))?
}
pub fn acquire_lock(path: &Path, timeout: Duration) -> Result<LockGuard> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).map_err(|e| AppError::io_at(parent, e))?;
}
let f = OpenOptions::new()
.create(true)
.read(true)
.write(true)
.truncate(false)
.open(path)
.map_err(|e| AppError::io_at(path, e))?;
let deadline = std::time::Instant::now() + timeout;
loop {
match f.try_lock_exclusive() {
Ok(()) => return Ok(LockGuard { file: f }),
Err(_) => {
if std::time::Instant::now() >= deadline {
return Err(AppError::Other(format!(
"cache lock timeout after {:?}",
timeout
)));
}
std::thread::sleep(Duration::from_millis(50));
}
}
}
}
pub struct LockGuard {
file: File,
}
impl Drop for LockGuard {
fn drop(&mut self) {
let _ = FileExt::unlock(&self.file);
}
}
pub fn atomic_write(path: &Path, bytes: &[u8]) -> Result<()> {
let dir = path.parent().ok_or_else(|| {
AppError::Other(format!(
"atomic_write: path has no parent: {}",
path.display()
))
})?;
fs::create_dir_all(dir).map_err(|e| AppError::io_at(dir, e))?;
let mut tmp = tempfile::Builder::new()
.prefix(".tmp.")
.tempfile_in(dir)
.map_err(|e| AppError::io_at(dir, e))?;
tmp.write_all(bytes)
.map_err(|e| AppError::io_at(tmp.path(), e))?;
tmp.as_file_mut()
.sync_all()
.map_err(|e| AppError::io_at(tmp.path(), e))?;
tmp.persist(path)
.map_err(|e| AppError::io_at(path, e.error))?;
Ok(())
}
fn xdg_cache_dir() -> Result<PathBuf> {
directories::BaseDirs::new()
.map(|b| b.cache_dir().to_path_buf())
.ok_or_else(|| AppError::Other("could not resolve XDG cache dir (no HOME?)".into()))
}
pub fn home_dir() -> Result<PathBuf> {
directories::BaseDirs::new()
.map(|b| b.home_dir().to_path_buf())
.ok_or_else(|| AppError::Other("could not resolve home directory (no HOME?)".into()))
}
#[cfg(test)]
pub(crate) fn closed_temp_file(name: &str, contents: Option<&str>) -> (tempfile::TempDir, PathBuf) {
let dir = tempfile::TempDir::new().unwrap();
let path = dir.path().join(name);
if let Some(c) = contents {
std::fs::write(&path, c).unwrap();
}
(dir, path)
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
fn fixture() -> (TempDir, Cache) {
let td = TempDir::new().unwrap();
let cache = Cache::at(td.path().join("anthropic"));
cache.ensure_dir().unwrap();
(td, cache)
}
#[test]
fn ensure_dir_is_idempotent() {
let (_td, cache) = fixture();
cache.ensure_dir().unwrap();
cache.ensure_dir().unwrap();
assert!(cache.dir().is_dir());
}
#[test]
fn write_then_read_round_trip() {
let (_td, cache) = fixture();
cache.write_payload(b"hello world").unwrap();
let got = cache.maybe_payload().unwrap();
assert_eq!(got.as_deref(), Some(&b"hello world"[..]));
}
#[test]
fn maybe_payload_returns_none_when_missing() {
let (_td, cache) = fixture();
assert!(cache.maybe_payload().unwrap().is_none());
}
#[test]
fn fresh_payload_respects_ttl() {
let (_td, cache) = fixture();
cache.write_payload(b"x").unwrap();
assert!(
cache
.fresh_payload(Duration::from_secs(10))
.unwrap()
.is_some()
);
assert!(
cache
.fresh_payload(Duration::from_secs(0))
.unwrap()
.is_none()
);
}
#[test]
fn write_clears_stale_marker_and_last_error() {
let (_td, cache) = fixture();
cache.mark_stale();
cache.write_last_error(429, "rate limited");
assert!(cache.is_stale());
assert!(cache.read_last_error().is_some());
cache.write_payload(b"fresh").unwrap();
assert!(!cache.is_stale());
assert!(cache.read_last_error().is_none());
}
#[test]
fn fallback_payload_refuses_a_payload_older_than_the_limit() {
let (_td, cache) = fixture();
cache.write_payload(b"old").unwrap();
std::thread::sleep(Duration::from_millis(60));
assert!(cache.maybe_payload().unwrap().is_some());
assert!(
cache
.fallback_payload(Duration::from_millis(5))
.unwrap()
.is_none()
);
assert_eq!(
cache.fallback_payload(MAX_STALE).unwrap().as_deref(),
Some(&b"old"[..])
);
}
#[test]
fn last_error_round_trip() {
let (_td, cache) = fixture();
cache.write_last_error(503, "service unavailable");
let (code, msg) = cache.read_last_error().unwrap();
assert_eq!(code, 503);
assert_eq!(msg, "service unavailable");
}
#[test]
fn last_error_with_empty_message_round_trips() {
let (_td, cache) = fixture();
cache.write_last_error(429, "");
let (code, msg) = cache.read_last_error().unwrap();
assert_eq!(code, 429);
assert_eq!(msg, "");
}
#[test]
fn last_error_round_trips_a_multi_line_message() {
let (_td, cache) = fixture();
let body = "{\n \"error\": \"quota exhausted\",\n \"retry_after\": 3600\n}";
cache.write_last_error(429, body);
let (code, msg) = cache.read_last_error().unwrap();
assert_eq!(code, 429);
assert_eq!(msg, body);
assert!(
msg.contains("quota exhausted"),
"message was truncated to its first line: {msg:?}"
);
}
#[test]
fn last_error_reads_files_written_by_the_previous_version() {
let (_td, cache) = fixture();
fs::write(cache.last_error_path(), "503\nservice unavailable").unwrap();
assert_eq!(
cache.read_last_error(),
Some((503, "service unavailable".into()))
);
fs::write(cache.last_error_path(), "429").unwrap();
assert_eq!(cache.read_last_error(), Some((429, String::new())));
fs::write(cache.last_error_path(), "not-a-code\nboom").unwrap();
assert!(cache.read_last_error().is_none());
}
#[test]
fn lock_serializes_concurrent_acquirers() {
let (_td, cache) = fixture();
let lock_path = cache.lock_path();
let _guard = acquire_lock(&lock_path, Duration::from_millis(500)).unwrap();
let res = acquire_lock(&lock_path, Duration::from_millis(100));
assert!(matches!(res, Err(AppError::Other(_))));
}
#[tokio::test(flavor = "current_thread")]
async fn async_lock_does_not_stall_the_runtime() {
let (_td, cache) = fixture();
let lock_path = cache.lock_path();
let _held = acquire_lock(&lock_path, Duration::from_millis(500)).unwrap();
let waiter = acquire_lock_async(&lock_path, Duration::from_millis(400));
let mut ticks = 0usize;
let ticker = async {
let mut iv = tokio::time::interval(Duration::from_millis(20));
iv.tick().await;
loop {
iv.tick().await;
ticks += 1;
}
};
tokio::select! {
res = waiter => {
assert!(matches!(res, Err(AppError::Other(_))));
}
_ = ticker => unreachable!("the ticker loops forever"),
}
assert!(
ticks > 1,
"runtime was starved while the lock was contended ({ticks} ticks)"
);
}
#[test]
fn atomic_write_creates_parent_dirs() {
let td = TempDir::new().unwrap();
let nested = td.path().join("a/b/c/file.txt");
atomic_write(&nested, b"abc").unwrap();
assert_eq!(fs::read(&nested).unwrap(), b"abc");
}
}