use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};
use serde_json::Value;
use thiserror::Error;
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct CacheKey(String);
impl CacheKey {
pub(crate) fn as_hex(&self) -> &str {
&self.0
}
}
#[derive(Debug, Error)]
pub enum CacheError {
#[error("could not create cache shard {0}: {1}")]
CreateShard(PathBuf, std::io::Error),
#[error("could not serialise cache entry {0}: {1}")]
Serialize(String, serde_json::Error),
#[error("could not write cache entry {0}: {1}")]
Write(PathBuf, std::io::Error),
#[error("could not read cache directory: {0}")]
Walk(std::io::Error),
#[error("could not stat cache entry {0}: {1}")]
Stat(PathBuf, std::io::Error),
#[error("could not remove cache entry {0}: {1}")]
Remove(PathBuf, std::io::Error),
}
#[derive(Debug, Clone)]
pub struct Cache {
root: PathBuf,
ttl: Duration,
max_bytes: u64,
}
impl Cache {
pub fn new(root: PathBuf, ttl_days: u64, max_bytes: u64) -> Self {
let _ = std::fs::create_dir_all(&root);
Self {
root,
ttl: Duration::from_secs(ttl_days.saturating_mul(86_400)),
max_bytes,
}
}
pub fn default_root() -> PathBuf {
if let Some(dirs) = directories::ProjectDirs::from("dev", "slb350", "drep") {
return dirs.cache_dir().to_path_buf();
}
PathBuf::from(".drep-cache")
}
pub fn key(
&self,
system_prompt: &str,
content: &str,
backend: &str,
model: &str,
request_shape: &str,
temperature: Option<f32>,
) -> CacheKey {
let mut hasher = blake3::Hasher::new();
write_field(&mut hasher, system_prompt.as_bytes());
write_field(&mut hasher, content.as_bytes());
write_field(&mut hasher, backend.as_bytes());
write_field(&mut hasher, model.as_bytes());
write_field(&mut hasher, request_shape.as_bytes());
let temp_str = match temperature {
Some(value) => format!("{value:?}"),
None => "unset".to_string(),
};
write_field(&mut hasher, temp_str.as_bytes());
CacheKey(hasher.finalize().to_hex().to_string())
}
pub fn get(&self, key: &CacheKey) -> Option<Value> {
let path = self.entry_path(key);
let meta = std::fs::metadata(&path).ok()?;
let mtime = meta.modified().ok()?;
let age = SystemTime::now()
.duration_since(mtime)
.unwrap_or(Duration::ZERO);
if age > self.ttl {
let _ = std::fs::remove_file(&path);
return None;
}
let bytes = std::fs::read(&path).ok()?;
serde_json::from_slice(&bytes).ok()
}
pub fn put(&self, key: &CacheKey, value: &Value) -> Result<(), CacheError> {
let path = self.entry_path(key);
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.map_err(|e| CacheError::CreateShard(parent.to_path_buf(), e))?;
}
let bytes = serde_json::to_vec(value)
.map_err(|e| CacheError::Serialize(key.as_hex().to_owned(), e))?;
std::fs::write(&path, &bytes).map_err(|e| CacheError::Write(path.clone(), e))?;
Ok(())
}
pub fn evict_if_needed(&self) -> Result<u64, CacheError> {
let mut entries = self.collect_entries()?;
let total: u64 = entries.iter().map(|e| e.size).sum();
if total <= self.max_bytes {
return Ok(0);
}
entries.sort_by_key(|e| e.mtime);
let mut current = total;
let mut freed = 0u64;
for entry in entries {
if current <= self.max_bytes {
break;
}
std::fs::remove_file(&entry.path)
.map_err(|e| CacheError::Remove(entry.path.clone(), e))?;
current = current.saturating_sub(entry.size);
freed = freed.saturating_add(entry.size);
}
Ok(freed)
}
pub(crate) fn entry_path(&self, key: &CacheKey) -> PathBuf {
let hex = key.as_hex();
let shard = &hex[..2];
self.root.join(shard).join(format!("{hex}.json"))
}
fn collect_entries(&self) -> Result<Vec<CacheEntry>, CacheError> {
let mut out = Vec::new();
let shards = std::fs::read_dir(&self.root).map_err(CacheError::Walk)?;
for shard in shards {
let shard = shard.map_err(CacheError::Walk)?;
let file_type = match shard.file_type() {
Ok(ft) => ft,
Err(_) => continue,
};
if !file_type.is_dir() {
continue;
}
if !is_shard_name(&shard.file_name()) {
continue;
}
let shard_path = shard.path();
let shard_name = shard.file_name();
let entries = match std::fs::read_dir(&shard_path) {
Ok(e) => e,
Err(_) => continue,
};
for entry in entries {
let entry = match entry {
Ok(e) => e,
Err(_) => continue,
};
if !is_entry_name(&shard_name, &entry.file_name()) {
continue;
}
let path = entry.path();
let file_type = match entry.file_type() {
Ok(file_type) => file_type,
Err(_) => continue,
};
if !file_type.is_file() {
continue;
}
let meta = match entry.metadata() {
Ok(m) => m,
Err(_) => continue,
};
if !meta.is_file() {
continue;
}
let mtime = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
let size = meta.len();
out.push(CacheEntry { path, mtime, size });
}
}
Ok(out)
}
}
impl Cache {
#[allow(dead_code)]
pub(crate) fn root(&self) -> &Path {
&self.root
}
}
struct CacheEntry {
path: PathBuf,
mtime: SystemTime,
size: u64,
}
fn is_shard_name(name: &std::ffi::OsStr) -> bool {
name.to_str().is_some_and(|name| is_lower_hex(name, 2))
}
fn is_entry_name(shard: &std::ffi::OsStr, name: &std::ffi::OsStr) -> bool {
let (Some(shard), Some(name)) = (shard.to_str(), name.to_str()) else {
return false;
};
let Some(digest) = name.strip_suffix(".json") else {
return false;
};
digest.starts_with(shard) && is_lower_hex(digest, 64)
}
fn is_lower_hex(value: &str, expected_len: usize) -> bool {
value.len() == expected_len
&& value
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}
fn write_field(hasher: &mut blake3::Hasher, bytes: &[u8]) {
let len = u64::try_from(bytes.len()).expect("prompt field longer than u64::MAX bytes");
hasher.update(&len.to_be_bytes());
hasher.update(bytes);
}
#[cfg(test)]
mod tests;