use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use dashmap::DashMap;
use notify::{RecommendedWatcher, RecursiveMode, Watcher};
use notify_debouncer_full::{new_debouncer, DebounceEventResult, Debouncer, FileIdMap};
use super::workspace::{SessionScope, WorkspaceBundle, WorkspaceLoader};
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
struct CacheKey {
root: PathBuf,
scope: SessionScope,
extras: Vec<String>,
}
#[derive(Debug)]
struct Entry {
bundle: Arc<WorkspaceBundle>,
loaded_at: Instant,
}
#[derive(Debug, Default)]
pub struct WorkspaceCacheMetrics {
pub hits: DashMap<String, AtomicU64>,
pub misses: DashMap<String, AtomicU64>,
pub invalidations: DashMap<String, AtomicU64>,
}
impl WorkspaceCacheMetrics {
fn bump(map: &DashMap<String, AtomicU64>, key: &str) {
if let Some(c) = map.get(key) {
c.value().fetch_add(1, Ordering::Relaxed);
return;
}
let entry = map
.entry(key.to_string())
.or_insert_with(|| AtomicU64::new(0));
entry.value().fetch_add(1, Ordering::Relaxed);
}
pub fn snapshot(&self) -> HashMap<String, (u64, u64, u64)> {
let mut out: HashMap<String, (u64, u64, u64)> = HashMap::new();
for kv in self.hits.iter() {
out.entry(kv.key().clone()).or_default().0 = kv.value().load(Ordering::Relaxed);
}
for kv in self.misses.iter() {
out.entry(kv.key().clone()).or_default().1 = kv.value().load(Ordering::Relaxed);
}
for kv in self.invalidations.iter() {
out.entry(kv.key().clone()).or_default().2 = kv.value().load(Ordering::Relaxed);
}
out
}
}
pub struct WorkspaceCache {
bundles: DashMap<CacheKey, Arc<Entry>>,
metrics: Arc<WorkspaceCacheMetrics>,
roots: Vec<PathBuf>,
max_age: Option<Duration>,
_watcher: Mutex<Option<Debouncer<RecommendedWatcher, FileIdMap>>>,
}
impl WorkspaceCache {
pub fn new(
roots: &[PathBuf],
debounce_ms: u32,
max_age_seconds: u32,
) -> anyhow::Result<Arc<Self>> {
let metrics = Arc::new(WorkspaceCacheMetrics::default());
let bundles: DashMap<CacheKey, Arc<Entry>> = DashMap::new();
let max_age = if max_age_seconds == 0 {
None
} else {
Some(Duration::from_secs(max_age_seconds as u64))
};
let cache = Arc::new(Self {
bundles,
metrics,
roots: roots.iter().map(|p| p.to_path_buf()).collect(),
max_age,
_watcher: Mutex::new(None),
});
cache.start_watcher(debounce_ms)?;
Ok(cache)
}
fn start_watcher(self: &Arc<Self>, debounce_ms: u32) -> anyhow::Result<()> {
if self.roots.is_empty() {
return Ok(());
}
let weak = Arc::downgrade(self);
let timeout = Duration::from_millis(debounce_ms.max(50) as u64);
let mut debouncer = new_debouncer(timeout, None, move |res: DebounceEventResult| {
let Some(this) = weak.upgrade() else {
return;
};
match res {
Ok(events) => {
let mut roots_to_drop: Vec<PathBuf> = Vec::new();
for ev in events {
for path in &ev.event.paths {
if path.extension().and_then(|e| e.to_str()) != Some("md") {
continue;
}
for root in &this.roots {
if path.starts_with(root) && !roots_to_drop.contains(root) {
roots_to_drop.push(root.clone());
}
}
}
}
for root in roots_to_drop {
this.invalidate_root(&root);
}
}
Err(errs) => {
for e in errs {
tracing::warn!(error = ?e, "workspace_cache: watcher error");
}
}
}
})?;
for root in &self.roots {
if let Err(e) = debouncer.watcher().watch(root, RecursiveMode::Recursive) {
tracing::warn!(
root = %root.display(),
error = %e,
"workspace_cache: failed to start watcher (cache will still serve hits, just no auto-invalidation)"
);
}
}
if let Ok(mut slot) = self._watcher.lock() {
*slot = Some(debouncer);
}
Ok(())
}
pub fn invalidate_root(&self, root: &Path) {
let key = root.display().to_string();
let before = self.bundles.len();
self.bundles.retain(|k, _| k.root != root);
let dropped = before - self.bundles.len();
if dropped > 0 {
WorkspaceCacheMetrics::bump(&self.metrics.invalidations, &key);
tracing::info!(
root = %root.display(),
dropped,
"workspace_cache: invalidated entries"
);
}
}
pub fn clear(&self) {
self.bundles.clear();
}
pub fn metrics(&self) -> Arc<WorkspaceCacheMetrics> {
Arc::clone(&self.metrics)
}
pub async fn get(
&self,
root: &Path,
scope: SessionScope,
extras: &[String],
) -> anyhow::Result<Arc<WorkspaceBundle>> {
let mut key_extras: Vec<String> = extras
.iter()
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
key_extras.sort();
let key = CacheKey {
root: root.to_path_buf(),
scope,
extras: key_extras,
};
let root_label = root.display().to_string();
if let Some(entry) = self.bundles.get(&key) {
if let Some(max_age) = self.max_age {
if entry.value().loaded_at.elapsed() > max_age {
drop(entry);
self.bundles.remove(&key);
WorkspaceCacheMetrics::bump(&self.metrics.invalidations, &root_label);
} else {
WorkspaceCacheMetrics::bump(&self.metrics.hits, &root_label);
return Ok(Arc::clone(&entry.value().bundle));
}
} else {
WorkspaceCacheMetrics::bump(&self.metrics.hits, &root_label);
return Ok(Arc::clone(&entry.value().bundle));
}
}
WorkspaceCacheMetrics::bump(&self.metrics.misses, &root_label);
let bundle = WorkspaceLoader::new(root)
.load_with_extras(scope, extras)
.await?;
let entry = Arc::new(Entry {
bundle: Arc::new(bundle),
loaded_at: Instant::now(),
});
let stored = self
.bundles
.entry(key)
.or_insert_with(|| Arc::clone(&entry))
.value()
.clone();
Ok(Arc::clone(&stored.bundle))
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
fn write(p: &Path, name: &str, text: &str) {
std::fs::write(p.join(name), text).unwrap();
}
#[tokio::test]
async fn cold_then_warm() {
let dir = tempdir().unwrap();
write(dir.path(), "IDENTITY.md", "- **Name:** Ana");
let cache = WorkspaceCache::new(&[dir.path().to_path_buf()], 100, 0).unwrap();
let a = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
let b = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
assert!(Arc::ptr_eq(&a, &b));
let snap = cache.metrics().snapshot();
let (hits, misses, _) = snap
.get(&dir.path().display().to_string())
.copied()
.unwrap_or_default();
assert_eq!(misses, 1);
assert_eq!(hits, 1);
}
#[tokio::test]
async fn different_scope_separate_entry() {
let dir = tempdir().unwrap();
write(dir.path(), "IDENTITY.md", "- **Name:** Ana");
let cache = WorkspaceCache::new(&[dir.path().to_path_buf()], 100, 0).unwrap();
let main = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
let shared = cache
.get(dir.path(), SessionScope::Shared, &[])
.await
.unwrap();
assert!(!Arc::ptr_eq(&main, &shared));
}
#[tokio::test]
async fn extras_order_does_not_split_cache() {
let dir = tempdir().unwrap();
write(dir.path(), "A.md", "a");
write(dir.path(), "B.md", "b");
let cache = WorkspaceCache::new(&[dir.path().to_path_buf()], 100, 0).unwrap();
let one = cache
.get(
dir.path(),
SessionScope::Main,
&["A.md".to_string(), "B.md".to_string()],
)
.await
.unwrap();
let two = cache
.get(
dir.path(),
SessionScope::Main,
&["B.md".to_string(), "A.md".to_string()],
)
.await
.unwrap();
assert!(Arc::ptr_eq(&one, &two));
}
#[tokio::test]
async fn manual_invalidate_drops_entries() {
let dir = tempdir().unwrap();
write(dir.path(), "IDENTITY.md", "- **Name:** Ana");
let cache = WorkspaceCache::new(&[dir.path().to_path_buf()], 100, 0).unwrap();
let _ = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
cache.invalidate_root(dir.path());
let snap = cache.metrics().snapshot();
let (_, _, inv) = snap
.get(&dir.path().display().to_string())
.copied()
.unwrap_or_default();
assert_eq!(inv, 1);
}
#[tokio::test]
async fn watcher_invalidates_after_md_write() {
let dir = tempdir().unwrap();
write(dir.path(), "IDENTITY.md", "- **Name:** Ana");
let cache = WorkspaceCache::new(&[dir.path().to_path_buf()], 80, 0).unwrap();
let a = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
std::fs::write(
dir.path().join("IDENTITY.md"),
"- **Name:** Ana\n- **Vibe:** updated",
)
.unwrap();
tokio::time::sleep(Duration::from_millis(500)).await;
let b = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
assert!(
!Arc::ptr_eq(&a, &b),
"expected cache miss after MD modification"
);
let snap = cache.metrics().snapshot();
let (_, _, inv) = snap
.get(&dir.path().display().to_string())
.copied()
.unwrap_or_default();
assert!(inv >= 1, "expected at least one invalidation, got {inv}");
}
#[tokio::test]
async fn max_age_forces_refresh() {
let dir = tempdir().unwrap();
write(dir.path(), "IDENTITY.md", "- **Name:** Ana");
let cache = WorkspaceCache::new(&[dir.path().to_path_buf()], 100, 1).unwrap();
let a = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(1100)).await;
let b = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
assert!(!Arc::ptr_eq(&a, &b));
}
#[tokio::test]
async fn non_md_file_ignored_by_watcher() {
let dir = tempdir().unwrap();
write(dir.path(), "IDENTITY.md", "- **Name:** Ana");
let cache = WorkspaceCache::new(&[dir.path().to_path_buf()], 80, 0).unwrap();
let a = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
std::fs::write(dir.path().join("note.txt"), "ignored").unwrap();
tokio::time::sleep(Duration::from_millis(400)).await;
let b = cache
.get(dir.path(), SessionScope::Main, &[])
.await
.unwrap();
assert!(Arc::ptr_eq(&a, &b));
}
}