use super::format::{get_serializer, Format, FormatError};
use super::index::{SessionIndex, SessionMeta};
use super::project::{ProjectIndex, ProjectMeta};
use crate::txlog::TxLog;
use std::fs::{self, File};
use std::io::BufReader;
use std::path::{Path, PathBuf};
use thiserror::Error;
#[derive(Debug, Error)]
pub enum StorageError {
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("Format error: {0}")]
Format(#[from] FormatError),
#[error("Session not found: {0}")]
SessionNotFound(String),
#[error("Project not found: {0}")]
ProjectNotFound(String),
#[error("Storage not initialized at {0}")]
NotInitialized(PathBuf),
#[error("Index error: {0}")]
Index(String),
}
pub type StorageResult<T> = Result<T, StorageError>;
#[derive(Debug)]
pub struct RyoStorage {
root: PathBuf,
sessions_dir: PathBuf,
projects_dir: PathBuf,
index_path: PathBuf,
project_index_path: PathBuf,
format: Format,
index: Option<SessionIndex>,
project_index: Option<ProjectIndex>,
}
impl RyoStorage {
pub const DIR_NAME: &'static str = ".ryo";
pub const SESSIONS_DIR: &'static str = "sessions";
pub const PROJECTS_DIR: &'static str = "projects";
pub const INDEX_FILE: &'static str = "index.json";
pub const PROJECT_INDEX_FILE: &'static str = "projects.json";
pub fn global() -> StorageResult<Self> {
let home = dirs::home_dir().ok_or_else(|| {
StorageError::Io(std::io::Error::new(
std::io::ErrorKind::NotFound,
"Could not find home directory",
))
})?;
Self::new(home.join(Self::DIR_NAME))
}
pub fn new(root: PathBuf) -> StorageResult<Self> {
let sessions_dir = root.join(Self::SESSIONS_DIR);
let projects_dir = root.join(Self::PROJECTS_DIR);
let index_path = root.join(Self::INDEX_FILE);
let project_index_path = root.join(Self::PROJECT_INDEX_FILE);
Ok(Self {
root,
sessions_dir,
projects_dir,
index_path,
project_index_path,
format: Format::default(),
index: None,
project_index: None,
})
}
pub fn with_format(mut self, format: Format) -> Self {
self.format = format;
self
}
pub fn root(&self) -> &Path {
&self.root
}
pub fn is_initialized(&self) -> bool {
self.root.exists() && self.sessions_dir.exists()
}
pub fn init(&self) -> StorageResult<()> {
fs::create_dir_all(&self.sessions_dir)?;
fs::create_dir_all(&self.projects_dir)?;
if !self.index_path.exists() {
let index = SessionIndex::new();
let json = serde_json::to_string_pretty(&index)
.map_err(|e| StorageError::Format(FormatError::Json(e)))?;
fs::write(&self.index_path, json)?;
}
if !self.project_index_path.exists() {
let index = ProjectIndex::new();
let json = serde_json::to_string_pretty(&index)
.map_err(|e| StorageError::Format(FormatError::Json(e)))?;
fs::write(&self.project_index_path, json)?;
}
Ok(())
}
pub fn ensure_init(&self) -> StorageResult<()> {
if !self.is_initialized() {
self.init()?;
}
Ok(())
}
pub fn dump(&mut self, log: &TxLog) -> StorageResult<String> {
self.ensure_init()?;
let session_id = log.session_id.clone();
let filename = self.session_filename(&session_id);
let path = self.sessions_dir.join(&filename);
let serializer = get_serializer(self.format);
let file = File::create(&path)?;
serializer.serialize_to_file(log, file)?;
self.add_to_index(log)?;
Ok(session_id)
}
pub fn load(&self, session_id: &str) -> StorageResult<TxLog> {
let formats_to_try = [self.format, Format::Json, Format::JsonCompact];
for format in formats_to_try {
let filename = format!("{}.txlog.{}", session_id, format.extension());
let path = self.sessions_dir.join(&filename);
if path.exists() {
let file = File::open(&path)?;
let reader = BufReader::new(file);
let serializer = get_serializer(format);
let log = serializer.deserialize_from_reader(reader)?;
return Ok(log);
}
}
Err(StorageError::SessionNotFound(session_id.to_string()))
}
pub fn exists(&self, session_id: &str) -> bool {
self.find_session_path(session_id).is_some()
}
pub fn delete(&mut self, session_id: &str) -> StorageResult<()> {
if let Some(path) = self.find_session_path(session_id) {
fs::remove_file(&path)?;
}
self.remove_from_index(session_id)?;
Ok(())
}
fn find_session_path(&self, session_id: &str) -> Option<PathBuf> {
for format in [Format::Json, Format::JsonCompact] {
let filename = format!("{}.txlog.{}", session_id, format.extension());
let path = self.sessions_dir.join(&filename);
if path.exists() {
return Some(path);
}
}
None
}
fn session_filename(&self, session_id: &str) -> String {
format!("{}.txlog.{}", session_id, self.format.extension())
}
pub fn index(&mut self) -> StorageResult<&SessionIndex> {
if self.index.is_none() {
self.load_index()?;
}
Ok(self
.index
.as_ref()
.expect("load_index() above sets self.index to Some"))
}
fn index_mut(&mut self) -> StorageResult<&mut SessionIndex> {
if self.index.is_none() {
self.load_index()?;
}
Ok(self
.index
.as_mut()
.expect("load_index() above sets self.index to Some"))
}
fn load_index(&mut self) -> StorageResult<()> {
if self.index_path.exists() {
let content = fs::read_to_string(&self.index_path)?;
let index: SessionIndex = serde_json::from_str(&content)
.map_err(|e| StorageError::Format(FormatError::Json(e)))?;
self.index = Some(index);
} else {
self.index = Some(SessionIndex::new());
}
Ok(())
}
fn save_index(&self) -> StorageResult<()> {
if let Some(ref index) = self.index {
let json = serde_json::to_string_pretty(index)
.map_err(|e| StorageError::Format(FormatError::Json(e)))?;
fs::write(&self.index_path, json)?;
}
Ok(())
}
fn add_to_index(&mut self, log: &TxLog) -> StorageResult<()> {
let meta = SessionMeta::from_log(log);
self.index_mut()?.add(meta);
self.save_index()?;
Ok(())
}
fn remove_from_index(&mut self, session_id: &str) -> StorageResult<()> {
self.index_mut()?.remove(session_id);
self.save_index()?;
Ok(())
}
pub fn list_sessions(&mut self) -> StorageResult<Vec<&SessionMeta>> {
Ok(self.index()?.list())
}
pub fn sessions_for_project(
&mut self,
project_path: &Path,
) -> StorageResult<Vec<&SessionMeta>> {
Ok(self.index()?.by_project(project_path))
}
pub fn latest_session(&mut self) -> StorageResult<Option<&SessionMeta>> {
Ok(self.index()?.latest())
}
pub fn latest_for_project(
&mut self,
project_path: &Path,
) -> StorageResult<Option<&SessionMeta>> {
Ok(self.index()?.latest_for_project(project_path))
}
pub fn cleanup(&mut self, keep_per_project: usize) -> StorageResult<usize> {
let to_delete = self.index_mut()?.cleanup(keep_per_project);
let count = to_delete.len();
for session_id in to_delete {
let filename = self.session_filename(&session_id);
let path = self.sessions_dir.join(&filename);
if path.exists() {
fs::remove_file(&path)?;
}
}
self.save_index()?;
Ok(count)
}
pub fn storage_size(&self) -> StorageResult<u64> {
let mut total = 0u64;
if self.sessions_dir.exists() {
for entry in fs::read_dir(&self.sessions_dir)? {
let entry = entry?;
if entry.file_type()?.is_file() {
total += entry.metadata()?.len();
}
}
}
if self.index_path.exists() {
total += fs::metadata(&self.index_path)?.len();
}
Ok(total)
}
pub fn project_index(&mut self) -> StorageResult<&ProjectIndex> {
if self.project_index.is_none() {
self.load_project_index()?;
}
Ok(self
.project_index
.as_ref()
.expect("load_project_index() above sets self.project_index to Some"))
}
fn project_index_mut(&mut self) -> StorageResult<&mut ProjectIndex> {
if self.project_index.is_none() {
self.load_project_index()?;
}
Ok(self
.project_index
.as_mut()
.expect("load_project_index() above sets self.project_index to Some"))
}
fn load_project_index(&mut self) -> StorageResult<()> {
if self.project_index_path.exists() {
let content = fs::read_to_string(&self.project_index_path)?;
let index: ProjectIndex = serde_json::from_str(&content)
.map_err(|e| StorageError::Format(FormatError::Json(e)))?;
self.project_index = Some(index);
} else {
self.project_index = Some(ProjectIndex::new());
}
Ok(())
}
fn save_project_index(&self) -> StorageResult<()> {
if let Some(ref index) = self.project_index {
let json = serde_json::to_string_pretty(index)
.map_err(|e| StorageError::Format(FormatError::Json(e)))?;
fs::write(&self.project_index_path, json)?;
}
Ok(())
}
pub fn register_project(&mut self, meta: ProjectMeta) -> StorageResult<String> {
self.ensure_init()?;
let project_id = meta.project_id.clone();
self.project_index_mut()?.add(meta);
self.save_project_index()?;
Ok(project_id)
}
pub fn unregister_project(&mut self, project_id: &str) -> StorageResult<Option<ProjectMeta>> {
let meta = self.project_index_mut()?.remove(project_id);
self.save_project_index()?;
Ok(meta)
}
pub fn get_project(&mut self, project_id: &str) -> StorageResult<Option<&ProjectMeta>> {
Ok(self.project_index()?.get(project_id))
}
pub fn get_project_by_path(&mut self, path: &Path) -> StorageResult<Option<&ProjectMeta>> {
Ok(self.project_index()?.get_by_path(path))
}
pub fn get_project_by_path_mut(
&mut self,
path: &Path,
) -> StorageResult<Option<&mut ProjectMeta>> {
let project_id = self
.project_index()?
.get_by_path(path)
.map(|p| p.project_id.clone());
if let Some(id) = project_id {
Ok(self.project_index_mut()?.get_mut(&id))
} else {
Ok(None)
}
}
pub fn save_projects(&self) -> StorageResult<()> {
self.save_project_index()
}
pub fn search_projects_by_name(&mut self, pattern: &str) -> StorageResult<Vec<&ProjectMeta>> {
Ok(self.project_index()?.search_by_name(pattern))
}
pub fn is_project_registered(&mut self, path: &Path) -> StorageResult<bool> {
Ok(self.project_index()?.contains_path(path))
}
pub fn list_projects(&mut self) -> StorageResult<Vec<&ProjectMeta>> {
Ok(self.project_index()?.list())
}
pub fn touch_project(&mut self, project_id: &str) -> StorageResult<()> {
if let Some(meta) = self.project_index_mut()?.get_mut(project_id) {
meta.touch();
self.save_project_index()?;
}
Ok(())
}
pub fn project_stats(&mut self) -> StorageResult<(usize, usize, usize)> {
let index = self.project_index()?;
Ok((index.count(), index.total_files(), index.total_lines()))
}
pub fn cleanup_dead_servers(&mut self) -> StorageResult<usize> {
let cleaned = self.project_index_mut()?.cleanup_dead_servers();
if cleaned > 0 {
self.save_project_index()?;
}
Ok(cleaned)
}
pub fn list_projects_with_cleanup(
&mut self,
cleanup: bool,
) -> StorageResult<Vec<&ProjectMeta>> {
if cleanup {
self.cleanup_dead_servers()?;
}
Ok(self.project_index()?.list())
}
pub const CACHE_DIR: &'static str = "cache";
pub fn cache_dir(&self) -> PathBuf {
self.root.join(Self::CACHE_DIR)
}
fn ensure_cache_dir(&self) -> StorageResult<()> {
let dir = self.cache_dir();
if !dir.exists() {
fs::create_dir_all(&dir)?;
}
Ok(())
}
pub fn save_graph_cache(&self, project_hash: &str, data: &[u8]) -> StorageResult<PathBuf> {
self.ensure_cache_dir()?;
let path = self.cache_dir().join(format!("{}.graph.bin", project_hash));
fs::write(&path, data)?;
Ok(path)
}
pub fn load_graph_cache(&self, project_hash: &str) -> StorageResult<Option<Vec<u8>>> {
let path = self.cache_dir().join(format!("{}.graph.bin", project_hash));
if !path.exists() {
return Ok(None);
}
let data = fs::read(&path)?;
Ok(Some(data))
}
pub fn delete_graph_cache(&self, project_hash: &str) -> StorageResult<()> {
let path = self.cache_dir().join(format!("{}.graph.bin", project_hash));
if path.exists() {
fs::remove_file(&path)?;
}
Ok(())
}
pub fn list_graph_caches(&self) -> StorageResult<Vec<String>> {
let dir = self.cache_dir();
if !dir.exists() {
return Ok(Vec::new());
}
let mut hashes = Vec::new();
for entry in fs::read_dir(&dir)? {
let entry = entry?;
let path = entry.path();
if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
if name.ends_with(".graph.bin") {
let hash = name.trim_end_matches(".graph.bin").to_string();
hashes.push(hash);
}
}
}
Ok(hashes)
}
pub fn cache_size(&self) -> StorageResult<u64> {
let dir = self.cache_dir();
if !dir.exists() {
return Ok(0);
}
let mut size = 0u64;
for entry in fs::read_dir(&dir)? {
let entry = entry?;
if let Ok(meta) = entry.metadata() {
size += meta.len();
}
}
Ok(size)
}
pub fn clear_graph_caches(&self) -> StorageResult<usize> {
let dir = self.cache_dir();
if !dir.exists() {
return Ok(0);
}
let mut count = 0;
for entry in fs::read_dir(&dir)? {
let entry = entry?;
let path = entry.path();
if path.extension().map(|e| e == "bin").unwrap_or(false) {
fs::remove_file(&path)?;
count += 1;
}
}
Ok(count)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::txlog::TxAction;
use tempfile::TempDir;
fn create_test_log(project: &str) -> TxLog {
let mut log = TxLog::with_project(project);
log.log(TxAction::GoalSet {
query: "test".to_string(),
intent_type: "test".to_string(),
confidence: 0.9,
});
log
}
#[test]
fn test_init_and_dump() {
let temp = TempDir::new().unwrap();
let mut storage = RyoStorage::new(temp.path().join(".ryo")).unwrap();
storage.init().unwrap();
assert!(storage.is_initialized());
let log = create_test_log("/test/project");
let session_id = storage.dump(&log).unwrap();
assert!(storage.exists(&session_id));
}
#[test]
fn test_load_session() {
let temp = TempDir::new().unwrap();
let mut storage = RyoStorage::new(temp.path().join(".ryo")).unwrap();
storage.init().unwrap();
let log = create_test_log("/test/project");
let session_id = storage.dump(&log).unwrap();
let loaded = storage.load(&session_id).unwrap();
assert_eq!(loaded.session_id, log.session_id);
assert_eq!(loaded.entries().len(), log.entries().len());
}
#[test]
fn test_session_index() {
let temp = TempDir::new().unwrap();
let mut storage = RyoStorage::new(temp.path().join(".ryo")).unwrap();
storage.init().unwrap();
let log1 = create_test_log("/project/a");
let log2 = create_test_log("/project/b");
let log3 = create_test_log("/project/a");
storage.dump(&log1).unwrap();
storage.dump(&log2).unwrap();
storage.dump(&log3).unwrap();
let all = storage.list_sessions().unwrap();
assert_eq!(all.len(), 3);
let proj_a = storage
.sessions_for_project(Path::new("/project/a"))
.unwrap();
assert_eq!(proj_a.len(), 2);
}
#[test]
fn test_delete_session() {
let temp = TempDir::new().unwrap();
let mut storage = RyoStorage::new(temp.path().join(".ryo")).unwrap();
storage.init().unwrap();
let log = create_test_log("/test/project");
let session_id = storage.dump(&log).unwrap();
assert!(storage.exists(&session_id));
storage.delete(&session_id).unwrap();
assert!(!storage.exists(&session_id));
}
}