use astrid_core::SessionId;
use astrid_core::dirs::AstridHome;
use std::path::{Path, PathBuf};
use tracing::{debug, info};
use crate::error::{RuntimeError, RuntimeResult};
use crate::session::{AgentSession, SerializableSession};
pub struct SessionStore {
sessions_dir: PathBuf,
dir_ensured: std::sync::atomic::AtomicBool,
}
impl SessionStore {
#[must_use]
pub fn new(sessions_dir: impl AsRef<Path>) -> Self {
let sessions_dir = sessions_dir.as_ref().to_path_buf();
let dir_exists = sessions_dir.is_dir();
Self {
sessions_dir,
dir_ensured: std::sync::atomic::AtomicBool::new(dir_exists),
}
}
#[must_use]
pub fn from_home(home: &AstridHome) -> Self {
Self::new(home.sessions_dir())
}
fn ensure_dir(&self) -> RuntimeResult<()> {
if self.dir_ensured.load(std::sync::atomic::Ordering::Relaxed) {
return Ok(());
}
std::fs::create_dir_all(&self.sessions_dir)?;
#[cfg(unix)]
{
use std::os::unix::fs::PermissionsExt;
let perms = std::fs::Permissions::from_mode(0o700);
if let Some(parent) = self.sessions_dir.parent() {
let _ = std::fs::set_permissions(parent, perms.clone());
}
let _ = std::fs::set_permissions(&self.sessions_dir, perms);
}
self.dir_ensured
.store(true, std::sync::atomic::Ordering::Relaxed);
Ok(())
}
fn session_path(&self, id: &SessionId) -> PathBuf {
self.sessions_dir.join(format!("{}.json", id.0))
}
pub fn save(&self, session: &AgentSession) -> RuntimeResult<()> {
self.ensure_dir()?;
let path = self.session_path(&session.id);
let serializable = SerializableSession::from(session);
let json = serde_json::to_string_pretty(&serializable)
.map_err(|e| RuntimeError::SerializationError(e.to_string()))?;
let temp_path = path.with_extension("json.tmp");
std::fs::write(&temp_path, &json)?;
std::fs::rename(&temp_path, &path).inspect_err(|_| {
let _ = std::fs::remove_file(&temp_path);
})?;
debug!(session_id = %session.id, path = ?path, "Session saved");
Ok(())
}
pub fn load(&self, id: &SessionId) -> RuntimeResult<Option<AgentSession>> {
let path = self.session_path(id);
if !path.exists() {
return Ok(None);
}
let json = std::fs::read_to_string(&path)?;
let serializable: SerializableSession = serde_json::from_str(&json)
.map_err(|e| RuntimeError::SerializationError(e.to_string()))?;
let session = serializable.to_session();
debug!(session_id = %id, "Session loaded");
Ok(Some(session))
}
pub fn load_by_str(&self, id: &str) -> RuntimeResult<Option<AgentSession>> {
let uuid =
uuid::Uuid::parse_str(id).map_err(|e| RuntimeError::StorageError(e.to_string()))?;
self.load(&SessionId::from_uuid(uuid))
}
pub fn delete(&self, id: &SessionId) -> RuntimeResult<()> {
let path = self.session_path(id);
if path.exists() {
std::fs::remove_file(&path)?;
info!(session_id = %id, "Session deleted");
}
Ok(())
}
pub fn list(&self) -> RuntimeResult<Vec<SessionId>> {
if !self.sessions_dir.is_dir() {
return Ok(Vec::new());
}
let mut sessions = Vec::new();
for entry in std::fs::read_dir(&self.sessions_dir)? {
let entry = entry?;
let path = entry.path();
if path.extension().is_some_and(|e| e == "json")
&& let Some(stem) = path.file_stem()
&& let Some(stem_str) = stem.to_str()
&& let Ok(uuid) = uuid::Uuid::parse_str(stem_str)
{
sessions.push(SessionId::from_uuid(uuid));
}
}
sessions.sort_by(|a, b| {
let path_a = self.session_path(a);
let path_b = self.session_path(b);
let time_a = std::fs::metadata(&path_a)
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH);
let time_b = std::fs::metadata(&path_b)
.and_then(|m| m.modified())
.unwrap_or(std::time::SystemTime::UNIX_EPOCH);
time_b.cmp(&time_a)
});
Ok(sessions)
}
pub fn list_with_metadata(&self) -> RuntimeResult<Vec<SessionSummary>> {
let ids = self.list()?;
let mut summaries = Vec::new();
for id in ids {
if let Ok(Some(session)) = self.load(&id) {
summaries.push(SessionSummary {
id: id.0.to_string(),
title: session.metadata.title.clone(),
created_at: session.created_at,
message_count: session.messages.len(),
token_count: session.token_count,
workspace_path: session.workspace_path.clone(),
});
}
}
Ok(summaries)
}
pub fn most_recent(&self) -> RuntimeResult<Option<AgentSession>> {
let ids = self.list()?;
if let Some(id) = ids.first() {
self.load(id)
} else {
Ok(None)
}
}
pub fn list_for_workspace(&self, workspace: &Path) -> RuntimeResult<Vec<SessionSummary>> {
let all = self.list_with_metadata()?;
Ok(all
.into_iter()
.filter(|s| s.workspace_path.as_deref().is_some_and(|p| p == workspace))
.collect())
}
pub fn cleanup_old(&self, max_age_days: i64) -> RuntimeResult<usize> {
#[allow(clippy::arithmetic_side_effects)]
let cutoff = chrono::Utc::now() - chrono::Duration::days(max_age_days);
let mut removed = 0usize;
for id in self.list()? {
if let Ok(Some(session)) = self.load(&id)
&& session.created_at < cutoff
&& self.delete(&id).is_ok()
{
removed = removed.saturating_add(1);
}
}
Ok(removed)
}
}
#[derive(Debug, Clone)]
pub struct SessionSummary {
pub id: String,
pub title: Option<String>,
pub created_at: chrono::DateTime<chrono::Utc>,
pub message_count: usize,
pub token_count: usize,
pub workspace_path: Option<PathBuf>,
}
impl SessionSummary {
#[must_use]
pub fn display_title(&self) -> String {
self.title.clone().unwrap_or_else(|| {
let short_id = &self.id[..8];
format!("Session {short_id}")
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_session_store() {
let temp_dir = tempfile::tempdir().unwrap();
let store = SessionStore::new(temp_dir.path());
let session = AgentSession::new([0u8; 8], "Test");
store.save(&session).unwrap();
let loaded = store.load(&session.id).unwrap().unwrap();
assert_eq!(loaded.system_prompt, session.system_prompt);
let ids = store.list().unwrap();
assert_eq!(ids.len(), 1);
store.delete(&session.id).unwrap();
assert!(store.load(&session.id).unwrap().is_none());
}
#[test]
fn test_session_store_lazy_dir_creation() {
let temp_dir = tempfile::tempdir().unwrap();
let sessions_path = temp_dir.path().join("lazy_sessions");
let store = SessionStore::new(&sessions_path);
assert!(!sessions_path.exists());
let ids = store.list().unwrap();
assert!(ids.is_empty());
let session = AgentSession::new([0u8; 8], "Test");
store.save(&session).unwrap();
assert!(sessions_path.exists());
}
#[test]
fn test_session_store_atomic_write() {
let temp_dir = tempfile::tempdir().unwrap();
let store = SessionStore::new(temp_dir.path());
let session = AgentSession::new([0u8; 8], "Test");
store.save(&session).unwrap();
let temp_path = temp_dir.path().join(format!("{}.json.tmp", session.id.0));
assert!(!temp_path.exists());
let real_path = temp_dir.path().join(format!("{}.json", session.id.0));
assert!(real_path.exists());
}
#[test]
fn test_session_store_from_home() {
let temp_dir = tempfile::tempdir().unwrap();
let home = AstridHome::from_path(temp_dir.path());
let store = SessionStore::from_home(&home);
let session = AgentSession::new([0u8; 8], "Test");
store.save(&session).unwrap();
let expected = temp_dir
.path()
.join("sessions")
.join(format!("{}.json", session.id.0));
assert!(expected.exists());
}
}