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

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

use crate::model::ChatMessage;
use anyhow::{anyhow, Context, Result};
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>,
    #[serde(default)]
    pub plan: Vec<crate::runtime::PlanStep>,
    #[serde(default)]
    pub pending_question: Option<serde_json::Value>,
    #[serde(default)]
    pub pending_tool: Option<String>,
    #[serde(default)]
    pub active_command: Option<serde_json::Value>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub last_verdict: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub last_message: Option<String>,
    /// Last page URL from optional Playwright browser tools (P3).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub last_browser_url: Option<String>,
}

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(),
            plan: vec![],
            pending_question: None,
            pending_tool: None,
            active_command: None,
            last_verdict: None,
            last_message: None,
            last_browser_url: None,
        }
    }

    /// Repair incomplete transcripts without replaying any operation.
    pub fn recover_interrupted(&mut self) {
        let mut repaired = Vec::new();
        let mut index = 0;
        while index < self.messages.len() {
            let message = self.messages[index].clone();
            let calls = message.tool_calls.clone().unwrap_or_default();
            repaired.push(message);
            index += 1;
            for call in calls {
                if self.messages.get(index).is_some_and(|m| {
                    m.role == crate::Role::Tool && m.tool_call_id.as_deref() == Some(&call.id)
                }) {
                    repaired.push(self.messages[index].clone());
                    index += 1;
                } else {
                    repaired.push(ChatMessage::tool(&call.id, serde_json::json!({
                        "error":"interrupted before a durable result was recorded; observe the current state before deciding what to do; do not replay automatically",
                        "outcome":"unknown", "was_active":self.pending_tool.as_deref()==Some(&call.id)
                    }).to_string()));
                }
            }
        }
        if let Some(command) = self.active_command.as_mut() {
            if command["running"] == true {
                command["running"] = serde_json::json!(false);
                command["outcome"] = serde_json::json!("unknown after process restart");
                repaired.push(ChatMessage::system(format!("Previous command is no longer a live session: {command}. Inspect saved output and current state before retrying; do not replay automatically.")));
            }
        }
        self.messages = repaired;
        self.pending_tool = None;
    }

    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> {
    if let Some(dir) = std::env::var_os("GROKBOY_SESSIONS_DIR").filter(|s| !s.is_empty()) {
        return Ok(PathBuf::from(dir));
    }
    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 mut session: Session = serde_json::from_slice(&data).context("parse session JSON")?;
    session.recover_interrupted();
    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 old_sessions_without_runtime_fields_remain_readable() {
        let session = Session::new(".");
        let mut value = serde_json::to_value(session).unwrap();
        for field in [
            "plan",
            "pending_question",
            "pending_tool",
            "active_command",
            "last_verdict",
            "last_message",
        ] {
            value.as_object_mut().unwrap().remove(field);
        }
        let old: Session = serde_json::from_value(value).unwrap();
        assert!(old.plan.is_empty());
        assert!(old.pending_question.is_none());
    }

    #[test]
    fn recovery_pairs_unknown_and_unstarted_calls_without_replaying() {
        let mut session = Session::new(".");
        let call = |id: &str| crate::ToolCall {
            id: id.into(),
            kind: "function".into(),
            function: crate::FunctionCall {
                name: "write_file".into(),
                arguments: "{}".into(),
            },
        };
        session.messages.push(ChatMessage::assistant_tool_calls(
            None,
            vec![call("a"), call("b")],
        ));
        session.pending_tool = Some("a".into());
        session.recover_interrupted();
        assert_eq!(session.messages.len(), 3);
        assert_eq!(session.messages[1].tool_call_id.as_deref(), Some("a"));
        assert_eq!(session.messages[2].tool_call_id.as_deref(), Some("b"));
        assert!(session.messages[1].text().contains("unknown"));
        session.recover_interrupted();
        assert_eq!(session.messages.len(), 3);
    }

    #[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());
    }
}