LazyBoy2/crates/grokboy-core/src/session.rs

//! Persist agent sessions under ~/.grokboy/sessions/.

use crate::model::ChatMessage;
use anyhow::{Context, Result, anyhow};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::{Path, PathBuf};
use uuid::Uuid;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
    pub id: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub cwd: PathBuf,
    pub messages: Vec<ChatMessage>,
}

impl Session {
    pub fn new(cwd: impl Into<PathBuf>) -> Self {
        let now = Utc::now();
        Self {
            id: Uuid::new_v4().to_string(),
            created_at: now,
            updated_at: now,
            cwd: cwd.into(),
            messages: Vec::new(),
        }
    }

    pub fn touch(&mut self) {
        self.updated_at = Utc::now();
    }

    pub fn push(&mut self, msg: ChatMessage) {
        self.messages.push(msg);
        self.touch();
    }
}

pub fn sessions_dir() -> Result<PathBuf> {
    let home = dirs_home().ok_or_else(|| anyhow!("cannot resolve home directory"))?;
    Ok(home.join(".grokboy").join("sessions"))
}

fn dirs_home() -> Option<PathBuf> {
    std::env::var_os("HOME")
        .or_else(|| std::env::var_os("USERPROFILE"))
        .map(PathBuf::from)
}

pub fn ensure_sessions_dir() -> Result<PathBuf> {
    let dir = sessions_dir()?;
    fs::create_dir_all(&dir).with_context(|| format!("create {}", dir.display()))?;
    Ok(dir)
}

pub fn session_path(id: &str) -> Result<PathBuf> {
    validate_session_id(id)?;
    Ok(ensure_sessions_dir()?.join(format!("{id}.json")))
}

fn validate_session_id(id: &str) -> Result<()> {
    if id.is_empty() || id.contains('/') || id.contains('\\') || id.contains("..") {
        return Err(anyhow!("invalid session id"));
    }
    Ok(())
}

pub fn save_session(session: &Session) -> Result<PathBuf> {
    let path = session_path(&session.id)?;
    save_session_to(session, &path)?;
    Ok(path)
}

pub fn save_session_to(session: &Session, path: &Path) -> Result<()> {
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
    }
    let tmp = path.with_extension("json.tmp");
    let data = serde_json::to_vec_pretty(session).context("serialize session")?;
    fs::write(&tmp, &data).with_context(|| format!("write {}", tmp.display()))?;
    fs::rename(&tmp, path).with_context(|| format!("rename to {}", path.display()))?;
    Ok(())
}

pub fn load_session(id: &str) -> Result<Session> {
    let path = session_path(id)?;
    load_session_from(&path)
}

pub fn load_session_from(path: &Path) -> Result<Session> {
    let data = fs::read(path).with_context(|| format!("read session {}", path.display()))?;
    let session: Session = serde_json::from_slice(&data).context("parse session JSON")?;
    Ok(session)
}

pub fn load_or_create(session_id: Option<&str>, cwd: &Path) -> Result<Session> {
    match session_id {
        Some(id) => load_session(id),
        None => Ok(Session::new(cwd)),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::ChatMessage;

    #[test]
    fn save_and_load_roundtrip() {
        let stamp = Uuid::new_v4();
        let dir = std::env::temp_dir().join(format!("grokboy-sess-{stamp}"));
        fs::create_dir_all(&dir).unwrap();
        let path = dir.join("sess.json");

        let mut s = Session::new(dir.join("proj"));
        s.push(ChatMessage::system("sys"));
        s.push(ChatMessage::user("hi"));
        save_session_to(&s, &path).unwrap();
        assert!(path.exists());

        let loaded = load_session_from(&path).unwrap();
        assert_eq!(loaded.id, s.id);
        assert_eq!(loaded.messages.len(), 2);
        assert_eq!(loaded.messages[1].text(), "hi");

        let _ = fs::remove_dir_all(&dir);
    }

    #[test]
    fn rejects_bad_session_id() {
        assert!(validate_session_id("../x").is_err());
        assert!(validate_session_id("a/b").is_err());
        assert!(validate_session_id("ok-id").is_ok());
    }
}