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)),
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionSummary {
    pub id: String,
    pub updated_at: DateTime<Utc>,
    pub preview: String,
    pub verdict: Option<String>,
}

pub fn list_sessions() -> Result<Vec<SessionSummary>> {
    let dir = ensure_sessions_dir()?;
    let mut rows = Vec::new();
    for entry in fs::read_dir(&dir).with_context(|| format!("read {}", dir.display()))? {
        let path = entry?.path();
        if path.extension().and_then(|s| s.to_str()) != Some("json") {
            continue;
        }
        let Ok(session) = load_session_from(&path) else {
            continue;
        };
        rows.push(SessionSummary {
            id: session.id.clone(),
            updated_at: session.updated_at,
            preview: session_preview(&session),
            verdict: session.last_verdict.clone(),
        });
    }
    rows.sort_by_key(|a| std::cmp::Reverse(a.updated_at));
    Ok(rows)
}

pub fn session_preview(session: &Session) -> String {
    let from_messages = session_preview_from_messages(&session.messages);
    if from_messages != "新對話" {
        from_messages
    } else {
        session
            .last_message
            .clone()
            .unwrap_or_else(|| "新對話".into())
            .chars()
            .take(80)
            .collect()
    }
}

pub fn session_preview_from_messages(messages: &[crate::ChatMessage]) -> String {
    messages
        .iter()
        .rev()
        .find_map(|m| match m.role {
            crate::Role::User if !m.text().starts_with("<system_reminder>") => {
                Some(m.text().trim().to_string())
            }
            _ => None,
        })
        .unwrap_or_else(|| "新對話".into())
        .chars()
        .take(80)
        .collect()
}

pub fn public_transcript(session: &Session) -> Vec<serde_json::Value> {
    public_transcript_from_messages(&session.messages)
}

pub fn public_transcript_from_messages(messages: &[crate::ChatMessage]) -> Vec<serde_json::Value> {
    use serde_json::json;
    let mut out = Vec::new();
    for msg in messages {
        match msg.role {
            crate::Role::User if !msg.text().starts_with("<system_reminder>") => {
                out.push(json!({"role": "user", "content": msg.text()}));
            }
            crate::Role::Assistant => {
                if let Some(calls) = &msg.tool_calls {
                    for call in calls {
                        if call.function.name != "send_message" {
                            continue;
                        }
                        if let Ok(args) = serde_json::from_str::<serde_json::Value>(&call.function.arguments)
                        {
                            if let Some(content) = args["content"]
                                .as_str()
                                .or(args["message"].as_str())
                                .filter(|s| !s.trim().is_empty())
                            {
                                out.push(json!({"role": "assistant", "content": content}));
                            }
                        }
                    }
                }
            }
            _ => {}
        }
    }
    out
}

#[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: "external_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 session_preview_uses_last_user_line() {
        let mut s = Session::new(".");
        s.push(ChatMessage::user("手機上也能用"));
        assert!(session_preview(&s).contains("手機"));
        assert_eq!(public_transcript(&s)[0]["role"], "user");
    }

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