huskies/server/src/agents/session_store.rs

//! Persistent session store — tracks the last Claude Code session_id per
//! (story_id, agent_name, model) triple so respawned agents can resume prior reasoning.
//!
//! The session_id is extracted from the `Done. Session: Some(<uuid>)` log entry
//! emitted at agent shutdown. When the same (story, agent, model) triple is
//! spawned again, the orchestrator passes `--resume <session_id>` so the new
//! session inherits the prior conversation context.
//!
//! Model is part of the key intentionally: resuming across models is not
//! supported (e.g. opus should not resume a sonnet session).

use std::collections::HashMap;
use std::path::Path;

/// Composite key for the session store: `{story_id}:{agent_name}:{model}`.
fn session_key(story_id: &str, agent_name: &str, model: &str) -> String {
    format!("{story_id}:{agent_name}:{model}")
}

/// Path to the persistent session store file.
fn store_path(project_root: &Path) -> std::path::PathBuf {
    project_root.join(".huskies/session_store.json")
}

/// Read the session store from disk. Returns an empty map if the file is
/// missing, empty, or corrupt.
fn read_store(project_root: &Path) -> HashMap<String, String> {
    let path = store_path(project_root);
    std::fs::read_to_string(path)
        .ok()
        .and_then(|s| serde_json::from_str(&s).ok())
        .unwrap_or_default()
}

/// Write the session store to disk. Silently ignores write errors (the store
/// is best-effort — a missed resume just means a fresh session starts).
fn write_store(project_root: &Path, data: &HashMap<String, String>) {
    let path = store_path(project_root);
    if let Ok(json) = serde_json::to_string_pretty(data) {
        let _ = std::fs::write(path, json);
    }
}

/// Record the session_id from a completed agent run, persisted to disk.
///
/// Called after an agent process exits with a valid session_id. The next
/// spawn of the same (story_id, agent_name, model) triple will find this
/// session and pass `--resume <session_id>`.
pub fn record_session(
    project_root: &Path,
    story_id: &str,
    agent_name: &str,
    model: &str,
    session_id: &str,
) {
    let key = session_key(story_id, agent_name, model);
    let mut data = read_store(project_root);
    data.insert(key, session_id.to_string());
    write_store(project_root, &data);
}

/// Look up the last session_id for a (story_id, agent_name, model) triple.
///
/// Returns `None` if no prior session exists (fresh story) or if the model
/// has changed (intentional — resuming across models is not supported).
pub fn lookup_session(
    project_root: &Path,
    story_id: &str,
    agent_name: &str,
    model: &str,
) -> Option<String> {
    let key = session_key(story_id, agent_name, model);
    read_store(project_root).get(&key).cloned()
}

/// Remove all session entries for a story (called when a story reaches done/archived).
#[cfg(test)]
pub fn remove_sessions_for_story(project_root: &Path, story_id: &str) {
    let mut data = read_store(project_root);
    let prefix = format!("{story_id}:");
    let before = data.len();
    data.retain(|k, _| !k.starts_with(&prefix));
    if data.len() < before {
        write_store(project_root, &data);
    }
}

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

    // ── AC1: record and lookup round-trip ─────────────────────────────────

    #[test]
    fn record_and_lookup_round_trip() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        std::fs::create_dir_all(root.join(".huskies")).unwrap();

        record_session(root, "42_story_foo", "coder-1", "sonnet", "sess-abc");

        let result = lookup_session(root, "42_story_foo", "coder-1", "sonnet");
        assert_eq!(result, Some("sess-abc".to_string()));
    }

    #[test]
    fn lookup_returns_none_for_unknown_story() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        std::fs::create_dir_all(root.join(".huskies")).unwrap();

        let result = lookup_session(root, "unknown_story", "coder-1", "sonnet");
        assert_eq!(result, None);
    }

    // ── AC3: model change semantics ───────────────────────────────────────

    /// When an operator escalates from sonnet to opus, the new opus spawn
    /// must NOT resume the prior sonnet session. The key includes the model,
    /// so a different model produces a cache miss.
    #[test]
    fn model_change_does_not_resume_prior_session() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        std::fs::create_dir_all(root.join(".huskies")).unwrap();

        // Record a sonnet session.
        record_session(root, "42_story_foo", "coder-1", "sonnet", "sess-sonnet");

        // Looking up with opus model returns None — no resume.
        let result = lookup_session(root, "42_story_foo", "coder-1", "opus");
        assert_eq!(result, None, "opus must not resume a sonnet session");

        // Looking up with sonnet still works.
        let result = lookup_session(root, "42_story_foo", "coder-1", "sonnet");
        assert_eq!(result, Some("sess-sonnet".to_string()));
    }

    // ── AC9: no prior session → fresh start ───────────────────────────────

    #[test]
    fn no_prior_session_returns_none() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        std::fs::create_dir_all(root.join(".huskies")).unwrap();

        // Empty store — every lookup is None.
        assert_eq!(
            lookup_session(root, "99_story_new", "coder-1", "sonnet"),
            None
        );
    }

    // ── AC1: persistence across "restarts" (re-read from disk) ────────────

    #[test]
    fn session_survives_store_reload() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        std::fs::create_dir_all(root.join(".huskies")).unwrap();

        record_session(root, "42_story_foo", "coder-1", "sonnet", "sess-persist");

        // Simulate restart: create a fresh store read.
        let result = lookup_session(root, "42_story_foo", "coder-1", "sonnet");
        assert_eq!(result, Some("sess-persist".to_string()));
    }

    #[test]
    fn record_overwrites_previous_session() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        std::fs::create_dir_all(root.join(".huskies")).unwrap();

        record_session(root, "42_story_foo", "coder-1", "sonnet", "sess-old");
        record_session(root, "42_story_foo", "coder-1", "sonnet", "sess-new");

        let result = lookup_session(root, "42_story_foo", "coder-1", "sonnet");
        assert_eq!(result, Some("sess-new".to_string()));
    }

    #[test]
    fn remove_sessions_for_story_cleans_up() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        std::fs::create_dir_all(root.join(".huskies")).unwrap();

        record_session(root, "42_story_foo", "coder-1", "sonnet", "sess-a");
        record_session(root, "42_story_foo", "coder-2", "opus", "sess-b");
        record_session(root, "99_story_bar", "coder-1", "sonnet", "sess-c");

        remove_sessions_for_story(root, "42_story_foo");

        assert_eq!(
            lookup_session(root, "42_story_foo", "coder-1", "sonnet"),
            None
        );
        assert_eq!(
            lookup_session(root, "42_story_foo", "coder-2", "opus"),
            None
        );
        // Other story is untouched.
        assert_eq!(
            lookup_session(root, "99_story_bar", "coder-1", "sonnet"),
            Some("sess-c".to_string())
        );
    }

    // ── AC1: corrupt/empty store file is handled gracefully ───────────────

    #[test]
    fn corrupt_store_file_returns_empty() {
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path();
        std::fs::create_dir_all(root.join(".huskies")).unwrap();
        std::fs::write(root.join(".huskies/session_store.json"), "NOT JSON").unwrap();

        let result = lookup_session(root, "42_story_foo", "coder-1", "sonnet");
        assert_eq!(result, None, "corrupt store should return None, not panic");
    }
}