use crate::memory_core::community::KnowledgeGap;
use crate::memory_core::palace::{Palace, PalaceId};
use crate::memory_core::retrieval::PalaceHandle;
use crate::memory_core::store::palace_store::PalaceStore;
use anyhow::{Context, Result};
use dashmap::DashMap;
use lru::LruCache;
use parking_lot::Mutex;
use std::num::NonZeroUsize;
use std::path::Path;
use std::sync::Arc;
pub const DEFAULT_MAX_OPEN_PALACES: usize = 64;
#[derive(Clone)]
pub struct PalaceRegistry {
handles: Arc<Mutex<LruCache<PalaceId, Arc<PalaceHandle>>>>,
gaps_cache: Arc<DashMap<PalaceId, Vec<KnowledgeGap>>>,
}
impl Default for PalaceRegistry {
fn default() -> Self {
Self::with_max_open(DEFAULT_MAX_OPEN_PALACES)
}
}
impl PalaceRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn with_max_open(max_open_palaces: usize) -> Self {
let cap = NonZeroUsize::new(max_open_palaces.max(1)).expect("max(1) is always nonzero");
Self {
handles: Arc::new(Mutex::new(LruCache::new(cap))),
gaps_cache: Arc::new(DashMap::new()),
}
}
pub fn register(&self, handle: PalaceHandle) {
let id = handle.id.clone();
let arc = Arc::new(handle);
let _evicted = {
let mut cache = self.handles.lock();
cache.put(id, arc)
};
}
pub fn register_arc(&self, handle: Arc<PalaceHandle>) {
let id = handle.id.clone();
let _evicted = {
let mut cache = self.handles.lock();
cache.put(id, handle)
};
}
pub fn get(&self, id: &PalaceId) -> Option<Arc<PalaceHandle>> {
let mut cache = self.handles.lock();
cache.get(id).cloned()
}
pub fn peek(&self, id: &PalaceId) -> Option<Arc<PalaceHandle>> {
let cache = self.handles.lock();
cache.peek(id).cloned()
}
pub fn list(&self) -> Vec<PalaceId> {
let cache = self.handles.lock();
cache.iter().map(|(k, _)| k.clone()).collect()
}
pub fn len(&self) -> usize {
self.handles.lock().len()
}
pub fn is_empty(&self) -> bool {
self.handles.lock().is_empty()
}
pub fn set_gaps(&self, palace_id: PalaceId, gaps: Vec<KnowledgeGap>) {
self.gaps_cache.insert(palace_id, gaps);
}
pub fn get_gaps(&self, palace_id: &PalaceId) -> Option<Vec<KnowledgeGap>> {
self.gaps_cache.get(palace_id).map(|r| r.value().clone())
}
pub fn clear_gaps(&self, palace_id: &PalaceId) {
self.gaps_cache.remove(palace_id);
}
pub fn remove(&self, palace_id: &PalaceId) {
let _evicted = {
let mut cache = self.handles.lock();
cache.pop(palace_id)
};
self.gaps_cache.remove(palace_id);
}
pub fn open_palace(&self, data_root: &Path, palace_id: &PalaceId) -> Result<Arc<PalaceHandle>> {
if let Some(h) = self.get(palace_id) {
return Ok(h);
}
let palace_dir = data_root.join(palace_id.as_str());
let palace = PalaceStore::load_palace(&palace_dir)
.with_context(|| format!("load palace metadata for {palace_id}"))?;
let handle = PalaceHandle::open(&palace)?;
self.register_arc(handle.clone());
Ok(handle)
}
pub fn create_palace(&self, data_root: &Path, mut palace: Palace) -> Result<Arc<PalaceHandle>> {
let palace_dir = data_root.join(palace.id.as_str());
palace.data_dir = palace_dir.clone();
std::fs::create_dir_all(&palace_dir)
.with_context(|| format!("create palace dir {}", palace_dir.display()))?;
PalaceStore::save_palace(&palace)
.with_context(|| format!("save palace metadata for {}", palace.id))?;
let handle = PalaceHandle::open(&palace)?;
self.register_arc(handle.clone());
Ok(handle)
}
pub fn list_palaces(data_root: &Path) -> Result<Vec<Palace>> {
PalaceStore::list_palaces(data_root)
.with_context(|| format!("list palaces under {}", data_root.display()))
}
pub fn open(data_root: &Path) -> Result<Self> {
std::fs::create_dir_all(data_root)
.with_context(|| format!("create registry root {}", data_root.display()))?;
let registry = Self::new();
let palaces = PalaceStore::list_palaces(data_root)
.with_context(|| format!("list palaces under {}", data_root.display()))?;
for palace in palaces {
match PalaceHandle::open(&palace) {
Ok(handle) => registry.register_arc(handle),
Err(e) => {
tracing::warn!(palace = %palace.id, "skipping palace during registry open: {e:#}");
}
}
}
Ok(registry)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::memory_core::retrieval::seed_shared_embedder_with_mock;
use crate::memory_core::store::{kg::KnowledgeGraph, vector::UsearchStore};
use tempfile::tempdir;
fn make_handle(id: &str, dir: &std::path::Path) -> PalaceHandle {
let vs = UsearchStore::new(dir.join(format!("{id}.usearch")), 384).unwrap();
let kg = KnowledgeGraph::open(&dir.join(format!("{id}.db"))).unwrap();
PalaceHandle::new(PalaceId::new(id), format!("Identity for {id}"), vs, kg)
}
#[test]
fn register_and_get_roundtrip() {
let dir = tempdir().unwrap();
let reg = PalaceRegistry::new();
reg.register(make_handle("alpha", dir.path()));
let h = reg.get(&PalaceId::new("alpha")).expect("registered");
assert_eq!(h.id.as_str(), "alpha");
}
#[test]
fn registry_remove_clears_cached_handle() {
let dir = tempdir().unwrap();
let reg = PalaceRegistry::new();
let id = PalaceId::new("doomed");
reg.register(make_handle("doomed", dir.path()));
reg.set_gaps(id.clone(), Vec::new());
assert!(reg.get(&id).is_some());
assert!(reg.get_gaps(&id).is_some());
reg.remove(&id);
assert!(reg.get(&id).is_none());
assert!(reg.get_gaps(&id).is_none());
reg.remove(&id);
}
#[test]
fn registry_create_and_open() {
use crate::memory_core::palace::Palace;
use chrono::Utc;
let dir = tempdir().unwrap();
let data_root = dir.path();
let palace = Palace {
id: PalaceId::new("alpha"),
name: "Alpha".to_string(),
description: Some("test".to_string()),
created_at: Utc::now(),
data_dir: data_root.join("alpha"),
};
{
let reg = PalaceRegistry::new();
let handle = reg
.create_palace(data_root, palace.clone())
.expect("create_palace");
assert_eq!(handle.id, PalaceId::new("alpha"));
crate::memory_core::store::palace_store::PalaceStore::save_identity(
&handle.id,
"I am Alpha",
handle.data_dir.as_ref().expect("data_dir set"),
)
.expect("save identity");
}
let reg2 = PalaceRegistry::new();
let handle2 = reg2
.open_palace(data_root, &PalaceId::new("alpha"))
.expect("open_palace");
assert_eq!(handle2.id, PalaceId::new("alpha"));
assert_eq!(handle2.identity, "I am Alpha");
let palaces = PalaceRegistry::list_palaces(data_root).unwrap();
assert_eq!(palaces.len(), 1);
assert_eq!(palaces[0].name, "Alpha");
}
#[tokio::test]
async fn palace_payloads_survive_registry_restart() {
seed_shared_embedder_with_mock();
use crate::memory_core::palace::{Palace, RoomType};
use chrono::Utc;
let dir = tempdir().unwrap();
let data_root = dir.path();
{
let registry = PalaceRegistry::open(data_root).unwrap();
let palace = Palace {
id: PalaceId::new("restart-test"),
name: "Restart".to_string(),
description: None,
created_at: Utc::now(),
data_dir: data_root.join("restart-test"),
};
let handle = registry.create_palace(data_root, palace).unwrap();
handle
.remember(
"the quokka is a small marsupial native to Western Australia".to_string(),
RoomType::Research,
vec!["wildlife".to_string()],
0.7,
)
.await
.expect("remember persists the drawer");
}
let registry = PalaceRegistry::open(data_root).unwrap();
assert_eq!(
registry.len(),
1,
"registry should have hydrated the persisted palace"
);
let handle = registry
.get(&PalaceId::new("restart-test"))
.expect("palace should be registered after open()");
let drawers = handle.drawers.read().clone();
assert!(
drawers
.iter()
.any(|d| d.content.contains("quokka") && d.tags.contains(&"wildlife".to_string())),
"persisted drawer content must survive restart; got {drawers:?}"
);
}
#[test]
fn gaps_cache_round_trip() {
use crate::memory_core::community::KnowledgeGap;
let reg = PalaceRegistry::new();
let pid = PalaceId::new("gap-cache");
assert!(reg.get_gaps(&pid).is_none());
let gaps = vec![KnowledgeGap {
entities: vec!["alpha".to_string(), "beta".to_string()],
internal_density: 0.1,
external_bridges: 1,
suggested_exploration: "Explore connections between alpha and beta".to_string(),
}];
reg.set_gaps(pid.clone(), gaps.clone());
let read = reg.get_gaps(&pid).expect("cached value");
assert_eq!(read.len(), 1);
assert_eq!(read[0].entities, gaps[0].entities);
assert!((read[0].internal_density - 0.1).abs() < 1e-6);
reg.clear_gaps(&pid);
assert!(reg.get_gaps(&pid).is_none());
}
#[test]
fn list_contains_all_registered() {
let dir = tempdir().unwrap();
let reg = PalaceRegistry::new();
reg.register(make_handle("a", dir.path()));
reg.register(make_handle("b", dir.path()));
let ids: Vec<_> = reg.list().into_iter().map(|p| p.0).collect();
assert_eq!(ids.len(), 2);
assert!(ids.contains(&"a".to_string()));
assert!(ids.contains(&"b".to_string()));
}
#[test]
fn lru_evicts_least_recently_used() {
let dir = tempdir().unwrap();
let reg = PalaceRegistry::with_max_open(2);
reg.register(make_handle("a", dir.path()));
reg.register(make_handle("b", dir.path()));
assert_eq!(reg.len(), 2, "two handles registered");
reg.register(make_handle("c", dir.path()));
assert_eq!(reg.len(), 2, "capacity-2 registry must stay at 2");
assert!(
reg.peek(&PalaceId::new("a")).is_none(),
"LRU handle 'a' must have been evicted"
);
assert!(
reg.peek(&PalaceId::new("b")).is_some(),
"MRU handle 'b' must survive"
);
assert!(
reg.peek(&PalaceId::new("c")).is_some(),
"newly inserted 'c' must be present"
);
}
#[test]
fn lru_get_promotes_to_mru() {
let dir = tempdir().unwrap();
let reg = PalaceRegistry::with_max_open(2);
reg.register(make_handle("a", dir.path()));
reg.register(make_handle("b", dir.path()));
let _ = reg.get(&PalaceId::new("a"));
reg.register(make_handle("c", dir.path()));
assert_eq!(reg.len(), 2);
assert!(
reg.peek(&PalaceId::new("b")).is_none(),
"'b' must be evicted — it was LRU after 'a' was promoted"
);
assert!(
reg.peek(&PalaceId::new("a")).is_some(),
"'a' must survive — it was promoted to MRU by get()"
);
assert!(
reg.peek(&PalaceId::new("c")).is_some(),
"'c' must be present"
);
}
#[test]
fn lru_evicted_handle_reopens() {
use crate::memory_core::palace::Palace;
use chrono::Utc;
let dir = tempdir().unwrap();
let data_root = dir.path();
let reg = PalaceRegistry::with_max_open(1);
let palace_a = Palace {
id: PalaceId::new("alpha"),
name: "Alpha".to_string(),
description: None,
created_at: Utc::now(),
data_dir: data_root.join("alpha"),
};
reg.create_palace(data_root, palace_a)
.expect("create alpha");
assert_eq!(reg.len(), 1, "'alpha' registered");
reg.register(make_handle("beta", data_root));
assert_eq!(reg.len(), 1, "capacity-1: only 'beta' remains");
assert!(
reg.peek(&PalaceId::new("alpha")).is_none(),
"'alpha' must have been evicted"
);
let reopened = reg
.open_palace(data_root, &PalaceId::new("alpha"))
.expect("open_palace after eviction must succeed");
assert_eq!(reopened.id, PalaceId::new("alpha"), "reopened id matches");
assert!(
reg.peek(&PalaceId::new("alpha")).is_some(),
"'alpha' must be back in the cache after reopen"
);
}
}